From ed8547018cefda1eb37d4afbadaf010e1a12183b Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 5 Jul 2026 08:30:06 -0400 Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=20perf:=20Persist=20HITL=20checkpoint?= =?UTF-8?q?s=20only=20on=20pause=20(lazy=20checkpointer)=20(#14024)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * ⚡ feat: Persist HITL checkpoints only on pause (skip clean-exit writes) With `durability: 'exit'` (set by the SDK whenever a checkpointer is active) LangGraph persists ONE checkpoint at the exit boundary on EVERY run — paused or not. So a non-paused HITL turn writes a dead checkpoint whose only fate is to be pruned by deleteAgentCheckpoint: pure write+delete churn on the common path, given HITL only ever resumes an *interrupt* checkpoint. `InterruptOnlyMongoSaver` (a MongoDBSaver subclass) persists only interrupt checkpoints and discards clean-exit ones, so a non-paused turn writes nothing. How it tells them apart (verified empirically against @langchain/langgraph, not docs): when a run interrupts, the runner calls `putWrites` with the `INTERRUPT` ("__interrupt__") channel for the checkpoint it's about to create, and that write's `config.checkpoint_id` equals the `checkpoint.id` of the `put` that immediately follows. A clean exit calls `put` with no preceding interrupt `putWrites`. So we record the checkpoint id of any interrupt `putWrites` and persist a `put` only when its `checkpoint.id` was so marked. Keying on the globally-unique checkpoint id (not thread_id) keeps this correct even when two runs race on the same conversation (the job-replacement scenario). Correctness is preserved end-to-end: interrupt checkpoints + their pending writes persist exactly as before (resume unchanged); clean checkpoints were only ever written-then-pruned, so not writing them is observationally equivalent. The eager prune stays as the backstop. Tests (mongodb-memory-server): a bare put() is discarded; an interrupt-seeded checkpoint is persisted with its __interrupt__ pending write; and an end-to-end real-graph run writes 0 checkpoints on a clean completion and a resumable one on interrupt. NOTE: a non-paused turn's deleteAgentCheckpoint now finds nothing to delete (a 0-match no-op) — a follow-up can skip that call entirely once the lingering-abandoned-pause cleanup role is reassigned to the TTL + expiry sweeper. * ⚡ feat: Drop the redundant clean-path checkpoint prune With the lazy checkpointer (InterruptOnlyMongoSaver) a non-paused turn no longer writes a clean-exit checkpoint, so the post-completion prune in chatCompletion's finally had nothing left to delete. It was also already redundant: every fresh turn runs a pre-run prune (`deleteAgentCheckpoint` before `processStream`) that clears any checkpoint orphaned by a prior abandoned pause — verified empirically that a lingering interrupt checkpoint WOULD otherwise poison a fresh turn (LangGraph continues the abandoned state + re-interrupts), and that the pre-run prune is what prevents it. The Mongo TTL remains the backstop, and the resume path still prunes after a successful finalize. Removing the clean-path prune also deletes its job-replacement race surface (round-17 F21): an older run's late finally can no longer delete a newer paused run's checkpoint, because there is no longer a clean-path prune to race. Dropped the now-dead F21 predicate test. Net per non-paused HITL turn: from {pre-run prune + checkpoint write + post-run prune} down to {pre-run prune} — no write, no post-completion delete. * 🛡️ fix: Anchor any pending-write checkpoint; stale-only eviction (Codex) Broaden the lazy saver's keep-rule from "interrupt-only" to "persist any checkpoint that carries pending writes" (renamed InterruptOnlyMongoSaver → LazyMongoSaver). This makes it robust to delta-channel graphs without changing behavior for LibreChat's graph: - K1 (P1): a delta-channel graph can write a synthetic PARENT/anchor checkpoint (no __interrupt__ mark) that the interrupt checkpoint then points at, with the delta writes stored under the parent id. The old rule discarded that parent, breaking delta-state resume. Now any checkpoint that received putWrites is persisted, so the anchor parent and its writes survive and resume can walk the chain. - K3 (P2): for the same reason, clean delta-write rows are no longer orphaned — their checkpoint is persisted alongside them. (For LibreChat's standard Annotation/messages graph a clean run makes no putWrites at all — verified empirically — so the common path still writes nothing and the optimization is unchanged.) - K2 (P2): the 1024 FIFO cap could evict a valid in-flight id whose put() was just behind Mongo I/O, mis-classifying its interrupt checkpoint as a clean exit. Replaced with time-based eviction: only ids older than 5 min (a put always follows its putWrites within ms) are swept; a recent in-flight id is never dropped, and the map grows rather than evict a valid id if nothing is stale. New integration test: a checkpoint anchored by a NON-interrupt write is persisted. Full agents/HITL suites green (108). * style(checkpointer): fix import order to satisfy sort-imports CI * fix(checkpointer): don't persist failed-turn (error-only) checkpoints LazyMongoSaver anchored on ANY pending write, so a non-paused turn that errors (LangGraph records an __error__ write then a put) was persisted and, with the clean-path prune removed, lingered until the next fresh-turn prune or the Mongo TTL. Anchor only on resumable writes — INTERRUPT or a real (non-__-prefixed) state/delta channel — so error/bookkeeping-only checkpoints are discarded at the source. Addresses Codex P3. Codex P2 (delta-stub parent orphan) is not reachable: the SDK graph uses standard Annotation/MessagesAnnotation channels (no DeltaChannel), and under durability:'exit' putWrites precedes put with a parentless boundary checkpoint — probe-confirmed against @langchain/langgraph@1.4. Documented the durability:'exit' invariant the saver depends on. Tests: error-only put discarded; e2e throwing graph persists 0 checkpoints. * fix(checkpointer): drop bookkeeping-only write batches, not just the checkpoint The prior fix stopped the failed-turn CHECKPOINT from persisting, but putWrites still forwarded the __error__ batch to MongoDBSaver.putWrites — writing a row to agent_checkpoint_writes whose parent checkpoint is then discarded. With the post-run deleteThread removed, that orphan row lingered until the Mongo TTL or the conversation's next pre-run prune. putWrites now drops a non-resumable (bookkeeping-only) batch entirely instead of forwarding it. Probed against a real MongoDBSaver (mongodb-memory-server): a throwing graph now leaves 0 checkpoints AND 0 write rows (was 0 + 1 orphan), while interrupt->resume is unaffected — the __interrupt__ write is resumable so it is still forwarded. Addresses Codex P2 (round 3). Tests: error-only put leaves no checkpoint and no write row; e2e throwing graph leaves both collections empty; new e2e interrupt->resume completes with the approval value. * fix(checkpointer): un-anchor a checkpoint whose putWrites failed; freshen comments Self-review findings on the converged PR: 1. LangGraph dispatches put() concurrently with putWrites (probe-confirmed on 1.4.5), and put() still completes when putWrites rejects — so a transient Mongo failure during the interrupt write could persist a checkpoint whose __interrupt__ row is missing (an unresumable phantom pause). putWrites now deletes the write anchor on rejection (best-effort) and rethrows, so that put() discards the checkpoint instead. The pre-recorded anchor stays where it is — recording after the await would drop slow-I/O interrupts on the success path, which the same probe showed is reachable. 2. Renamed leftovers: two comments still said InterruptOnlyMongoSaver; the class is LazyMongoSaver. 3. Documented why the pre-run prune is deliberately unconditional per HITL turn (any cheaper gate can go stale across replicas and skip the prune exactly when an orphaned interrupt exists). Test: failed putWrites → subsequent put persists nothing (14/14 green). * fix(checkpointer): bookkeeping write batches follow their checkpoint's fate The round-3 rule dropped bookkeeping-only putWrites batches (__error__/ __resume__/__no_writes__) unconditionally — batch-scoped, when the decision must be checkpoint-scoped. Probe-confirmed (langgraph 1.4.5, durability:'exit'): a Send fan-out that pauses on one sibling records the completed siblings as pure __no_writes__ batches on the RETAINED interrupt checkpoint; dropping those markers makes resume re-execute the completed siblings (side effects measured twice). Addresses Codex M2 (P2). putWrites now PARKS a bookkeeping-only batch in memory until the checkpoint's fate is known: forwarded when the checkpoint is anchored (or was just persisted — put is dispatched concurrently), dropped when put discards it. Net: an errored turn still leaves nothing durable (0 checkpoints, 0 write rows), and a retained checkpoint stores byte-for-byte what a plain MongoDBSaver would. Codex M1 (__resume__ lost on re-pause) did not reproduce: the re-pause emits [__interrupt__,__resume__] as ONE batch (anchored, forwarded whole) and a second resume on a rebuilt graph replays both answers correctly — but the fate-scoped buffering now covers a lone __resume__ batch in any ordering too. Tests: bookkeeping preserved on a retained checkpoint in either arrival order; e2e Send-sibling pause/resume with side-effect counters (was {a:2,c:2} under the drop rule, now {a:1,c:1}); error-only turn still leaves both collections empty. 16/16 green. --- .../agents/__tests__/jobReplacement.spec.js | 31 +- api/server/controllers/agents/client.js | 40 +-- package-lock.json | 2 + packages/api/package.json | 2 + .../agents/checkpointer.integration.spec.ts | 317 +++++++++++++++++- packages/api/src/agents/checkpointer.ts | 232 ++++++++++++- 6 files changed, 561 insertions(+), 63 deletions(-) diff --git a/api/server/controllers/agents/__tests__/jobReplacement.spec.js b/api/server/controllers/agents/__tests__/jobReplacement.spec.js index c4247fcf30..ec70b66854 100644 --- a/api/server/controllers/agents/__tests__/jobReplacement.spec.js +++ b/api/server/controllers/agents/__tests__/jobReplacement.spec.js @@ -337,32 +337,11 @@ describe('HITL Terminal-Side-Effect Guards', () => { }); }); - describe('F21 — checkpoint prune is skipped when the generation was replaced', () => { - // Mirrors client.js chatCompletion finally: prune only when NOT replaced. - const shouldPrune = async ({ resumableStreamId, jobCreatedAt }) => { - let replaced = false; - if (resumableStreamId && jobCreatedAt != null) { - const liveJob = await mockGenerationJobManager.getJob(resumableStreamId); - replaced = !liveJob || liveJob.createdAt !== jobCreatedAt; - } - return !replaced; - }; - - it('skips the prune when a newer job replaced this one', async () => { - mockGenerationJobManager.getJob.mockResolvedValue({ createdAt: 2000 }); - expect(await shouldPrune({ resumableStreamId: 'c1', jobCreatedAt: 1000 })).toBe(false); - }); - - it('prunes when this is still the live job', async () => { - mockGenerationJobManager.getJob.mockResolvedValue({ createdAt: 1000 }); - expect(await shouldPrune({ resumableStreamId: 'c1', jobCreatedAt: 1000 })).toBe(true); - }); - - it('prunes without a lookup when there is no resumable stream id', async () => { - expect(await shouldPrune({ resumableStreamId: undefined, jobCreatedAt: 1000 })).toBe(true); - expect(mockGenerationJobManager.getJob).not.toHaveBeenCalled(); - }); - }); + // (Removed: F21 — the chatCompletion clean-path checkpoint prune + its job-replacement + // guard no longer exist. The lazy checkpointer never writes a clean-exit checkpoint, so + // there is nothing to prune after a non-paused turn; the pre-run prune (before + // processStream) clears any orphaned interrupt checkpoint instead. See + // checkpointer.ts LazyMongoSaver and client.js chatCompletion.) describe('F24 — resume catch-path terminal writes are skipped when replaced', () => { // Mirrors resume.js: stillLive gate around emitError/completeJob/deleteAgentCheckpoint. diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index f9432ae821..8b639f01a8 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -1615,6 +1615,10 @@ class AgentClient extends BaseClient { // conversation (one that expired or was aborted while paused) so this fresh // turn starts clean instead of rehydrating a stale interrupt — thread_id is // the stable conversationId. No-op when HITL is off or nothing is orphaned. + // Deliberately UNCONDITIONAL per HITL turn: any cheaper gate (job metadata, + // a Redis flag) can go stale across replicas/restarts and skip the prune + // exactly when an orphan exists, while these are two indexed, usually-empty + // deleteMany ops — correctness over a micro-optimization. if (streamId && isHITLEnabled(agentsEConfig?.toolApproval)) { await deleteAgentCheckpoint(this.conversationId, agentsEConfig?.checkpointer); } @@ -1743,35 +1747,13 @@ class AgentClient extends BaseClient { this._resolveRun = null; } - // HITL: a turn that completed (or errored) without pausing leaves a dead - // checkpoint. thread_id is the conversationId — stable across turns — so it - // MUST be pruned before the next turn, or LangGraph would resume this turn's - // state instead of starting fresh. Skip when paused (the checkpoint is needed - // to resume) or when HITL is off (none was written). The Mongo TTL is the backstop. - const agentsEConfig = appConfig?.endpoints?.[EModelEndpoint.agents]; - if (!this.pendingApproval && isHITLEnabled(agentsEConfig?.toolApproval)) { - try { - // Job-replacement guard: only prune if THIS generation is still the live job. - // A newer request can replace this one on the same conversationId; if this - // (older) run's finally lands after the newer run paused, pruning by - // conversationId would delete the NEWER run's checkpoint and break its /resume. - const resumableStreamId = this.options.req?._resumableStreamId; - let replaced = false; - if (resumableStreamId && this.jobCreatedAt != null) { - const liveJob = await GenerationJobManager.getJobStore().getJob(resumableStreamId); - replaced = !liveJob || liveJob.createdAt !== this.jobCreatedAt; - } - if (replaced) { - logger.debug('[AgentClient] Skipping checkpoint prune — job was replaced', { - streamId: resumableStreamId, - }); - } else { - await deleteAgentCheckpoint(this.conversationId, agentsEConfig?.checkpointer); - } - } catch (err) { - logger.warn('[AgentClient] Failed to prune checkpoint after completion', err); - } - } + // HITL: a non-paused turn deliberately prunes nothing here. The lazy checkpointer + // (LazyMongoSaver) never persists a clean-exit checkpoint, so there is + // nothing this turn left to delete. A checkpoint orphaned by a PRIOR abandoned pause + // is cleared by the pre-run prune (before processStream) on the next turn, with the + // Mongo TTL as the backstop. Dropping this post-completion prune also removes its + // job-replacement race: an older run's late finally can no longer delete a newer + // paused run's checkpoint, because there is no longer a clean-path prune to race. run = null; config = null; diff --git a/package-lock.json b/package-lock.json index 09f9455322..6efb68eb16 100644 --- a/package-lock.json +++ b/package-lock.json @@ -42345,12 +42345,14 @@ "version": "1.7.34", "license": "ISC", "dependencies": { + "@langchain/langgraph-checkpoint": "^1.1.2", "@langchain/langgraph-checkpoint-mongodb": "^1.4.0" }, "devDependencies": { "@babel/preset-env": "^7.29.5", "@babel/preset-react": "^7.18.6", "@babel/preset-typescript": "^7.21.0", + "@langchain/langgraph": "^1.4.5", "@rollup/plugin-alias": "^5.1.0", "@rollup/plugin-commonjs": "^29.0.0", "@rollup/plugin-json": "^6.1.0", diff --git a/packages/api/package.json b/packages/api/package.json index 0d663ccd05..079f0d5464 100644 --- a/packages/api/package.json +++ b/packages/api/package.json @@ -56,6 +56,7 @@ "@babel/preset-env": "^7.29.5", "@babel/preset-react": "^7.18.6", "@babel/preset-typescript": "^7.21.0", + "@langchain/langgraph": "^1.4.5", "@rollup/plugin-alias": "^5.1.0", "@rollup/plugin-commonjs": "^29.0.0", "@rollup/plugin-json": "^6.1.0", @@ -165,6 +166,7 @@ "zod": "^3.22.4" }, "dependencies": { + "@langchain/langgraph-checkpoint": "^1.1.2", "@langchain/langgraph-checkpoint-mongodb": "^1.4.0" } } diff --git a/packages/api/src/agents/checkpointer.integration.spec.ts b/packages/api/src/agents/checkpointer.integration.spec.ts index ae10502608..e11451c54e 100644 --- a/packages/api/src/agents/checkpointer.integration.spec.ts +++ b/packages/api/src/agents/checkpointer.integration.spec.ts @@ -1,6 +1,7 @@ import mongoose from 'mongoose'; import { MongoMemoryServer } from 'mongodb-memory-server'; -import { emptyCheckpoint } from '@langchain/langgraph-checkpoint'; +import { MongoDBSaver } from '@langchain/langgraph-checkpoint-mongodb'; +import { emptyCheckpoint, ERROR, INTERRUPT } from '@langchain/langgraph-checkpoint'; import { getAgentCheckpointer, deleteAgentCheckpoint, @@ -28,6 +29,22 @@ const readConfig = (threadId: string) => ({ configurable: { thread_id: threadId, checkpoint_ns: '' }, }); +/** + * Persist a checkpoint the way LangGraph does on a PAUSE: an interrupt `putWrites` on the + * `INTERRUPT` channel for the checkpoint id, then `put` of that checkpoint. The lazy saver + * persists only checkpoints seeded this way (a bare `put` is a clean exit → discarded). + */ +async function seedInterruptCheckpoint(saver: MongoDBSaver, threadId: string) { + const { config, checkpoint, metadata } = putArgs(threadId); + await saver.putWrites( + { configurable: { thread_id: threadId, checkpoint_ns: '', checkpoint_id: checkpoint.id } }, + [[INTERRUPT, 'approve?']], + 'task-1', + ); + await saver.put(config, checkpoint, metadata); + return checkpoint; +} + let mongoServer: MongoMemoryServer; beforeAll(async () => { @@ -75,8 +92,7 @@ describe('checkpointer (mongodb-memory-server integration)', () => { expect(saver).toBeDefined(); const threadId = `convo-${new mongoose.Types.ObjectId().toString()}`; - const { config, checkpoint, metadata } = putArgs(threadId); - await saver!.put(config, checkpoint, metadata); + await seedInterruptCheckpoint(saver!, threadId); // The checkpoint is durably readable before pruning… expect(await saver!.getTuple(readConfig(threadId))).toBeDefined(); @@ -92,10 +108,8 @@ describe('checkpointer (mongodb-memory-server integration)', () => { const threadA = `convo-${new mongoose.Types.ObjectId().toString()}`; const threadB = `convo-${new mongoose.Types.ObjectId().toString()}`; - const a = putArgs(threadA); - const b = putArgs(threadB); - await saver!.put(a.config, a.checkpoint, a.metadata); - await saver!.put(b.config, b.checkpoint, b.metadata); + await seedInterruptCheckpoint(saver!, threadA); + await seedInterruptCheckpoint(saver!, threadB); await deleteAgentCheckpoint(threadA, MONGO_CFG); @@ -107,3 +121,292 @@ describe('checkpointer (mongodb-memory-server integration)', () => { await expect(deleteAgentCheckpoint(undefined, MONGO_CFG)).resolves.toBeUndefined(); }); }); + +describe('LazyMongoSaver (lazy persistence — mongodb-memory-server)', () => { + it('does NOT persist a clean-exit checkpoint (a bare put with no pending writes)', async () => { + const saver = await getAgentCheckpointer(MONGO_CFG); + const threadId = `convo-${new mongoose.Types.ObjectId().toString()}`; + const { config, checkpoint, metadata } = putArgs(threadId); + + // A non-paused run's exit put — no preceding putWrites. + await saver!.put(config, checkpoint, metadata); + + expect(await saver!.getTuple(readConfig(threadId))).toBeUndefined(); + const count = await mongoose.connection + .db!.collection('agent_checkpoints') + .countDocuments({ thread_id: threadId }); + expect(count).toBe(0); + }); + + it('persists a checkpoint anchored by a NON-interrupt write (delta-channel safety)', async () => { + // K1/K3: a delta-channel graph can putWrites on a checkpoint that an interrupt + // checkpoint then depends on — even without the __interrupt__ marker. A write on a + // real (non-`__`-prefixed) channel must anchor its checkpoint so resume can walk the chain. + const saver = await getAgentCheckpointer(MONGO_CFG); + const threadId = `convo-${new mongoose.Types.ObjectId().toString()}`; + const { config, checkpoint, metadata } = putArgs(threadId); + await saver!.putWrites( + { configurable: { thread_id: threadId, checkpoint_ns: '', checkpoint_id: checkpoint.id } }, + [['some_delta_channel', { msgs: ['delta'] }]], + 'task-1', + ); + await saver!.put(config, checkpoint, metadata); + + expect(await saver!.getTuple(readConfig(threadId))).toBeDefined(); + }); + + it('does NOT persist an error-only checkpoint OR its write row (failed non-paused turn — no leak)', async () => { + // A turn that throws before any pause records a pending write on the `__error__` + // bookkeeping channel, then a `put` (probe-confirmed against @langchain/langgraph). That + // checkpoint is never HITL-resumable, so the lazy saver must leave NOTHING durable: the + // bookkeeping batch is PARKED (not forwarded) until the checkpoint's fate is known, and the + // discarding `put` drops both — no dead checkpoint, no orphan row in the writes collection. + const saver = await getAgentCheckpointer(MONGO_CFG); + const threadId = `convo-${new mongoose.Types.ObjectId().toString()}`; + const { config, checkpoint, metadata } = putArgs(threadId); + await saver!.putWrites( + { configurable: { thread_id: threadId, checkpoint_ns: '', checkpoint_id: checkpoint.id } }, + [[ERROR, 'boom']], // '__error__' — bookkeeping channel, not resumable + 'task-1', + ); + await saver!.put(config, checkpoint, metadata); + + expect(await saver!.getTuple(readConfig(threadId))).toBeUndefined(); + const db = mongoose.connection.db!; + expect(await db.collection('agent_checkpoints').countDocuments({ thread_id: threadId })).toBe( + 0, + ); + expect( + await db.collection('agent_checkpoint_writes').countDocuments({ thread_id: threadId }), + ).toBe(0); + }); + + it('preserves bookkeeping writes on a RETAINED checkpoint, in either arrival order', async () => { + // Codex M2 (probe-confirmed): when one Send-sibling interrupts, siblings that completed + // without state updates are recorded as `__no_writes__` pending writes on the SAME retained + // checkpoint. Those markers must persist — dropping them makes LangGraph re-execute the + // completed siblings on resume (duplicated side effects). The saver parks bookkeeping + // batches until the checkpoint is anchored, so both orders must end durable. + const NO_WRITES = '__no_writes__'; // langgraph constants.NO_WRITES (runner marker) + const saver = await getAgentCheckpointer(MONGO_CFG); + + const runOrder = async (bookkeepingFirst: boolean) => { + const threadId = `convo-${new mongoose.Types.ObjectId().toString()}`; + const { config, checkpoint, metadata } = putArgs(threadId); + const writeCfg = { + configurable: { thread_id: threadId, checkpoint_ns: '', checkpoint_id: checkpoint.id }, + }; + const bookkeeping = () => saver!.putWrites(writeCfg, [[NO_WRITES, null]], 'task-sibling'); + const anchor = () => saver!.putWrites(writeCfg, [[INTERRUPT, 'approve?']], 'task-gate'); + if (bookkeepingFirst) { + await bookkeeping(); + await anchor(); + } else { + await anchor(); + await bookkeeping(); + } + await saver!.put(config, checkpoint, metadata); + + const tuple = await saver!.getTuple(readConfig(threadId)); + expect(tuple).toBeDefined(); + const channels = (tuple?.pendingWrites ?? []).map((w) => w[1]); + expect(channels).toContain(INTERRUPT); + expect(channels).toContain(NO_WRITES); + }; + + await runOrder(true); // parked, then flushed by the anchoring batch + await runOrder(false); // forwarded directly (checkpoint already anchored) + }); + + it('un-anchors a checkpoint whose putWrites failed (no phantom pause persisted)', async () => { + // A transient Mongo failure while writing the interrupt batch must not leave a persisted + // checkpoint with MISSING pending writes: LangGraph dispatches the matching `put` + // concurrently with `putWrites` (probe-confirmed), so the failed batch's anchor is removed + // on rejection and that `put` discards the checkpoint instead of saving a phantom pause. + const saver = await getAgentCheckpointer(MONGO_CFG); + const threadId = `convo-${new mongoose.Types.ObjectId().toString()}`; + const { config, checkpoint, metadata } = putArgs(threadId); + + const spy = jest + .spyOn(MongoDBSaver.prototype, 'putWrites') + .mockRejectedValueOnce(new Error('transient mongo failure')); + try { + await expect( + saver!.putWrites( + { + configurable: { thread_id: threadId, checkpoint_ns: '', checkpoint_id: checkpoint.id }, + }, + [[INTERRUPT, 'approve?']], + 'task-1', + ), + ).rejects.toThrow('transient mongo failure'); + } finally { + spy.mockRestore(); + } + + // The `put` LangGraph issues for that checkpoint finds no anchor → discarded. + await saver!.put(config, checkpoint, metadata); + + expect(await saver!.getTuple(readConfig(threadId))).toBeUndefined(); + const count = await mongoose.connection + .db!.collection('agent_checkpoints') + .countDocuments({ thread_id: threadId }); + expect(count).toBe(0); + }); + + it('persists an interrupt checkpoint and carries its __interrupt__ pending write', async () => { + const saver = await getAgentCheckpointer(MONGO_CFG); + const threadId = `convo-${new mongoose.Types.ObjectId().toString()}`; + await seedInterruptCheckpoint(saver!, threadId); + + const tuple = await saver!.getTuple(readConfig(threadId)); + expect(tuple).toBeDefined(); + // pendingWrites entries are [taskId, channel, value]; the interrupt is on INTERRUPT. + expect((tuple?.pendingWrites ?? []).some((w) => w[1] === INTERRUPT)).toBe(true); + }); + + it('end-to-end: a real graph writes nothing on a clean run, a checkpoint on interrupt', async () => { + const { StateGraph, START, END, interrupt, Annotation } = await import('@langchain/langgraph'); + const saver = await getAgentCheckpointer(MONGO_CFG); + const coll = mongoose.connection.db!.collection('agent_checkpoints'); + + const State = Annotation.Root({ x: Annotation }); + const build = (withInterrupt: boolean) => + new StateGraph(State) + .addNode('a', () => { + if (withInterrupt) { + interrupt('approve?'); + } + return { x: 'done' }; + }) + .addEdge(START, 'a') + .addEdge('a', END) + // version skew between checkpoint-mongodb's BaseCheckpointSaver and langgraph's. + .compile({ checkpointer: saver as never }); + + // Clean run → nothing persisted. + const tClean = `convo-${new mongoose.Types.ObjectId().toString()}`; + await build(false).invoke( + { x: 'start' }, + { configurable: { thread_id: tClean }, durability: 'exit' }, + ); + expect(await coll.countDocuments({ thread_id: tClean })).toBe(0); + + // Interrupt run → a checkpoint is persisted and the graph reports a pending pause. + const tPause = `convo-${new mongoose.Types.ObjectId().toString()}`; + const pauseGraph = build(true); + await pauseGraph.invoke( + { x: 'start' }, + { configurable: { thread_id: tPause }, durability: 'exit' }, + ); + expect(await coll.countDocuments({ thread_id: tPause })).toBeGreaterThan(0); + const state = await pauseGraph.getState({ configurable: { thread_id: tPause } }); + expect(state.next.length).toBeGreaterThan(0); // the interrupted node is still pending → resumable + }); + + it('end-to-end: a real graph that THROWS before pausing persists no checkpoint', async () => { + // F2: a failed non-paused turn records an `__error__` pending write + a put. The lazy + // saver must discard it so a conversation that errors (and is never retried) leaves nothing + // durable behind — the clean-path prune that used to catch this was removed. + const { StateGraph, START, END, Annotation } = await import('@langchain/langgraph'); + const saver = await getAgentCheckpointer(MONGO_CFG); + const coll = mongoose.connection.db!.collection('agent_checkpoints'); + + const State = Annotation.Root({ x: Annotation }); + const boomGraph = new StateGraph(State) + .addNode('a', () => { + throw new Error('boom'); + }) + .addEdge(START, 'a') + .addEdge('a', END) + .compile({ checkpointer: saver as never }); + + const tErr = `convo-${new mongoose.Types.ObjectId().toString()}`; + await expect( + boomGraph.invoke({ x: 'start' }, { configurable: { thread_id: tErr }, durability: 'exit' }), + ).rejects.toThrow('boom'); + + // Nothing durable: neither the checkpoint nor an orphan row in the writes collection. + expect(await coll.countDocuments({ thread_id: tErr })).toBe(0); + const writesColl = mongoose.connection.db!.collection('agent_checkpoint_writes'); + expect(await writesColl.countDocuments({ thread_id: tErr })).toBe(0); + }); + + it('end-to-end: an interrupt persists, then resumes to completion with the approval value', async () => { + // Guards the `putWrites` change: the `__interrupt__` write must still be forwarded (it is + // resumable) so a paused run rehydrates and the resume value flows in. Mirrors the real HITL + // round-trip across a fresh `invoke` on the same thread_id. + const { StateGraph, START, END, interrupt, Annotation, Command } = await import( + '@langchain/langgraph' + ); + const saver = await getAgentCheckpointer(MONGO_CFG); + + const State = Annotation.Root({ approved: Annotation }); + const graph = new StateGraph(State) + .addNode('gate', () => ({ approved: interrupt('approve?') })) + .addNode('done', () => ({})) + .addEdge(START, 'gate') + .addEdge('gate', 'done') + .addEdge('done', END) + .compile({ checkpointer: saver as never }); + + const thread = `convo-${new mongoose.Types.ObjectId().toString()}`; + const cfg = { configurable: { thread_id: thread }, durability: 'exit' as const }; + + // Pause at the interrupt. + await graph.invoke({ approved: null }, cfg); + const paused = await graph.getState(cfg); + expect(paused.next.length).toBeGreaterThan(0); + + // Resume with the approval — the paused run rehydrates from the durable interrupt checkpoint. + const out = await graph.invoke(new Command({ resume: 'YES' }), cfg); + expect(out.approved).toBe('YES'); + }); + + it('end-to-end: completed Send-siblings are NOT re-executed after a pause (no duplicate side effects)', async () => { + // Codex M2 end-to-end: a Send fan-out where 'b' pauses for approval while 'a'/'c' complete + // with side effects but NO state writes (→ `__no_writes__` markers on the retained + // checkpoint). On resume — through a REBUILT graph, as resume.js rebuilds the Run — the + // completed siblings must not run again. Before the buffering fix this probe measured + // effects {a:2, c:2}: the dropped markers made LangGraph re-execute both siblings. + const { StateGraph, START, END, interrupt, Annotation, Send, Command } = await import( + '@langchain/langgraph' + ); + const saver = await getAgentCheckpointer(MONGO_CFG); + + const effects: Record = {}; + const State = Annotation.Root({ + items: Annotation({ reducer: (a: string[], b: string[]) => (a ?? []).concat(b ?? []) }), + results: Annotation({ reducer: (a: string[], b: string[]) => (a ?? []).concat(b ?? []) }), + }); + const build = () => + new StateGraph(State) + .addNode('fan', () => ({})) + .addNode('work', (s: { item?: string }) => { + if (s.item === 'b') { + return { results: [`b:${interrupt('approve b?')}`] }; + } + effects[s.item!] = (effects[s.item!] ?? 0) + 1; // side effect, no state update + return {}; // no channel writes → langgraph records a `__no_writes__` marker + }) + .addConditionalEdges( + 'fan', + (s) => s.items.map((i: string) => new Send('work', { item: i })), + ['work'], + ) + .addEdge(START, 'fan') + .addEdge('work', END) + .compile({ checkpointer: saver as never }); + + const thread = `convo-${new mongoose.Types.ObjectId().toString()}`; + const cfg = { configurable: { thread_id: thread }, durability: 'exit' as const }; + + await build().invoke({ items: ['a', 'b', 'c'], results: [] }, cfg); // pauses on 'b' + expect(effects).toEqual({ a: 1, c: 1 }); + + // Resume on a REBUILT graph sharing the durable saver (mirrors resume.js). + const out = await build().invoke(new Command({ resume: 'YES' }), cfg); + expect(out.results).toEqual(['b:YES']); + expect(effects).toEqual({ a: 1, c: 1 }); // siblings did NOT re-execute + }); +}); diff --git a/packages/api/src/agents/checkpointer.ts b/packages/api/src/agents/checkpointer.ts index 6a99a616aa..2a24810e03 100644 --- a/packages/api/src/agents/checkpointer.ts +++ b/packages/api/src/agents/checkpointer.ts @@ -1,7 +1,10 @@ import mongoose from 'mongoose'; import { logger } from '@librechat/data-schemas'; +import { INTERRUPT } from '@langchain/langgraph-checkpoint'; import { MongoDBSaver } from '@langchain/langgraph-checkpoint-mongodb'; +import type { Checkpoint, CheckpointMetadata, PendingWrite } from '@langchain/langgraph-checkpoint'; import type { TCheckpointerConfig } from 'librechat-data-provider'; +import type { RunnableConfig } from '@langchain/core/runnables'; /** * Durable checkpointing for human-in-the-loop (HITL) resume. @@ -22,6 +25,233 @@ import type { TCheckpointerConfig } from 'librechat-data-provider'; * prunes a thread's checkpoints eagerly on every terminal transition. */ +/** + * Soft size threshold that triggers a sweep of STALE write-anchor ids. The map normally + * holds a handful (each id is consumed by the `put` that immediately follows its + * `putWrites`); this only bounds a slow leak from a process that dies in that window. + */ +const WRITE_ANCHOR_SWEEP_THRESHOLD = 1024; + +/** + * A write-anchor id is considered stale once this much wall-clock has passed without its + * matching `put` — a `put` always follows its `putWrites` within the same exit sequence + * (milliseconds), so anything this old is from a crashed run, never a valid in-flight id. + * Generous on purpose: we would rather keep a tracked id slightly too long than evict a + * valid one and mis-classify its (possibly slow-I/O) interrupt `put` as a clean exit. + */ +const WRITE_ANCHOR_STALE_MS = 5 * 60 * 1000; + +/** + * Does a pending-write batch make its checkpoint worth persisting (ANCHOR it)? True if it + * carries an interrupt (the HITL pause that resume targets) or any real state/delta channel + * write (a value a later checkpoint's resume depends on). False for pure bookkeeping batches — + * `__error__` (a failed, non-paused turn), `__no_writes__` (a task that completed without state + * updates), a lone `__resume__`, `__scheduled__` — which never justify keeping a checkpoint on + * their own. A false verdict does NOT mean the batch is discarded: bookkeeping rows are still + * required when the checkpoint is retained (see the buffering in `LazyMongoSaver.putWrites`); + * this predicate only decides anchoring. + * + * `INTERRUPT` is the one `__`-prefixed channel that IS anchor-worthy; every other `__…__` + * channel is langgraph bookkeeping. Constants verified against `@langchain/langgraph`. + */ +function hasResumableWrite(writes: PendingWrite[]): boolean { + return (writes ?? []).some(([channel]) => { + const name = String(channel); + return name === INTERRUPT || !name.startsWith('__'); + }); +} + +/** + * A `MongoDBSaver` that persists ONLY checkpoints carrying a {@link hasResumableWrite resumable} + * pending write — an interrupt (a HITL pause) or a real-channel/delta anchor — and discards both + * the no-write checkpoint LangGraph writes on a CLEAN exit and the bookkeeping-only checkpoint of + * a failed (non-paused) turn. + * + * **Why.** With `durability: 'exit'` (set by the SDK whenever a checkpointer is active) the + * graph persists exactly one checkpoint at the exit boundary on EVERY run — paused or not. + * A non-paused turn therefore writes a dead checkpoint whose only fate is to be pruned by + * {@link deleteAgentCheckpoint}. HITL only ever resumes a checkpoint that has pending writes, + * so the clean (write-less) exit checkpoint is pure write+delete churn on the common path. + * This saver skips it. + * + * **How it tells them apart** (verified empirically with throwaway runnable probes against + * `@langchain/langgraph@1.4`, not source-reading): under `durability: 'exit'` LangGraph + * calls `putWrites` for a checkpoint BEFORE the `put` that creates it, with `config.checkpoint_id` + * equal to that `put`'s `checkpoint.id`. An interrupt records an `INTERRUPT` ("__interrupt__") + * write; a delta-channel graph records its delta writes on a real (non-`__`-prefixed) channel. + * A CLEAN exit produces a checkpoint with NO pending writes. So we record the checkpoint id of + * each `putWrites` that carries a {@link hasResumableWrite resumable} write and persist a `put` + * only when its `checkpoint.id` was so marked — which keeps interrupt checkpoints AND any + * real-channel/delta anchors (resume can walk the chain), while dropping the write-less clean + * exit. Keying on the globally-unique checkpoint id (NOT thread_id) stays correct even when two + * runs race on the same conversation (`thread_id`). + * + * **Bookkeeping batches follow their checkpoint's fate.** Only a {@link hasResumableWrite + * resumable} batch ANCHORS a checkpoint (justifies persisting it); a bookkeeping-only batch + * (`__error__`, `__no_writes__`, a lone `__resume__`, …) never does — but whether its ROWS matter + * depends on whether the checkpoint survives, which `put` decides later. Probe-confirmed both + * ways: a failed non-paused turn emits `putWrites([__error__])` + `put` — persisting either half + * would leak (an orphan row or a dead checkpoint) — while a paused Send fan-out records completed + * siblings as `__no_writes__` markers on the RETAINED interrupt checkpoint, and dropping those + * re-executes the siblings on resume (duplicated side effects). So bookkeeping batches are + * BUFFERED in memory until the fate is known: forwarded once the checkpoint is anchored (or was + * just persisted), dropped when its `put` discards it. Net effect: an errored turn still leaves + * NOTHING durable (0 checkpoints, 0 write rows), and a retained checkpoint keeps EVERY pending + * write LangGraph recorded for it — byte-for-byte what a plain `MongoDBSaver` would store. + * + * For LibreChat's agent graph (standard `Annotation`/`MessagesAnnotation` channels, no + * `DeltaChannel` — grep-confirmed in `@librechat/agents`) a clean run makes no `putWrites` at all, + * so this is effectively interrupt-only and the common path writes nothing; the broader + * real-channel rule just keeps it honest for delta graphs. + * + * **Invariant.** Correctness depends on `durability: 'exit'` (which the SDK sets whenever a + * checkpointer is active): exactly one parentless boundary checkpoint per run, with its + * `putWrites` ordered before its `put`. Under per-step durability LangGraph instead emits + * `put`-before-`putWrites` for chained checkpoints — the anchor would arrive too late and a + * checkpoint could be wrongly discarded. The SDK never runs HITL that way; if that ever changes, + * this saver must be revisited (a parent-based guard is NOT viable — a resumed turn's clean + * completion is itself a parented, write-less checkpoint that we correctly discard). + * + * **Correctness.** Checkpoints with resumable writes (interrupt + real-channel/delta anchor) and + * the writes themselves persist exactly as before, so resume is unchanged. The write-less clean + * checkpoint (and the now-discarded error-only checkpoint) was only ever written-then-pruned, so + * not writing it is observationally equivalent; the pre-run prune + Mongo TTL remain the + * backstops. `getTuple`/`list`/`deleteThread`/`setup` are inherited. + */ +/** A bookkeeping-only pending-write batch held until its checkpoint's fate is decided. */ +interface BufferedWriteBatch { + at: number; + batches: Array<{ config: RunnableConfig; writes: PendingWrite[]; taskId: string }>; +} + +export class LazyMongoSaver extends MongoDBSaver { + /** checkpoint id → time the resumable `putWrites` anchoring it arrived; consumed by `put`. */ + private readonly writeAnchorIds = new Map(); + /** checkpoint id → time its anchored `put` persisted it, so a bookkeeping batch that lands + * after the `put` (concurrent dispatch) is forwarded instead of buffered forever. */ + private readonly persistedIds = new Map(); + /** checkpoint id → bookkeeping batches parked until the checkpoint persists or is discarded. */ + private readonly bufferedBookkeeping = new Map(); + + override async putWrites( + config: RunnableConfig, + writes: PendingWrite[], + taskId: string, + ): Promise { + const checkpointId = config.configurable?.checkpoint_id as string | undefined; + if (!checkpointId) { + // No checkpoint id to tie a fate to — forward untouched (the base saver's contract). + return super.putWrites(config, writes, taskId); + } + if (!hasResumableWrite(writes)) { + // A bookkeeping-only batch (`__error__` from a failed turn, a completed Send-sibling's + // `__no_writes__` marker, a lone `__resume__`, …). It must NOT anchor the checkpoint, + // but its rows follow the checkpoint's fate: required on a RETAINED checkpoint + // (probe-confirmed — dropping a sibling's `__no_writes__` marker re-executes the + // sibling on resume), an orphan on a discarded one. Forward when the fate is already + // known to be "persist"; otherwise buffer until an anchoring batch or `put` decides. + if (this.writeAnchorIds.has(checkpointId) || this.persistedIds.has(checkpointId)) { + return super.putWrites(config, writes, taskId); + } + const buffered = this.bufferedBookkeeping.get(checkpointId); + if (buffered) { + buffered.batches.push({ config, writes, taskId }); + } else { + sweepStale(this.bufferedBookkeeping, (b) => b.at); + this.bufferedBookkeeping.set(checkpointId, { + at: Date.now(), + batches: [{ config, writes, taskId }], + }); + } + return; + } + // A resumable batch — an interrupt (a HITL pause) or a real state/delta channel a later + // checkpoint depends on — anchors the checkpoint so its `put` persists it. Keyed on the + // globally-unique checkpoint id so concurrent runs on the same `thread_id` can't + // cross-consume anchors. The anchor is recorded BEFORE the awaited super call on purpose: + // LangGraph dispatches the matching `put` concurrently with `putWrites` (probe-confirmed), + // so recording after the await could let a slow-I/O interrupt `put` miss its anchor and be + // 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)), + ); + } + return await super.putWrites(config, writes, taskId); + } catch (err) { + // The write batch never landed — best-effort un-anchor so the concurrent `put` doesn't + // persist a checkpoint whose pending writes are missing (an unresumable phantom pause). + // If `put` already consumed the anchor, the thrown error still fails the run and the + // pre-run prune / Mongo TTL reclaim the orphan. + this.writeAnchorIds.delete(checkpointId); + throw err; + } + } + + override async put( + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: CheckpointMetadata, + ): Promise { + if (this.writeAnchorIds.delete(checkpoint.id)) { + // 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. + sweepStale(this.persistedIds, (t) => t); + this.persistedIds.set(checkpoint.id, Date.now()); + return super.put(config, checkpoint, metadata); + } + // 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 + // it — this is what keeps a failed turn from leaving orphan rows in the writes + // collection. Return the config LangGraph expects (pointing at the checkpoint it believes + // was saved) so the run finishes normally; nothing durable is written. + this.bufferedBookkeeping.delete(checkpoint.id); + return { + ...config, + configurable: { + ...config.configurable, + checkpoint_id: checkpoint.id, + }, + }; + } + + /** + * Track a checkpoint id whose `put` must persist it. Evicts ONLY genuinely-stale ids + * (older than {@link WRITE_ANCHOR_STALE_MS}, i.e. from a crashed run whose `put` never + * landed) — never a recent in-flight id — so a slow-I/O interrupt `put` is never + * mis-classified as a clean exit. If nothing is stale the map is allowed to grow rather + * than drop a valid id; the next sweep reclaims the crashed ones. + */ + private recordWriteAnchor(checkpointId: string): void { + sweepStale(this.writeAnchorIds, (t) => t); + this.writeAnchorIds.set(checkpointId, Date.now()); + } +} + +/** + * Evict genuinely-stale entries from a fate-tracking map once it is crowded + * ({@link WRITE_ANCHOR_SWEEP_THRESHOLD}). Entries from a crashed run (older than + * {@link WRITE_ANCHOR_STALE_MS}) are reclaimed; recent in-flight entries never are. + */ +function sweepStale(map: Map, timeOf: (value: T) => number): void { + if (map.size < WRITE_ANCHOR_SWEEP_THRESHOLD) { + return; + } + const now = Date.now(); + for (const [id, value] of map) { + if (now - timeOf(value) > WRITE_ANCHOR_STALE_MS) { + map.delete(id); + } + } +} + /** Default approval window and checkpoint TTL: 24h. */ export const DEFAULT_CHECKPOINT_TTL_SECONDS = 86400; @@ -104,7 +334,7 @@ async function buildMongoSaver( resolved: ResolvedCheckpointerConfig, ): Promise { try { - const saver = new MongoDBSaver({ + const saver = new LazyMongoSaver({ // mongoose vends the live MongoClient; reuse it instead of opening a second // connection. The driver type is structurally identical but resolves to a // different `mongodb` copy than checkpoint-mongodb's, hence the cast.