diff --git a/apps/realtime/src/handlers/file-doc-store.test.ts b/apps/realtime/src/handlers/file-doc-store.test.ts index f5159c10cbc..1904f1eb31a 100644 --- a/apps/realtime/src/handlers/file-doc-store.test.ts +++ b/apps/realtime/src/handlers/file-doc-store.test.ts @@ -128,6 +128,31 @@ vi.mock('redis', () => ({ createClient: () => makeClient() })) import { FileDocStore, REDIS_AGENT_ORIGIN, REDIS_ORIGIN } from '@/handlers/file-doc-store' const REDIS_URL = 'redis://fake' + +interface StoreRoomInternals { + lastId: string + pendingDeltas: Map + realEdited: boolean + publishes: number + compactRetryAfter: number + doc: Y.Doc + seededObserved: boolean +} + +interface FileDocStoreInternals { + rooms: Map + applyEntry(room: StoreRoomInternals, id: string, message: Record): void + appendUpdate(name: string, update: Uint8Array, agent?: boolean): Promise + write: { xTrim: (...args: unknown[]) => Promise } + maybeCompact(name: string, force?: boolean): Promise +} + +/** Reaches the private state these tests assert on, without `any`. */ +function internals(store: object): FileDocStoreInternals { + return store as unknown as FileDocStoreInternals +} + +const COMPACT_THRESHOLD_ENTRIES = 400 const NAME = 'workspace-file-doc:file-1' function docWithText(text: string): Y.Doc { @@ -338,14 +363,16 @@ describe('FileDocStore', () => { const a = await newStore() // This task has integrated only up to entry 400 (all no-ops) — its local doc is empty and lags the // two peer entries. Inject that lagging room directly (a real edit was integrated → realEdited). - ;(a as any).rooms.set(NAME, { + internals(a).rooms.set(NAME, { doc: new Y.Doc(), lastId: '400-0', publishes: 0, + compactRetryAfter: 0, + pendingDeltas: new Map(), seededObserved: true, realEdited: true, }) - await (a as any).maybeCompact(NAME) + await internals(a).maybeCompact(NAME) // A fresh catch-up must still reconstruct the peer content — compaction must not have trimmed 401/402. const doc = new Y.Doc() @@ -391,16 +418,241 @@ describe('FileDocStore', () => { const a = await newStore() const doc = new Y.Doc() await a.attachRoom(NAME, doc) - const room = (a as any).rooms.get(NAME) + const room = internals(a).rooms.get(NAME)! expect(room.realEdited).toBe(false) // Kick off a real (non-agent) append but do NOT await it: realEdited must already be true before the // xAdd/expire awaits resolve, so any compaction racing on the awaits sees the real edit. - const pending = (a as any).appendUpdate(NAME, updateFor('real user edit')) + const pending = internals(a).appendUpdate(NAME, updateFor('real user edit')) expect(room.realEdited).toBe(true) await pending doc.destroy() }) + it('compacts on appended bytes, before the entry threshold is anywhere near reached', async () => { + const streamKey = `filedoc:stream:${NAME}` + const a = await newStore() + const doc = new Y.Doc() + await a.attachRoom(NAME, doc) + + // A handful of large pastes: far below COMPACT_THRESHOLD entries, far above the byte ceiling. + // Before bytes were counted this stream held tens of megabytes and never compacted. + const updates: Uint8Array[] = [] + doc.on('update', (u: Uint8Array) => updates.push(u)) + for (let i = 0; i < 4; i++) { + doc.getText('body').insert(0, 'x'.repeat(3 * 1024 * 1024)) + } + for (const update of updates) { + await a.publishAndWait(NAME, update) + } + + await vi.waitFor( + () => { + const stream = state.backing!.streams.get(streamKey)! + expect(stream.length).toBeLessThan(COMPACT_THRESHOLD_ENTRIES) + expect(stream.some((entry) => entry.message.s === '1')).toBe(true) + }, + { timeout: 5000 } + ) + + // Compaction must be lossless: the whole document is still reconstructable from what remains. + const rebuilt = new Y.Doc() + Y.applyUpdate(rebuilt, (await a.getStreamState(NAME))!) + expect(rebuilt.getText('body').length).toBe(4 * 3 * 1024 * 1024) + rebuilt.destroy() + doc.destroy() + }) + + it('does not re-compact on every publish once the document itself exceeds the byte ceiling', async () => { + const streamKey = `filedoc:stream:${NAME}` + const a = await newStore() + const doc = new Y.Doc() + await a.attachRoom(NAME, doc) + + const updates: Uint8Array[] = [] + doc.on('update', (u: Uint8Array) => updates.push(u)) + // Grow the document past the byte ceiling so its own snapshot exceeds it, then keep editing. + // Counting the snapshot as appended bytes would leave the threshold permanently breached and + // force a full snapshot append per keystroke — the amplification the threshold exists to stop. + doc.getText('body').insert(0, 'x'.repeat(12 * 1024 * 1024)) + for (let i = 0; i < 30; i++) doc.getText('body').insert(0, 'tiny') + for (const update of updates) { + await a.publishAndWait(NAME, update) + } + await vi.waitFor(() => { + const stream = state.backing!.streams.get(streamKey)! + expect(stream.some((entry) => entry.message.s === '1')).toBe(true) + }) + + const snapshots = state + .backing!.streams.get(streamKey)! + .filter((entry) => entry.message.s === '1').length + expect(snapshots).toBeLessThanOrEqual(2) + + const rebuilt = new Y.Doc() + Y.applyUpdate(rebuilt, (await a.getStreamState(NAME))!) + expect(rebuilt.getText('body').toString().startsWith('tiny')).toBe(true) + expect(rebuilt.getText('body').length).toBe(12 * 1024 * 1024 + 30 * 4) + rebuilt.destroy() + doc.destroy() + }) + + it('keeps the byte trigger armed when compaction fails', async () => { + const a = await newStore() + const doc = new Y.Doc() + await a.attachRoom(NAME, doc) + const room = internals(a).rooms.get(NAME)! + room.pendingDeltas = new Map([['1-0', 9 * 1024 * 1024]]) + room.realEdited = true + + const write = internals(a).write + const original = write.xTrim.bind(write) + write.xTrim = async () => { + throw new Error('redis blip') + } + await internals(a).maybeCompact(NAME, true) + + // A failed fold must not disarm the trigger — otherwise the stream stays oversized until + // this task happens to append another full threshold's worth of deltas. + expect([...room.pendingDeltas]).toEqual([['1-0', 9 * 1024 * 1024]]) + + // But it must not retry immediately either: the snapshot XADD lands before the XTRIM, so a + // persistent trim failure would append a full-document snapshot on every attempt. + const snapshotsAfterFailure = state.backing!.streams.get(`filedoc:stream:${NAME}`)?.length ?? 0 + await internals(a).maybeCompact(NAME, true) + await internals(a).maybeCompact(NAME, true) + expect(state.backing!.streams.get(`filedoc:stream:${NAME}`)?.length ?? 0).toBe( + snapshotsAfterFailure + ) + + write.xTrim = original + doc.destroy() + }) + + it('keeps counting deltas the trim retained because they sit past the fold boundary', async () => { + const a = await newStore() + const doc = new Y.Doc() + await a.attachRoom(NAME, doc) + const room = internals(a).rooms.get(NAME)! + room.realEdited = true + // The tailer has integrated up to 5-0, so `MINID 5-0` retains both 5-0 (the boundary is + // INCLUSIVE) and 9-0. Their bytes are still in Redis, and dropping them would disarm the + // byte trigger while the stream kept growing. + room.lastId = '5-0' + room.pendingDeltas = new Map([ + ['3-0', 4 * 1024 * 1024], + ['5-0', 6 * 1024 * 1024], + ['9-0', 7 * 1024 * 1024], + ]) + + await internals(a).maybeCompact(NAME, true) + + expect([...room.pendingDeltas]).toEqual([ + ['5-0', 6 * 1024 * 1024], + ['9-0', 7 * 1024 * 1024], + ]) + doc.destroy() + }) + + it('adopts accounting for a stream it takes over, and folds it if already over the ceiling', async () => { + const streamKey = `filedoc:stream:${NAME}` + // A stream left behind by a previous task: two entries, so far under the entry threshold, and + // far over the byte ceiling. A fresh room starting from an empty ledger would never fold it, + // while its own heartbeat kept refreshing the TTL. + const seedDoc = new Y.Doc() + const updates: Uint8Array[] = [] + seedDoc.on('update', (u: Uint8Array) => updates.push(u)) + seedDoc.getText('body').insert(0, 'x'.repeat(9 * 1024 * 1024)) + seedDoc.getText('body').insert(0, 'tail') + state.backing!.streams.set( + streamKey, + updates.map((update, index) => ({ + id: `${index + 1}-0`, + message: { u: Buffer.from(update).toString('base64') }, + })) + ) + state.backing!.seq = updates.length + + const a = await newStore() + const doc = new Y.Doc() + await a.attachRoom(NAME, doc) + + // Either marker counts as a fold: this room only ever replayed entries, so it never observed + // a real edit and its snapshot is stamped as an agent frame (the no-persist guarantee). + await vi.waitFor(() => { + const stream = state.backing!.streams.get(streamKey)! + expect(stream.some((entry) => entry.message.s === '1' || entry.message.a === '1')).toBe(true) + }) + + // Lossless: the adopted content survives the fold it triggered. + const rebuilt = new Y.Doc() + Y.applyUpdate(rebuilt, (await a.getStreamState(NAME))!) + expect(rebuilt.getText('body').length).toBe(9 * 1024 * 1024 + 4) + rebuilt.destroy() + doc.destroy() + seedDoc.destroy() + }) + + it('counts agent preview deltas, which share a marker with an agent-only snapshot', async () => { + const streamKey = `filedoc:stream:${NAME}` + // Agent preview frames are the LARGE ones — a copilot file edit re-serialising a document is + // what filled Redis. They carry the same marker as a fold of an agent-only stream, so keying + // exclusion on that marker would drop exactly the payloads this bound exists for. + const seedDoc = new Y.Doc() + const updates: Uint8Array[] = [] + seedDoc.on('update', (u: Uint8Array) => updates.push(u)) + seedDoc.getText('body').insert(0, 'x'.repeat(9 * 1024 * 1024)) + seedDoc.getText('body').insert(0, 'tail') + state.backing!.streams.set( + streamKey, + updates.map((update, index) => ({ + id: `${index + 1}-0`, + message: { u: Buffer.from(update).toString('base64'), a: '1' }, + })) + ) + state.backing!.seq = updates.length + + const a = await newStore() + const doc = new Y.Doc() + await a.attachRoom(NAME, doc) + + await vi.waitFor(() => { + const stream = state.backing!.streams.get(streamKey)! + expect(stream.some((entry) => entry.message.c === '1')).toBe(true) + }) + + const rebuilt = new Y.Doc() + Y.applyUpdate(rebuilt, (await a.getStreamState(NAME))!) + expect(rebuilt.getText('body').length).toBe(9 * 1024 * 1024 + 4) + rebuilt.destroy() + doc.destroy() + seedDoc.destroy() + }) + + it("never counts a fold's own output, so a large document cannot arm the trigger against itself", async () => { + const a = await newStore() + const doc = new Y.Doc() + await a.attachRoom(NAME, doc) + const room = internals(a).rooms.get(NAME)! + + internals(a).applyEntry(room, '7-0', { u: 'x'.repeat(9 * 1024 * 1024), a: '1', c: '1' }) + + expect(room.pendingDeltas.has('7-0')).toBe(false) + doc.destroy() + }) + + it('counts a delta published by a peer task, which this room only ever tails', async () => { + const a = await newStore() + const doc = new Y.Doc() + await a.attachRoom(NAME, doc) + const room = internals(a).rooms.get(NAME)! + + // Never published locally, so publish-side accounting would miss it entirely. + internals(a).applyEntry(room, '4-0', { u: 'x'.repeat(1024) }) + + expect(room.pendingDeltas.get('4-0')).toBe(1024) + doc.destroy() + }) + it('stamps a compaction snapshot of an agent-ONLY stream as an agent frame (never persisted)', async () => { const streamKey = `filedoc:stream:${NAME}` const noop = Buffer.from(Y.encodeStateAsUpdate(new Y.Doc())).toString('base64') @@ -414,14 +666,16 @@ describe('FileDocStore', () => { state.backing!.seq = 400 const a = await newStore() - ;(a as any).rooms.set(NAME, { + internals(a).rooms.set(NAME, { doc: agentDoc, lastId: '400-0', publishes: 0, + compactRetryAfter: 0, + pendingDeltas: new Map(), seededObserved: true, realEdited: false, }) - await (a as any).maybeCompact(NAME) + await internals(a).maybeCompact(NAME) // The snapshot must carry the AGENT marker, NOT the snapshot marker, so a peer catch-up applies it as // REDIS_AGENT_ORIGIN and never marks the doc edited — the no-persist guarantee survives compaction. @@ -580,21 +834,25 @@ describe('FileDocStore', () => { const b = await newStore() const docA = new Y.Doc() Y.applyUpdate(docA, peerUpdates[0]) // A integrated up to 401 - ;(a as any).rooms.set(NAME, { + internals(a).rooms.set(NAME, { doc: docA, lastId: '401-0', publishes: 0, + compactRetryAfter: 0, + pendingDeltas: new Map(), seededObserved: true, realEdited: true, }) - ;(b as any).rooms.set(NAME, { + internals(b).rooms.set(NAME, { doc: new Y.Doc(), lastId: '400-0', publishes: 0, + compactRetryAfter: 0, + pendingDeltas: new Map(), seededObserved: true, realEdited: true, }) - await Promise.all([(a as any).maybeCompact(NAME), (b as any).maybeCompact(NAME)]) + await Promise.all([internals(a).maybeCompact(NAME), internals(b).maybeCompact(NAME)]) const doc = new Y.Doc() Y.applyUpdate(doc, (await a.getStreamState(NAME))!) diff --git a/apps/realtime/src/handlers/file-doc-store.ts b/apps/realtime/src/handlers/file-doc-store.ts index 537f7f4db12..43332963feb 100644 --- a/apps/realtime/src/handlers/file-doc-store.ts +++ b/apps/realtime/src/handlers/file-doc-store.ts @@ -123,6 +123,20 @@ const SNAPSHOT_FIELD = 's' /** Marks a stream entry as an AGENT-STREAMED preview frame, so the tailer applies it with * {@link REDIS_AGENT_ORIGIN} (never marks the doc edited). Present only on agent-frame entries. */ const AGENT_FIELD = 'a' +/** + * Marks a stream entry as the OUTPUT of a compaction, for byte accounting only. + * + * {@link SNAPSHOT_FIELD} cannot serve this purpose: a fold of an agent-only stream is stamped + * {@link AGENT_FIELD} instead, so it is indistinguishable from an ordinary agent preview frame — + * and those are the large ones. Excluding both markers would drop preview deltas from accounting; + * excluding neither would count a snapshot as something a fold can reclaim, arming the trigger + * against its own output. A separate field settles it without touching origin selection, which + * must keep treating an agent-only fold as an agent frame to preserve the no-persist guarantee. + * + * Entries written before this field existed carry no marker and are counted as deltas. That + * over-arms by at most one fold, which then trims them. + */ +const COMPACTION_FIELD = 'c' /** Sentinel token a DISABLED store returns from a lock acquire, so single-replica callers proceed * without special-casing; {@link FileDocStore.releaseLock} treats it as a no-op. Not a real UUID, so it @@ -140,11 +154,37 @@ const IDLE_POLL_MS = 250 const READ_COUNT = 200 /** Compact a stream once it exceeds this many entries (snapshot + trim). */ const COMPACT_THRESHOLD = 400 +/** + * Compact a stream once its appended deltas exceed this many bytes, whichever comes first. + * + * The entry threshold alone bounds how many entries a stream holds and says nothing about + * how large each one is: one pasted block is a single entry carrying megabytes, so a stream + * can sit at a few dozen entries and hundreds of megabytes and never reach + * {@link COMPACT_THRESHOLD} before its TTL. Folding by bytes as well keeps a stream's cost + * proportional to its document rather than to the size of the edits that produced it. + * + * Compaction is the only safe way to shrink one of these streams: a task attaching later + * replays every entry to rebuild the doc, so dropping the oldest entries — what a native + * `MAXLEN` retention bound would do — loses edits outright. A snapshot folds them first. + * + * Measured over deltas appended since the last fold, never over the resulting snapshot, so a + * stream settles at roughly one document snapshot plus this much churn. + */ +const COMPACT_BYTES_THRESHOLD = 8 * 1024 * 1024 /** Check whether compaction is due only every Nth local publish, to avoid an XLEN per keystroke. */ const COMPACT_CHECK_EVERY = 64 /** Compaction critical section (snapshot + xAdd + xTrim) is fast; a generous TTL covers a slow Redis * round-trip without risking expiry mid-compact. Released via compare-and-delete regardless. */ const COMPACT_LOCK_TTL_MS = 10_000 +/** + * Quiet period after a failed fold before another may be forced. + * + * A failed fold deliberately leaves the trigger armed so the bytes are not forgotten, but the + * snapshot `XADD` lands before the `XTRIM` — so if the trim is what failed, retrying immediately + * appends another full-document snapshot each time, turning a Redis blip into exactly the write + * amplification the threshold exists to prevent. The entry-count path is unaffected. + */ +const COMPACT_RETRY_COOLDOWN_MS = 30_000 /** Retry a failed stream append this many times before giving up, so a transient Redis blip doesn't * silently drop an edit from the shared log (which no peer would then ever see). */ const PUBLISH_MAX_RETRIES = 3 @@ -169,6 +209,24 @@ const READER_ERROR_LOG_EVERY = 20 const streamKey = (name: string) => `${STREAM_PREFIX}${name}` +/** + * Unfolded delta bytes a compaction could actually reclaim right now. + * + * Only entries strictly before `room.lastId` count. A fold trims with `MINID upTo`, which is + * inclusive, so everything from `upTo` onward survives it; counting those would re-arm the trigger + * the moment a fold finished and force a full snapshot append per publish that reclaims nothing. + * They stay in `pendingDeltas` and start counting once the tailer has moved past them. + */ +function foldableDeltaBytes(room: StoreRoom): number { + let bytes = 0 + for (const [id, deltaBytes] of room.pendingDeltas) { + // Strictly before the boundary: MINID is inclusive, so the entry AT `lastId` survives the + // trim and folding cannot reclaim it. + if (isAfterStreamId(room.lastId, id)) bytes += deltaBytes + } + return bytes +} + /** * Decode one stream entry's base64 Yjs update and apply it to `doc`. A malformed entry is logged and * SKIPPED — never thrown — so one bad frame can neither wedge the tailer nor abort a headless @@ -218,6 +276,25 @@ interface StoreRoom { lastId: string /** Local publish count, to pace compaction checks. */ publishes: number + /** Epoch ms before which no forced fold is attempted, after one failed. */ + compactRetryAfter: number + /** + * Unfolded delta bytes in the shared stream, by entry id. + * + * Recorded in {@link FileDocStore.applyEntry}, so it covers EVERY entry this room's tailer + * observes — this task's own appends, a peer task's, and one published with no room attached + * anywhere. Accounting on publish instead would see only this task's writes. + * + * Keyed by id rather than summed, because a fold trims to `room.lastId` and retains anything + * from that boundary on. Those bytes are still in Redis, so dropping them would disarm the + * trigger while the stream kept growing; entries are removed only once an `XTRIM` provably + * removed them. + * + * Excludes what a fold produces (see {@link COMPACTION_FIELD}) — a snapshot is a function of + * document size rather than edit volume, and counting one would make a large document breach + * the threshold permanently. + */ + pendingDeltas: Map /** Set once the doc has been observed seeded, so the seed transition itself is never mistaken for an * edit (mirrors the relay's `seededObserved`). */ seededObserved: boolean @@ -299,6 +376,8 @@ export class FileDocStore { doc, lastId: '0', publishes: 0, + compactRetryAfter: 0, + pendingDeltas: new Map(), seededObserved: false, realEdited: false, } @@ -331,6 +410,9 @@ export class FileDocStore { this.applyEntry(room, entry.id, entry.message) } await this.write.expire(streamKey(name), STREAM_TTL_SEC) + // A stream taken over may already be past the ceiling, and nothing else re-checks until the + // next local publish — which a read-only participant never makes. + if (foldableDeltaBytes(room) >= COMPACT_BYTES_THRESHOLD) void this.maybeCompact(name, true) } catch (error) { logger.warn(`FileDocStore catch-up failed for ${name}`, { error: getErrorMessage(error) }) } @@ -379,7 +461,14 @@ export class FileDocStore { } await this.write.expire(streamKey(name), STREAM_TTL_SEC).catch(() => {}) const room = this.rooms.get(name) - if (room && ++room.publishes % COMPACT_CHECK_EVERY === 0) void this.maybeCompact(name) + if (!room) return + // Bytes are checked every publish: one entry can cross the ceiling on its own, so pacing this + // check the way the entry count is paced would let a stream sit far over the ceiling for up to + // COMPACT_CHECK_EVERY more appends. The check itself is a local sum over unfolded entries. + const overBytes = foldableDeltaBytes(room) >= COMPACT_BYTES_THRESHOLD + if (overBytes || ++room.publishes % COMPACT_CHECK_EVERY === 0) { + void this.maybeCompact(name, overBytes) + } } /** @@ -639,6 +728,12 @@ export class FileDocStore { private applyEntry(room: StoreRoom, id: string, message: Record): void { room.lastId = id + // Account for every entry the tailer sees, whoever wrote it — this is the only point that + // observes peer and roomless appends. A fold's own output is excluded so it cannot arm the + // trigger against itself. + if (!message[COMPACTION_FIELD]) { + room.pendingDeltas.set(id, message[UPDATE_FIELD]?.length ?? 0) + } // A compaction snapshot folds seed + edits into one frame; stamp it so the relay's edit-tracker // treats a fresh catch-up from it as edited (a snapshot only exists once real edits accumulated). An // agent-streamed preview frame is stamped separately so the tracker NEVER marks it edited. @@ -691,6 +786,11 @@ export class FileDocStore { // but wasteful re-delivery). The new room caught itself up via xRange already. if (!room || room !== snapshot.get(name)) continue for (const entry of stream.messages) this.applyEntry(room, entry.id, entry.message) + // Foldability is decided by `lastId`, which only the tailer advances — so a burst of + // large edits followed by silence would otherwise sit unfolded until the next publish + // happened to re-evaluate the trigger. Re-check it where the boundary actually moved. + if (foldableDeltaBytes(room) >= COMPACT_BYTES_THRESHOLD) + void this.maybeCompact(name, true) } } catch (error) { if (!this.running) break @@ -734,12 +834,13 @@ export class FileDocStore { * only one task compacts a given stream at a time (concurrent snapshot+trim would race). Trims only up * to what the snapshot provably contains — never un-integrated peer entries (see below). */ - private async maybeCompact(name: string): Promise { + private async maybeCompact(name: string, force = false): Promise { if (!this.write) return const room = this.rooms.get(name) if (!room) return try { - if ((await this.write.xLen(streamKey(name))) < COMPACT_THRESHOLD) return + if (force && Date.now() < room.compactRetryAfter) return + if (!force && (await this.write.xLen(streamKey(name))) < COMPACT_THRESHOLD) return const key = `${COMPACT_LOCK_PREFIX}${name}` const token = await this.acquireLock(key, COMPACT_LOCK_TTL_MS) if (!token) return @@ -760,14 +861,22 @@ export class FileDocStore { await this.write.xAdd(streamKey(name), '*', { [UPDATE_FIELD]: snapshot, [marker]: '1', + [COMPACTION_FIELD]: '1', }) // MINID keeps entries with id >= upTo: the snapshot, any un-integrated peer entries, and // `upTo` itself (redundant with the snapshot, harmless); it drops only the folded older deltas. await this.write.xTrim(streamKey(name), 'MINID', upTo) + // Drop exactly what the trim removed, which `MINID upTo` being INCLUSIVE makes `id < upTo` + // — the entry at the boundary survives, and its bytes are still in Redis. Run after the + // trim, so a failed fold leaves the ledger intact and the trigger armed. + for (const id of room.pendingDeltas.keys()) { + if (isAfterStreamId(upTo, id)) room.pendingDeltas.delete(id) + } } finally { await this.releaseLock(key, token) } } catch (error) { + room.compactRetryAfter = Date.now() + COMPACT_RETRY_COOLDOWN_MS logger.warn(`FileDocStore compaction failed for ${name}`, { error: getErrorMessage(error) }) } } diff --git a/apps/sim/lib/copilot/request/lifecycle/start.ts b/apps/sim/lib/copilot/request/lifecycle/start.ts index 2d7943f74ad..d9f477a8404 100644 --- a/apps/sim/lib/copilot/request/lifecycle/start.ts +++ b/apps/sim/lib/copilot/request/lifecycle/start.ts @@ -116,7 +116,7 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS const abortController = new AbortController() registerActiveStream(streamId, abortController) - const publisher = new StreamWriter({ streamId, chatId, requestId }) + const publisher = new StreamWriter({ streamId, chatId, requestId, userId }) // Declared at function scope (same rationale as `cancelReason` below) so the // leak backstop in the orchestration's outer finally can always reach them: diff --git a/apps/sim/lib/copilot/request/session/buffer.test.ts b/apps/sim/lib/copilot/request/session/buffer.test.ts index e0fec738227..a0807556995 100644 --- a/apps/sim/lib/copilot/request/session/buffer.test.ts +++ b/apps/sim/lib/copilot/request/session/buffer.test.ts @@ -9,6 +9,7 @@ import { MothershipStreamV1TextChannel, } from '@/lib/copilot/generated/mothership-stream-v1' import { createEvent } from '@/lib/copilot/request/session/event' +import { getRedisBudgetLimits } from '@/lib/core/redis/byte-budget.server' type StoredEnvelope = { score: number @@ -64,6 +65,34 @@ const createRedisStub = () => { return Promise.resolve('OK') }), get: vi.fn().mockImplementation((key: string) => Promise.resolve(values.get(key) ?? null)), + /** + * Stands in for `APPEND_EVENTS_SCRIPT`. It reproduces the script's observable + * effects — dedupe, zadd, rank-trim, seq — so the read-path tests still exercise + * real data, and exposes `budgetRefusal` so the refusal branch can be driven + * without reimplementing the budget arithmetic here. + */ + budgetRefusal: null as null | [number, string, number], + eval: vi.fn().mockImplementation((...args: unknown[]) => { + const numKeys = Number(args[1]) + const keys = args.slice(2, 2 + numKeys) as string[] + const argv = args.slice(2 + numKeys) as Array + + if (api.budgetRefusal) return Promise.resolve(api.budgetRefusal) + + const [eventsKey, seqKey] = keys + const eventLimit = Number(argv[1]) + const lastSeq = String(argv[5]) + const entries = sortedSets.get(eventsKey) ?? [] + for (let i = 6; i < argv.length; i += 2) { + const score = Number(argv[i]) + const value = String(argv[i + 1]) + if (!entries.some((entry) => entry.value === value)) entries.push({ score, value }) + } + entries.sort((a, b) => a.score - b.score) + sortedSets.set(eventsKey, entries.slice(Math.max(0, entries.length - eventLimit))) + values.set(seqKey, lastSeq) + return Promise.resolve([1]) + }), pipeline: vi.fn().mockImplementation(() => { const operations: Array<() => Promise> = [] const pipeline = { @@ -103,11 +132,24 @@ let mockRedis: ReturnType import { allocateCursor, appendEvent, + appendEvents, clearBuffer, readEvents, scheduleBufferCleanup, } from '@/lib/copilot/request/session/buffer' +async function makeEnvelope(text: string) { + const cursor = await allocateCursor('stream-1') + return createEvent({ + streamId: 'stream-1', + cursor: cursor.cursor, + seq: cursor.seq, + requestId: 'req-1', + type: MothershipStreamV1EventType.text, + payload: { channel: MothershipStreamV1TextChannel.assistant, text }, + }) +} + describe('mothership-stream-outbox', () => { beforeEach(() => { mockRedis = createRedisStub() @@ -161,11 +203,84 @@ describe('mothership-stream-outbox', () => { }) ) - expect(mockRedis.zremrangebyrank).toHaveBeenCalledWith( - 'mothership_stream:stream-1:events', - 0, - -100_001 - ) + // KEYS: [events, seq, budgetOwner]; ARGV follows. + const [, numKeys, eventsKey, seqKey, ownerKey, ...argv] = mockRedis.eval.mock.calls[0] + expect(numKeys).toBe(3) + expect(eventsKey).toBe('mothership_stream:stream-1:events') + expect(seqKey).toBe('mothership_stream:stream-1:seq') + expect(ownerKey).toBe('execution:redis-budget:copilot_stream:stream-1') + // ARGV: [ttl, eventLimit, ownerLimit, userLimit, budgetTtl, lastSeq, ...zaddArgs] + expect(argv[1]).toBe(100_000) + }) + + /** + * The stream's replay copy is charged to a budget, and a refusal is reported rather + * than thrown: `flush()` rethrows what it is handed, and that throw reaches the + * error-path finalize, which would reject a response stream whose bytes the user + * already received. + */ + it('reports a budget refusal instead of throwing', async () => { + const cursor = await allocateCursor('stream-1') + mockRedis.budgetRefusal = [0, 'owner_redis_bytes', 40_000_000] + + const result = await appendEvents([ + createEvent({ + streamId: 'stream-1', + cursor: cursor.cursor, + seq: cursor.seq, + requestId: 'req-1', + type: MothershipStreamV1EventType.text, + payload: { channel: MothershipStreamV1TextChannel.assistant, text: 'hello' }, + }), + ]) + + expect(result.persisted).toBe(false) + if (!result.persisted) { + expect(result.refusal.resource).toBe('owner_redis_bytes') + expect(result.refusal.currentBytes).toBe(40_000_000) + } + }) + + it('refuses a batch past the single-write ceiling without reaching Redis', async () => { + const cursor = await allocateCursor('stream-1') + + const result = await appendEvents([ + createEvent({ + streamId: 'stream-1', + cursor: cursor.cursor, + seq: cursor.seq, + requestId: 'req-1', + type: MothershipStreamV1EventType.text, + payload: { + channel: MothershipStreamV1TextChannel.assistant, + text: 'x'.repeat(2 * 1024 * 1024), + }, + }), + ]) + + expect(result.persisted).toBe(false) + expect(mockRedis.eval).not.toHaveBeenCalled() + }) + + it('charges the user ceiling only when a user is in scope', async () => { + const cursor = await allocateCursor('stream-1') + const envelope = createEvent({ + streamId: 'stream-1', + cursor: cursor.cursor, + seq: cursor.seq, + requestId: 'req-1', + type: MothershipStreamV1EventType.text, + payload: { channel: MothershipStreamV1TextChannel.assistant, text: 'hello' }, + }) + + await appendEvents([envelope], { streamId: 'stream-1' }) + expect(mockRedis.eval.mock.calls[0][1]).toBe(3) + expect(mockRedis.eval.mock.calls[0][4]).toBe('execution:redis-budget:copilot_stream:stream-1') + + mockRedis.eval.mockClear() + await appendEvents([envelope], { streamId: 'stream-1', userId: 'user-1' }) + expect(mockRedis.eval.mock.calls[0][1]).toBe(4) + expect(mockRedis.eval.mock.calls[0][5]).toBe('execution:redis-budget:user:user-1') }) it('clears persisted stream state during teardown cleanup', async () => { @@ -245,4 +360,65 @@ describe('mothership-stream-outbox', () => { expect(replayed).toHaveLength(1) expect(replayed[0]?.payload.text).toBe('hello') }) + + it('splits an oversized batch instead of refusing it', async () => { + const limits = getRedisBudgetLimits('copilot_stream') + // Individually writable frames that collectively exceed the per-write ceiling. Refusing the + // whole batch would stop replay persistence for the rest of the stream over a batching artefact. + const envelopes = await Promise.all( + Array.from({ length: 3 }, () => + makeEnvelope('x'.repeat(Math.floor(limits.maxSingleWriteBytes * 0.45))) + ) + ) + + const result = await appendEvents(envelopes, { streamId: 'stream-1' }) + + expect(result.persisted).toBe(true) + expect(mockRedis.eval).toHaveBeenCalledTimes(2) + }) + + it('refuses a single frame that can never land, without splitting', async () => { + const limits = getRedisBudgetLimits('copilot_stream') + const oversized = await makeEnvelope('x'.repeat(limits.maxSingleWriteBytes + 10)) + const result = await appendEvents([oversized], { streamId: 'stream-1' }) + + expect(result.persisted).toBe(false) + expect(mockRedis.eval).not.toHaveBeenCalled() + }) + + it('measures the ceiling in UTF-8 bytes, not UTF-16 units', async () => { + const limits = getRedisBudgetLimits('copilot_stream') + // Each astral char is 2 UTF-16 units but 4 UTF-8 bytes, so `String.length` under-reports by 2x + // and would call this batch writable when Redis will not. + const chars = Math.floor(limits.maxSingleWriteBytes / 3) + const astral = await makeEnvelope('\u{1D306}'.repeat(chars)) + expect(JSON.stringify(astral).length).toBeLessThan(limits.maxSingleWriteBytes) + + const result = await appendEvents([astral], { streamId: 'stream-1' }) + + expect(result.persisted).toBe(false) + expect(mockRedis.eval).not.toHaveBeenCalled() + }) + + it('drops the owner counter together with the buffer it accounts for', async () => { + // The buffer keys are deleted rather than expired, so a counter left behind would refuse a + // retry that reuses the same streamId against bytes that no longer exist anywhere. One + // script, so a concurrent append cannot land between the delete and the release and keep + // its events stored with its reservation already erased. + await clearBuffer('stream-1') + + // One variadic DEL: a single atomic command, so no script is needed for the counter to go + // with the data it accounts for. + expect(mockRedis.del).toHaveBeenCalledTimes(1) + expect(mockRedis.del.mock.calls[0]).toContain('execution:redis-budget:copilot_stream:stream-1') + }) + + it('never touches the shared user counter when clearing a buffer', async () => { + // An owner id is not proof of who wrote the bytes, so crediting the user counter here would + // let anyone who can name a stream decrement a ceiling they never charged. + await clearBuffer('stream-1') + + const keys = mockRedis.del.mock.calls[0] as string[] + expect(keys.some((key) => key.includes('redis-budget:user:'))).toBe(false) + }) }) diff --git a/apps/sim/lib/copilot/request/session/buffer.ts b/apps/sim/lib/copilot/request/session/buffer.ts index 871bdeb8d25..f928a2971f2 100644 --- a/apps/sim/lib/copilot/request/session/buffer.ts +++ b/apps/sim/lib/copilot/request/session/buffer.ts @@ -3,6 +3,14 @@ import { toError } from '@sim/utils/errors' import { sleep } from '@sim/utils/helpers' import { env, envNumber } from '@/lib/core/config/env' import { getRedisClient } from '@/lib/core/config/redis' +import { + getRedisBudgetKeys, + getRedisBudgetLimits, + logRedisBudgetRefusal, + parseRedisBudgetRefusal, + type RedisBudgetRefusal, + renderRedisBudgetLua, +} from '@/lib/core/redis/byte-budget.server' import { type PersistedStreamEventEnvelope, parsePersistedStreamEventEnvelopeJson, @@ -99,8 +107,27 @@ export async function resetBuffer(streamId: string): Promise { } export async function clearBuffer(streamId: string, operation = 'clear_outbox'): Promise { + /* + The owner counter is deleted WITH the data it accounts for. These keys are deleted rather + than expired, so a counter left behind would refuse a retry reusing the same streamId + against bytes that no longer exist; dropping it in a second round trip would be its own + hole, since a concurrent append landing between the two would keep its events stored with + its reservation already erased. One variadic DEL is a single atomic command, so no script + is needed to get that. + + The shared user counter is deliberately untouched: an owner id is not proof of who wrote + the bytes, so crediting it here would let anyone able to name a stream decrement a ceiling + they never charged — and a counter driven down grants writes rather than denying them. Its + fixed window settles it instead, over-counting in the safe direction meanwhile. + */ + const [ownerBudgetKey] = getRedisBudgetKeys({ kind: 'copilot_stream', id: streamId }) await withRedisRetry({ operation, streamId }, async (redis) => { - await redis.del(getEventsKey(streamId), getSeqKey(streamId), getAbortKey(streamId)) + await redis.del( + getEventsKey(streamId), + getSeqKey(streamId), + getAbortKey(streamId), + ownerBudgetKey + ) }) } @@ -125,38 +152,197 @@ export async function scheduleBufferCleanup( } } +/** + * Appends a batch, trims the ring, refreshes both TTLs and charges the net bytes to + * the stream's budget — in one script, so the reservation and the write it pays for + * commit together. + * + * Entries already present are skipped when counting, which makes the script + * idempotent: `withRedisRetry` may run it up to three times, and a retry after a + * partial failure must not charge the same bytes twice. + * + * KEYS: [events, seq, budgetOwner, budgetUser?] + * ARGV: [ttlSeconds, eventLimit, ownerLimit, userLimit, budgetTtlSeconds, lastSeq, + * score, member, ...] + * Returns {1} on success, or {0, resource, currentBytes} when the budget refuses. + */ +const APPEND_EVENTS_SCRIPT = ` +local ttl_seconds = tonumber(ARGV[1]) +local event_limit = tonumber(ARGV[2]) +local owner_limit = tonumber(ARGV[3]) +local user_limit = tonumber(ARGV[4]) +local budget_ttl_seconds = tonumber(ARGV[5]) +local last_seq = ARGV[6] + +local new_count = 0 +local new_bytes = 0 +local new_members = {} +for i = 7, #ARGV, 2 do + local member = ARGV[i + 1] + if not redis.call('ZSCORE', KEYS[1], member) then + new_count = new_count + 1 + new_bytes = new_bytes + string.len(member) + table.insert(new_members, member) + end +end + +local current_count = redis.call('ZCARD', KEYS[1]) +local prune_count = current_count + new_count - event_limit +if prune_count < 0 then + prune_count = 0 +end +local existing_prune_count = math.min(prune_count, current_count) +local pruned_bytes = 0 +if existing_prune_count > 0 then + local pruned = redis.call('ZRANGE', KEYS[1], 0, existing_prune_count - 1) + for _, member in ipairs(pruned) do + pruned_bytes = pruned_bytes + string.len(member) + end +end +for i = 1, prune_count - existing_prune_count do + local member = new_members[i] + if member then + pruned_bytes = pruned_bytes + string.len(member) + end +end + +local net_bytes = new_bytes - pruned_bytes +${renderRedisBudgetLua(2)} + +for i = 7, #ARGV, 2 do + redis.call('ZADD', KEYS[1], ARGV[i], ARGV[i + 1]) +end +redis.call('ZREMRANGEBYRANK', KEYS[1], 0, -event_limit - 1) +redis.call('EXPIRE', KEYS[1], ttl_seconds) +redis.call('SET', KEYS[2], last_seq, 'EX', ttl_seconds) +return {1} +` + +/** What a stream is charged against. `userId` adds the cross-stream user ceiling. */ +export interface StreamBudgetScope { + streamId: string + userId?: string +} + +export type AppendEventsResult = + | { persisted: true } + | { persisted: false; refusal: RedisBudgetRefusal } + +/** + * Persists a batch for replay. + * + * A refusal is returned, never thrown. A throw here reaches + * `finalizeStream`'s second flush, which runs inside the error handler and so + * escapes to reject the response stream — a stream that has already delivered every + * byte to the user would end in an error because its *replay copy* did not fit. + * Refusing to persist costs a resume; throwing costs the turn. + */ export async function appendEvents( - envelopes: PersistedStreamEventEnvelope[] -): Promise { + envelopes: PersistedStreamEventEnvelope[], + scope?: StreamBudgetScope +): Promise { if (envelopes.length === 0) { - return envelopes + return { persisted: true } } - const streamId = envelopes[0].stream.streamId + const streamId = scope?.streamId ?? envelopes[0].stream.streamId const config = getStreamConfig() + const limits = getRedisBudgetLimits('copilot_stream') + const budgetScope = { + kind: 'copilot_stream' as const, + id: streamId, + ...(scope?.userId ? { userId: scope.userId } : {}), + } + const budgetKeys = getRedisBudgetKeys(budgetScope) + /* + A counter must never expire before the data it accounts for: the next write would then + see zero reserved and let the stream grow by another full ceiling. `COPILOT_STREAM_TTL_SECONDS` + is configurable and defaults to exactly the budget window, so raising it would otherwise + break that invariant silently. + */ + const budgetTtlSeconds = Math.max(limits.ttlSeconds, config.ttlSeconds) + + /* + Redis measures a member in UTF-8 bytes, so the ceiling has to be measured the same + way — `String.length` counts UTF-16 units and under-reports every non-ASCII frame, + which would let a batch past a check the Lua then applies differently. + */ + const members = envelopes.map((envelope) => { + const member = JSON.stringify(envelope) + return { seq: envelope.seq, member, bytes: Buffer.byteLength(member, 'utf8') } + }) + + /* + Split on the per-write ceiling rather than refusing the whole batch: a flush carries + whatever accumulated since the last one, so an ordinary run of large frames can exceed + the ceiling collectively while every frame is individually writable. Refusing that + batch would stop replay persistence for the rest of the stream over a batching + artefact. Chunks are written in sequence order, so the stored cursor stays monotonic. + */ + const chunks: Array<{ members: typeof members; bytes: number }> = [] + for (const entry of members) { + const last = chunks[chunks.length - 1] + if (!last || last.bytes + entry.bytes > limits.maxSingleWriteBytes) { + chunks.push({ members: [entry], bytes: entry.bytes }) + } else { + last.members.push(entry) + last.bytes += entry.bytes + } + } + + for (const chunk of chunks) { + /* + A single frame past the ceiling can never land, and retrying it would stall every + later batch behind it. Refuse it the same way the budget would. + */ + if (chunk.bytes > limits.maxSingleWriteBytes) { + const refusal: RedisBudgetRefusal = { + resource: 'owner_redis_bytes', + currentBytes: 0, + limitBytes: limits.maxSingleWriteBytes, + attemptedBytes: chunk.bytes, + } + logRedisBudgetRefusal(refusal, { operation: 'append_event', scope: budgetScope, logger }) + return { persisted: false, refusal } + } - await withRedisRetry({ operation: 'append_event', streamId }, async (redis) => { - const key = getEventsKey(streamId) - const seqKey = getSeqKey(streamId) - const pipeline = redis.pipeline() const zaddArgs: Array = [] - for (const envelope of envelopes) { - zaddArgs.push(envelope.seq, JSON.stringify(envelope)) + for (const entry of chunk.members) { + zaddArgs.push(entry.seq, entry.member) } - pipeline.zadd(key, ...(zaddArgs as [number, string, ...Array])) - pipeline.zremrangebyrank(key, 0, -config.eventLimit - 1) - pipeline.expire(key, config.ttlSeconds) - pipeline.set(seqKey, String(envelopes[envelopes.length - 1].seq), 'EX', config.ttlSeconds) - await pipeline.exec() - }) - return envelopes + const result = await withRedisRetry({ operation: 'append_event', streamId }, async (redis) => + redis.eval( + APPEND_EVENTS_SCRIPT, + 2 + budgetKeys.length, + getEventsKey(streamId), + getSeqKey(streamId), + ...budgetKeys, + config.ttlSeconds, + config.eventLimit, + limits.maxOwnerBytes, + limits.maxUserBytes, + budgetTtlSeconds, + String(chunk.members[chunk.members.length - 1].seq), + ...zaddArgs + ) + ) + + const refusal = parseRedisBudgetRefusal(result, chunk.bytes, limits) + if (refusal) { + logRedisBudgetRefusal(refusal, { operation: 'append_event', scope: budgetScope, logger }) + return { persisted: false, refusal } + } + } + + return { persisted: true } } export async function appendEvent( - envelope: PersistedStreamEventEnvelope + envelope: PersistedStreamEventEnvelope, + scope?: StreamBudgetScope ): Promise { - await appendEvents([envelope]) + await appendEvents([envelope], scope) return envelope } diff --git a/apps/sim/lib/copilot/request/session/writer.test.ts b/apps/sim/lib/copilot/request/session/writer.test.ts index 719a22f978c..62a594988d7 100644 --- a/apps/sim/lib/copilot/request/session/writer.test.ts +++ b/apps/sim/lib/copilot/request/session/writer.test.ts @@ -27,14 +27,16 @@ describe('StreamWriter', () => { beforeEach(() => { vi.clearAllMocks() vi.useRealTimers() + // The buffer reports a refusal rather than throwing, so every persist resolves. + appendEvents.mockResolvedValue({ persisted: true }) }) it('enqueues before persistence completes and flushes pending writes on close', async () => { let releasePersist: (() => void) | null = null appendEvents.mockImplementation( () => - new Promise((resolve) => { - releasePersist = resolve + new Promise<{ persisted: true }>((resolve) => { + releasePersist = () => resolve({ persisted: true }) }) ) @@ -86,7 +88,7 @@ describe('StreamWriter', () => { const persistedSeqs: number[] = [] appendEvents.mockImplementation(async (envelopes) => { persistedSeqs.push(...envelopes.map((envelope) => envelope.seq)) - return envelopes + return { persisted: true } }) const writer = new StreamWriter({ @@ -119,10 +121,10 @@ describe('StreamWriter', () => { await writer.close() expect(persistedSeqs).toEqual([1, 2]) - expect(appendEvents).toHaveBeenCalledWith([ - expect.objectContaining({ seq: 1 }), - expect.objectContaining({ seq: 2 }), - ]) + expect(appendEvents).toHaveBeenCalledWith( + [expect.objectContaining({ seq: 1 }), expect.objectContaining({ seq: 2 })], + { streamId: 'stream-1' } + ) expect(chunks[0]).toContain('"seq":1') expect(chunks[1]).toContain('"seq":2') }) @@ -149,7 +151,7 @@ describe('StreamWriter', () => { }) it('persists synthetic preview events alongside contract events', async () => { - appendEvents.mockResolvedValue([]) + appendEvents.mockResolvedValue({ persisted: true }) const writer = new StreamWriter({ streamId: 'stream-1', @@ -176,15 +178,18 @@ describe('StreamWriter', () => { await writer.flush() expect(chunks[0]).toContain('"previewPhase":"file_preview_start"') - expect(appendEvents).toHaveBeenCalledWith([ - expect.objectContaining({ - type: MothershipStreamV1EventType.tool, - payload: expect.objectContaining({ - toolCallId: 'preview-1', - previewPhase: 'file_preview_start', + expect(appendEvents).toHaveBeenCalledWith( + [ + expect.objectContaining({ + type: MothershipStreamV1EventType.tool, + payload: expect.objectContaining({ + toolCallId: 'preview-1', + previewPhase: 'file_preview_start', + }), }), - }), - ]) + ], + { streamId: 'stream-1' } + ) }) /** @@ -197,7 +202,7 @@ describe('StreamWriter', () => { * failed and never advances past a gap the buffer does not have. */ it('persists an envelope whose delivery failed, and stops enqueuing after', async () => { - appendEvents.mockResolvedValue(undefined) + appendEvents.mockResolvedValue({ persisted: true }) const writer = new StreamWriter({ streamId: 'stream-gap', @@ -238,4 +243,137 @@ describe('StreamWriter', () => { // The failed enqueue disconnects; nothing is pushed at the dead controller again. expect(enqueueCalls).toBe(1) }) + + /** + * A refused write is not a fault. + * + * `flush()` rethrows whatever it is handed, and that throw reaches the error-path + * `finalizeStream`, which runs inside the catch and so escapes to reject the + * response stream. A turn whose bytes the user already received must not end in an + * error because its replay copy did not fit — so the buffer stops and the turn + * finishes. + */ + it('stops persisting on a budget refusal without failing the stream', async () => { + appendEvents.mockResolvedValue({ + persisted: false, + refusal: { + resource: 'owner_redis_bytes', + currentBytes: 33_000_000, + limitBytes: 32 * 1024 * 1024, + attemptedBytes: 4_096, + }, + }) + + const writer = new StreamWriter({ + streamId: 'stream-budget', + chatId: 'chat-budget', + requestId: 'req-budget', + userId: 'user-budget', + }) + + const chunks: string[] = [] + writer.attach({ + enqueue: vi.fn((value: Uint8Array) => { + chunks.push(decodeChunk(value)) + }), + close: vi.fn(), + } as unknown as ReadableStreamDefaultController) + + writer.publish({ + type: MothershipStreamV1EventType.text, + payload: { channel: MothershipStreamV1TextChannel.assistant, text: 'one' }, + } as StreamEvent) + + await expect(writer.flush()).resolves.toBeUndefined() + expect(writer.persistenceStopped).toBe(true) + + // The turn keeps streaming; only the replay copy stopped. + appendEvents.mockClear() + writer.publish({ + type: MothershipStreamV1EventType.text, + payload: { channel: MothershipStreamV1TextChannel.assistant, text: 'two' }, + } as StreamEvent) + await writer.flush() + + expect(appendEvents).not.toHaveBeenCalled() + expect(chunks.join('')).toContain('"text":"two"') + }) + + it('charges the replay buffer to the user when one is known', async () => { + const writer = new StreamWriter({ + streamId: 'stream-1', + chatId: 'chat-1', + requestId: 'req-1', + userId: 'user-7', + }) + writer.attach({ + enqueue: vi.fn(), + close: vi.fn(), + } as unknown as ReadableStreamDefaultController) + + writer.publish({ + type: MothershipStreamV1EventType.text, + payload: { channel: MothershipStreamV1TextChannel.assistant, text: 'x' }, + } as StreamEvent) + await writer.flush() + + expect(appendEvents).toHaveBeenCalledWith(expect.any(Array), { + streamId: 'stream-1', + userId: 'user-7', + }) + }) + + it('does not persist a batch queued while an earlier append was already refusing', async () => { + vi.useFakeTimers() + let releaseFirst: () => void = () => {} + appendEvents + .mockImplementationOnce( + () => + new Promise((resolve) => { + releaseFirst = () => + resolve({ + persisted: false, + refusal: { + resource: 'owner_redis_bytes', + currentBytes: 1, + limitBytes: 1, + attemptedBytes: 1, + }, + }) + }) + ) + .mockResolvedValue({ persisted: true }) + + const writer = new StreamWriter({ + streamId: 'stream-1', + chatId: 'chat-1', + requestId: 'req-1', + }) + const controller = { + enqueue: vi.fn(), + close: vi.fn(), + } as unknown as ReadableStreamDefaultController + writer.attach(controller) + + await writer.publish({ + type: MothershipStreamV1EventType.text, + payload: { channel: MothershipStreamV1TextChannel.assistant, text: 'one' }, + }) + await vi.advanceTimersByTimeAsync(15) + + // Queued while the first append is still in flight, so the enqueue-time check cannot see the + // refusal about to latch. Persisting it would leave replay holding a later event but not the + // refused one — a hole a resuming client cannot detect. + await writer.publish({ + type: MothershipStreamV1EventType.text, + payload: { channel: MothershipStreamV1TextChannel.assistant, text: 'two' }, + }) + await vi.advanceTimersByTimeAsync(15) + + releaseFirst() + await writer.close() + + expect(writer.persistenceStopped).toBe(true) + expect(appendEvents).toHaveBeenCalledTimes(1) + }) }) diff --git a/apps/sim/lib/copilot/request/session/writer.ts b/apps/sim/lib/copilot/request/session/writer.ts index 7ccabf83dd3..8699b790c71 100644 --- a/apps/sim/lib/copilot/request/session/writer.ts +++ b/apps/sim/lib/copilot/request/session/writer.ts @@ -17,12 +17,18 @@ export interface StreamWriterOptions { streamId: string chatId?: string requestId: string + /** Charges this stream's replay buffer to the user's cross-stream byte ceiling. */ + userId?: string keepaliveMs?: number } +/** Result used when the soft stop is already latched, so no further append is attempted. */ +const PERSISTENCE_ALREADY_STOPPED = { persisted: true } as const + export class StreamWriter { private readonly streamId: string private readonly chatId: string | undefined + private readonly userId: string | undefined private requestId: string private readonly keepaliveMs: number private readonly flushIntervalMs: number @@ -33,6 +39,7 @@ export class StreamWriter { private flushTimer: ReturnType | null = null private _clientDisconnected = false private _sawComplete = false + private _persistenceStopped = false private nextSeq = 0 private pendingEnvelopes: PersistedStreamEventEnvelope[] = [] private persistenceTail: Promise = Promise.resolve() @@ -41,6 +48,7 @@ export class StreamWriter { constructor(options: StreamWriterOptions) { this.streamId = options.streamId this.chatId = options.chatId + this.userId = options.userId this.requestId = options.requestId this.keepaliveMs = options.keepaliveMs ?? DEFAULT_KEEPALIVE_MS this.flushIntervalMs = DEFAULT_PERSIST_FLUSH_INTERVAL_MS @@ -56,6 +64,14 @@ export class StreamWriter { return this._sawComplete } + /** + * The replay buffer stopped accepting writes because this stream exhausted its byte + * budget. Live delivery is unaffected; only a resume would come back short. + */ + get persistenceStopped(): boolean { + return this._persistenceStopped + } + updateRequestId(id: string): void { this.requestId = id } @@ -151,6 +167,9 @@ export class StreamWriter { } private queuePersistence(envelope: PersistedStreamEventEnvelope): void { + // Once the budget has refused, every later batch would be refused too; stop + // paying for the round trip. + if (this._persistenceStopped) return this.pendingEnvelopes.push(envelope) if (this.pendingEnvelopes.length >= this.flushMaxBatch) { this.flushPendingPersistence() @@ -174,9 +193,40 @@ export class StreamWriter { this.pendingEnvelopes = [] this.persistenceTail = this.persistenceTail .catch(() => undefined) - .then(() => appendEvents(batch)) - .then(() => { + .then(() => + /* + Re-checked here, not only at enqueue: a batch queued while an earlier append was + in flight would otherwise land after that append had already stopped persistence, + leaving a replay that holds later events but not the refused ones — a hole a + resuming client cannot detect. + */ + this._persistenceStopped + ? PERSISTENCE_ALREADY_STOPPED + : appendEvents(batch, { + streamId: this.streamId, + ...(this.userId ? { userId: this.userId } : {}), + }) + ) + .then((result) => { this.lastPersistenceError = null + if (!result.persisted) { + /* + A budget refusal is deliberate, not a fault: it is left out of + `lastPersistenceError` so `flush()` does not rethrow it. That throw would + reach `finalizeStream`'s error-path flush and reject the response stream, + ending a turn whose bytes the user already has. Stop persisting instead — + a resume comes back short, which the caller can see. + */ + this._persistenceStopped = true + logger.warn('Stream replay buffer stopped: byte budget exhausted', { + streamId: this.streamId, + requestId: this.requestId, + resource: result.refusal.resource, + attemptedBytes: result.refusal.attemptedBytes, + currentBytes: result.refusal.currentBytes, + limitBytes: result.refusal.limitBytes, + }) + } }) .catch((error) => { this.lastPersistenceError = toError(error) diff --git a/apps/sim/lib/core/redis/byte-budget.server.test.ts b/apps/sim/lib/core/redis/byte-budget.server.test.ts new file mode 100644 index 00000000000..514cd5a44b4 --- /dev/null +++ b/apps/sim/lib/core/redis/byte-budget.server.test.ts @@ -0,0 +1,58 @@ +/** + * @vitest-environment node + */ +import { describe, expect, it } from 'vitest' +import { getRedisBudgetKeys, getRedisBudgetLimits } from '@/lib/core/redis/byte-budget.server' + +describe('getRedisBudgetKeys', () => { + it('charges the owner only when no user is in scope', () => { + expect(getRedisBudgetKeys({ kind: 'execution', id: 'exec-1' })).toEqual([ + 'execution:redis-budget:execution:exec-1', + ]) + }) + + it('charges the owner and the user when a user is in scope', () => { + expect(getRedisBudgetKeys({ kind: 'execution', id: 'exec-1', userId: 'user-1' })).toEqual([ + 'execution:redis-budget:execution:exec-1', + 'execution:redis-budget:user:user-1', + ]) + }) + + /** + * These keys are shared with counters written before this module existed, so the + * layout is a wire contract: a change here strands every counter in flight. + */ + it('separates owner kinds without disturbing the execution key layout', () => { + expect( + getRedisBudgetKeys({ kind: 'copilot_stream', id: 'stream-1', userId: 'user-1' }) + ).toEqual([ + 'execution:redis-budget:copilot_stream:stream-1', + 'execution:redis-budget:user:user-1', + ]) + }) + + it('shares one user ceiling across owner kinds', () => { + const [, executionUserKey] = getRedisBudgetKeys({ + kind: 'execution', + id: 'exec-1', + userId: 'user-1', + }) + const [, streamUserKey] = getRedisBudgetKeys({ + kind: 'copilot_stream', + id: 'stream-1', + userId: 'user-1', + }) + expect(streamUserKey).toBe(executionUserKey) + }) +}) + +describe('getRedisBudgetLimits', () => { + it('preserves the ceilings the execution buffer has always enforced', () => { + expect(getRedisBudgetLimits('execution')).toEqual({ + maxSingleWriteBytes: 8 * 1024 * 1024, + maxOwnerBytes: 64 * 1024 * 1024, + maxUserBytes: 256 * 1024 * 1024, + ttlSeconds: 60 * 60, + }) + }) +}) diff --git a/apps/sim/lib/core/redis/byte-budget.server.ts b/apps/sim/lib/core/redis/byte-budget.server.ts new file mode 100644 index 00000000000..5728d69960a --- /dev/null +++ b/apps/sim/lib/core/redis/byte-budget.server.ts @@ -0,0 +1,222 @@ +import type { Logger } from '@sim/logger' + +/** + * Per-owner byte accounting for shared Redis. + * + * Redis has no per-tenant memory limit — the documented way to get one is to meter + * in the application, which is what this does. It is the generalization of the + * budget the execution event buffer has enforced since it was written, which the + * copilot stream buffer now shares rather than inventing a bound of its own. + * + * A quota is the right bound for a buffer whose contents must stay contiguous: the + * copilot replay chain and an execution's event history are read from a cursor, so + * the write that would breach the ceiling is refused and the buffer stops growing. + * A live-update feed is bounded differently — see `lib/realtime/event-log.ts`, whose + * readers already handle a prune by refetching, so it drops oldest-first instead. + * + * The unit is bytes, deliberately. An entry cap bounds cardinality and says nothing + * about size, so a key holding a few hundred entries of a few hundred KB passes an + * entry cap of any value while holding hundreds of megabytes. That is how a copilot + * file-edit stream reached gigabytes under a 100,000-entry cap. + * + * Values too large to store belong in blob storage behind a reference — see + * `lib/execution/payloads/large-value-ref.ts`. This module is the other half: it + * bounds the aggregate once each value is already small enough to keep. + */ + +/** + * Historical prefix, kept verbatim. + * + * It reads as execution-scoped because executions were the first owner. Renaming it + * would orphan every counter in flight at deploy for no behavioural gain, and the + * kind segment below already disambiguates. + */ +const REDIS_BUDGET_PREFIX = 'execution:redis-budget:' + +/** What a budget is charged to. One counter per owner, plus one per user across owners. */ +export type RedisBudgetOwnerKind = 'execution' | 'copilot_stream' + +export interface RedisBudgetScope { + kind: RedisBudgetOwnerKind + /** The owner's id — an execution id, a stream id, a table id. */ + id: string + /** + * Charges the write to a second, user-wide counter as well. Omitted where the + * writer has no user in scope; the owner counter still applies. + */ + userId?: string +} + +export interface RedisBudgetLimits { + maxSingleWriteBytes: number + maxOwnerBytes: number + maxUserBytes: number + ttlSeconds: number +} + +/** + * Window applied to both counters, extended differently on purpose. + * + * An owner counter accounts for data refreshed on the same schedule as the counter + * itself, so sliding its TTL on every write keeps the counter and the bytes it + * represents in step. + * + * A user counter aggregates across every owner that user writes to. Sliding it on + * each write would keep it alive indefinitely for anyone who stays active while the + * data underneath it keeps expiring — so the counter would accrue bytes Redis has + * already dropped and eventually pin the user at their ceiling until they went a full + * window without writing. User counters therefore get a fixed window: set on + * creation, never extended. + */ +const REDIS_BUDGET_TTL_SECONDS = 60 * 60 + +const LIMITS: Record> = { + /** Unchanged from what the execution event buffer has always enforced. */ + execution: { + maxSingleWriteBytes: 8 * 1024 * 1024, + maxOwnerBytes: 64 * 1024 * 1024, + maxUserBytes: 256 * 1024 * 1024, + }, + /** + * A copilot turn streams text and tool frames, not payloads — a single frame past + * 1 MB is already pathological. The owner ceiling is what a long agentic session + * may retain for replay across its whole hour. + */ + copilot_stream: { + maxSingleWriteBytes: 1 * 1024 * 1024, + maxOwnerBytes: 32 * 1024 * 1024, + maxUserBytes: 128 * 1024 * 1024, + }, +} + +export function getRedisBudgetLimits(kind: RedisBudgetOwnerKind): RedisBudgetLimits { + return { ...LIMITS[kind], ttlSeconds: REDIS_BUDGET_TTL_SECONDS } +} + +/** + * The counter keys a write is charged to, owner first. + * + * Callers append these to their script's `KEYS` **last** and pass the number of keys + * that precede them, which is what lets {@link renderRedisBudgetLua} address them + * without every script agreeing on a fixed layout. + */ +export function getRedisBudgetKeys(scope: RedisBudgetScope): string[] { + const keys = [`${REDIS_BUDGET_PREFIX}${scope.kind}:${scope.id}`] + if (scope.userId) { + keys.push(`${REDIS_BUDGET_PREFIX}user:${scope.userId}`) + } + return keys +} + +export interface RedisBudgetRefusal { + resource: 'owner_redis_bytes' | 'user_redis_bytes' + currentBytes: number + limitBytes: number + attemptedBytes: number +} + +/** + * Lua that reserves or releases `net_bytes` against the caller's budget keys. + * + * Rendered into the caller's own script so the reservation and the write it pays for + * commit together — a budget checked in a separate round trip is a budget two + * concurrent writers can both pass. + * + * Contract for the caller's script: + * - budget keys are the **last** one or two entries of `KEYS`, in the order + * {@link getRedisBudgetKeys} returns them + * - `baseKeyCount` is how many keys precede them + * - before including this fragment, define `net_bytes` (may be negative, for bytes + * the same write releases by trimming), `owner_limit`, `user_limit` and + * `budget_ttl_seconds` + * - on refusal the fragment `return`s, so include it before the write it guards + */ +export function renderRedisBudgetLua(baseKeyCount: number): string { + const ownerKey = `KEYS[${baseKeyCount + 1}]` + const userKey = `KEYS[${baseKeyCount + 2}]` + const hasUserKey = `#KEYS >= ${baseKeyCount + 2}` + + return ` +if net_bytes > 0 then + local owner_current = tonumber(redis.call('GET', ${ownerKey}) or '0') + if owner_limit > 0 and owner_current + net_bytes > owner_limit then + return {0, 'owner_redis_bytes', owner_current} + end + if ${hasUserKey} then + local user_current = tonumber(redis.call('GET', ${userKey}) or '0') + if user_limit > 0 and user_current + net_bytes > user_limit then + return {0, 'user_redis_bytes', user_current} + end + end + redis.call('INCRBY', ${ownerKey}, net_bytes) + redis.call('EXPIRE', ${ownerKey}, budget_ttl_seconds) + if ${hasUserKey} then + redis.call('INCRBY', ${userKey}, net_bytes) + if redis.call('TTL', ${userKey}) < 0 then + redis.call('EXPIRE', ${userKey}, budget_ttl_seconds) + end + end +elseif net_bytes < 0 then + local release_bytes = -net_bytes + local owner_next = redis.call('DECRBY', ${ownerKey}, release_bytes) + if owner_next <= 0 then + redis.call('DEL', ${ownerKey}) + else + redis.call('EXPIRE', ${ownerKey}, budget_ttl_seconds) + end + if ${hasUserKey} then + local user_next = redis.call('DECRBY', ${userKey}, release_bytes) + if user_next <= 0 then + redis.call('DEL', ${userKey}) + elseif redis.call('TTL', ${userKey}) < 0 then + redis.call('EXPIRE', ${userKey}, budget_ttl_seconds) + end + end +else + if redis.call('EXISTS', ${ownerKey}) == 1 then + redis.call('EXPIRE', ${ownerKey}, budget_ttl_seconds) + end + if ${hasUserKey} and redis.call('EXISTS', ${userKey}) == 1 and redis.call('TTL', ${userKey}) < 0 then + redis.call('EXPIRE', ${userKey}, budget_ttl_seconds) + end +end +` +} + +/** Parses the `{0, resource, current}` refusal a guarded script returns. */ +export function parseRedisBudgetRefusal( + result: unknown, + attemptedBytes: number, + limits: RedisBudgetLimits +): RedisBudgetRefusal | null { + if (!Array.isArray(result) || result[0] !== 0) return null + const resource = result[1] === 'user_redis_bytes' ? 'user_redis_bytes' : 'owner_redis_bytes' + return { + resource, + currentBytes: Number(result[2] ?? 0), + limitBytes: resource === 'user_redis_bytes' ? limits.maxUserBytes : limits.maxOwnerBytes, + attemptedBytes, + } +} + +export interface RedisBudgetLogContext { + operation: string + scope: RedisBudgetScope + logger?: Logger +} + +/** One place that decides how a refusal is reported, so every writer reports it alike. */ +export function logRedisBudgetRefusal( + refusal: RedisBudgetRefusal, + context: RedisBudgetLogContext +): void { + context.logger?.warn('Redis byte budget refused a write', { + operation: context.operation, + ownerKind: context.scope.kind, + ownerId: context.scope.id, + resource: refusal.resource, + attemptedBytes: refusal.attemptedBytes, + currentBytes: refusal.currentBytes, + limitBytes: refusal.limitBytes, + }) +} diff --git a/apps/sim/lib/execution/event-buffer.ts b/apps/sim/lib/execution/event-buffer.ts index 834edf8ac4f..abd1561f36e 100644 --- a/apps/sim/lib/execution/event-buffer.ts +++ b/apps/sim/lib/execution/event-buffer.ts @@ -4,6 +4,7 @@ import { randomInt } from '@sim/utils/random' import { getConfiguredCacheProvider } from '@/lib/core/config/env-capabilities.server' import { getRedisClient } from '@/lib/core/config/redis' import { getMaxExecutionTimeout } from '@/lib/core/execution-limits' +import { getRedisBudgetKeys, getRedisBudgetLimits } from '@/lib/core/redis/byte-budget.server' import { getExecutionSignalChannel, publishLocalExecutionSignal, @@ -11,11 +12,6 @@ import { import { LARGE_VALUE_THRESHOLD_BYTES } from '@/lib/execution/payloads/large-value-ref' import { compactExecutionPayload } from '@/lib/execution/payloads/serializer' import type { LargeValueStoreContext } from '@/lib/execution/payloads/store' -import { - type ExecutionRedisBudgetReservation, - getExecutionRedisBudgetKeys, - getExecutionRedisBudgetLimits, -} from '@/lib/execution/redis-budget.server' import { ExecutionResourceLimitError, isExecutionResourceLimitError, @@ -47,9 +43,9 @@ const MAX_ACTIVE_BLOCK_SNAPSHOT_BYTES = 256 * 1024 * run has actually buffered its way into the danger zone, so a short run keeps * full-fidelity output and a runaway one stops accumulating. */ -const EXECUTION_EVENT_OFFLOAD_PRESSURE_BYTES = getExecutionRedisBudgetLimits().maxExecutionBytes / 2 +const EXECUTION_EVENT_OFFLOAD_PRESSURE_BYTES = getRedisBudgetLimits('execution').maxOwnerBytes / 2 const EXECUTION_EVENT_PRESSURE_VALUE_BYTES = - getExecutionRedisBudgetLimits().maxExecutionBytes / EVENT_LIMIT + getRedisBudgetLimits('execution').maxOwnerBytes / EVENT_LIMIT const ACTIVE_META_ATTEMPTS = 3 const FINALIZE_FLUSH_ATTEMPTS = 2 const FLUSH_EVENTS_SCRIPT = ` @@ -562,15 +558,7 @@ export async function resetExecutionStreamBuffer(executionId: string): Promise ({}))) as Record const userId = typeof meta.userId === 'string' ? meta.userId : undefined - const budgetReservation: ExecutionRedisBudgetReservation = { - executionId, - userId, - category: 'event_buffer', - operation: 'reset_events', - bytes: 0, - logger, - } - const budgetKeys = getExecutionRedisBudgetKeys(budgetReservation) + const budgetKeys = getRedisBudgetKeys({ kind: 'execution', id: executionId, userId }) await redis.eval( RESET_STREAM_SCRIPT, 2 + budgetKeys.length, @@ -580,7 +568,7 @@ export async function resetExecutionStreamBuffer(executionId: string): Promise limits.maxSingleWriteBytes) { // A single entry above the cap can never be written; dropping it is the // only way the rest of the buffer makes progress. @@ -1068,7 +1048,11 @@ export function createExecutionEventWriter( limitBytes: limits.maxSingleWriteBytes, }) } - const budgetKeys = getExecutionRedisBudgetKeys(budgetReservation) + const budgetKeys = getRedisBudgetKeys({ + kind: 'execution', + id: executionId, + userId: context.userId, + }) const flushResult = getFlushScriptResult( await redis.eval( FLUSH_EVENTS_SCRIPT, @@ -1083,7 +1067,7 @@ export function createExecutionEventWriter( new Date().toISOString(), chunkTerminalStatus ?? '', batchBytes, - limits.maxExecutionBytes, + limits.maxOwnerBytes, limits.maxUserBytes, limits.ttlSeconds, ...zaddArgs @@ -1100,7 +1084,7 @@ export function createExecutionEventWriter( limitBytes: flushResult.resource === 'user_redis_bytes' ? limits.maxUserBytes - : limits.maxExecutionBytes, + : limits.maxOwnerBytes, }) } consecutiveFlushFailures = 0 diff --git a/apps/sim/lib/execution/redis-budget.server.test.ts b/apps/sim/lib/execution/redis-budget.server.test.ts deleted file mode 100644 index f9456fc3fae..00000000000 --- a/apps/sim/lib/execution/redis-budget.server.test.ts +++ /dev/null @@ -1,28 +0,0 @@ -/** - * @vitest-environment node - */ -import { describe, expect, it } from 'vitest' -import { getExecutionRedisBudgetKeys } from '@/lib/execution/redis-budget.server' - -describe('getExecutionRedisBudgetKeys', () => { - it('scopes the reservation to the execution, and to the user when one is known', () => { - expect( - getExecutionRedisBudgetKeys({ - executionId: 'exec-1', - category: 'event_buffer', - operation: 'write_events', - bytes: 1, - }) - ).toEqual(['execution:redis-budget:execution:exec-1']) - - expect( - getExecutionRedisBudgetKeys({ - executionId: 'exec-1', - userId: 'user-1', - category: 'event_buffer', - operation: 'write_events', - bytes: 1, - }) - ).toEqual(['execution:redis-budget:execution:exec-1', 'execution:redis-budget:user:user-1']) - }) -}) diff --git a/apps/sim/lib/execution/redis-budget.server.ts b/apps/sim/lib/execution/redis-budget.server.ts deleted file mode 100644 index daf5822fb4c..00000000000 --- a/apps/sim/lib/execution/redis-budget.server.ts +++ /dev/null @@ -1,56 +0,0 @@ -import type { Logger } from '@sim/logger' - -const REDIS_BUDGET_PREFIX = 'execution:redis-budget:' -const MAX_SINGLE_REDIS_WRITE_BYTES = 8 * 1024 * 1024 -const MAX_EXECUTION_REDIS_BYTES = 64 * 1024 * 1024 -const MAX_USER_REDIS_BYTES = 256 * 1024 * 1024 - -/** - * Window applied to both budget keys, but extended differently on purpose by - * every Lua script that enforces them — `FLUSH_EVENTS_SCRIPT` and - * `RESET_STREAM_SCRIPT` in `event-buffer.ts`, and the base64 cache pair in - * `lib/uploads/utils/user-file-base64.server.ts`. - * - * An execution key accounts for data that is refreshed on the same schedule as - * the key itself, so sliding its TTL on every write keeps the counter and the - * bytes it represents in step. - * - * A user key aggregates across every execution that user runs. Sliding its TTL - * on each write keeps it alive indefinitely for any user who stays active, - * while the per-execution data it accounts for keeps expiring underneath it — - * so the counter accrues bytes Redis has already dropped and eventually pins - * the user at their ceiling until they go a full TTL without writing. User - * keys therefore get a fixed window: the TTL is set when the key is created - * and never extended. - */ -const REDIS_BUDGET_TTL_SECONDS = 60 * 60 - -export type ExecutionRedisBudgetCategory = 'event_buffer' | 'base64_cache' - -export interface ExecutionRedisBudgetReservation { - executionId: string - userId?: string - category: ExecutionRedisBudgetCategory - bytes: number - operation: string - logger?: Logger -} - -export function getExecutionRedisBudgetLimits() { - return { - maxSingleWriteBytes: MAX_SINGLE_REDIS_WRITE_BYTES, - maxExecutionBytes: MAX_EXECUTION_REDIS_BYTES, - maxUserBytes: MAX_USER_REDIS_BYTES, - ttlSeconds: REDIS_BUDGET_TTL_SECONDS, - } -} - -export function getExecutionRedisBudgetKeys( - reservation: ExecutionRedisBudgetReservation -): string[] { - const keys = [`${REDIS_BUDGET_PREFIX}execution:${reservation.executionId}`] - if (reservation.userId) { - keys.push(`${REDIS_BUDGET_PREFIX}user:${reservation.userId}`) - } - return keys -} diff --git a/apps/sim/lib/realtime/event-log.test.ts b/apps/sim/lib/realtime/event-log.test.ts index b8159b108cf..e49de7d78b4 100644 --- a/apps/sim/lib/realtime/event-log.test.ts +++ b/apps/sim/lib/realtime/event-log.test.ts @@ -3,7 +3,8 @@ */ import { beforeEach, describe, expect, it, vi } from 'vitest' -const { mockEnv } = vi.hoisted(() => ({ +const { mockEnv, mockRedisClient } = vi.hoisted(() => ({ + mockRedisClient: { current: null as { eval: ReturnType } | null }, mockEnv: { REDIS_URL: undefined as string | undefined, REDIS_TLS_SERVERNAME: undefined as string | undefined, @@ -11,7 +12,7 @@ const { mockEnv } = vi.hoisted(() => ({ })) vi.mock('@/lib/core/config/env', () => ({ env: mockEnv })) -vi.mock('@/lib/core/config/redis', () => ({ getRedisClient: () => null })) +vi.mock('@/lib/core/config/redis', () => ({ getRedisClient: () => mockRedisClient.current })) import { appendEvent, @@ -28,7 +29,13 @@ interface TestEntry extends EventLogEntry { value: string } -const config: EventLogConfig = { prefix: 'test:stream:', ttlSeconds: 3600, cap: 3, readChunk: 500 } +const config: EventLogConfig = { + prefix: 'test:stream:', + ttlSeconds: 3600, + cap: 3, + maxBytes: 0, + readChunk: 500, +} function serializerFor(streamId: string, value: string) { return { @@ -42,6 +49,7 @@ describe('event-log (memory fallback)', () => { beforeEach(() => { mockEnv.REDIS_URL = undefined mockEnv.REDIS_TLS_SERVERNAME = undefined + mockRedisClient.current = null resetEventLogMemoryForTesting() }) @@ -113,3 +121,70 @@ describe('event-log (memory fallback)', () => { ) }) }) + +describe('event-log byte ceiling', () => { + beforeEach(() => { + mockEnv.REDIS_URL = undefined + mockEnv.REDIS_TLS_SERVERNAME = undefined + mockRedisClient.current = null + resetEventLogMemoryForTesting() + }) + + /** + * The entry cap is what let one writer hold hundreds of megabytes: `cap` bounds how + * many entries a stream keeps and nothing about how large each one is. + */ + it('drops oldest entries once the buffer exceeds maxBytes, under the entry cap', async () => { + const bounded: EventLogConfig = { ...config, cap: 1000, maxBytes: 400 } + const big = 'x'.repeat(150) + + for (let i = 0; i < 6; i++) { + await appendEvent(bounded, 's1', serializerFor('s1', big)) + } + + const fromStart = await readEventsSince(bounded, 's1', 0) + expect(fromStart.status).toBe('pruned') + const earliest = fromStart.status === 'pruned' ? fromStart.earliestEventId : undefined + expect(earliest).toBeGreaterThan(1) + + const retained = await readEventsSince(bounded, 's1', (earliest as number) - 1) + expect(retained.status).toBe('ok') + const events = retained.status === 'ok' ? retained.events : [] + expect(events.at(-1)?.eventId).toBe(6) + const bytes = events.reduce((total, e) => total + JSON.stringify(e).length, 0) + expect(bytes).toBeLessThanOrEqual(400) + }) + + it('keeps the newest entry even when it alone exceeds maxBytes', async () => { + const bounded: EventLogConfig = { ...config, cap: 1000, maxBytes: 10 } + await appendEvent(bounded, 's1', serializerFor('s1', 'a')) + await appendEvent(bounded, 's1', serializerFor('s1', 'b'.repeat(500))) + + const result = await readEventsSince(bounded, 's1', 1) + expect(result.status).toBe('ok') + const events = result.status === 'ok' ? result.events : [] + expect(events).toHaveLength(1) + expect(events[0]?.eventId).toBe(2) + }) + + it('leaves the buffer unbounded by bytes when maxBytes is 0', async () => { + const unbounded: EventLogConfig = { ...config, cap: 1000, maxBytes: 0 } + for (let i = 0; i < 5; i++) { + await appendEvent(unbounded, 's1', serializerFor('s1', 'x'.repeat(500))) + } + const result = await readEventsSince(unbounded, 's1', 0) + expect(result.status).toBe('ok') + expect(result.status === 'ok' ? result.events : []).toHaveLength(5) + }) + + it('passes the ceiling to the Redis script', async () => { + const evalFn = vi.fn().mockResolvedValue(1) + mockRedisClient.current = { eval: evalFn } + mockEnv.REDIS_URL = 'redis://localhost:6379' + + await appendEvent({ ...config, maxBytes: 4096 }, 's1', serializerFor('s1', 'a')) + + expect(evalFn).toHaveBeenCalledTimes(1) + expect(evalFn.mock.calls[0]?.at(-1)).toBe(4096) + }) +}) diff --git a/apps/sim/lib/realtime/event-log.ts b/apps/sim/lib/realtime/event-log.ts index 5916c4b6820..f002b6a1c89 100644 --- a/apps/sim/lib/realtime/event-log.ts +++ b/apps/sim/lib/realtime/event-log.ts @@ -23,26 +23,66 @@ const logger = createLogger('EventLog') /** * Atomic append: INCR the seq counter to mint a new eventId, splice it into the - * adapter-supplied entry JSON, ZADD it, refresh TTLs, trim to cap, and record the - * resulting earliestEventId in meta — one round-trip. Without atomicity a slow - * reader could observe the trim before the meta update and miss the prune signal. + * adapter-supplied entry JSON, ZADD it, refresh TTLs, trim, and record the resulting + * earliestEventId in meta — one round-trip. Without atomicity a slow reader could + * observe the trim before the meta update and miss the prune signal. + * + * The buffer is bounded twice: to `cap` entries, and to `maxBytes`. The entry bound + * alone bounds cardinality and says nothing about size — an entry here carries a + * cell's outputs, which a dispatch resends cumulatively, so `cap` entries of a few + * hundred KB is gigabytes for one table. Both trims drop the oldest, which is the + * behaviour readers already handle: `earliestEventId` moves, `readEventsSince` + * returns `pruned`, and the client refetches and resumes from latest. + * + * The running total is kept in meta rather than summed per append, and both keys + * share a TTL so the counter cannot outlive the bytes it counts. * * KEYS: [events, seq, meta] - * ARGV: [ttlSec, cap, updatedAtIso, entryPrefix, entrySuffix] + * ARGV: [ttlSec, cap, updatedAtIso, entryPrefix, entrySuffix, maxBytes] * The new eventId is spliced between prefix/suffix to form the entry JSON. * Returns the new eventId. */ const APPEND_EVENT_SCRIPT = ` +local ttl_seconds = tonumber(ARGV[1]) +local cap = tonumber(ARGV[2]) +local max_bytes = tonumber(ARGV[6]) + local eventId = redis.call('INCR', KEYS[2]) local entry = ARGV[4] .. eventId .. ARGV[5] redis.call('ZADD', KEYS[1], eventId, entry) -redis.call('EXPIRE', KEYS[1], tonumber(ARGV[1])) -redis.call('EXPIRE', KEYS[2], tonumber(ARGV[1])) -redis.call('ZREMRANGEBYRANK', KEYS[1], 0, -tonumber(ARGV[2]) - 1) +redis.call('EXPIRE', KEYS[1], ttl_seconds) +redis.call('EXPIRE', KEYS[2], ttl_seconds) + +local total = tonumber(redis.call('HGET', KEYS[3], 'bytes') or '0') + string.len(entry) + +local over = redis.call('ZCARD', KEYS[1]) - cap +if over > 0 then + local dropped = redis.call('ZRANGE', KEYS[1], 0, over - 1) + for _, member in ipairs(dropped) do + total = total - string.len(member) + end + redis.call('ZREMRANGEBYRANK', KEYS[1], 0, over - 1) +end + +while max_bytes > 0 and total > max_bytes and redis.call('ZCARD', KEYS[1]) > 1 do + local oldest_member = redis.call('ZRANGE', KEYS[1], 0, 0) + if not oldest_member[1] then break end + total = total - string.len(oldest_member[1]) + redis.call('ZREMRANGEBYRANK', KEYS[1], 0, 0) +end +if total < 0 then total = 0 end +-- Self-correct: the counter is an accumulator, so an independently evicted events key would leave +-- it over-reporting forever and pin the buffer at a single entry. Whenever the buffer is down to one +-- entry its exact size is known, so drift cannot outlive a trim. +if redis.call('ZCARD', KEYS[1]) == 1 then + local only = redis.call('ZRANGE', KEYS[1], 0, 0) + if only[1] then total = string.len(only[1]) end +end + local oldest = redis.call('ZRANGE', KEYS[1], 0, 0, 'WITHSCORES') if oldest[2] then - redis.call('HSET', KEYS[3], 'earliestEventId', tostring(math.floor(tonumber(oldest[2]))), 'updatedAt', ARGV[3]) - redis.call('EXPIRE', KEYS[3], tonumber(ARGV[1])) + redis.call('HSET', KEYS[3], 'earliestEventId', tostring(math.floor(tonumber(oldest[2]))), 'bytes', tostring(total), 'updatedAt', ARGV[3]) + redis.call('EXPIRE', KEYS[3], ttl_seconds) end return eventId ` @@ -57,6 +97,13 @@ export interface EventLogConfig { prefix: string ttlSeconds: number cap: number + /** + * Byte ceiling for one stream's buffer. Entries are dropped oldest-first until the + * buffer fits, exactly as `cap` does — an entry cap bounds how many entries a key + * holds and nothing about how large each one is, which is how a key of a few + * hundred entries reaches hundreds of megabytes. + */ + maxBytes: number /** Max entries returned by one read; the SSE route drains in chunks. */ readChunk: number } @@ -149,8 +196,19 @@ export async function appendEvent( stream.events.push(entry) if (stream.events.length > config.cap) { stream.events = stream.events.slice(-config.cap) - stream.earliestEventId = stream.events[0]?.eventId } + if (config.maxBytes > 0) { + // UTF-8 bytes, so this path bounds a stream identically to the Lua's `string.len`; + // `String.length` counts UTF-16 units and under-reports every non-ASCII event. + const entryBytes = (event: EventLogEntry) => + Buffer.byteLength(JSON.stringify(event), 'utf8') + let bytes = stream.events.reduce((total, event) => total + entryBytes(event), 0) + while (bytes > config.maxBytes && stream.events.length > 1) { + bytes -= entryBytes(stream.events[0]) + stream.events = stream.events.slice(1) + } + } + stream.earliestEventId = stream.events[0]?.eventId stream.expiresAt = Date.now() + config.ttlSeconds * 1000 return entry } catch (error) { @@ -174,7 +232,8 @@ export async function appendEvent( config.cap, new Date().toISOString(), serializer.entryPrefix, - serializer.entrySuffix + serializer.entrySuffix, + config.maxBytes ) const eventId = typeof result === 'number' ? result : Number(result) if (!Number.isFinite(eventId)) return null diff --git a/apps/sim/lib/table/events.ts b/apps/sim/lib/table/events.ts index 034b20b5db8..d1c507ec60b 100644 --- a/apps/sim/lib/table/events.ts +++ b/apps/sim/lib/table/events.ts @@ -25,6 +25,16 @@ import { export const TABLE_EVENT_TTL_SECONDS = 60 * 60 // 1 hour export const TABLE_EVENT_CAP = 5000 +/** + * Byte ceiling for one table's buffer. + * + * An event carries a cell's outputs, and a dispatch across many rows emits one per + * cell, so `TABLE_EVENT_CAP` entries says nothing about the bytes they hold — a table + * of large text cells reaches hundreds of megabytes well inside the entry cap. 32 MB + * is far above what an interactive dispatch buffers in its TTL and far below what one + * table may cost a shared Redis. + */ +export const TABLE_EVENT_MAX_BYTES = 32 * 1024 * 1024 /** Max events returned by a single read; the SSE route drains in chunks. */ export const TABLE_EVENT_READ_CHUNK = 500 @@ -33,6 +43,7 @@ const TABLE_EVENT_LOG: EventLogConfig = { prefix: 'table:stream:', ttlSeconds: TABLE_EVENT_TTL_SECONDS, cap: TABLE_EVENT_CAP, + maxBytes: TABLE_EVENT_MAX_BYTES, readChunk: TABLE_EVENT_READ_CHUNK, } diff --git a/apps/sim/lib/uploads/utils/user-file-base64.server.ts b/apps/sim/lib/uploads/utils/user-file-base64.server.ts index 87921a31836..a1398e3a176 100644 --- a/apps/sim/lib/uploads/utils/user-file-base64.server.ts +++ b/apps/sim/lib/uploads/utils/user-file-base64.server.ts @@ -3,6 +3,7 @@ import type { Logger } from '@sim/logger' import { createLogger } from '@sim/logger' import { isPlainRecord } from '@sim/utils/object' import { getRedisClient } from '@/lib/core/config/redis' +import { getRedisBudgetKeys, getRedisBudgetLimits } from '@/lib/core/redis/byte-budget.server' import { isUserFileWithMetadata } from '@/lib/core/utils/user-file' import { recordMaterializedAccessKeys } from '@/lib/execution/payloads/access-keys' import { @@ -19,11 +20,6 @@ import { readUserFileContentWithContributors, } from '@/lib/execution/payloads/materialization.server' import { materializeLargeValueRef } from '@/lib/execution/payloads/store' -import { - type ExecutionRedisBudgetReservation, - getExecutionRedisBudgetKeys, - getExecutionRedisBudgetLimits, -} from '@/lib/execution/redis-budget.server' import { ExecutionResourceLimitError } from '@/lib/execution/resource-errors' import type { WorkspaceFileSecretProvenanceIdentity } from '@/lib/uploads/contexts/workspace/workspace-file-secret-provenance' import { isGeneratedDocumentSourceType } from '@/lib/uploads/utils/file-utils' @@ -254,7 +250,7 @@ function createBase64Cache(options: Base64HydrationOptions, logger: Logger): Bas return } - const limits = getExecutionRedisBudgetLimits() + const limits = getRedisBudgetLimits('execution') if (valueBytes > limits.maxSingleWriteBytes) { logSkippedCacheWrite( logger, @@ -269,15 +265,11 @@ function createBase64Cache(options: Base64HydrationOptions, logger: Logger): Bas return } const cacheTtlSeconds = Math.max(ttlSeconds, limits.ttlSeconds) - const budgetReservation: ExecutionRedisBudgetReservation = { - executionId, + const budgetKeys = getRedisBudgetKeys({ + kind: 'execution', + id: executionId, userId: options.userId, - category: 'base64_cache', - operation: 'set_base64_cache', - bytes: valueBytes, - logger, - } - const budgetKeys = getExecutionRedisBudgetKeys(budgetReservation) + }) const result = (await redis.eval( SET_BASE64_CACHE_SCRIPT, 2 + budgetKeys.length, @@ -289,7 +281,7 @@ function createBase64Cache(options: Base64HydrationOptions, logger: Logger): Bas getFileCacheKey(file), serializeBudgetEntry({ bytes: valueBytes, userId: options.userId }), valueBytes, - limits.maxExecutionBytes, + limits.maxOwnerBytes, limits.maxUserBytes, limits.ttlSeconds )) as [number, string, number | string | null] @@ -305,7 +297,7 @@ function createBase64Cache(options: Base64HydrationOptions, logger: Logger): Bas attemptedBytes: valueBytes, currentBytes: Number(current ?? 0), limitBytes: - resource === 'user_redis_bytes' ? limits.maxUserBytes : limits.maxExecutionBytes, + resource === 'user_redis_bytes' ? limits.maxUserBytes : limits.maxOwnerBytes, }) ) } @@ -379,15 +371,12 @@ async function cleanupBudgetEntry( rawEntry: string, entry: Base64BudgetEntry ): Promise<{ claimed: boolean; deletedCount: number }> { - const limits = getExecutionRedisBudgetLimits() - const budgetReservation: ExecutionRedisBudgetReservation = { - executionId, + const limits = getRedisBudgetLimits('execution') + const budgetKeys = getRedisBudgetKeys({ + kind: 'execution', + id: executionId, userId: entry.userId, - category: 'base64_cache', - operation: 'cleanup_base64_cache', - bytes: entry.bytes, - } - const budgetKeys = getExecutionRedisBudgetKeys(budgetReservation) + }) const result = (await redis.eval( CLEANUP_BASE64_CACHE_ENTRY_SCRIPT, 2 + budgetKeys.length,