From e5fa0395b6ce73d7ab7b065dc64649a31ac4b170 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 7 Jul 2026 13:38:41 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=9B=A1=EF=B8=8F=20fix:=20Guard=20HITL=20c?= =?UTF-8?q?heckpoint=20size=20against=20MongoDB=2016MB=20limit?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A LangGraph HITL checkpoint embeds the whole serialized message history in a single BSON document, so a large conversation (inlined base64 media, big tool outputs, long history) can serialize past MongoDB's 16MB document ceiling. `MongoDBSaver.put` would then throw a raw `BSONObjectTooLarge` at pause time and the pause would be lost with no legible error. `LazyMongoSaver` now measures the serialized checkpoint on the persist path (rare HITL pauses only — the clean-exit common path is untouched): debug-logs the size, warns past a soft 8MB threshold, and throws a typed `CheckpointTooLargeError` before the doomed write past a 15MB hard limit (16MB minus headroom for the document's other fields). Thresholds are overridable via the constructor for testing. Adds integration coverage (real serde + mongodb-memory-server) for the under-threshold, soft-warn, and hard-reject cases. --- .../agents/checkpointer.integration.spec.ts | 88 ++++++++++++++ packages/api/src/agents/checkpointer.ts | 111 ++++++++++++++++++ 2 files changed, 199 insertions(+) diff --git a/packages/api/src/agents/checkpointer.integration.spec.ts b/packages/api/src/agents/checkpointer.integration.spec.ts index cc67b96817..100e948458 100644 --- a/packages/api/src/agents/checkpointer.integration.spec.ts +++ b/packages/api/src/agents/checkpointer.integration.spec.ts @@ -1,4 +1,5 @@ import mongoose from 'mongoose'; +import { logger } from '@librechat/data-schemas'; import { MongoMemoryServer } from 'mongodb-memory-server'; import { MongoDBSaver } from '@langchain/langgraph-checkpoint-mongodb'; import { emptyCheckpoint, ERROR, INTERRUPT } from '@langchain/langgraph-checkpoint'; @@ -6,6 +7,8 @@ import { getAgentCheckpointer, deleteAgentCheckpoint, deleteAgentCheckpoints, + LazyMongoSaver, + CheckpointTooLargeError, __resetCheckpointerForTests, } from './checkpointer'; @@ -446,3 +449,88 @@ describe('LazyMongoSaver (lazy persistence — mongodb-memory-server)', () => { expect(effects).toEqual({ a: 1, c: 1 }); // siblings did NOT re-execute }); }); + +describe('LazyMongoSaver checkpoint size guard (mongodb-memory-server integration)', () => { + // Guards the single-document ceiling: a checkpoint embeds the whole message history, so a + // large conversation can serialize past MongoDB's 16 MB limit. The guard measures the + // serialized state on the persist path, WARNS past a soft threshold, and REJECTS past a hard + // limit BEFORE the write — a typed CheckpointTooLargeError instead of a raw BSONObjectTooLarge. + // Thresholds are shrunk here so payloads stay tiny while exercising the real serde + Mongo. + const clientForSaver = () => + // mongoose vends the live MongoClient; the driver type resolves to a different `mongodb` copy + // than checkpoint-mongodb's, so the cast mirrors buildMongoSaver in checkpointer.ts. + mongoose.connection.getClient() as unknown as ConstructorParameters< + typeof MongoDBSaver + >[0]['client']; + + const makeSaver = (overrides?: { warnBytes?: number; hardLimitBytes?: number }) => + new LazyMongoSaver({ + client: clientForSaver(), + checkpointCollectionName: 'agent_checkpoints', + checkpointWritesCollectionName: 'agent_checkpoint_writes', + ttl: 3600, + ...overrides, + }); + + /** Seed an interrupt anchor for a checkpoint whose serialized state is inflated to ~`payloadBytes`. */ + async function seedSizedInterrupt(saver: LazyMongoSaver, threadId: string, payloadBytes: number) { + const checkpoint = emptyCheckpoint(); + checkpoint.channel_values = { messages: 'x'.repeat(payloadBytes) }; + 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: '' } }; + const metadata = { source: 'input' as const, step: -1, writes: null, parents: {} }; + return { checkpoint, config, metadata }; + } + + afterEach(() => { + jest.restoreAllMocks(); + }); + + it('persists a checkpoint under the soft threshold', async () => { + const saver = makeSaver({ warnBytes: 5_000, hardLimitBytes: 50_000 }); + await saver.setup(); + const threadId = `convo-${new mongoose.Types.ObjectId().toString()}`; + const { checkpoint, config, metadata } = await seedSizedInterrupt(saver, threadId, 100); + + await saver.put(config, checkpoint, metadata); + + expect(await saver.getTuple(readConfig(threadId))).toBeDefined(); + }); + + it('persists but WARNS when a checkpoint crosses the soft threshold', async () => { + const warnSpy = jest.spyOn(logger, 'warn').mockImplementation(() => logger); + const saver = makeSaver({ warnBytes: 500, hardLimitBytes: 50_000 }); + await saver.setup(); + const threadId = `convo-${new mongoose.Types.ObjectId().toString()}`; + const { checkpoint, config, metadata } = await seedSizedInterrupt(saver, threadId, 2_000); + + await saver.put(config, checkpoint, metadata); + + expect(await saver.getTuple(readConfig(threadId))).toBeDefined(); + expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining('soft threshold')); + }); + + it('REFUSES to persist and throws CheckpointTooLargeError over the hard limit', async () => { + 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, config, metadata } = await seedSizedInterrupt(saver, threadId, 10_000); + + await expect(saver.put(config, checkpoint, metadata)).rejects.toBeInstanceOf( + CheckpointTooLargeError, + ); + + // Nothing durable was written for that thread — getTuple reads the checkpoint document. + expect(await saver.getTuple(readConfig(threadId))).toBeUndefined(); + expect( + await mongoose.connection.db!.collection('agent_checkpoints').countDocuments({ + thread_id: threadId, + }), + ).toBe(0); + }); +}); diff --git a/packages/api/src/agents/checkpointer.ts b/packages/api/src/agents/checkpointer.ts index 79b6954346..0491617f45 100644 --- a/packages/api/src/agents/checkpointer.ts +++ b/packages/api/src/agents/checkpointer.ts @@ -124,6 +124,68 @@ interface BufferedWriteBatch { batches: Array<{ config: RunnableConfig; writes: PendingWrite[]; taskId: string }>; } +/** + * MongoDB's hard per-document ceiling. A checkpoint whose serialized state pushes its + * document past this cannot be stored — the driver throws `BSONObjectTooLarge` (code 10334). + */ +const MAX_BSON_DOCUMENT_BYTES = 16 * 1024 * 1024; + +/** + * Headroom reserved below {@link MAX_BSON_DOCUMENT_BYTES} for a checkpoint document's + * non-state fields (ids, metadata, `metadata_search`, BSON framing). The serialized + * `checkpoint` blob dominates the document; this margin covers everything else so the guard + * rejects before Mongo does — with a legible error instead of a raw driver failure. + */ +const CHECKPOINT_SIZE_HEADROOM_BYTES = 1024 * 1024; + +/** + * Reject a checkpoint whose serialized state exceeds this. The pause is unrecoverable either + * way (the document can't be written), so failing here as a typed {@link CheckpointTooLargeError} + * turns an opaque `BSONObjectTooLarge` crash into an actionable one. + */ +export const CHECKPOINT_HARD_LIMIT_BYTES = MAX_BSON_DOCUMENT_BYTES - CHECKPOINT_SIZE_HEADROOM_BYTES; + +/** + * Warn once a persisted checkpoint crosses this soft threshold (~50% of the ceiling), so a + * conversation's checkpoint growth is visible in logs well before it reaches the hard limit. + */ +export const CHECKPOINT_WARN_BYTES = 8 * 1024 * 1024; + +/** + * A HITL checkpoint whose serialized state exceeds {@link CHECKPOINT_HARD_LIMIT_BYTES} — more + * than MongoDB can hold in a single document. Thrown BEFORE the doomed write so the run fails + * with a clear, typed message instead of a raw driver `BSONObjectTooLarge`. The pause cannot be + * persisted regardless of how it is handled upstream; a durable resume is impossible for this turn. + */ +export class CheckpointTooLargeError extends Error { + readonly code = 'CHECKPOINT_TOO_LARGE'; + constructor( + readonly bytes: number, + readonly limit: number, + readonly threadId?: string, + ) { + const mb = (n: number): string => (n / 1024 / 1024).toFixed(1); + super( + `Checkpoint state is ${mb(bytes)} MB, over the ${mb(limit)} MB limit for a durable pause. ` + + 'This conversation carries too much state to pause for input — large tool outputs or ' + + 'inlined media are the usual cause. Start a new conversation or reduce context.', + ); + this.name = 'CheckpointTooLargeError'; + } +} + +/** + * Construction options for {@link LazyMongoSaver}: the base saver options plus optional + * size-guard overrides. The overrides default to the module thresholds and exist so tests can + * exercise the guard at small sizes; production always uses the defaults. + */ +export type LazyMongoSaverOptions = ConstructorParameters[0] & { + /** Soft warn threshold in bytes. Defaults to {@link CHECKPOINT_WARN_BYTES}. */ + warnBytes?: number; + /** Hard reject limit in bytes. Defaults to {@link CHECKPOINT_HARD_LIMIT_BYTES}. */ + hardLimitBytes?: number; +}; + export class LazyMongoSaver extends MongoDBSaver { /** checkpoint id → time the resumable `putWrites` anchoring it arrived; consumed by `put`. */ private readonly writeAnchorIds = new Map(); @@ -133,6 +195,18 @@ export class LazyMongoSaver extends MongoDBSaver { /** checkpoint id → bookkeeping batches parked until the checkpoint persists or is discarded. */ private readonly bufferedBookkeeping = new Map(); + /** Soft threshold (bytes) past which a persisted checkpoint is warned about. */ + private readonly warnBytes: number; + /** Hard limit (bytes) past which a checkpoint is refused with {@link CheckpointTooLargeError}. */ + private readonly hardLimitBytes: number; + + constructor(options: LazyMongoSaverOptions) { + const { warnBytes, hardLimitBytes, ...mongoOptions } = options; + super(mongoOptions); + this.warnBytes = warnBytes ?? CHECKPOINT_WARN_BYTES; + this.hardLimitBytes = hardLimitBytes ?? CHECKPOINT_HARD_LIMIT_BYTES; + } + override async putWrites( config: RunnableConfig, writes: PendingWrite[], @@ -203,6 +277,7 @@ 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); sweepStale(this.persistedIds, (t) => t); this.persistedIds.set(checkpoint.id, Date.now()); return super.put(config, checkpoint, metadata); @@ -233,6 +308,42 @@ export class LazyMongoSaver extends MongoDBSaver { sweepStale(this.writeAnchorIds, (t) => t); this.writeAnchorIds.set(checkpointId, Date.now()); } + + /** + * 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 + * {@link hardLimitBytes} — BEFORE the write, so an oversize pause fails legibly rather than as a + * raw `BSONObjectTooLarge`. Serializes with the same `serde` the base `put` uses, so the measured + * bytes match what would be stored. The extra serialization runs only on the (rare) HITL pause + * path — never the clean-exit common path, which is discarded before reaching here. + */ + private async assertCheckpointFitsDocument( + config: RunnableConfig, + checkpoint: Checkpoint, + ): Promise { + const [, serialized] = await this.serde.dumpsTyped(checkpoint); + const bytes = serialized.byteLength; + const threadId = config.configurable?.thread_id as string | undefined; + const mb = (n: number): string => (n / 1024 / 1024).toFixed(1); + if (bytes > this.hardLimitBytes) { + // The anchoring write row was already persisted by `putWrites`; the pre-run prune and Mongo + // TTL reclaim it. Drop any parked bookkeeping so it doesn't linger in memory. + this.bufferedBookkeeping.delete(checkpoint.id); + logger.error( + `[checkpointer] HITL checkpoint for thread ${threadId ?? 'unknown'} is ${mb(bytes)} MB, over the ${mb(this.hardLimitBytes)} MB durable-pause limit; refusing the write (a document past 16 MB cannot be stored in MongoDB).`, + ); + throw new CheckpointTooLargeError(bytes, this.hardLimitBytes, threadId); + } + if (bytes >= this.warnBytes) { + logger.warn( + `[checkpointer] HITL checkpoint for thread ${threadId ?? 'unknown'} is ${mb(bytes)} MB, past the ${mb(this.warnBytes)} MB soft threshold (hard limit ${mb(this.hardLimitBytes)} MB) — approaching MongoDB's single-document ceiling.`, + ); + return; + } + logger.debug( + `[checkpointer] Persisting HITL checkpoint for thread ${threadId ?? 'unknown'}: ${bytes} bytes`, + ); + } } /**