diff --git a/.changeset/fix-mongodb-checkpoint-put-writes-special-channel-idx.md b/.changeset/fix-mongodb-checkpoint-put-writes-special-channel-idx.md new file mode 100644 index 0000000000..72b4d23284 --- /dev/null +++ b/.changeset/fix-mongodb-checkpoint-put-writes-special-channel-idx.md @@ -0,0 +1,5 @@ +--- +"@langchain/langgraph-checkpoint-mongodb": patch +--- + +fix: `MongoDBSaver.putWrites` now honors `WRITES_IDX_MAP`, pinning special channels (`__error__`, `__scheduled__`, `__interrupt__`, `__resume__`) to fixed negative indices instead of the call-local ordinal. Previously a mixed `putWrites([[...regular...], [INTERRUPT, …]], taskId)` placed the INTERRUPT at a positive idx that could collide with a regular write at the same `(task_id, idx)`, and the unconditional `$set` upsert silently overwrote whichever row landed there first. The conflict-resolution clause now matches the Postgres / SQLite (TS and Python) checkpointers: `$set` only when every channel is a special one, `$setOnInsert` otherwise. diff --git a/libs/checkpoint-mongodb/src/checkpoint.ts b/libs/checkpoint-mongodb/src/checkpoint.ts index 07f9419039..17c61704f0 100644 --- a/libs/checkpoint-mongodb/src/checkpoint.ts +++ b/libs/checkpoint-mongodb/src/checkpoint.ts @@ -9,6 +9,7 @@ import { type PendingWrite, type CheckpointMetadata, CheckpointPendingWrite, + WRITES_IDX_MAP, } from "@langchain/langgraph-checkpoint"; export type MongoDBSaverParams = { @@ -376,6 +377,17 @@ export class MongoDBSaver extends BaseCheckpointSaver { { required: true } ); + // Conflict resolution matches the Python MongoDB checkpointer and the + // langgraph-checkpoint contract (BaseCheckpointSaver.put_writes): + // - When every write targets a special channel (ERROR / SCHEDULED / + // INTERRUPT / RESUME, each pinned to a negative `idx` by + // WRITES_IDX_MAP), we $set so e.g. INTERRUPT can be overwritten on + // RESUME. + // - Otherwise we $setOnInsert so a regular write from one task can + // never clobber a regular write that another concurrent task already + // stored at the same (task_id, idx). + const allSpecial = writes.every(([channel]) => channel in WRITES_IDX_MAP); + const operations = await Promise.all( writes.map(async ([channel, value], idx) => { const upsertQuery = { @@ -383,18 +395,29 @@ export class MongoDBSaver extends BaseCheckpointSaver { checkpoint_ns, checkpoint_id, task_id: taskId, - idx, + // Special channels are stored at fixed negative indices so they + // never collide with regular per-step writes (whose `idx` is the + // ordinal within `writes`). + idx: WRITES_IDX_MAP[channel] ?? idx, }; const [type, serializedValue] = await this.serde.dumpsTyped(value); + const fields = { channel, type, value: serializedValue }; return { updateOne: { filter: upsertQuery, - update: { - $set: { channel, type, value: serializedValue }, - ...this.timestampOp, - }, + update: allSpecial + ? { $set: fields, ...this.timestampOp } + : { + // Insert-or-ignore: the row's channel/value AND its + // upserted_at timestamp must only be written on first insert + // so a no-op update against a peer task's existing row + // doesn't bump that row's "last modified" stamp. + $setOnInsert: this.enableTimestamps + ? { ...fields, upserted_at: new Date() } + : fields, + }, upsert: true, }, }; diff --git a/libs/checkpoint-mongodb/src/tests/checkpoints.test.ts b/libs/checkpoint-mongodb/src/tests/checkpoints.test.ts index 64a62bd35d..1448e7e785 100644 --- a/libs/checkpoint-mongodb/src/tests/checkpoints.test.ts +++ b/libs/checkpoint-mongodb/src/tests/checkpoints.test.ts @@ -491,6 +491,91 @@ describe("MongoDBSaver", () => { ).rejects.toThrow('Invalid configurable.thread_id: expected a string'); }); + it("should pin special channels to fixed negative indices and switch to $setOnInsert for regular writes", async () => { + // Capture the bulkWrite operations so we can assert on their structure + // without standing up a real MongoDB. + const bulkWriteCalls: unknown[][] = []; + const writesCollection = { + find: vi.fn(() => ({ + sort: vi.fn(() => ({ + limit: vi.fn(() => ({ + toArray: vi.fn(() => Promise.resolve([])), + })), + })), + })), + bulkWrite: vi.fn((ops: unknown[]) => { + bulkWriteCalls.push(ops); + return Promise.resolve({}); + }), + }; + const client = { + appendMetadata: vi.fn(), + db: vi.fn(() => ({ + collection: vi.fn(() => writesCollection), + })), + }; + const saver = new MongoDBSaver({ + client: client as unknown as MongoClient, + }); + const config = { + configurable: { + thread_id: "t", + checkpoint_ns: "", + checkpoint_id: "c", + }, + }; + + // Mix of regular writes and a special-channel write in the same call. + // Regular writes must NOT push the special-channel onto a positive idx. + await saver.putWrites( + config, + [ + ["foo", "v_foo"], + ["bar", "v_bar"], + ["__interrupt__", "paused"], + ], + "task_A" + ); + + const ops = bulkWriteCalls[bulkWriteCalls.length - 1] as Array<{ + updateOne: { + filter: { idx: number }; + update: Record; + }; + }>; + const indices = ops.map((op) => op.updateOne.filter.idx).sort( + (a, b) => a - b + ); + // foo (idx 0), bar (idx 1), __interrupt__ (idx -3 via WRITES_IDX_MAP) + expect(indices).toEqual([-3, 0, 1]); + + // Mixed batch must go through $setOnInsert so a peer task's row at + // (task_A, idx=0) can't be silently overwritten. + for (const op of ops) { + expect(op.updateOne.update).toHaveProperty("$setOnInsert"); + expect(op.updateOne.update).not.toHaveProperty("$set"); + } + + // A separate call where every write is a special channel must go + // through $set so e.g. INTERRUPT → RESUME state transitions overwrite. + await saver.putWrites( + config, + [["__resume__", "carry_on"]], + "task_A" + ); + const specialOnly = bulkWriteCalls[bulkWriteCalls.length - 1] as Array<{ + updateOne: { + filter: { idx: number }; + update: Record; + }; + }>; + expect(specialOnly[0].updateOne.filter.idx).toBe(-4); // RESUME + expect(specialOnly[0].updateOne.update).toHaveProperty("$set"); + expect(specialOnly[0].updateOne.update).not.toHaveProperty( + "$setOnInsert" + ); + }); + it("should reject non-string thread_id in putWrites", async () => { const client = createMockClient(); const saver = new MongoDBSaver({