Skip to content

Commit 4af6e71

Browse files
waleedlatif1claude
andcommitted
fix(redis): make buffer cleanup atomic and match MINID's inclusive boundary
Deleting a copilot buffer and releasing its reservation were two round trips, so a concurrent append landing between them kept its events stored with its reservation already erased. Both now run in one script, composed from a rendered release fragment the same way the reservation is. `XTRIM MINID upTo` is inclusive — it keeps the entry whose id equals the boundary. Accounting treated that entry as folded, so its bytes stopped counting while they were still in Redis, and a large paste landing exactly on the boundary could leave the trigger disarmed. Both directions now use the same strict/inclusive split: only entries strictly before the boundary arm the trigger, and only those are dropped once a trim removes them. Replaces the tests' `any` casts with a typed accessor, per the repository's TypeScript rule. This immediately caught injected test rooms missing `pendingDeltas`, which made compaction throw into its catch while the assertions still passed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 parent 8e68c31 commit 4af6e71

5 files changed

Lines changed: 126 additions & 55 deletions

File tree

‎apps/realtime/src/handlers/file-doc-store.test.ts‎

Lines changed: 48 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,28 @@ vi.mock('redis', () => ({ createClient: () => makeClient() }))
128128
import { FileDocStore, REDIS_AGENT_ORIGIN, REDIS_ORIGIN } from '@/handlers/file-doc-store'
129129

130130
const REDIS_URL = 'redis://fake'
131+
132+
interface StoreRoomInternals {
133+
lastId: string
134+
pendingDeltas: Array<{ id: string; bytes: number }>
135+
realEdited: boolean
136+
publishes: number
137+
doc: Y.Doc
138+
seededObserved: boolean
139+
}
140+
141+
interface FileDocStoreInternals {
142+
rooms: Map<string, StoreRoomInternals>
143+
appendUpdate(name: string, update: Uint8Array, agent?: boolean): Promise<void>
144+
write: { xTrim: (...args: unknown[]) => Promise<unknown> }
145+
maybeCompact(name: string, force?: boolean): Promise<void>
146+
}
147+
148+
/** Reaches the private state these tests assert on, without `any`. */
149+
function internals(store: object): FileDocStoreInternals {
150+
return store as unknown as FileDocStoreInternals
151+
}
152+
131153
const COMPACT_THRESHOLD_ENTRIES = 400
132154
const NAME = 'workspace-file-doc:file-1'
133155

