From f5e86166c70962370db1ca66bdf2a42f322b882d Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 8 Jul 2026 11:18:39 -0400 Subject: [PATCH] =?UTF-8?q?fix:=20Codex=20review=20=E2=80=94=20include=20m?= =?UTF-8?q?etadata=20in=20the=20size=20guard=20+=20flush=20parked=20bookke?= =?UTF-8?q?eping?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P1 (lost bookkeeping): `put` consumes the write anchor, then AWAITS assertCheckpointFitsDocument (checkpoint serialization). A bookkeeping-only putWrites dispatched in that window sees neither the anchor nor persistedIds, so it parks — and `put` never flushed it, dropping the marker (e.g. a completed Send-sibling's __no_writes__) so a resume re-executed the sibling. Extract flushBufferedBookkeeping (shared with the anchoring putWrites) and call it after super.put in the persist path. P2 (metadata ignored by guard): MongoDBSaver.put stores the serialized checkpoint AND metadata (plus metadata_search) in the SAME document, but the guard measured only the checkpoint — a just-under-limit checkpoint with large metadata fell through to a raw BSONObjectTooLarge. Measure checkpoint + metadata; the fixed headroom now only covers metadata_search/ids/framing. Two integration regressions added (both fail without the fix, pass with it): metadata-pushes-over-the-ceiling, and flush-during-the-serialization-window. --- .../agents/checkpointer.integration.spec.ts | 75 +++++++++++++++++++ packages/api/src/agents/checkpointer.ts | 51 +++++++++---- 2 files changed, 113 insertions(+), 13 deletions(-) diff --git a/packages/api/src/agents/checkpointer.integration.spec.ts b/packages/api/src/agents/checkpointer.integration.spec.ts index 100e948458..1eb5488a80 100644 --- a/packages/api/src/agents/checkpointer.integration.spec.ts +++ b/packages/api/src/agents/checkpointer.integration.spec.ts @@ -533,4 +533,79 @@ describe('LazyMongoSaver checkpoint size guard (mongodb-memory-server integratio }), ).toBe(0); }); + + it('counts metadata toward the ceiling — refuses when the checkpoint is under but metadata pushes over', async () => { + // MongoDBSaver stores the serialized checkpoint AND metadata in one document, so a + // just-under-limit checkpoint with large metadata still overflows. The guard must catch it + // as a typed CheckpointTooLargeError, not let it fall through to a raw BSONObjectTooLarge. + jest.spyOn(logger, 'error').mockImplementation(() => logger); + const saver = makeSaver({ warnBytes: 500, hardLimitBytes: 2_000 }); + await saver.setup(); + const threadId = `convo-${new mongoose.Types.ObjectId().toString()}`; + + const checkpoint = emptyCheckpoint(); + checkpoint.channel_values = { messages: 'x'.repeat(400) }; // checkpoint alone well UNDER 2 KB + await saver.putWrites( + { configurable: { thread_id: threadId, checkpoint_ns: '', checkpoint_id: checkpoint.id } }, + [[INTERRUPT, 'approve?']], + 'task-1', + ); + const config = { configurable: { thread_id: threadId, checkpoint_ns: '' } }; + // Metadata alone pushes checkpoint + metadata over the 2 KB hard limit. + const metadata = { + source: 'input' as const, + step: -1, + writes: { payload: 'y'.repeat(4_000) }, + parents: {}, + }; + + await expect(saver.put(config, checkpoint, metadata)).rejects.toBeInstanceOf( + CheckpointTooLargeError, + ); + expect(await saver.getTuple(readConfig(threadId))).toBeUndefined(); + }); + + it('flushes bookkeeping parked during the size-serialization window (not dropped on resume)', async () => { + // `put` consumes the write anchor, then AWAITS `assertCheckpointFitsDocument` (serialization). + // A bookkeeping-only putWrites dispatched in that window sees no anchor and no persisted + // marker, so it parks — and must be flushed once the checkpoint persists, or a resume + // re-executes the completed sibling. + const NO_WRITES = '__no_writes__'; + const saver = makeSaver({ warnBytes: 5_000, hardLimitBytes: 50_000 }); + await saver.setup(); + const threadId = `convo-${new mongoose.Types.ObjectId().toString()}`; + const checkpoint = emptyCheckpoint(); + const writeCfg = { + configurable: { thread_id: threadId, checkpoint_ns: '', checkpoint_id: checkpoint.id }, + }; + await saver.putWrites(writeCfg, [[INTERRUPT, 'approve?']], 'task-gate'); + + // Pause the checkpoint serialization inside `put` so the bookkeeping batch lands mid-window. + const serde = (saver as unknown as { serde: { dumpsTyped: (v: unknown) => unknown } }).serde; + const realDumps = serde.dumpsTyped.bind(serde); + let releaseGate!: () => void; + const gate = new Promise((resolve) => { + releaseGate = resolve; + }); + let paused = false; + jest.spyOn(serde, 'dumpsTyped').mockImplementation(async (value: unknown) => { + if (!paused) { + paused = true; + await gate; // hold inside assertCheckpointFitsDocument (anchor consumed, not yet persisted) + } + return realDumps(value); + }); + + const config = { configurable: { thread_id: threadId, checkpoint_ns: '' } }; + const metadata = { source: 'input' as const, step: -1, writes: null, parents: {} }; + const putPromise = saver.put(config, checkpoint, metadata); // suspends at the gate + await saver.putWrites(writeCfg, [[NO_WRITES, null]], 'task-sibling'); // parks in the window + releaseGate(); + await putPromise; + + const tuple = await saver.getTuple(readConfig(threadId)); + const channels = (tuple?.pendingWrites ?? []).map((w) => w[1]); + expect(channels).toContain(INTERRUPT); + expect(channels).toContain(NO_WRITES); // flushed, not dropped + }); }); diff --git a/packages/api/src/agents/checkpointer.ts b/packages/api/src/agents/checkpointer.ts index 54c4d089a1..2304fd8988 100644 --- a/packages/api/src/agents/checkpointer.ts +++ b/packages/api/src/agents/checkpointer.ts @@ -249,15 +249,9 @@ export class LazyMongoSaver extends MongoDBSaver { // wrongly discarded. this.recordWriteAnchor(checkpointId); try { - const buffered = this.bufferedBookkeeping.get(checkpointId); - this.bufferedBookkeeping.delete(checkpointId); - if (buffered) { - // The checkpoint's fate is now "persist" — flush the bookkeeping batches that - // arrived before this anchor so the stored pending writes are complete. - await Promise.all( - buffered.batches.map((b) => super.putWrites(b.config, b.writes, b.taskId)), - ); - } + // The checkpoint's fate is now "persist" — flush the bookkeeping batches that + // arrived before this anchor so the stored pending writes are complete. + await this.flushBufferedBookkeeping(checkpointId); return await super.putWrites(config, writes, taskId); } catch (err) { // The write batch never landed — best-effort un-anchor so the concurrent `put` doesn't @@ -278,10 +272,17 @@ export class LazyMongoSaver extends MongoDBSaver { // Carries a resumable write (interrupt / real-channel delta anchor) — persist so resume // can read it, and remember the id briefly so any bookkeeping batch dispatched after // this `put` is forwarded rather than parked. - await this.assertCheckpointFitsDocument(config, checkpoint); + await this.assertCheckpointFitsDocument(config, checkpoint, metadata); sweepStale(this.persistedIds, (t) => t); this.persistedIds.set(checkpoint.id, Date.now()); - return super.put(config, checkpoint, metadata); + const persisted = await super.put(config, checkpoint, metadata); + // `assertCheckpointFitsDocument` awaits a (potentially slow) serialization AFTER the + // anchor was consumed above but BEFORE `persistedIds` was set — a bookkeeping-only + // `putWrites` dispatched in that window sees neither marker and parks its batch. Flush + // it now that the checkpoint is persisted; without this the marker is dropped and a + // resume can re-execute already-completed work. + await this.flushBufferedBookkeeping(checkpoint.id); + return persisted; } // No resumable writes ⇒ a clean exit (a non-paused completion, a resumed turn's clean // finish, or an error-only turn): discard, and drop the parked bookkeeping batches with @@ -310,6 +311,22 @@ export class LazyMongoSaver extends MongoDBSaver { this.writeAnchorIds.set(checkpointId, Date.now()); } + /** + * Forward the bookkeeping batches parked for `checkpointId` while its fate was undecided, + * now that the checkpoint is being persisted. Snapshot-and-delete before awaiting so a batch + * that arrives afterwards can't be double-forwarded — by then the anchor/persisted marker is + * set, so it forwards directly instead of parking. Shared by the anchoring `putWrites` and by + * `put` (for a batch parked during the size-check serialization window). + */ + private async flushBufferedBookkeeping(checkpointId: string): Promise { + const buffered = this.bufferedBookkeeping.get(checkpointId); + if (!buffered) { + return; + } + this.bufferedBookkeeping.delete(checkpointId); + await Promise.all(buffered.batches.map((b) => super.putWrites(b.config, b.writes, b.taskId))); + } + /** * Measure the checkpoint's serialized size on the persist path and act on it: `debug`-log it, * `warn` past {@link warnBytes}, and throw {@link CheckpointTooLargeError} past @@ -321,9 +338,17 @@ export class LazyMongoSaver extends MongoDBSaver { private async assertCheckpointFitsDocument( config: RunnableConfig, checkpoint: Checkpoint, + metadata: CheckpointMetadata, ): Promise { - const [, serialized] = await this.serde.dumpsTyped(checkpoint); - const bytes = serialized.byteLength; + // `MongoDBSaver.put` stores the serialized checkpoint AND the serialized metadata + // (plus a small raw `metadata_search` subset) in the SAME `agent_checkpoints` + // document, so BOTH count toward the 16 MB ceiling. Measuring only the checkpoint + // let a just-under-limit checkpoint with large metadata fall through to a raw + // `BSONObjectTooLarge`; the headroom now only has to cover `metadata_search`, ids + // and BSON framing. + const [, serializedCheckpoint] = await this.serde.dumpsTyped(checkpoint); + const [, serializedMetadata] = await this.serde.dumpsTyped(metadata); + const bytes = serializedCheckpoint.byteLength + serializedMetadata.byteLength; const threadId = config.configurable?.thread_id as string | undefined; const mb = (n: number): string => (n / 1024 / 1024).toFixed(1); if (bytes > this.hardLimitBytes) {