fix: Codex review — include metadata in the size guard + flush parked bookkeeping

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.

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<void>((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
+  });
 });

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<void> {
+    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<void> {
-    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) {