@@ -339,14 +361,15 @@ describe('FileDocStore', () => {
339361
const a = await newStore()
340362
// This task has integrated only up to entry 400 (all no-ops) — its local doc is empty and lags the
341363
// two peer entries. Inject that lagging room directly (a real edit was integrated → realEdited).
342-
;(a as any).rooms.set(NAME, {
364+
internals(a).rooms.set(NAME, {
343365
doc: new Y.Doc(),
344366
lastId: '400-0',
345367
publishes: 0,
368+
pendingDeltas: [],
346369
seededObserved: true,
347370
realEdited: true,
348371
})
349-
await (a as any).maybeCompact(NAME)
372+
await internals(a).maybeCompact(NAME)
350373

351374
// A fresh catch-up must still reconstruct the peer content — compaction must not have trimmed 401/402.
352375
const doc = new Y.Doc()
@@ -392,11 +415,11 @@ describe('FileDocStore', () => {
392415
const a = await newStore()
393416
const doc = new Y.Doc()
394417
await a.attachRoom(NAME, doc)
395-
const room = (a as any).rooms.get(NAME)
418+
const room = internals(a).rooms.get(NAME)!
396419
expect(room.realEdited).toBe(false)
397420
// Kick off a real (non-agent) append but do NOT await it: realEdited must already be true before the
398421
// xAdd/expire awaits resolve, so any compaction racing on the awaits sees the real edit.
399-
const pending = (a as any).appendUpdate(NAME, updateFor('real user edit'))
422+
const pending = internals(a).appendUpdate(NAME, updateFor('real user edit'))
400423
expect(room.realEdited).toBe(true)
401424
await pending
402425
doc.destroy()
@@ -474,16 +497,16 @@ describe('FileDocStore', () => {
474497
const a = await newStore()
475498
const doc = new Y.Doc()
476499
await a.attachRoom(NAME, doc)
477-
const room = (a as any).rooms.get(NAME)
500+
const room = internals(a).rooms.get(NAME)!
478501
room.pendingDeltas = [{ id: '1-0', bytes: 9 * 1024 * 1024 }]
479502
room.realEdited = true
480503

481-
const write = (a as any).write
504+
const write = internals(a).write
482505
const original = write.xTrim.bind(write)
483506
write.xTrim = async () => {
484507
throw new Error('redis blip')
485508
}
486-
await (a as any).maybeCompact(NAME, true)
509+
await internals(a).maybeCompact(NAME, true)
487510
write.xTrim = original
488511

489512
// A failed fold must not disarm the trigger — otherwise the stream stays oversized until
@@ -496,19 +519,24 @@ describe('FileDocStore', () => {
496519
const a = await newStore()
497520
const doc = new Y.Doc()
498521
await a.attachRoom(NAME, doc)
499-
const room = (a as any).rooms.get(NAME)
522+
const room = internals(a).rooms.get(NAME)!
500523
room.realEdited = true
501-
// The tailer has integrated up to 5-0, so MINID retains 9-0. Its bytes are still in Redis,
502-
// and dropping them would disarm the byte trigger while the stream kept growing.
524+
// The tailer has integrated up to 5-0, so `MINID 5-0` retains both 5-0 (the boundary is
525+
// INCLUSIVE) and 9-0. Their bytes are still in Redis, and dropping them would disarm the
526+
// byte trigger while the stream kept growing.
503527
room.lastId = '5-0'
504528
room.pendingDeltas = [
505529
{ id: '3-0', bytes: 4 * 1024 * 1024 },
530+
{ id: '5-0', bytes: 6 * 1024 * 1024 },
506531
{ id: '9-0', bytes: 7 * 1024 * 1024 },
507532
]
508533

509-
await (a as any).maybeCompact(NAME, true)
534+
await internals(a).maybeCompact(NAME, true)
510535

511-
expect(room.pendingDeltas).toEqual([{ id: '9-0', bytes: 7 * 1024 * 1024 }])
536+
expect(room.pendingDeltas).toEqual([
537+
{ id: '5-0', bytes: 6 * 1024 * 1024 },
538+
{ id: '9-0', bytes: 7 * 1024 * 1024 },
539+
])
512540
doc.destroy()
513541
})
514542

@@ -525,14 +553,15 @@ describe('FileDocStore', () => {
525553
state.backing!.seq = 400
526554

527555
const a = await newStore()
528-
;(a as any).rooms.set(NAME, {
556+
internals(a).rooms.set(NAME, {
529557
doc: agentDoc,
530558
lastId: '400-0',
531559
publishes: 0,
560+
pendingDeltas: [],
532561
seededObserved: true,
533562
realEdited: false,
534563
})
535-
await (a as any).maybeCompact(NAME)
564+
await internals(a).maybeCompact(NAME)
536565

537566
// The snapshot must carry the AGENT marker, NOT the snapshot marker, so a peer catch-up applies it as
538567
// REDIS_AGENT_ORIGIN and never marks the doc edited — the no-persist guarantee survives compaction.
@@ -691,21 +720,23 @@ describe('FileDocStore', () => {
691720
const b = await newStore()
692721
const docA = new Y.Doc()
693722
Y.applyUpdate(docA, peerUpdates[0]) // A integrated up to 401
694-
;(a as any).rooms.set(NAME, {
723+
internals(a).rooms.set(NAME, {
695724
doc: docA,
696725
lastId: '401-0',
697726
publishes: 0,
727+
pendingDeltas: [],
698728
seededObserved: true,
699729
realEdited: true,
700730
})
701-
;(b as any).rooms.set(NAME, {
731+
internals(b).rooms.set(NAME, {
702732
doc: new Y.Doc(),
703733
lastId: '400-0',
704734
publishes: 0,
735+
pendingDeltas: [],
705736
seededObserved: true,
706737
realEdited: true,
707738
})
708-
await Promise.all([(a as any).maybeCompact(NAME), (b as any).maybeCompact(NAME)])
739+
await Promise.all([internals(a).maybeCompact(NAME), internals(b).maybeCompact(NAME)])
709740

710741
const doc = new Y.Doc()
711742
Y.applyUpdate(doc, (await a.getStreamState(NAME))!)

‎apps/realtime/src/handlers/file-doc-store.ts‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -189,15 +189,17 @@ const streamKey = (name: string) => `${STREAM_PREFIX}${name}`
189189
/**
190190
* Unfolded delta bytes a compaction could actually reclaim right now.
191191
*
192-
* Only entries at or before `room.lastId` count: a fold trims to that boundary, so bytes past it
193-
* would survive the trim and re-arm the trigger immediately, forcing a full snapshot append per
194-
* publish that reclaims nothing. They stay in `pendingDeltas` and start counting once the tailer
195-
* has integrated them.
192+
* Only entries strictly before `room.lastId` count. A fold trims with `MINID upTo`, which is
193+
* inclusive, so everything from `upTo` onward survives it; counting those would re-arm the trigger
194+
* the moment a fold finished and force a full snapshot append per publish that reclaims nothing.
195+
* They stay in `pendingDeltas` and start counting once the tailer has moved past them.
196196
*/
197197
function foldableDeltaBytes(room: StoreRoom): number {
198198
let bytes = 0
199199
for (const delta of room.pendingDeltas) {
200-
if (!isAfterStreamId(delta.id, room.lastId)) bytes += delta.bytes
200+
// Strictly before the boundary: MINID is inclusive, so the entry AT `lastId` survives the
201+
// trim and folding cannot reclaim it.
202+
if (isAfterStreamId(room.lastId, delta.id)) bytes += delta.bytes
201203
}
202204
return bytes
203205
}
@@ -835,7 +837,9 @@ export class FileDocStore {
835837
// MINID and its bytes are still in Redis, so it stays counted; dropping it would disarm the
836838
// trigger while the stream kept growing. Done after the trim, so a failed fold changes
837839
// nothing and leaves the trigger armed.
838-
room.pendingDeltas = room.pendingDeltas.filter((delta) => isAfterStreamId(delta.id, upTo))
840+
// `MINID upTo` is inclusive — it keeps the entry whose id EQUALS `upTo`, so that entry's
841+
// bytes are still in Redis and must stay counted. Keeps exactly what survived the trim.
842+
room.pendingDeltas = room.pendingDeltas.filter((delta) => !isAfterStreamId(upTo, delta.id))
839843
} finally {
840844
await this.releaseLock(key, token)
841845
}

‎apps/sim/lib/copilot/request/session/buffer.test.ts‎

Lines changed: 21 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -66,16 +66,29 @@ const createRedisStub = () => {
6666
}),
6767
get: vi.fn().mockImplementation((key: string) => Promise.resolve(values.get(key) ?? null)),
6868
/**
69-
* Stands in for `APPEND_EVENTS_SCRIPT`. It reproduces the script's observable
69+
* Stands in for both Lua scripts, dispatching on the leading `DEL` that only
70+
* `CLEAR_BUFFER_SCRIPT` has. It reproduces their observable
7071
* effects — dedupe, zadd, rank-trim, seq — so the read-path tests still exercise
7172
* real data, and exposes `budgetRefusal` so the refusal branch can be driven
7273
* without reimplementing the budget arithmetic here.
7374
*/
7475
budgetRefusal: null as null | [number, string, number],
7576
eval: vi.fn().mockImplementation((...args: unknown[]) => {
77+
const script = String(args[0])
7678
const numKeys = Number(args[1])
7779
const keys = args.slice(2, 2 + numKeys) as string[]
7880
const argv = args.slice(2 + numKeys) as Array<string | number>
81+
82+
// CLEAR_BUFFER_SCRIPT is the only one that opens with a DEL.
83+
if (script.trimStart().startsWith("redis.call('DEL'")) {
84+
for (const key of keys) {
85+
values.delete(key)
86+
sortedSets.delete(key)
87+
counters.delete(key)
88+
}
89+
return Promise.resolve(1)
90+
}
91+
7992
if (api.budgetRefusal) return Promise.resolve(api.budgetRefusal)
8093

8194
const [eventsKey, seqKey] = keys
@@ -404,18 +417,19 @@ describe('mothership-stream-outbox', () => {
404417
// retry that reuses the same streamId against bytes that no longer exist anywhere.
405418
await clearBuffer('stream-1', 'clear_outbox', { streamId: 'stream-1', userId: 'user-1' })
406419

407-
expect(mockRedis.del).toHaveBeenCalled()
420+
// One script, so a concurrent append cannot land between the delete and the release and
421+
// keep its events stored with its reservation already erased.
408422
const evalCall = mockRedis.eval.mock.calls.at(-1)
409-
expect(evalCall?.[1]).toBe(2)
410-
expect(evalCall?.[2]).toBe('execution:redis-budget:copilot_stream:stream-1')
411-
expect(evalCall?.[3]).toBe('execution:redis-budget:user:user-1')
423+
expect(evalCall?.[1]).toBe(5)
424+
expect(evalCall?.[5]).toBe('execution:redis-budget:copilot_stream:stream-1')
425+
expect(evalCall?.[6]).toBe('execution:redis-budget:user:user-1')
412426
})
413427

414428
it('releases only the owner counter when no user is in scope', async () => {
415429
await clearBuffer('stream-1')
416430

417431
const evalCall = mockRedis.eval.mock.calls.at(-1)
418-
expect(evalCall?.[1]).toBe(1)
419-
expect(evalCall?.[2]).toBe('execution:redis-budget:copilot_stream:stream-1')
432+
expect(evalCall?.[1]).toBe(4)
433+
expect(evalCall?.[5]).toBe('execution:redis-budget:copilot_stream:stream-1')
420434
})
421435
})

‎apps/sim/lib/copilot/request/session/buffer.ts‎

Lines changed: 28 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,9 @@ import {
88
getRedisBudgetLimits,
99
logRedisBudgetRefusal,
1010
parseRedisBudgetRefusal,
11-
REDIS_BUDGET_RELEASE_SCRIPT,
1211
type RedisBudgetRefusal,
1312
renderRedisBudgetLua,
13+
renderRedisBudgetReleaseLua,
1414
} from '@/lib/core/redis/byte-budget.server'
1515
import {
1616
type PersistedStreamEventEnvelope,
@@ -103,6 +103,13 @@ export async function allocateCursor(streamId: string): Promise<{
103103
return { seq, cursor: String(seq) }
104104
}
105105

106+
/** Deletes a stream's buffer and releases its reservation together. KEYS: [events, seq, abort, budget...]. */
107+
const CLEAR_BUFFER_SCRIPT = `
108+
redis.call('DEL', KEYS[1], KEYS[2], KEYS[3])
109+
${renderRedisBudgetReleaseLua(3)}
110+
return 1
111+
`
112+
106113
export async function resetBuffer(streamId: string, scope?: StreamBudgetScope): Promise<void> {
107114
await clearBuffer(streamId, 'reset_outbox', scope)
108115
}
@@ -112,19 +119,27 @@ export async function clearBuffer(
112119
operation = 'clear_outbox',
113120
scope?: StreamBudgetScope
114121
): Promise<void> {
122+
/*
123+
Delete and release in ONE script. The counter outlives the data it accounts for
124+
unless it is released here — these keys are deleted rather than expired, so a retry
125+
reusing the same streamId would be refused against bytes that no longer exist. Doing
126+
it in a second round trip would be its own hole: a concurrent append landing between
127+
the two would keep its events stored with its reservation already erased.
128+
*/
129+
const budgetKeys = getRedisBudgetKeys({
130+
kind: 'copilot_stream',
131+
id: streamId,
132+
...(scope?.userId ? { userId: scope.userId } : {}),
133+
})
115134
await withRedisRetry({ operation, streamId }, async (redis) => {
116-
await redis.del(getEventsKey(streamId), getSeqKey(streamId), getAbortKey(streamId))
117-
/*
118-
The counter outlives the data it accounts for unless it is released here: the keys
119-
above are deleted rather than expired, so without this a retry reusing the same
120-
streamId would be refused against bytes that no longer exist anywhere.
121-
*/
122-
const budgetKeys = getRedisBudgetKeys({
123-
kind: 'copilot_stream',
124-
id: streamId,
125-
...(scope?.userId ? { userId: scope.userId } : {}),
126-
})
127-
await redis.eval(REDIS_BUDGET_RELEASE_SCRIPT, budgetKeys.length, ...budgetKeys)
135+
await redis.eval(
136+
CLEAR_BUFFER_SCRIPT,
137+
3 + budgetKeys.length,
138+
getEventsKey(streamId),
139+
getSeqKey(streamId),
140+
getAbortKey(streamId),
141+
...budgetKeys
142+
)
128143
})
129144
}
130145

‎apps/sim/lib/core/redis/byte-budget.server.ts‎

Lines changed: 19 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -202,25 +202,32 @@ end
202202
}
203203

204204
/**
205-
* Releases an owner's whole reservation when its data is deleted rather than expired.
205+
* Lua that releases an owner's whole reservation, for data that is deleted rather than
206+
* left to expire.
206207
*
207-
* The owner counter is dropped and the user counter credited by exactly what the owner
208-
* held, in one script — crediting the user from a separately read value would let a
209-
* concurrent write land in between and be released twice.
208+
* Rendered into the caller's own script, the same way {@link renderRedisBudgetLua} is, so
209+
* the release commits together with the delete it accounts for. Releasing in a second
210+
* round trip would let a concurrent write land in between and keep its bytes stored with
211+
* its reservation already erased.
210212
*
211-
* KEYS: [ownerKey] or [ownerKey, userKey], as {@link getRedisBudgetKeys} returns them.
213+
* Contract: budget keys are the **last** one or two entries of `KEYS`, in the order
214+
* {@link getRedisBudgetKeys} returns them, and `baseKeyCount` is how many precede them.
212215
*/
213-
export const REDIS_BUDGET_RELEASE_SCRIPT = `
214-
local owner_bytes = tonumber(redis.call('GET', KEYS[1]) or '0')
215-
redis.call('DEL', KEYS[1])
216-
if #KEYS >= 2 and owner_bytes > 0 then
217-
local user_next = redis.call('DECRBY', KEYS[2], owner_bytes)
216+
export function renderRedisBudgetReleaseLua(baseKeyCount: number): string {
217+
const ownerKey = `KEYS[${baseKeyCount + 1}]`
218+
const userKey = `KEYS[${baseKeyCount + 2}]`
219+
220+
return `
221+
local owner_bytes = tonumber(redis.call('GET', ${ownerKey}) or '0')
222+
redis.call('DEL', ${ownerKey})
223+
if #KEYS >= ${baseKeyCount + 2} and owner_bytes > 0 then
224+
local user_next = redis.call('DECRBY', ${userKey}, owner_bytes)
218225
if user_next <= 0 then
219-
redis.call('DEL', KEYS[2])
226+
redis.call('DEL', ${userKey})
220227
end
221228
end
222-
return owner_bytes
223229
`
230+
}
224231

225232
/** Parses the `{0, resource, current}` refusal a guarded script returns. */
226233
export function parseRedisBudgetRefusal(

0 commit comments

Comments
 (0)