From 8a46ca070c72bc0ff51ad37c1ac97d582d181227 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 2 Sep 2026 20:29:43 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=A7=8A=20feat:=20Persist=20Context=20Fadi?= =?UTF-8?q?ng=20Tier=20Across=20Agent=20Runs=20(#15505)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * ๐ŸงŠ feat: Persist Context Fading Tier Across Agent Runs Carry the pruner's latched context-fading tier from `@librechat/agents` through `contextMeta` the way `calibrationRatio` already travels, so a historical tool result keeps the same truncated bytes from one run to the next and Anthropic's prefix-based prompt cache survives across turns (danny-avila/agents#497). - Add `IAgentFadingTier` to the message and conversation `contextMeta` types, schemas and the data-provider zod schema - Add `isAgentFadingTier`, `resolvePersistableFadingTier` and `resolveRunContextMeta` in `packages/api`; a tier is persisted once it carries information (masking active or a budget below the window), and a latched tier is kept even at a neutral calibration ratio - Seed `createRun({ fadingTier })` from the parent response's contextMeta regardless of encoding, since caps are character-based; the field rides the existing forward-compatible `runConfig` variable and older SDK versions ignore it, as does `Run.getFadingTier` being absent - Accept a valid `fading` field in event actor context and reject a malformed one, matching the calibration checks - Capture run context meta through one module-level helper for both completion paths instead of two inline copies Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐ŸงŠ fix: Seed Fading Tier Across HITL Resume, Validate at the DB Layer Address review findings on fading tier persistence: - Leave the "is this tier worth persisting" decision to the SDK, whose `Run.getFadingTier()` now returns only informative tiers; comparing against the client's window misclassified untouched conversations whenever a reserve ratio was configured - Carry `contextMeta` through the HITL pause projection (staged approval, job metadata, resume state) and seed the rebuilt client from it, so a paused turn resumes with the same tier and calibration instead of re-deriving a shallower one and rewriting the prefix - Share run seeding between `chatCompletion` and `resumeCompletion`, and make capture tolerate partial runs and bare resume contexts - Validate `contextMeta.fading` in `commitAgentEventActorState` with a shared `isAgentFadingTier` guard in data-schemas, so a malformed tier is rejected at the DB layer instead of silently cold-starting the actor - Reuse `IAgentEventActorContextMeta` for `IMessage.contextMeta` instead of extending an inline duplicate Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐ŸงŠ test: Cover Resume Seeding From Paused Context Meta Assert the resume controller seeds the rebuilt client from the context meta captured at the pause before rebuilding the run, and move the capture helper's JSDoc back above its function. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐Ÿงน style: Sort Imports in Fading Tier Plumbing Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐ŸงŠ fix: Persist Fading Tier Through Redis Jobs and Branches Address Codex and Copilot review findings on fading tier persistence: - Deserialize `contextMeta` in `RedisJobStore`, whose explicit read mapper rebuilt every other pause field but this one, so Redis-backed deployments resumed with no tier to seed after every HITL pause - Copy server-private `contextMeta` onto the assistant message created by `POST /api/messages/branch`, so the next turn seeds its pruner the same way it would from the parallel source response - Constrain the `fading` subdocuments in both Mongoose schemas (version enum, positive budget, required fields) to match the zod schema - Throw a distinct "fading tier is invalid" error from both context meta validators instead of reporting a calibration failure Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐ŸงŠ fix: Carry Context Meta Through Stop and Re-Pause Addresses the second Codex pass on the fading-tier persistence: - A re-pause whose resumed segment has nothing persistable now clears the job's `contextMeta` in the same transition instead of leaving the first pause's calibration and tier for the next resume to seed from. - The client publishes the run's live calibration and fading tier onto the job after each pre-invoke context snapshot (deduplicated on value), and the Stop path copies `jobData.contextMeta` onto the stopped response, so a follow-up from a stopped turn keeps the prompt prefix stable. - Mid-run captures (HITL pause, Stop) read the graph's live state; `Run` refreshes its own getters only after the stream settles, so the pause used to persist the seeds it started from. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐ŸงŠ fix: Fence Context Meta Publication Ahead of Stop A Stop handled by another replica reads the job before the owner is signaled, so a fire-and-forget publish could lose the race, and a run stopped before its first context snapshot never published the tier it inherited. The client now publishes the inherited seed before the run streams or resumes, and the context-usage handler awaits each snapshot's publish so the write lands before the model call it describes begins. Failures still only log; the pre-run publish is optional-called so bare resume contexts keep working. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐ŸงŠ fix: Read Back Context Meta After the Abort Claim The run owner publishes its context meta ahead of each model call and awaits the write, so a publish landing between the abort's initial job read and its terminal claim describes the call whose partial output the abort snapshot carries. abortJob now reads the same-epoch job back after the claim and content re-read and takes its context meta, so the stopped response persists the tier that produced its bytes. The refresh is best-effort and logs on failure like the content refresh beside it. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐ŸงŠ fix: Share In-Flight Context Meta Writes and Clear Neutral Live State Two snapshot callbacks in a parallel-agent run can capture the same tier before the first write settles; the publisher now caches the in-flight promise per serialized value so the second caller awaits the same durable write instead of treating an uncommitted value as published. A live snapshot with nothing left to persist after an earlier publish now writes a neutral record (ratio 1, no tier), since a running job's fields cannot be deleted through the metadata writer, so a Stop no longer persists an earlier non-neutral state onto the response. The pre-run seed publish is unchanged and is told apart from live snapshots by a flag from the context-usage sink. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐ŸงŠ chore: Drop Unused Rest Parameter in Context Meta Test Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐ŸงŠ feat: Persist Per-Agent Fading Tiers as Compact Context Meta Reconciles the host with @librechat/agents at 9aacad1a, where graph history stays canonical and each Run derives a provider-only projection from the latched tiers. LibreChat therefore persists only compact state: the calibration ratio, the default agent's tier, and now the per-agent tiers from Run.getFadingTiers(), stored as validated entries so agent IDs never become MongoDB field names and restored onto a null-prototype record for RunConfig.fadingTiers. Truncated messages, canonical tool content, projection provenance and per-message truncation state are never persisted; the SDK's latched flag is stripped on capture. The per-agent map travels everywhere the single tier already did: response contextMeta, event-actor checkpoint state, HITL pause and resume job metadata in both stores, the Stop path, and branch creation. A shared Mongoose definition keeps the message and conversation schemas identical, the zod schema validates entries, and commitAgentEventActorState rejects malformed maps. Comments no longer describe graph messages as truncated; only the provider projection is. Tests cover: tool results and inputs reach the persisted content parts at full size; contextMeta carries exactly the compact tier and calibration fields; a persisted snapshot round-trips into the seeded RunConfig unchanged and prototype-safe; per-agent tiers survive pause, Redis job round-trips, and event-actor context validation. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐ŸงŠ fix: Restore Event Actor Types and Publish Seeds Before Abortable Setup The previous commit's doc rewording in the conversation types dropped the event-actor interfaces that precede the fading tier, which broke the declaration bundle in CI; they are restored unchanged. The inherited context meta is now published onto the job as soon as the parent's state is loaded and awaited at the top of chatCompletion and resumeCompletion, ahead of run creation and the other abortable setup stages, so a Stop during setup still persists the parent's tiers onto the stopped response. The partial response saved when every subscriber disconnects now carries the job's context meta like the Stop and pause paths do. The branch route keeps the source's context meta in the saved message but strips it from the HTTP response, matching client reads. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐ŸงŠ fix: Serialize Distinct Context Meta Publications Overlapping snapshots from parallel agents could carry different tiers; each started its own write and both job stores keep whichever write finished last, so an older snapshot could overwrite a newer one. Each distinct publication is now issued after the previous one settles, in call order, so the newest snapshot is what a Stop reads back. Equal values still share one write. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐ŸงŠ fix: Project Root Message Reads and Keep Inherited Meta on Setup Failure The root GET /api/messages paths returned raw documents: the single-message query read called getMessages without a projection and the cursor page performed an unprojected find. Both now apply CLIENT_MESSAGE_SELECT, with getMessagesByCursor accepting a select option, so the persisted fading tiers never reach the client through the paginated endpoint; the search-hydration read projects the same way. A turn that loads a parent's context meta but fails before its run exists used to replace the inherited meta with the capture of a missing run, so the persisted error response lost the last valid tiers. Both finalizers now keep the inherited meta when no run was created while a created run's neutral state may still clear it. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐ŸงŠ refactor: Move Context Meta Publication Into packages/api The publication coordinator (ordering, deduplication, in-flight sharing, failure state) now lives in packages/api as createContextMetaPublisher, with selectRunContextMetaToPublish deciding what a publication carries; the legacy client keeps only the wiring to the job metadata writer. A transiently failing write is retried with a short backoff before the publication is reported and forgotten, so a snapshot's model call no longer proceeds on the first rejection while the job holds the previous record. The SDK swallows handler errors, so retrying is the only way to strengthen the fence without failing the turn. A terminal save now supplies contextMeta: null when the finished run has nothing to carry, and saveMessage unsets a previously stored value in that case, so a disconnect snapshot's record cannot outlive the run that completed neutrally. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐ŸงŠ fix: Widen the saveMessage Method Type for the Null Unset The exported method interface still typed contextMeta as the message field, so the spec exercising the null unset failed typechecking; it now mirrors the implementation signature. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐ŸงŠ fix: Publish Every Agent's Snapshot, Sanitize Search Hits, Unbound Agent IDs The context-usage handler now awaits the run's context meta publish for every ON_CONTEXT_USAGE event, hidden sequential agents included, while the visible-snapshot bookkeeping and the client emission keep their gate; a Stop before the next visible snapshot finds the tier a hidden agent's model call latched. Search results are built from the hydrated hit minus its server-private contextMeta, since Meilisearch hydration projects every schema field. The per-agent tier entry guard no longer caps the agent ID length: ephemeral agent IDs encode endpoint, model and sender without a bound, and dropping such an entry would let that agent alone re-derive its tier on the next turn. Also formats the publication spec so the static checks pass. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐ŸงŠ fix: Unset Context Meta Atomically, Keep Publication State After a Failed Write - `saveMessage` folds the `contextMeta` unset into the same update that persists the terminal response, on both the plain and the provenance-merge paths, so a failure between two writes can no longer leave a completed row carrying a disconnect snapshot's state. - `createContextMetaPublisher` tracks whether any record has committed separately from the retryable latest publication, so `hasPublished` stays true after a later publication exhausts its retries and a following neutral snapshot still overwrites the earlier record on the job. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐ŸงŠ fix: Carry Context Meta Through the Resumable Stop Route and Neutral Resumed Completions - The resumable Stop route copies the refreshed `jobData.contextMeta` onto the terminal assistant message it persists, as the abort middleware already does, so a stopped resumable response keeps the tiers that produced its partial output. - A resumed run that completes with neutral context state now saves `contextMeta: null`, which unsets the paused segment's record instead of letting an omitted field keep it on the completed row. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐ŸงŠ fix: Trust Only Server-Authored Context Meta, Unset It on a Neutral Stop, Read It From the Job on Disconnect - `POST /api/messages/:conversationId` and conversation imports drop any client-supplied `contextMeta`, and the run seeds its calibration and fading tiers only from a server-authored parent response, so a forged tier can never shape a provider projection. - Both Stop paths save `contextMeta: jobData.contextMeta ?? null`, so a job that re-paused with neutral state unsets what an earlier pause stored on the row instead of leaving it behind. - The disconnect partial save reads the same-epoch job record for the run's published context meta, since the client-facing resume snapshot never carries server-private state. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐Ÿงช test: Mock the Job Store in the Resume Metadata Spec The disconnect partial save now reads the same-epoch job record through `GenerationJobManager.getJobStore()`, so the resume metadata spec's manager mock provides it as the other disconnect spec already does. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐ŸงŠ fix: Keep Context Meta Off the POST Message Response `POST /api/messages/:conversationId` returned the saved row verbatim, so a client write against an existing server-authored message echoed the row's stored context meta. The route now projects the row through the same client projection the branch route uses, which drops the server-private field. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF * ๐Ÿ“ฆ chore: Bump @librechat/agents to 3.7.16 for Persisted Fading Tiers 3.7.16 is the first release that carries danny-avila/agents#497, so `Run.getFadingTier()` / `getFadingTiers()` now report the latched tier this branch persists and `RunConfig.fadingTier` / `fadingTiers` seed the next run from it. The dependency set is unchanged from 3.7.15; only the package entry moves. Adds an SDK-backed round trip to the fading spec: a first pruner latches an informative tier, `resolveRunContextMeta` reduces it to the persisted shape, and a second pruner seeded from that shape reproduces the first run's projected tool-exchange bytes under instruction and calibration drift, while an unseeded pruner under the same drift relaxes its tier and rewrites them. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WNbgMeR2aZDyhsqAy4L8gF --------- Co-authored-by: Claude --- api/package.json | 2 +- .../agents/__tests__/callbacks.spec.js | 119 +++ .../request.partialDisconnect.spec.js | 39 +- .../__tests__/request.resumeMetadata.spec.js | 4 + .../agents/__tests__/resume.spec.js | 49 ++ api/server/controllers/agents/callbacks.js | 43 +- api/server/controllers/agents/client.js | 189 ++++- api/server/controllers/agents/client.test.js | 704 +++++++++++++++++- api/server/controllers/agents/request.js | 15 +- api/server/controllers/agents/resume.js | 14 +- api/server/middleware/abortMiddleware.js | 5 + api/server/middleware/abortMiddleware.spec.js | 41 + .../routes/__tests__/messages-get.spec.js | 106 +++ .../routes/agents/__tests__/abort.spec.js | 69 ++ api/server/routes/agents/index.js | 5 + api/server/routes/messages.js | 43 +- .../services/Endpoints/agents/initialize.js | 3 +- api/server/utils/import/importers.js | 5 +- api/server/utils/import/importers.spec.js | 27 + package-lock.json | 10 +- packages/api/package.json | 2 +- packages/api/src/agents/fading.spec.ts | 259 +++++++ packages/api/src/agents/fading.ts | 100 +++ packages/api/src/agents/index.ts | 2 + packages/api/src/agents/publication.spec.ts | 201 +++++ packages/api/src/agents/publication.ts | 141 ++++ packages/api/src/agents/run.ts | 21 +- packages/api/src/stream/ApprovalLifecycle.ts | 13 +- .../api/src/stream/GenerationJobManager.ts | 17 + .../stream/__tests__/RedisJobStore.spec.ts | 12 + .../stream/__tests__/abortContextMeta.spec.ts | 79 ++ .../stream/__tests__/pendingAction.spec.ts | 38 + .../stream/implementations/RedisJobStore.ts | 3 + .../api/src/stream/interfaces/IJobStore.ts | 8 +- packages/api/src/stream/metadata.ts | 3 + packages/api/src/types/stream.ts | 7 +- packages/data-provider/src/schemas.ts | 16 + .../src/methods/conversation.spec.ts | 33 + .../data-schemas/src/methods/conversation.ts | 10 + .../data-schemas/src/methods/message.spec.ts | 71 ++ packages/data-schemas/src/methods/message.ts | 63 +- packages/data-schemas/src/schema/convo.ts | 2 + packages/data-schemas/src/schema/fading.ts | 35 + packages/data-schemas/src/schema/message.ts | 2 + packages/data-schemas/src/types/convo.ts | 28 + packages/data-schemas/src/types/message.ts | 6 +- .../data-schemas/src/utils/fading.spec.ts | 44 ++ packages/data-schemas/src/utils/fading.ts | 47 ++ packages/data-schemas/src/utils/index.ts | 1 + 49 files changed, 2657 insertions(+), 99 deletions(-) create mode 100644 packages/api/src/agents/fading.spec.ts create mode 100644 packages/api/src/agents/fading.ts create mode 100644 packages/api/src/agents/publication.spec.ts create mode 100644 packages/api/src/agents/publication.ts create mode 100644 packages/api/src/stream/__tests__/abortContextMeta.spec.ts create mode 100644 packages/data-schemas/src/schema/fading.ts create mode 100644 packages/data-schemas/src/utils/fading.spec.ts create mode 100644 packages/data-schemas/src/utils/fading.ts diff --git a/api/package.json b/api/package.json index 2a45e638c4..d9745c14c0 100644 --- a/api/package.json +++ b/api/package.json @@ -46,7 +46,7 @@ "@azure/storage-blob": "^12.30.0", "@google/genai": "^2.8.0", "@keyv/redis": "5.1.6", - "@librechat/agents": "^3.7.15", + "@librechat/agents": "^3.7.16", "@librechat/api": "*", "@librechat/data-schemas": "*", "@microsoft/microsoft-graph-client": "^3.0.7", diff --git a/api/server/controllers/agents/__tests__/callbacks.spec.js b/api/server/controllers/agents/__tests__/callbacks.spec.js index 383d1d083c..34bab04619 100644 --- a/api/server/controllers/agents/__tests__/callbacks.spec.js +++ b/api/server/controllers/agents/__tests__/callbacks.spec.js @@ -124,6 +124,125 @@ describe('resumable event generation fencing', () => { ); }); + it('keeps a tool result and its input at full size in the content parts LibreChat persists', async () => { + const { GraphEvents, createContentAggregator } = jest.requireActual('@librechat/agents'); + const { getDefaultHandlers } = require('../callbacks'); + /** Far beyond any cap the SDK's provider-only projection would apply. */ + const output = 'r'.repeat(600_000); + const args = { query: 'q'.repeat(120_000) }; + const { contentParts, stepMap, aggregateContent } = createContentAggregator(); + const handlers = getDefaultHandlers({ + res: { write: jest.fn() }, + aggregateContent, + contentParts, + stepMap, + toolEndCallback: jest.fn(), + collectedUsage: [], + streamId: 'conversation-1', + jobCreatedAt: 1234, + }); + const step = { + id: 'step-full-size', + index: 0, + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'call-full-size', name: 'fetch', args: JSON.stringify(args) }], + }, + }; + + await handlers[GraphEvents.ON_RUN_STEP].handle(GraphEvents.ON_RUN_STEP, step); + await handlers[GraphEvents.ON_RUN_STEP_COMPLETED].handle(GraphEvents.ON_RUN_STEP_COMPLETED, { + result: { + id: step.id, + index: 0, + tool_call: { id: 'call-full-size', name: 'fetch', args, output }, + }, + }); + + expect(contentParts).toHaveLength(1); + expect(contentParts[0].tool_call.output).toBe(output); + expect(contentParts[0].tool_call.args).toEqual(args); + expect(JSON.stringify(contentParts[0])).not.toContain('[truncated:'); + }); + + it('publishes a hidden sequential agent snapshot without recording or forwarding it', async () => { + const { GenerationJobManager } = require('@librechat/api'); + const { GraphEvents } = jest.requireActual('@librechat/agents'); + const { getDefaultHandlers } = require('../callbacks'); + const onSnapshot = jest.fn(async () => undefined); + const contextUsageSink = { latest: null, count: 0, onSnapshot }; + const handlers = getDefaultHandlers({ + res: { write: jest.fn() }, + aggregateContent: jest.fn(), + toolEndCallback: jest.fn(), + collectedUsage: [], + streamId: 'conversation-1', + jobCreatedAt: 1234, + contextUsageSink, + usageEmitSink: [], + }); + + await handlers[GraphEvents.ON_CONTEXT_USAGE].handle( + GraphEvents.ON_CONTEXT_USAGE, + { contextBudget: 1000, remainingContextTokens: 400 }, + { + hide_sequential_outputs: true, + last_agent_id: 'agent-final', + langgraph_node: 'agent-intermediate', + }, + ); + + expect(onSnapshot).toHaveBeenCalledTimes(1); + expect(contextUsageSink).toMatchObject({ latest: null, count: 0 }); + expect(GenerationJobManager.emitChunk).not.toHaveBeenCalled(); + }); + + it('records a context snapshot and notifies the sink before forwarding it', async () => { + const { GenerationJobManager } = require('@librechat/api'); + const { GraphEvents } = jest.requireActual('@librechat/agents'); + const { getDefaultHandlers } = require('../callbacks'); + let releaseSnapshot; + const onSnapshot = jest.fn( + () => + new Promise((resolve) => { + releaseSnapshot = resolve; + }), + ); + const contextUsageSink = { latest: null, count: 0, onSnapshot }; + const usageEmitSink = [{ input_tokens: 10 }]; + const data = { contextBudget: 1000, remainingContextTokens: 400 }; + const handlers = getDefaultHandlers({ + res: { write: jest.fn() }, + aggregateContent: jest.fn(), + toolEndCallback: jest.fn(), + collectedUsage: [], + streamId: 'conversation-1', + jobCreatedAt: 1234, + contextUsageSink, + usageEmitSink, + }); + + let settled = false; + const handled = handlers[GraphEvents.ON_CONTEXT_USAGE] + .handle(GraphEvents.ON_CONTEXT_USAGE, data, { hide_sequential_outputs: false }) + .then(() => { + settled = true; + }); + await Promise.resolve(); + + expect(contextUsageSink).toMatchObject({ latest: data, count: 1, latestUsageIndex: 1 }); + expect(onSnapshot).toHaveBeenCalledTimes(1); + expect(settled).toBe(false); + releaseSnapshot(); + await handled; + expect(settled).toBe(true); + expect(GenerationJobManager.emitChunk).toHaveBeenCalledWith( + 'conversation-1', + { event: GraphEvents.ON_CONTEXT_USAGE, data }, + { expectedCreatedAt: 1234 }, + ); + }); + it('resolves MCP identity from a function-shaped root tool call', async () => { const { GraphEvents } = jest.requireActual('@librechat/agents'); const { getDefaultHandlers } = require('../callbacks'); diff --git a/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js b/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js index 85caf062cb..29970e0f30 100644 --- a/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js +++ b/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js @@ -28,6 +28,7 @@ const mockGenerationJobManager = { beginProviderExecution: jest.fn(), markProviderExecutionDrained: jest.fn(), getResumeState: jest.fn(), + getJobStore: jest.fn(), updateMetadata: jest.fn(), claimGeneration: jest.fn(), releaseGeneration: jest.fn(), @@ -142,8 +143,20 @@ describe('ResumableAgentController tenant context', () => { * Drives the controller far enough to register the `allSubscribersLeft` handler, * fires it, and returns the tenant context that was active during `saveMessage`. */ - const firePartialDisconnect = async (user) => { + const partialContextMeta = { + calibrationRatio: 1.2, + encoding: 'claude', + fading: { v: 1, budgetTokens: 50_000, masked: true }, + }; + + const firePartialDisconnect = async ( + user, + jobRecord = { createdAt: 1000, contextMeta: partialContextMeta }, + ) => { let allSubscribersLeftHandler; + mockGenerationJobManager.getJobStore.mockReturnValue({ + getJob: jest.fn().mockResolvedValue(jobRecord), + }); mockGenerationJobManager.createJob.mockResolvedValue({ createdAt: 1000, metadata: { @@ -202,6 +215,30 @@ describe('ResumableAgentController tenant context', () => { return tenantSeenBySave; }; + it('carries the context meta the run published onto the partial response saved on disconnect', async () => { + await firePartialDisconnect({ id: 'user-123', tenantId: 'tenant-a' }); + + expect(mockSaveMessage).toHaveBeenCalledWith( + expect.any(Object), + expect.objectContaining({ + messageId: 'response-message', + unfinished: true, + contextMeta: partialContextMeta, + }), + expect.any(Object), + ); + }); + + it('leaves context meta off the partial response when the job record belongs to another epoch', async () => { + await firePartialDisconnect( + { id: 'user-123', tenantId: 'tenant-a' }, + { createdAt: 2000, contextMeta: partialContextMeta }, + ); + + const [, savedMessage] = mockSaveMessage.mock.calls[0]; + expect(savedMessage).not.toHaveProperty('contextMeta'); + }); + it('restores the authenticated tenant before saving a partial response on disconnect', async () => { const tenantSeenBySave = await firePartialDisconnect({ id: 'user-123', tenantId: 'tenant-a' }); diff --git a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js index a9e2e6475e..fd49e3bfd4 100644 --- a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js +++ b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js @@ -24,6 +24,7 @@ const mockGenerationJobManager = { markProviderExecutionDrained: jest.fn(), failPausePersistence: jest.fn(), getResumeState: jest.fn(), + getJobStore: jest.fn(), updateMetadata: jest.fn(), persistAgentEventDetachedTerminalEvidence: jest.fn(), claimGeneration: jest.fn(), @@ -444,6 +445,9 @@ describe('ResumableAgentController resume metadata', () => { emitter: { on: jest.fn() }, }); mockGenerationJobManager.getResumeState.mockResolvedValue(null); + mockGenerationJobManager.getJobStore.mockReturnValue({ + getJob: jest.fn().mockResolvedValue(null), + }); mockGenerationJobManager.getJob.mockResolvedValue(undefined); mockGenerationJobManager.updateMetadata.mockResolvedValue(undefined); mockGenerationJobManager.isRedis = false; diff --git a/api/server/controllers/agents/__tests__/resume.spec.js b/api/server/controllers/agents/__tests__/resume.spec.js index 2845c5b683..fcf040f763 100644 --- a/api/server/controllers/agents/__tests__/resume.spec.js +++ b/api/server/controllers/agents/__tests__/resume.spec.js @@ -281,6 +281,7 @@ function makeClient(overrides = {}) { pendingApproval: false, buildResponseMetadata: jest.fn(() => null), resumeCompletion: jest.fn().mockResolvedValue(undefined), + seedContextMeta: jest.fn(), ...overrides, }; } @@ -2677,6 +2678,26 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { ); }); + it('seeds the rebuilt client from the context meta captured at the pause', async () => { + const contextMeta = { + calibrationRatio: 1.25, + encoding: 'claude', + fading: { v: 1, budgetTokens: 50_000, masked: true }, + }; + mockGenerationJobManager.getJob.mockResolvedValue( + makeToolApprovalJob({ metadata: { contextMeta } }), + ); + await post(approveBody()); + await settled; + await flush(); + + const client = await mockInitializeClient.mock.results[0].value.then((r) => r.client); + expect(client.seedContextMeta).toHaveBeenCalledWith(contextMeta); + expect(client.seedContextMeta.mock.invocationCallOrder[0]).toBeLessThan( + client.resumeCompletion.mock.invocationCallOrder[0], + ); + }); + it('reuses the persisted MCP identity for edited and overridden turns', async () => { const persistedMCPRequestBody = { messageId: RESPONSE_MSG_ID, @@ -3206,6 +3227,34 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { ); }); + it('unsets the paused row context meta when the resumed run completes neutrally', async () => { + mockGenerationJobManager.getJob.mockResolvedValue( + makeToolApprovalJob({ + metadata: { + contextMeta: { + calibrationRatio: 1.2, + encoding: 'claude', + fading: { v: 1, budgetTokens: 50_000, masked: true }, + }, + }, + }), + ); + mockInitializeClient.mockResolvedValue({ + client: makeClient({ contextMeta: undefined }), + userMCPAuthMap: {}, + }); + + await post(approveBody()); + await settled; + await flush(); + + expect(mockSaveMessage).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ contextMeta: null }), + expect.objectContaining({ context: expect.stringContaining('resumed response end') }), + ); + }); + it('carries manualSkills/alwaysAppliedSkills onto the resumed requestMessage', async () => { const job = makeToolApprovalJob(); job.metadata.userMessage.manualSkills = ['skill-a']; diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js index 582c9fbaef..c97a196aad 100644 --- a/api/server/controllers/agents/callbacks.js +++ b/api/server/controllers/agents/callbacks.js @@ -806,26 +806,31 @@ function getDefaultHandlers({ * @param {GraphRunnableConfig['configurable']} [metadata] The runnable metadata. */ handle: async (event, data, metadata) => { - if ( + const visible = checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node) || - !metadata?.hide_sequential_outputs - ) { - /** Capture the latest visible snapshot (last-wins) and how many usage - * events preceded it BEFORE awaiting the emit. `emitEvent` can yield - * (resumable SSE / Redis publish); with parallel runs active this - * call's own primary usage could land in `usageEmitSink` during that - * yield, pushing `latestUsageIndex` past the very event that proves the - * snapshot completed โ€” the save path would then slice it away and drop - * a valid breakdown. The recorded index lets the save path persist only - * when a PRIMARY usage follows this snapshot (the snapshot's call - * actually invoked the model); a summarization detour emits a snapshot - * whose only following usage is tagged `summarization`, which a plain - * snapshot-count would over-count and wrongly drop. */ - if (contextUsageSink) { - contextUsageSink.latest = data; - contextUsageSink.count = (contextUsageSink.count ?? 0) + 1; - contextUsageSink.latestUsageIndex = usageEmitSink?.length ?? 0; - } + !metadata?.hide_sequential_outputs; + /** Capture the latest visible snapshot (last-wins) and how many usage + * events preceded it BEFORE awaiting the emit. `emitEvent` can yield + * (resumable SSE / Redis publish); with parallel runs active this + * call's own primary usage could land in `usageEmitSink` during that + * yield, pushing `latestUsageIndex` past the very event that proves the + * snapshot completed โ€” the save path would then slice it away and drop + * a valid breakdown. The recorded index lets the save path persist only + * when a PRIMARY usage follows this snapshot (the snapshot's call + * actually invoked the model); a summarization detour emits a snapshot + * whose only following usage is tagged `summarization`, which a plain + * snapshot-count would over-count and wrongly drop. */ + if (visible && contextUsageSink) { + contextUsageSink.latest = data; + contextUsageSink.count = (contextUsageSink.count ?? 0) + 1; + contextUsageSink.latestUsageIndex = usageEmitSink?.length ?? 0; + } + /** Every agent's snapshot publishes the run's context meta, hidden + * sequential agents included: their model calls latch tiers too, and a + * Stop before the next visible snapshot must find them on the job. Awaited + * so the write lands before the model call it describes begins. */ + await contextUsageSink?.onSnapshot?.(); + if (visible) { await emitForJob({ event, data }); } }, diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index 34f3e254af..5a2e93ba68 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -142,6 +142,12 @@ const { createCompactionSemanticIndexProjection, restoreCompactionSemanticIndexSnapshot, MAX_AGENT_CONTEXT_SKILLS, + isAgentFadingTier, + isAgentFadingTierEntries, + resolveRunContextMeta, + resolveRunFadingTiers, + createContextMetaPublisher, + selectRunContextMetaToPublish, } = require('@librechat/api'); const { Run, @@ -210,7 +216,7 @@ function normalizeEventActorContextMeta(contextMeta) { if (contextMeta == null) { return undefined; } - const { calibrationRatio, encoding } = contextMeta; + const { calibrationRatio, encoding, fading, fadingTiers } = contextMeta; if ( !Number.isFinite(calibrationRatio) || calibrationRatio < 0.5 || @@ -222,7 +228,64 @@ function normalizeEventActorContextMeta(contextMeta) { ) { throw new RangeError('Event actor context calibration is invalid'); } - return { calibrationRatio, ...(encoding == null ? {} : { encoding }) }; + if (fading != null && !isAgentFadingTier(fading)) { + throw new RangeError('Event actor context fading tier is invalid'); + } + if (fadingTiers != null && !isAgentFadingTierEntries(fadingTiers)) { + throw new RangeError('Event actor context fading tiers are invalid'); + } + return { + calibrationRatio, + ...(encoding == null ? {} : { encoding }), + ...(fading == null ? {} : { fading }), + ...(fadingTiers == null ? {} : { fadingTiers }), + }; +} + +/** + * Seeds for a new run from the previous run's contextMeta: the calibration + * ratio when the tokenizer encoding still matches, and the fading tiers, which + * are character-based and so seed regardless of encoding. The default agent's + * tier and the per-agent map are both passed; the SDK restores each agent from + * its own entry and falls back to the default tier for the first agent. + */ +function resolveRunSeeds(client) { + const prevMeta = client.contextMeta; + if (prevMeta == null) { + return {}; + } + const currentEncoding = client.getEncoding(); + const encodingMatch = prevMeta.encoding === currentEncoding; + const calibrationRatio = + encodingMatch && prevMeta.calibrationRatio > 0 ? prevMeta.calibrationRatio : undefined; + const fadingTier = isAgentFadingTier(prevMeta.fading) ? prevMeta.fading : undefined; + const fadingTiers = resolveRunFadingTiers(prevMeta.fadingTiers); + logger.debug( + `[AgentClient] contextMeta from parent: ratio=${prevMeta.calibrationRatio}, encoding=${prevMeta.encoding}, current=${currentEncoding}, seeded=${calibrationRatio ?? 'none'}, fading=${fadingTier ? `${fadingTier.budgetTokens}/${fadingTier.masked}` : 'none'}, agents=${fadingTiers ? Object.keys(fadingTiers).length : 0}`, + ); + return { calibrationRatio, fadingTier, fadingTiers }; +} + +/** + * Captures the compact context state of a run for persistence on the response + * message: calibration plus the latched fading tiers, never message content. + * Called from `finally`, so values survive an abort. The tier getters are + * optional so SDK versions without them persist calibration alone, and they + * already return only tiers that carry information; the encoding is only + * resolved when there is something to persist. + */ +function captureRunContextMeta(client) { + const run = client.run; + /** `Run` refreshes its own getters only after `processStream` settles, so a + * capture taken mid-run (a HITL pause, a Stop) reads the live graph state. */ + const graph = run?.Graph; + const source = graph ?? run; + return resolveRunContextMeta({ + calibrationRatio: source?.getCalibrationRatio?.() ?? 0, + fadingTier: source?.getFadingTier?.(), + fadingTiers: source?.getFadingTiers?.(), + getEncoding: () => client.getEncoding(), + }); } function getLatestEventActorSummary(contentParts) { @@ -356,6 +419,9 @@ class AgentClient extends BaseClient { * ON_CONTEXT_USAGE handler; persisted on `metadata.contextUsage`. * @type {{ latest: import('librechat-data-provider').TContextUsageEvent | null } | undefined} */ this.contextUsageSink = contextUsageSink; + if (this.contextUsageSink != null) { + this.contextUsageSink.onSnapshot = () => this.publishRunContextMeta({ live: true }); + } /** Every emitted `on_token_usage` payload for this response (primary, * summarization, sequential, and subagent); aggregated into the rollup * persisted on `metadata.usage`. @@ -1866,6 +1932,56 @@ class AgentClient extends BaseClient { }; } + /** + * Seeds context meta captured at a pause (or from a parent response) into a + * rebuilt client. Malformed values are dropped rather than trusted. + * @param {unknown} contextMeta + */ + /** + * Publishes the run's compact context state onto the job: the inherited seed + * before the run starts, then the live state after each pre-invoke context + * snapshot. A Stop or disconnect persists the response from job data alone, on + * whichever replica handles it, so callers await the publish ahead of the + * model call it describes. Ordering, deduplication, retries and failure + * handling live in the `packages/api` publisher; this is only the wiring. + * @param {{ live?: boolean }} [options] `live` marks a snapshot from the running + * graph; the pre-run call publishes the inherited seed instead. + * @returns {Promise} + */ + publishRunContextMeta({ live = false } = {}) { + const streamId = this.options?.req?._resumableStreamId; + if (!streamId) { + return Promise.resolve(); + } + this.contextMetaPublisher ??= createContextMetaPublisher({ + write: (contextMeta) => + GenerationJobManager.updateMetadata(streamId, { contextMeta }, this.jobCreatedAt), + onFailure: (err) => + logger.warn( + `[AgentClient] Failed to publish context meta for ${streamId}`, + getSafeErrorMetadata(err), + ), + }); + const contextMeta = selectRunContextMetaToPublish({ + live, + captured: captureRunContextMeta(this), + inherited: this.contextMeta, + hasPublished: this.contextMetaPublisher.hasPublished, + getEncoding: () => this.getEncoding(), + }); + return contextMeta == null ? Promise.resolve() : this.contextMetaPublisher.publish(contextMeta); + } + + seedContextMeta(contextMeta) { + try { + this.contextMeta = normalizeEventActorContextMeta(contextMeta); + } catch (err) { + logger.warn('[AgentClient] Ignoring malformed context meta', getSafeErrorMetadata(err)); + this.contextMeta = undefined; + } + void this.publishRunContextMeta?.(); + } + async loadHistory(conversationId, parentMessageId = null) { if (this.eventActorContinuation === 'warm') { logger.debug('[AgentClient] Skipping durable history for compatible event actor', { @@ -2425,8 +2541,17 @@ class AgentClient extends BaseClient { * last is the current user message). Seeds the pruner's calibration EMA for this run. */ const parentResponse = orderedMessages.length >= 2 ? orderedMessages[orderedMessages.length - 2] : undefined; - if (parentResponse?.contextMeta && !parentResponse.isCreatedByUser) { + /** Only a server-authored response may seed the run: a client-submitted + * row carries no trusted calibration or fading state. */ + if ( + parentResponse?.contextMeta && + !parentResponse.isCreatedByUser && + parentResponse.isUserSubmitted !== true + ) { this.contextMeta = parentResponse.contextMeta; + /** Start the seed publish as soon as the parent's state is known: a Stop + * during the rest of setup must already find it on the job. */ + void this.publishRunContextMeta?.(); } const result = { @@ -3515,6 +3640,7 @@ class AgentClient extends BaseClient { ...(staged.compactionSemanticIndex == null ? {} : { compactionSemanticIndex: staged.compactionSemanticIndex }), + ...(staged.contextMeta == null ? {} : { contextMeta: staged.contextMeta }), persistencePending: true, ...(eventActorSuspension == null ? {} @@ -3760,6 +3886,10 @@ class AgentClient extends BaseClient { compactionSemanticIndex: createCompactionSemanticIndexProjection( this.compactionSemanticIndexSnapshot, ), + // Calibration and fading state at the pause, so the resumed segment seeds + // its rebuilt pruner from the same tiers and its provider projection of + // history keeps the same bytes. + contextMeta: captureRunContextMeta({ run, getEncoding: () => this.getEncoding() }), }; if (this.eventActorInvocationId != null) { return; @@ -3770,6 +3900,9 @@ class AgentClient extends BaseClient { } async chatCompletion({ payload, userMCPAuthMap, abortController = null }) { + /** The inherited state is on the job before any abortable setup begins; a + * publish already started while loading history is simply awaited. */ + await this.publishRunContextMeta?.(); /** @type {Partial} */ let config; /** @type {ReturnType} */ @@ -4143,18 +4276,7 @@ class AgentClient extends BaseClient { memoryPromise = this.runMemory(memoryMessages); } - /** Seed calibration state from previous run if encoding matches */ - const currentEncoding = this.getEncoding(); - const prevMeta = this.contextMeta; - const encodingMatch = prevMeta?.encoding === currentEncoding; - const calibrationRatio = - encodingMatch && prevMeta?.calibrationRatio > 0 ? prevMeta.calibrationRatio : undefined; - - if (prevMeta) { - logger.debug( - `[AgentClient] contextMeta from parent: ratio=${prevMeta.calibrationRatio}, encoding=${prevMeta.encoding}, current=${currentEncoding}, seeded=${calibrationRatio ?? 'none'}`, - ); - } + const { calibrationRatio, fadingTier, fadingTiers } = resolveRunSeeds(this); const streamId = this.options.req?._resumableStreamId; // HITL: establish an empty checkpoint barrier for THIS immutable generation @@ -4260,6 +4382,8 @@ class AgentClient extends BaseClient { : { compactionSemanticIndex: continuationCompactionSemanticIndex }), initialSessions, calibrationRatio, + fadingTier, + fadingTiers, runId: this.responseMessageId, signal: abortController.signal, /** The phase wrapper stays outermost: it claims and offsets the @@ -4328,6 +4452,8 @@ class AgentClient extends BaseClient { if (this.activityLabelsMarkedPromise != null) { await this.activityLabelsMarkedPromise; } + /** The inherited tier must be on the job before any Stop can read it. */ + await this.publishRunContextMeta?.(); try { const invocationMessages = this.eventActorContinuation === 'warm' ? messages.slice(-1) : messages; @@ -4479,17 +4605,10 @@ class AgentClient extends BaseClient { * the failure; retain that model-visible state for actor reconciliation. */ this.eventActorSummary = getLatestEventActorSummary(this.contentParts) ?? this.eventActorSummary; - /** Capture calibration state from the run for persistence on the response message. - * Runs in finally so values are captured even on abort. */ - const ratio = this.run?.getCalibrationRatio() ?? 0; - if (ratio > 0 && ratio !== 1) { - this.contextMeta = { - calibrationRatio: Math.round(ratio * 1000) / 1000, - encoding: this.getEncoding(), - }; - } else { - this.contextMeta = undefined; - } + /** A run that never came to exist has no state of its own: keep the + * inherited meta so the persisted error response still seeds the next + * turn. A created run's neutral state may still clear it. */ + this.contextMeta = this.run == null ? this.contextMeta : captureRunContextMeta(this); this.finalizeSubagentContent(); this.stampMcpServerIdentities(); @@ -4591,6 +4710,9 @@ class AgentClient extends BaseClient { activityPhaseSnapshot, compactionSemanticIndex, }) { + /** The seeded state is on the job before the run is rebuilt, so a Stop + * during rebuild still persists it onto the stopped response. */ + await this.publishRunContextMeta?.(); /** @type {Partial} */ let config; /** @type {ReturnType} */ @@ -4791,6 +4913,7 @@ class AgentClient extends BaseClient { // rebuilt model binding. Undefined/empty for non-deferred turns is a no-op. discoveredToolNames, initialSessions, + ...resolveRunSeeds(this), runId: this.responseMessageId, signal: abortController.signal, // The rebuilt graph numbers content indices from 0, but the aggregator was @@ -4849,6 +4972,7 @@ class AgentClient extends BaseClient { if (this.activityLabelsMarkedPromise != null) { await this.activityLabelsMarkedPromise; } + await this.publishRunContextMeta?.(); try { await run.resume( resumeValue, @@ -4926,15 +5050,10 @@ class AgentClient extends BaseClient { } finally { this.eventActorSummary = getLatestEventActorSummary(this.contentParts) ?? this.eventActorSummary; - const ratio = this.run?.getCalibrationRatio() ?? 0; - if (ratio > 0 && ratio !== 1) { - this.contextMeta = { - calibrationRatio: Math.round(ratio * 1000) / 1000, - encoding: this.getEncoding(), - }; - } else { - this.contextMeta = undefined; - } + /** A run that never came to exist has no state of its own: keep the + * inherited meta so the persisted error response still seeds the next + * turn. A created run's neutral state may still clear it. */ + this.contextMeta = this.run == null ? this.contextMeta : captureRunContextMeta(this); this.finalizeSubagentContent(); this.stampMcpServerIdentities(); diff --git a/api/server/controllers/agents/client.test.js b/api/server/controllers/agents/client.test.js index 0eb7d3d1e7..1b351ad011 100644 --- a/api/server/controllers/agents/client.test.js +++ b/api/server/controllers/agents/client.test.js @@ -921,6 +921,570 @@ describe('AgentClient - interrupt discovery persistence', () => { }); }); + it('makes the run context meta durable when the run pauses', async () => { + const streamId = 'conversation-context-meta-pause'; + const job = await GenerationJobManager.createJob(streamId, 'user-123', streamId); + const client = new AgentClient({ + req: { + user: { id: 'user-123' }, + body: { endpoint: EModelEndpoint.agents, agent_id: 'agent-123' }, + config: { endpoints: { [EModelEndpoint.agents]: {} } }, + }, + res: {}, + agent: { + id: 'agent-123', + endpoint: EModelEndpoint.openAI, + provider: EModelEndpoint.openAI, + model_parameters: { model: 'gpt-4' }, + }, + contentParts: [], + collectedUsage: [], + artifactPromises: [], + }); + client.conversationId = streamId; + client.responseMessageId = 'response-context-meta-pause'; + client.jobCreatedAt = job.createdAt; + const fading = { v: 1, budgetTokens: 50_000, masked: true }; + + await client.handleRunInterrupt( + { + getInterrupt: () => ({ + interruptId: 'ask-interrupt', + threadId: streamId, + payload: { + type: 'ask_user_question', + question: { question: 'Proceed?' }, + }, + }), + getDiscoveredTools: () => [], + getRunMessages: () => [], + getCalibrationRatio: () => 1.2, + getFadingTier: () => ({ ...fading, latched: true }), + getFadingTiers: () => ({ + 'agent-123': { ...fading, latched: true }, + 'agent-worker': { v: 1, budgetTokens: 8_000, masked: false, latched: true }, + }), + }, + streamId, + ); + + const paused = await GenerationJobManager.getJob(streamId); + expect(paused?.status).toBe('requires_action'); + /** Only the compact tiers travel: the SDK's provenance flag is stripped. */ + expect(paused?.metadata.contextMeta).toEqual({ + calibrationRatio: 1.2, + encoding: client.getEncoding(), + fading, + fadingTiers: [ + { agentId: 'agent-123', v: 1, budgetTokens: 50_000, masked: true }, + { agentId: 'agent-worker', v: 1, budgetTokens: 8_000, masked: false }, + ], + }); + }); + + it('captures the live graph state at a pause instead of the run seeds', async () => { + const streamId = 'conversation-context-meta-live-pause'; + const job = await GenerationJobManager.createJob(streamId, 'user-123', streamId); + const client = new AgentClient({ + req: { + user: { id: 'user-123' }, + body: { endpoint: EModelEndpoint.agents, agent_id: 'agent-123' }, + config: { endpoints: { [EModelEndpoint.agents]: {} } }, + }, + res: {}, + agent: { + id: 'agent-123', + endpoint: EModelEndpoint.openAI, + provider: EModelEndpoint.openAI, + model_parameters: { model: 'gpt-4' }, + }, + contentParts: [], + collectedUsage: [], + artifactPromises: [], + }); + client.conversationId = streamId; + client.responseMessageId = 'response-context-meta-live-pause'; + client.jobCreatedAt = job.createdAt; + const seeded = { v: 1, budgetTokens: 50_000, masked: false }; + const live = { v: 1, budgetTokens: 25_000, masked: true }; + + await client.handleRunInterrupt( + { + getInterrupt: () => ({ + interruptId: 'ask-interrupt', + threadId: streamId, + payload: { + type: 'ask_user_question', + question: { question: 'Proceed?' }, + }, + }), + getDiscoveredTools: () => [], + getRunMessages: () => [], + getCalibrationRatio: () => 1, + getFadingTier: () => seeded, + Graph: { + getCalibrationRatio: () => 1.4, + getFadingTier: () => live, + }, + }, + streamId, + ); + + const paused = await GenerationJobManager.getJob(streamId); + expect(paused?.metadata.contextMeta).toEqual({ + calibrationRatio: 1.4, + encoding: client.getEncoding(), + fading: live, + }); + }); + + it('publishes the live context meta onto the job after each context snapshot', async () => { + const streamId = 'conversation-context-meta-publish'; + const job = await GenerationJobManager.createJob(streamId, 'user-123', streamId); + const contextUsageSink = { latest: null, count: 0 }; + const client = new AgentClient({ + req: { + user: { id: 'user-123' }, + body: { endpoint: EModelEndpoint.agents, agent_id: 'agent-123' }, + config: { endpoints: { [EModelEndpoint.agents]: {} } }, + _resumableStreamId: streamId, + }, + res: {}, + agent: { + id: 'agent-123', + endpoint: EModelEndpoint.openAI, + provider: EModelEndpoint.openAI, + model_parameters: { model: 'gpt-4' }, + }, + contentParts: [], + collectedUsage: [], + artifactPromises: [], + jobCreatedAt: job.createdAt, + contextUsageSink, + }); + const seed = { calibrationRatio: 1.05, encoding: client.getEncoding() }; + client.contextMeta = seed; + const updateMetadata = jest.spyOn(GenerationJobManager, 'updateMetadata'); + + /** Before the run exists the inherited seed is what a Stop must carry. */ + await client.publishRunContextMeta(); + await expect(GenerationJobManager.getJob(streamId)).resolves.toMatchObject({ + metadata: { contextMeta: seed }, + }); + + const tier = { v: 1, budgetTokens: 25_000, masked: true }; + let ratio = 1.1; + client.run = { + Graph: { + getCalibrationRatio: () => ratio, + getFadingTier: () => tier, + }, + }; + + expect(typeof contextUsageSink.onSnapshot).toBe('function'); + await contextUsageSink.onSnapshot(); + await contextUsageSink.onSnapshot(); + ratio = 1.3; + await contextUsageSink.onSnapshot(); + + expect(updateMetadata).toHaveBeenCalledTimes(3); + expect(updateMetadata).toHaveBeenLastCalledWith( + streamId, + { + contextMeta: { + calibrationRatio: 1.3, + encoding: client.getEncoding(), + fading: tier, + }, + }, + job.createdAt, + ); + const running = await GenerationJobManager.getJob(streamId); + expect(running?.metadata.contextMeta).toEqual({ + calibrationRatio: 1.3, + encoding: client.getEncoding(), + fading: tier, + }); + updateMetadata.mockRestore(); + }); + + it('publishes a neutral record when live state stops carrying anything', async () => { + const streamId = 'conversation-context-meta-neutral'; + const job = await GenerationJobManager.createJob(streamId, 'user-123', streamId); + const contextUsageSink = { latest: null, count: 0 }; + const client = new AgentClient({ + req: { + user: { id: 'user-123' }, + body: { endpoint: EModelEndpoint.agents, agent_id: 'agent-123' }, + config: { endpoints: { [EModelEndpoint.agents]: {} } }, + _resumableStreamId: streamId, + }, + res: {}, + agent: { + id: 'agent-123', + endpoint: EModelEndpoint.openAI, + provider: EModelEndpoint.openAI, + model_parameters: { model: 'gpt-4' }, + }, + contentParts: [], + collectedUsage: [], + artifactPromises: [], + jobCreatedAt: job.createdAt, + contextUsageSink, + }); + let ratio = 1; + let tier; + client.run = { + Graph: { + getCalibrationRatio: () => ratio, + getFadingTier: () => tier, + }, + }; + const updateMetadata = jest.spyOn(GenerationJobManager, 'updateMetadata'); + + /** A fresh conversation's first neutral snapshot has nothing to clear. */ + await contextUsageSink.onSnapshot(); + expect(updateMetadata).not.toHaveBeenCalled(); + await expect(GenerationJobManager.getJob(streamId)).resolves.toMatchObject({ + metadata: expect.not.objectContaining({ contextMeta: expect.anything() }), + }); + + ratio = 1.25; + await contextUsageSink.onSnapshot(); + await expect(GenerationJobManager.getJob(streamId)).resolves.toMatchObject({ + metadata: { contextMeta: { calibrationRatio: 1.25, encoding: client.getEncoding() } }, + }); + + ratio = 1; + await contextUsageSink.onSnapshot(); + await contextUsageSink.onSnapshot(); + expect(updateMetadata).toHaveBeenCalledTimes(2); + await expect(GenerationJobManager.getJob(streamId)).resolves.toMatchObject({ + metadata: { contextMeta: { calibrationRatio: 1, encoding: client.getEncoding() } }, + }); + updateMetadata.mockRestore(); + }); + + it('shares one in-flight write between equal concurrent snapshots', async () => { + const streamId = 'conversation-context-meta-inflight'; + const job = await GenerationJobManager.createJob(streamId, 'user-123', streamId); + const contextUsageSink = { latest: null, count: 0 }; + const client = new AgentClient({ + req: { + user: { id: 'user-123' }, + body: { endpoint: EModelEndpoint.agents, agent_id: 'agent-123' }, + config: { endpoints: { [EModelEndpoint.agents]: {} } }, + _resumableStreamId: streamId, + }, + res: {}, + agent: { + id: 'agent-123', + endpoint: EModelEndpoint.openAI, + provider: EModelEndpoint.openAI, + model_parameters: { model: 'gpt-4' }, + }, + contentParts: [], + collectedUsage: [], + artifactPromises: [], + jobCreatedAt: job.createdAt, + contextUsageSink, + }); + const tier = { v: 1, budgetTokens: 25_000, masked: true }; + client.run = { + Graph: { + getCalibrationRatio: () => 1.2, + getFadingTier: () => tier, + }, + }; + let settled = false; + let releaseWrite; + const updateMetadata = jest + .spyOn(GenerationJobManager, 'updateMetadata') + .mockImplementationOnce(async () => { + await new Promise((resolve) => { + releaseWrite = resolve; + }); + settled = true; + }); + + const first = contextUsageSink.onSnapshot(); + const second = contextUsageSink.onSnapshot(); + let secondSettled = false; + void second.then(() => { + secondSettled = true; + }); + await Promise.resolve(); + expect(secondSettled).toBe(false); + expect(updateMetadata).toHaveBeenCalledTimes(1); + + releaseWrite(); + await Promise.all([first, second]); + expect(settled).toBe(true); + expect(secondSettled).toBe(true); + updateMetadata.mockRestore(); + }); + + it('serializes distinct concurrent snapshots so the newest value wins', async () => { + const streamId = 'conversation-context-meta-ordered'; + const job = await GenerationJobManager.createJob(streamId, 'user-123', streamId); + const contextUsageSink = { latest: null, count: 0 }; + const client = new AgentClient({ + req: { + user: { id: 'user-123' }, + body: { endpoint: EModelEndpoint.agents, agent_id: 'agent-123' }, + config: { endpoints: { [EModelEndpoint.agents]: {} } }, + _resumableStreamId: streamId, + }, + res: {}, + agent: { + id: 'agent-123', + endpoint: EModelEndpoint.openAI, + provider: EModelEndpoint.openAI, + model_parameters: { model: 'gpt-4' }, + }, + contentParts: [], + collectedUsage: [], + artifactPromises: [], + jobCreatedAt: job.createdAt, + contextUsageSink, + }); + let tier = { v: 1, budgetTokens: 50_000, masked: false }; + client.run = { + Graph: { + getCalibrationRatio: () => 1.2, + getFadingTier: () => tier, + }, + }; + let releaseFirstWrite; + const originalUpdateMetadata = Object.getPrototypeOf(GenerationJobManager).updateMetadata; + const updateMetadata = jest.spyOn(GenerationJobManager, 'updateMetadata'); + updateMetadata.mockImplementationOnce(async (...args) => { + await new Promise((resolve) => { + releaseFirstWrite = resolve; + }); + return originalUpdateMetadata.apply(GenerationJobManager, args); + }); + + const older = contextUsageSink.onSnapshot(); + tier = { v: 1, budgetTokens: 25_000, masked: true }; + const newer = contextUsageSink.onSnapshot(); + await Promise.resolve(); + /** The newer write waits for the older one instead of racing it. */ + expect(updateMetadata).toHaveBeenCalledTimes(1); + + releaseFirstWrite(); + await Promise.all([older, newer]); + expect(updateMetadata).toHaveBeenCalledTimes(2); + await expect(GenerationJobManager.getJob(streamId)).resolves.toMatchObject({ + metadata: { + contextMeta: { + calibrationRatio: 1.2, + encoding: client.getEncoding(), + fading: { v: 1, budgetTokens: 25_000, masked: true }, + }, + }, + }); + updateMetadata.mockRestore(); + }); + + it('publishes the inherited context meta before the resumed run continues', async () => { + const streamId = 'conversation-context-meta-resume-seed'; + const job = await GenerationJobManager.createJob(streamId, 'user-123', streamId); + const client = new AgentClient({ + req: { + user: { id: 'user-123' }, + body: { endpoint: EModelEndpoint.agents, agent_id: 'agent-123' }, + config: { endpoints: { [EModelEndpoint.agents]: {} } }, + _resumableStreamId: streamId, + }, + res: {}, + agent: { + id: 'agent-123', + endpoint: EModelEndpoint.openAI, + provider: EModelEndpoint.openAI, + model_parameters: { model: 'gpt-4' }, + }, + contentParts: [], + collectedUsage: [], + artifactPromises: [], + jobCreatedAt: job.createdAt, + }); + const seed = { + calibrationRatio: 1.15, + encoding: client.getEncoding(), + fading: { v: 1, budgetTokens: 30_000, masked: true }, + }; + client.seedContextMeta(seed); + let metaWhenResumed; + let metaWhenCreated; + const resume = jest.fn(async () => { + metaWhenResumed = (await GenerationJobManager.getJob(streamId))?.metadata.contextMeta; + }); + mockCreateRun.mockImplementationOnce(async () => { + metaWhenCreated = (await GenerationJobManager.getJob(streamId))?.metadata.contextMeta; + return { + Graph: null, + resume, + processStream: jest.fn().mockResolvedValue(), + getCalibrationRatio: jest.fn(() => 0), + getInterrupt: jest.fn(() => undefined), + }; + }); + client.conversationId = streamId; + client.responseMessageId = 'response-context-meta-resume-seed'; + client.recordCollectedUsage = jest.fn().mockResolvedValue(); + + await client.resumeCompletion({ + resumeValue: { decisions: [] }, + streamId, + checkpointNamespace: 'resume-seed', + }); + + expect(resume).toHaveBeenCalledTimes(1); + expect(metaWhenCreated).toEqual(seed); + expect(metaWhenResumed).toEqual(seed); + }); + + it('publishes the inherited context meta before a fresh run streams', async () => { + const streamId = 'conversation-context-meta-stream-seed'; + const job = await GenerationJobManager.createJob(streamId, 'user-123', streamId); + const client = new AgentClient({ + req: { + user: { id: 'user-123' }, + body: { endpoint: EModelEndpoint.agents, agent_id: 'agent-123' }, + config: { endpoints: { [EModelEndpoint.agents]: {} } }, + _resumableStreamId: streamId, + }, + res: {}, + agent: { + id: 'agent-123', + endpoint: EModelEndpoint.openAI, + provider: EModelEndpoint.openAI, + model_parameters: { model: 'gpt-4' }, + }, + contentParts: [], + collectedUsage: [], + artifactPromises: [], + jobCreatedAt: job.createdAt, + }); + const seed = { + calibrationRatio: 1.15, + encoding: client.getEncoding(), + fading: { v: 1, budgetTokens: 30_000, masked: true }, + }; + client.contextMeta = seed; + let metaWhenStreamed; + let metaWhenCreated; + const processStream = jest.fn(async () => { + metaWhenStreamed = (await GenerationJobManager.getJob(streamId))?.metadata.contextMeta; + }); + /** Run creation is the first abortable setup stage; the seed must precede it. */ + mockCreateRun.mockImplementationOnce(async () => { + metaWhenCreated = (await GenerationJobManager.getJob(streamId))?.metadata.contextMeta; + return { + Graph: null, + processStream, + getCalibrationRatio: jest.fn(() => 0), + getInterrupt: jest.fn(() => undefined), + }; + }); + client.conversationId = streamId; + client.responseMessageId = 'response-context-meta-stream-seed'; + client.parentMessageId = 'parent-context-meta-stream-seed'; + client.recordCollectedUsage = jest.fn().mockResolvedValue(); + + await client.chatCompletion({ payload: [] }); + + expect(processStream).toHaveBeenCalledTimes(1); + expect(metaWhenCreated).toEqual(seed); + expect(metaWhenStreamed).toEqual(seed); + }); + + it('keeps the inherited context meta when a fresh run fails before it exists', async () => { + const streamId = 'conversation-context-meta-setup-failure'; + const job = await GenerationJobManager.createJob(streamId, 'user-123', streamId); + const client = new AgentClient({ + req: { + user: { id: 'user-123' }, + body: { endpoint: EModelEndpoint.agents, agent_id: 'agent-123' }, + config: { endpoints: { [EModelEndpoint.agents]: {} } }, + _resumableStreamId: streamId, + }, + res: {}, + agent: { + id: 'agent-123', + endpoint: EModelEndpoint.openAI, + provider: EModelEndpoint.openAI, + model_parameters: { model: 'gpt-4' }, + }, + contentParts: [], + collectedUsage: [], + artifactPromises: [], + jobCreatedAt: job.createdAt, + }); + const seed = { + calibrationRatio: 1.15, + encoding: client.getEncoding(), + fading: { v: 1, budgetTokens: 30_000, masked: true }, + }; + client.contextMeta = seed; + mockCreateRun.mockRejectedValueOnce(new Error('run creation failed')); + client.conversationId = streamId; + client.responseMessageId = 'response-context-meta-setup-failure'; + client.parentMessageId = 'parent-context-meta-setup-failure'; + client.recordCollectedUsage = jest.fn().mockResolvedValue(); + + await client.chatCompletion({ payload: [] }).catch(() => undefined); + + expect(client.run).toBeUndefined(); + expect(client.contextMeta).toEqual(seed); + }); + + it('keeps the inherited context meta when a resumed run fails to rebuild', async () => { + const streamId = 'conversation-context-meta-resume-failure'; + const job = await GenerationJobManager.createJob(streamId, 'user-123', streamId); + const client = new AgentClient({ + req: { + user: { id: 'user-123' }, + body: { endpoint: EModelEndpoint.agents, agent_id: 'agent-123' }, + config: { endpoints: { [EModelEndpoint.agents]: {} } }, + _resumableStreamId: streamId, + }, + res: {}, + agent: { + id: 'agent-123', + endpoint: EModelEndpoint.openAI, + provider: EModelEndpoint.openAI, + model_parameters: { model: 'gpt-4' }, + }, + contentParts: [], + collectedUsage: [], + artifactPromises: [], + jobCreatedAt: job.createdAt, + }); + const seed = { + calibrationRatio: 1.15, + encoding: client.getEncoding(), + fading: { v: 1, budgetTokens: 30_000, masked: true }, + }; + client.seedContextMeta(seed); + mockCreateRun.mockRejectedValueOnce(new Error('rebuild failed')); + client.conversationId = streamId; + client.responseMessageId = 'response-context-meta-resume-failure'; + client.recordCollectedUsage = jest.fn().mockResolvedValue(); + + await client + .resumeCompletion({ + resumeValue: { decisions: [] }, + streamId, + checkpointNamespace: 'resume-failure', + }) + .catch(() => undefined); + + expect(client.run).toBeUndefined(); + expect(client.contextMeta).toEqual(seed); + }); + it('caps an event-bound pause at the inherited binding deadline', async () => { const now = Date.now(); const streamId = 'conversation-event-bound-pause'; @@ -2236,7 +2800,15 @@ describe('AgentClient - startup telemetry', () => { client.eventActorContinuation = 'warm'; client.eventActorDiscoveredToolNames = ['deferred_tool']; client.eventActorSummary = { text: 'summary of earlier turns', tokenCount: 40 }; - client.contextMeta = { calibrationRatio: 1.25, encoding: client.getEncoding() }; + client.contextMeta = { + calibrationRatio: 1.25, + encoding: client.getEncoding(), + fading: { v: 1, budgetTokens: 20_000, masked: true }, + fadingTiers: [ + { agentId: 'agent-123', v: 1, budgetTokens: 20_000, masked: true }, + { agentId: 'agent-worker', v: 1, budgetTokens: 8_000, masked: false }, + ], + }; client.compactionSemanticIndexSnapshot = baseCompactionSemanticIndexSnapshot; client.recordCollectedUsage = jest.fn().mockResolvedValue(); @@ -2255,9 +2827,17 @@ describe('AgentClient - startup telemetry', () => { indexTokenCountMap: {}, initialSummary: { text: 'summary of earlier turns', tokenCount: 40 }, calibrationRatio: 1.25, + fadingTier: { v: 1, budgetTokens: 20_000, masked: true }, + fadingTiers: { + 'agent-123': { v: 1, budgetTokens: 20_000, masked: true }, + 'agent-worker': { v: 1, budgetTokens: 8_000, masked: false }, + }, compactionSemanticIndex: evolvedCompactionSemanticIndexSnapshot.entries, }), ); + /** The seeded map must be prototype-safe: a null-prototype record built from own keys. */ + const seededTiers = mockCreateRun.mock.calls.at(-1)[0].fadingTiers; + expect(Object.getPrototypeOf(seededTiers)).toBeNull(); expect(mockFormatAgentMessages.mock.calls[0][4]).toEqual( expect.objectContaining({ compactionSemanticIndex: expect.objectContaining({ @@ -4516,6 +5096,43 @@ describe('AgentClient - titleConvo', () => { ).resolves.toEqual(expect.objectContaining({ prompt: expect.any(Array) })); }); + it.each([ + ['seeds the run from a server-authored parent response', undefined, true], + ['ignores context meta on a client-submitted parent response', true, false], + ])('%s', async (_label, isUserSubmitted, expectSeed) => { + const contextMeta = { + calibrationRatio: 1.2, + encoding: 'claude', + fading: { v: 1, budgetTokens: 50_000, masked: true }, + }; + const parentResponse = { + messageId: 'assistant-seed', + parentMessageId: null, + sender: 'Assistant', + role: 'assistant', + isCreatedByUser: false, + ...(isUserSubmitted != null && { isUserSubmitted }), + text: 'Prior response', + contextMeta, + }; + const userMessage = { + messageId: 'user-next', + parentMessageId: 'assistant-seed', + sender: 'User', + role: 'user', + isCreatedByUser: true, + text: 'Next question', + }; + client.setModelBoundStoredMessages([parentResponse, userMessage]); + client.contextMeta = undefined; + client.publishRunContextMeta = jest.fn().mockResolvedValue(undefined); + + await client.buildMessages([parentResponse, userMessage], 'user-next', {}, {}); + + expect(client.contextMeta).toEqual(expectSeed ? contextMeta : undefined); + expect(client.publishRunContextMeta).toHaveBeenCalledTimes(expectSeed ? 1 : 0); + }); + it('preserves persisted source identity through both agent formatting passes', async () => { mockReq.config.filters = { messages: { @@ -7656,3 +8273,88 @@ describe('AgentClient - resumeCompletion content protection', () => { }); }); }); + +describe('fading tier context meta', () => { + const fading = { v: 1, budgetTokens: 20_000, masked: true }; + const fadingTiers = [ + { agentId: 'agent-123', v: 1, budgetTokens: 20_000, masked: true }, + { agentId: 'agent-worker', v: 1, budgetTokens: 8_000, masked: false }, + ]; + + it('carries valid fading tiers through event actor context', async () => { + const client = Object.create(AgentClient.prototype); + client.options = { req: { config: {} } }; + const contextMeta = { calibrationRatio: 1.25, encoding: 'o200k_base', fading, fadingTiers }; + client.getEventActorContext = jest.fn().mockResolvedValue({ + fingerprint: 'fingerprint-1', + skillManifest: [], + discoveredToolNames: [], + contextMeta, + }); + + await expect( + client.prepareEventActorContext({ + contextFingerprint: 'fingerprint-1', + skillManifest: [], + discoveredToolNames: [], + contextMeta, + }), + ).resolves.toMatchObject({ contextMeta }); + expect(client.contextMeta).toEqual(contextMeta); + }); + + it('rejects event actor context carrying a malformed fading tier', async () => { + const client = Object.create(AgentClient.prototype); + client.options = { req: { config: {} } }; + client.getEventActorContext = jest.fn(); + + await expect( + client.prepareEventActorContext({ + contextFingerprint: 'fingerprint-1', + skillManifest: [], + discoveredToolNames: [], + contextMeta: { calibrationRatio: 1.25, fading: { v: 1, budgetTokens: -5, masked: true } }, + }), + ).resolves.toBeUndefined(); + expect(client.getEventActorContext).not.toHaveBeenCalled(); + }); + + it('rejects event actor context carrying malformed per-agent tiers', async () => { + const client = Object.create(AgentClient.prototype); + client.options = { req: { config: {} } }; + client.getEventActorContext = jest.fn(); + + await expect( + client.prepareEventActorContext({ + contextFingerprint: 'fingerprint-1', + skillManifest: [], + discoveredToolNames: [], + contextMeta: { + calibrationRatio: 1.25, + fadingTiers: [{ agentId: '', v: 1, budgetTokens: 20_000, masked: true }], + }, + }), + ).resolves.toBeUndefined(); + expect(client.getEventActorContext).not.toHaveBeenCalled(); + }); +}); + +describe('seedContextMeta', () => { + it('seeds valid context meta and drops a malformed fading tier', () => { + const client = Object.create(AgentClient.prototype); + const contextMeta = { + calibrationRatio: 1.25, + encoding: 'claude', + fading: { v: 1, budgetTokens: 20_000, masked: true }, + }; + + client.seedContextMeta(contextMeta); + expect(client.contextMeta).toEqual(contextMeta); + + client.seedContextMeta({ calibrationRatio: 1.25, fading: { v: 1, budgetTokens: 0 } }); + expect(client.contextMeta).toBeUndefined(); + + client.seedContextMeta(undefined); + expect(client.contextMeta).toBeUndefined(); + }); +}); diff --git a/api/server/controllers/agents/request.js b/api/server/controllers/agents/request.js index 8e8880fc52..7d9a22a908 100644 --- a/api/server/controllers/agents/request.js +++ b/api/server/controllers/agents/request.js @@ -1679,7 +1679,10 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit return; } - const resumeState = await GenerationJobManager.getResumeState(streamId, jobCreatedAt); + const [resumeState, jobRecord] = await Promise.all([ + GenerationJobManager.getResumeState(streamId, jobCreatedAt), + GenerationJobManager.getJobStore().getJob(streamId), + ]); if (!resumeState?.userMessage) { logger.debug('[ResumableAgentController] No user message to save partial response for'); return; @@ -1687,6 +1690,13 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit partialResponseSaved = true; const responseConversationId = resumeState.conversationId || conversationId; + /** The run publishes its calibration and fading tiers onto the job; a + * partial response saved on disconnect must carry them like the Stop and + * pause paths do, or a turn continued from it re-derives its provider + * projection of history and loses the cached prefix. The same-epoch job + * record is the source, since the client-facing resume snapshot never + * carries server-private state. */ + const contextMeta = jobRecord?.createdAt === jobCreatedAt ? jobRecord.contextMeta : undefined; try { const partialMessage = { @@ -1702,6 +1712,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit endpoint: endpointOption.endpoint, iconURL: resumeState.iconURL || endpointIconURL, model: resumeState.model || responseModel, + ...(contextMeta != null && { contextMeta }), }; if (req.body?.agent_id) { @@ -2722,6 +2733,8 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit reqCtx, { ...response, + /** A neutral finish unsets what a disconnect snapshot may have stored. */ + contextMeta: response.contextMeta ?? null, user: userId, unfinished: responseIsUnfinished, /** Distinguishes "ran out of steps" from a user stop, so the client can diff --git a/api/server/controllers/agents/resume.js b/api/server/controllers/agents/resume.js index 5d767a7c69..2edec4f3ac 100644 --- a/api/server/controllers/agents/resume.js +++ b/api/server/controllers/agents/resume.js @@ -500,11 +500,12 @@ async function finalizeResumedTurn({ if (Object.keys(responseMetadata).length > 0) { responseMessage.metadata = responseMetadata; } - // Carry the resumed run's context-window calibration (BaseClient.sendMessage persists - // this on the response). Without it, the NEXT turn can't seed its pruner from this - // run and falls back to uncalibrated token accounting. - if (client?.contextMeta != null) { - responseMessage.contextMeta = client.contextMeta; + // Carry the resumed run's compact context meta (calibration and fading tiers), as + // BaseClient.sendMessage persists it on the response. Without it, the NEXT turn can't + // seed its pruner from this run. A neutral finish unsets what the paused segment + // stored on this row, since an omitted field would otherwise survive the save. + if (client != null) { + responseMessage.contextMeta = client.contextMeta ?? null; } // Win terminal ownership BEFORE the outcome-defining response write. Stop @@ -1797,6 +1798,9 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) client.checkpointNamespace = checkpointNamespace; client.responseMessageId = job.metadata.responseMessageId; client.parentMessageId = job.metadata.userMessage?.messageId ?? Constants.NO_PARENT; + // Seed the rebuilt pruner from the tier and calibration captured at the pause, so the + // resumed segment keeps historical tool results byte-identical to the paused one. + client.seedContextMeta?.(job.metadata?.contextMeta); if (client.contentParts) { GenerationJobManager.setContentParts(streamId, client.contentParts, job.createdAt); } diff --git a/api/server/middleware/abortMiddleware.js b/api/server/middleware/abortMiddleware.js index 9add302e37..1f78c5eb54 100644 --- a/api/server/middleware/abortMiddleware.js +++ b/api/server/middleware/abortMiddleware.js @@ -114,6 +114,11 @@ async function abortMessage(req, res) { error: false, isCreatedByUser: false, tokenCount: completionTokens, + /** The run publishes its calibration and fading tiers onto the job as it + * goes; a stopped response must carry them or the next turn re-derives its + * provider projection of history from scratch and loses the cached prefix. + * A job with none unsets what an earlier pause stored on this row. */ + ...(jobData != null && { contextMeta: jobData.contextMeta ?? null }), }; /** Persist the usage/cost rollup + context breakdown for the stopped response diff --git a/api/server/middleware/abortMiddleware.spec.js b/api/server/middleware/abortMiddleware.spec.js index 276760162b..8a64feae89 100644 --- a/api/server/middleware/abortMiddleware.spec.js +++ b/api/server/middleware/abortMiddleware.spec.js @@ -403,6 +403,47 @@ describe('abortMiddleware - transactions config', () => { ); }); + it('carries the context meta the run published onto the job into the stopped response', async () => { + const contextMeta = { + calibrationRatio: 1.2, + encoding: 'claude', + fading: { v: 1, budgetTokens: 50_000, masked: true }, + }; + GenerationJobManager.abortJob.mockResolvedValue({ + success: true, + jobData: { ...buildJobData(), contextMeta }, + content: [], + text: 'partial', + collectedUsage: [], + }); + const res = buildRes(); + + await handleAbort()(buildReq(), res); + + expect(db.saveMessage).toHaveBeenCalledWith( + expect.any(Object), + expect.objectContaining({ messageId: 'msg-123', contextMeta }), + expect.any(Object), + ); + const finalEvent = JSON.parse(res.send.mock.calls[0][0]); + expect(finalEvent.responseMessage.contextMeta).toEqual(contextMeta); + }); + + it('unsets context meta on a stopped response when the job carries none', async () => { + GenerationJobManager.abortJob.mockResolvedValue({ + success: true, + jobData: buildJobData(), + content: [], + text: 'partial', + collectedUsage: [], + }); + + await handleAbort()(buildReq(), buildRes()); + + const [, savedMessage] = db.saveMessage.mock.calls[0]; + expect(savedMessage.contextMeta).toBeNull(); + }); + it('resolves the config from req and forwards it on the token-count fallback path', async () => { GenerationJobManager.abortJob.mockResolvedValue({ success: true, diff --git a/api/server/routes/__tests__/messages-get.spec.js b/api/server/routes/__tests__/messages-get.spec.js index 08dbc6954b..d1b397d6b9 100644 --- a/api/server/routes/__tests__/messages-get.spec.js +++ b/api/server/routes/__tests__/messages-get.spec.js @@ -318,6 +318,47 @@ describe('message route conversation ownership filters', () => { ]); }); + it.each([ + { + name: 'carries server-private context meta onto the branch', + contextMeta: { + calibrationRatio: 1.25, + encoding: 'claude', + fading: { v: 1, budgetTokens: 50_000, masked: true }, + }, + }, + { name: 'leaves context meta absent when the source has none', contextMeta: undefined }, + ])('$name', async ({ contextMeta }) => { + getMessage.mockResolvedValue({ + messageId: 'source-message', + conversationId: 'convo-1', + parentMessageId: 'parent-1', + isCreatedByUser: false, + ...(contextMeta == null ? {} : { contextMeta }), + content: [ + { type: 'text', text: 'Different agent content', agentId: 'agent-2' }, + { type: 'text', text: 'Assistant content', agentId: 'agent-1' }, + ], + }); + saveMessage.mockImplementation(async (_ctx, message) => message); + + const response = await request(app).post('/api/messages/branch').send({ + messageId: 'source-message', + agentId: 'agent-1', + }); + + expect(response.status).toBe(201); + const savedMessage = saveMessage.mock.calls[0][1]; + if (contextMeta == null) { + expect(savedMessage).not.toHaveProperty('contextMeta'); + } else { + expect(savedMessage.contextMeta).toEqual(contextMeta); + } + /** Server-private state stays in the database; normal reads strip it, so must this response. */ + expect(response.body).not.toHaveProperty('contextMeta'); + expect(response.body.messageId).toBe(savedMessage.messageId); + }); + it('does not hydrate branch files for an inert file policy', async () => { getMessage.mockResolvedValue({ messageId: 'source-message', @@ -401,6 +442,12 @@ describe('message route conversation ownership filters', () => { isTemporary: false, files: [{ file_id: 'file-1' }], user: authenticatedUserId, + /** An existing server-authored row keeps its stored state through the update. */ + contextMeta: { + calibrationRatio: 1.2, + encoding: 'claude', + fading: { v: 1, budgetTokens: 50_000, masked: true }, + }, }; saveMessage.mockResolvedValue(savedMessage); @@ -420,6 +467,11 @@ describe('message route conversation ownership filters', () => { userSubmittedMessageFieldPaths: [ { path: '/forged/model/output', field: 'decision_reason' }, ], + contextMeta: { + calibrationRatio: 1, + encoding: 'claude', + fading: { v: 1, budgetTokens: 1, masked: true }, + }, }); expect(response.status).toBe(201); @@ -437,6 +489,9 @@ describe('message route conversation ownership filters', () => { expect(saveMessage.mock.calls[0][1].conversationId).not.toBe(bodyConversationId); expect(saveMessage.mock.calls[0][1]).not.toHaveProperty('userSubmittedPaths'); expect(saveMessage.mock.calls[0][1]).not.toHaveProperty('userSubmittedMessageFieldPaths'); + expect(saveMessage.mock.calls[0][1]).not.toHaveProperty('contextMeta'); + expect(response.body.messageId).toBe(savedMessage.messageId); + expect(response.body).not.toHaveProperty('contextMeta'); expect(saveConvo).toHaveBeenCalledWith( expect.objectContaining({ userId: authenticatedUserId }), { @@ -513,6 +568,36 @@ describe('message route conversation ownership filters', () => { }, ); + it('strips server-private context meta from hydrated search hits', async () => { + searchMessages.mockResolvedValue({ + hits: [ + { + messageId: 'hit-1', + conversationId: 'convo-1', + text: 'needle in a haystack', + contextMeta: { + calibrationRatio: 1.2, + encoding: 'claude', + fading: { v: 1, budgetTokens: 50_000, masked: true }, + }, + }, + ], + }); + getConvosQueried.mockResolvedValue({ + convoMap: { 'convo-1': { title: 'Found', model: 'gpt-4' } }, + }); + getMessages.mockResolvedValue([ + { messageId: 'hit-1', isCreatedByUser: false, endpoint: 'agents', iconURL: null }, + ]); + + const response = await request(app).get('/api/messages?search=needle'); + + expect(response.status).toBe(200); + expect(response.body.messages).toHaveLength(1); + expect(response.body.messages[0]).toMatchObject({ messageId: 'hit-1', title: 'Found' }); + expect(response.body.messages[0]).not.toHaveProperty('contextMeta'); + }); + it('returns indistinguishable not-found responses for child and missing query reads', async () => { getConvoOwnership.mockResolvedValueOnce({ user: authenticatedUserId, @@ -616,6 +701,27 @@ describe('message route conversation ownership filters', () => { expect.objectContaining({ user: { id: authenticatedUserId } }), 'active-convo', ); + /** The paginated read projects like every other client read, so persisted + * `contextMeta` never reaches the client through this endpoint. */ + expect(getMessagesByCursor).toHaveBeenCalledWith( + { conversationId: 'active-convo', user: authenticatedUserId }, + expect.objectContaining({ select: CLIENT_MESSAGE_SELECT }), + ); + }); + + it('projects a single-message query read like every other client read', async () => { + getConvoOwnership.mockResolvedValue({ conversationId: 'convo-1' }); + getMessages.mockResolvedValue([{ messageId: 'message-1', conversationId: 'convo-1' }]); + + const response = await request(app).get( + '/api/messages?conversationId=convo-1&messageId=message-1', + ); + + expect(response.status).toBe(200); + expect(getMessages).toHaveBeenCalledWith( + { conversationId: 'convo-1', messageId: 'message-1', user: authenticatedUserId }, + CLIENT_MESSAGE_SELECT, + ); }); it('should start conversation message reads before validation resolves', async () => { diff --git a/api/server/routes/agents/__tests__/abort.spec.js b/api/server/routes/agents/__tests__/abort.spec.js index b99d0fb59a..3d3c748b06 100644 --- a/api/server/routes/agents/__tests__/abort.spec.js +++ b/api/server/routes/agents/__tests__/abort.spec.js @@ -424,6 +424,75 @@ describe('Agent Abort Endpoint', () => { ); }); + it('carries the context meta the run published onto the stopped response', async () => { + const contextMeta = { + calibrationRatio: 1.2, + encoding: 'claude', + fading: { v: 1, budgetTokens: 50_000, masked: true }, + fadingTiers: [{ agentId: 'agent-a', v: 1, budgetTokens: 50_000, masked: true }], + }; + mockGenerationJobManager.getJob.mockResolvedValue({ + metadata: { userId: 'test-user-123' }, + }); + const abortResult = { + success: true, + jobData: { + userMessage: { messageId: 'user-msg-123' }, + responseMessageId: 'response-msg-456', + conversationId: 'test-stream-123', + contextMeta, + }, + content: [{ type: 'text', text: 'Partial response...' }], + text: 'Partial response...', + }; + mockGenerationJobManager.abortJob.mockImplementation(async (_streamId, options) => { + await options.beforePublish(abortResult); + return abortResult; + }); + + const response = await request(app) + .post('/api/agents/chat/abort') + .send({ conversationId: 'test-stream-123' }); + + expect(response.status).toBe(200); + expect(mockSaveMessage).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ messageId: 'response-msg-456', contextMeta }), + expect.objectContaining({ + context: 'api/server/routes/agents/index.js - abort endpoint', + }), + ); + }); + + it('unsets context meta on the stopped response when the job has none', async () => { + mockGenerationJobManager.getJob.mockResolvedValue({ + metadata: { userId: 'test-user-123' }, + }); + const abortResult = { + success: true, + jobData: { + userMessage: { messageId: 'user-msg-123' }, + responseMessageId: 'response-msg-456', + conversationId: 'test-stream-123', + }, + content: [{ type: 'text', text: 'Partial response...' }], + text: 'Partial response...', + }; + mockGenerationJobManager.abortJob.mockImplementation(async (_streamId, options) => { + await options.beforePublish(abortResult); + return abortResult; + }); + + await request(app) + .post('/api/agents/chat/abort') + .send({ conversationId: 'test-stream-123' }); + + const [, savedResponse] = mockSaveMessage.mock.calls.find( + ([, message]) => message.messageId === 'response-msg-456', + ); + expect(savedResponse.contextMeta).toBeNull(); + }); + it('saves the aborted partial as temporary from job metadata, not the request body', async () => { const jobStreamId = 'test-stream-123'; diff --git a/api/server/routes/agents/index.js b/api/server/routes/agents/index.js index 3ffe49835f..46c50cd88f 100644 --- a/api/server/routes/agents/index.js +++ b/api/server/routes/agents/index.js @@ -794,6 +794,11 @@ router.post('/chat/abort', configMiddleware, async (req, res, next) => { jobData.userSubmittedMessageFieldPaths.length > 0 && { userSubmittedMessageFieldPaths: jobData.userSubmittedMessageFieldPaths, }), + /** The run published its compact context meta onto the job ahead + * of each model call; the stopped response must carry it so the + * next turn seeds the same tiers. A job with none unsets what an + * earlier pause stored on this row, since omission would keep it. */ + contextMeta: jobData.contextMeta ?? null, user: userId, }; diff --git a/api/server/routes/messages.js b/api/server/routes/messages.js index f4b52eab77..306cdd2621 100644 --- a/api/server/routes/messages.js +++ b/api/server/routes/messages.js @@ -101,11 +101,12 @@ router.get('/', async (req, res) => { let scopedMessageRead; if (typeof conversationId === 'string') { const ownershipRead = db.getConvoOwnership(user, conversationId); + /** Client-facing reads never expose server-private fields such as `contextMeta`. */ const messageRead = messageId - ? db.getMessages({ conversationId, messageId, user }) + ? db.getMessages({ conversationId, messageId, user }, CLIENT_MESSAGE_SELECT) : db.getMessagesByCursor( { conversationId, user }, - { sortField, sortOrder, limit: pageSize, cursor }, + { sortField, sortOrder, limit: pageSize, cursor, select: CLIENT_MESSAGE_SELECT }, ); scopedMessageRead = Promise.resolve(messageRead).then( (value) => ({ ok: true, value }), @@ -161,10 +162,13 @@ router.get('/', async (req, res) => { } } - const dbMessages = await db.getMessages({ - user, - messageId: { $in: messageIds }, - }); + const dbMessages = await db.getMessages( + { + user, + messageId: { $in: messageIds }, + }, + CLIENT_MESSAGE_SELECT, + ); const dbMessageMap = {}; for (const dbMessage of dbMessages) { @@ -175,9 +179,12 @@ router.get('/', async (req, res) => { for (const message of cleanedMessages) { const convo = result.convoMap[message.conversationId]; const dbMessage = dbMessageMap[message.messageId]; + /** Search hydrates every schema field; server-private state never leaves. */ + const publicHit = { ...message }; + delete publicHit.contextMeta; activeMessages.push({ - ...message, + ...publicHit, title: convo.title, conversationId: message.conversationId, model: convo.model, @@ -209,6 +216,18 @@ router.get('/', async (req, res) => { * @param {string} req.body.agentId - The agentId to filter content by * @returns {TMessage} The newly created branch message */ +/** + * Projects a saved row for a client response. The context meta stays in the + * database for the next turn; no client-facing read or write response carries it. + * @param {TMessage} message + * @returns {TMessage} + */ +function toClientMessage(message) { + const clientMessage = { ...message }; + delete clientMessage.contextMeta; + return clientMessage; +} + router.post('/branch', configMiddleware, async (req, res) => { try { const { messageId, agentId } = req.body; @@ -298,6 +317,9 @@ router.post('/branch', configMiddleware, async (req, res) => { endpoint: sourceMessage.endpoint, sender: sourceMessage.sender, iconURL: sourceMessage.iconURL, + // Server-private context meta (calibration and fading tier) travels with the branch so + // the next turn seeds its pruner the same way it would from the source response. + ...(sourceMessage.contextMeta != null && { contextMeta: sourceMessage.contextMeta }), ...(typeof sourceMessage.isUserSubmitted === 'boolean' && { isUserSubmitted: sourceMessage.isUserSubmitted, }), @@ -339,7 +361,7 @@ router.post('/branch', configMiddleware, async (req, res) => { return res.status(500).json({ error: 'Failed to save branch message' }); } - res.status(201).json(savedMessage); + res.status(201).json(toClientMessage(savedMessage)); } catch (error) { if (isContentFilterError(error)) { return res.status(error.statusCode).json(error.body); @@ -491,6 +513,9 @@ router.post('/:conversationId', storedMessageMutationMiddleware, async (req, res delete message.isUserSubmitted; delete message.userSubmittedPaths; delete message.userSubmittedMessageFieldPaths; + /** Server-private run state: a client-authored row must never seed a run's + * calibration or fading tiers, so the field only ever comes from the server. */ + delete message.contextMeta; const reqCtx = { userId: req?.user?.id, isTemporary: req?.body?.isTemporary, @@ -514,7 +539,7 @@ router.post('/:conversationId', storedMessageMutationMiddleware, async (req, res context: 'POST /api/messages/:conversationId', ...(savedMessage._id != null ? { appendMessageIds: [savedMessage._id] } : {}), }); - res.status(201).json(savedMessage); + res.status(201).json(toClientMessage(savedMessage)); } catch (error) { logger.error('Error saving message:', error); res.status(500).json({ error: 'Internal server error' }); diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js index e5c3cad454..764d1f0277 100644 --- a/api/server/services/Endpoints/agents/initialize.js +++ b/api/server/services/Endpoints/agents/initialize.js @@ -495,7 +495,8 @@ const initializeClient = async ({ /** Latest visible context snapshot + every emitted usage payload for this * response, captured by the handlers and persisted on the response message's * metadata so the breakdown and branch/total cost survive a reload. - * @type {{ latest: import('librechat-data-provider').TContextUsageEvent | null, count: number }} */ + * `onSnapshot` is installed by the client to publish run state after each snapshot. + * @type {{ latest: import('librechat-data-provider').TContextUsageEvent | null, count: number, onSnapshot?: () => void }} */ const contextUsageSink = { latest: null, count: 0 }; /** @type {Array} */ const usageEmitSink = []; diff --git a/api/server/utils/import/importers.js b/api/server/utils/import/importers.js index ec63649959..181796c339 100644 --- a/api/server/utils/import/importers.js +++ b/api/server/utils/import/importers.js @@ -59,8 +59,11 @@ function sanitizeImportedMessage(message) { const text = castPersistedImportedText(message.text); const content = normalizeImportedArray(message.content); const attachments = normalizeImportedArray(message.attachments); + const importable = { ...message }; + /** Server-private run state never comes from an import. */ + delete importable.contextMeta; return { - ...message, + ...importable, isUserSubmitted: true, ...(text !== message.text && { text }), ...(sanitizeTextMarkers && diff --git a/api/server/utils/import/importers.spec.js b/api/server/utils/import/importers.spec.js index 3fcc529977..0554c4a0c8 100644 --- a/api/server/utils/import/importers.spec.js +++ b/api/server/utils/import/importers.spec.js @@ -966,6 +966,33 @@ describe('importLibreChatConvo', () => { ]); }); + it('drops server-private context meta from imported messages', async () => { + const message = { + messageId: 'message-1', + parentMessageId: Constants.NO_PARENT, + text: 'Imported response', + isCreatedByUser: false, + contextMeta: { + calibrationRatio: 1, + encoding: 'claude', + fading: { v: 1, budgetTokens: 1, masked: true }, + }, + }; + const jsonData = { + conversationId: 'context-meta-import', + title: 'Context meta import', + recursive: false, + messages: [message], + }; + const importBatchBuilder = new ImportBatchBuilder('user-123'); + + const importer = getImporter(jsonData); + await importer(jsonData, 'user-123', () => importBatchBuilder); + + expect(importBatchBuilder.messages[0]).not.toHaveProperty('contextMeta'); + expect(importBatchBuilder.messages[0].isUserSubmitted).toBe(true); + }); + it('sanitizes singleton content and attachment fields before Mongoose array casting', async () => { const message = { messageId: 'message-1', diff --git a/package-lock.json b/package-lock.json index 3701b90dd2..5461761a97 100644 --- a/package-lock.json +++ b/package-lock.json @@ -63,7 +63,7 @@ "@azure/storage-blob": "^12.30.0", "@google/genai": "^2.8.0", "@keyv/redis": "5.1.6", - "@librechat/agents": "^3.7.15", + "@librechat/agents": "^3.7.16", "@librechat/api": "*", "@librechat/data-schemas": "*", "@microsoft/microsoft-graph-client": "^3.0.7", @@ -10633,9 +10633,9 @@ } }, "node_modules/@librechat/agents": { - "version": "3.7.15", - "resolved": "https://registry.npmjs.org/@librechat/agents/-/agents-3.7.15.tgz", - "integrity": "sha512-4huKqYzZ7mOav0NdtzId0adT4YE00bRfx+FaPtfzAeZxUryKXwRPmq8/La77tH7Ljf+AnC00pdJHUcPezuBoHQ==", + "version": "3.7.16", + "resolved": "https://registry.npmjs.org/@librechat/agents/-/agents-3.7.16.tgz", + "integrity": "sha512-mhtvBeBQ6/LR03QQuQdatTd3nPOYzZpxXtxOo75CzkwqdL4m7Gg+5rPhRjV9rOD8ZZqUoXJfWh1pNPEz9XQjeg==", "license": "MIT", "dependencies": { "@anthropic-ai/sdk": "^0.115.0", @@ -42639,7 +42639,7 @@ "@azure/storage-blob": "^12.30.0", "@google/genai": "^2.8.0", "@keyv/redis": "5.1.6", - "@librechat/agents": "^3.7.15", + "@librechat/agents": "^3.7.16", "@librechat/data-schemas": "*", "@modelcontextprotocol/sdk": "^1.30.0", "@opentelemetry/api": "^1.9.0", diff --git a/packages/api/package.json b/packages/api/package.json index 30a7024c92..8ccef59653 100644 --- a/packages/api/package.json +++ b/packages/api/package.json @@ -113,7 +113,7 @@ "@azure/storage-blob": "^12.30.0", "@google/genai": "^2.8.0", "@keyv/redis": "5.1.6", - "@librechat/agents": "^3.7.15", + "@librechat/agents": "^3.7.16", "@librechat/data-schemas": "*", "@modelcontextprotocol/sdk": "^1.30.0", "@opentelemetry/api": "^1.9.0", diff --git a/packages/api/src/agents/fading.spec.ts b/packages/api/src/agents/fading.spec.ts new file mode 100644 index 0000000000..37bb283411 --- /dev/null +++ b/packages/api/src/agents/fading.spec.ts @@ -0,0 +1,259 @@ +import { AIMessage, HumanMessage, ToolMessage } from '@langchain/core/messages'; +import { createPruneMessages, resolveFadingCaps, seedFadingTier } from '@librechat/agents'; +import type { IAgentFadingTier } from '@librechat/data-schemas'; +import type { BaseMessage } from '@langchain/core/messages'; +import { + isAgentFadingTier, + resolvePersistableFadingTier, + resolvePersistableFadingTiers, + resolveRunContextMeta, + resolveRunFadingTiers, +} from './fading'; + +describe('isAgentFadingTier', () => { + it('accepts a well-formed tier and rejects everything else', () => { + expect(isAgentFadingTier({ v: 1, budgetTokens: 20_000, masked: true })).toBe(true); + expect(isAgentFadingTier({ v: 2, budgetTokens: 20_000, masked: true })).toBe(false); + expect(isAgentFadingTier({ v: 1, budgetTokens: 0, masked: true })).toBe(false); + expect(isAgentFadingTier({ v: 1, budgetTokens: 20_000, masked: 'yes' })).toBe(false); + expect(isAgentFadingTier({ v: 1, budgetTokens: Number.NaN, masked: false })).toBe(false); + expect(isAgentFadingTier(null)).toBe(false); + expect(isAgentFadingTier(undefined)).toBe(false); + }); +}); + +describe('resolvePersistableFadingTier', () => { + it('strips a valid tier to its fields and drops invalid input', () => { + expect( + resolvePersistableFadingTier({ v: 1, budgetTokens: 50_000, masked: true, extra: 1 }), + ).toEqual({ v: 1, budgetTokens: 50_000, masked: true }); + expect(resolvePersistableFadingTier(undefined)).toBeUndefined(); + expect(resolvePersistableFadingTier({ v: 1, budgetTokens: -1, masked: true })).toBeUndefined(); + }); +}); + +describe('resolvePersistableFadingTiers', () => { + it('keeps only valid tiers, stripped to the compact shape, from own keys', () => { + const snapshot = Object.fromEntries([ + ['agent-a', { v: 1, budgetTokens: 20_000, masked: true, latched: true }], + ['agent-b', { v: 1, budgetTokens: 0, masked: false }], + ['__proto__', { v: 1, budgetTokens: 10_000, masked: false, latched: true }], + ]); + expect(resolvePersistableFadingTiers(snapshot)).toEqual([ + { agentId: 'agent-a', v: 1, budgetTokens: 20_000, masked: true }, + { agentId: '__proto__', v: 1, budgetTokens: 10_000, masked: false }, + ]); + }); + + it('ignores inherited keys and yields nothing for an empty or invalid snapshot', () => { + const inherited = Object.create({ ghost: { v: 1, budgetTokens: 20_000, masked: true } }); + expect(resolvePersistableFadingTiers(inherited)).toBeUndefined(); + expect(resolvePersistableFadingTiers({})).toBeUndefined(); + expect(resolvePersistableFadingTiers(null)).toBeUndefined(); + expect(resolvePersistableFadingTiers('agent-a')).toBeUndefined(); + }); +}); + +describe('resolveRunFadingTiers', () => { + it('rebuilds a prototype-safe record from persisted entries', () => { + const tiers = resolveRunFadingTiers([ + { agentId: 'agent-a', v: 1, budgetTokens: 20_000, masked: true }, + { agentId: '__proto__', v: 1, budgetTokens: 10_000, masked: false }, + ]); + + expect(tiers).toBeDefined(); + expect(Object.getPrototypeOf(tiers)).toBeNull(); + expect(Object.keys(tiers ?? {})).toEqual(['agent-a', '__proto__']); + expect(Object.prototype.hasOwnProperty.call(tiers, '__proto__')).toBe(true); + expect(tiers?.['agent-a']).toEqual({ v: 1, budgetTokens: 20_000, masked: true }); + expect('budgetTokens' in {}).toBe(false); + }); + + it('rejects duplicate, malformed, or empty entry lists', () => { + expect( + resolveRunFadingTiers([ + { agentId: 'agent-a', v: 1, budgetTokens: 20_000, masked: true }, + { agentId: 'agent-a', v: 1, budgetTokens: 10_000, masked: true }, + ]), + ).toBeUndefined(); + expect( + resolveRunFadingTiers([{ agentId: '', v: 1, budgetTokens: 1, masked: true }]), + ).toBeUndefined(); + expect(resolveRunFadingTiers([])).toBeUndefined(); + expect(resolveRunFadingTiers(undefined)).toBeUndefined(); + }); + + it('round-trips a run snapshot through the persisted entries unchanged', () => { + const snapshot = { + 'agent-a': { v: 1 as const, budgetTokens: 20_000, masked: true, latched: true as const }, + 'agent-b': { v: 1 as const, budgetTokens: 5_000, masked: false, latched: true as const }, + }; + const entries = resolvePersistableFadingTiers(snapshot); + const restored = resolveRunFadingTiers(entries); + expect(restored).toEqual({ + 'agent-a': { v: 1, budgetTokens: 20_000, masked: true }, + 'agent-b': { v: 1, budgetTokens: 5_000, masked: false }, + }); + }); +}); + +describe('resolveRunContextMeta', () => { + const fading = { v: 1, budgetTokens: 20_000, masked: true }; + const getEncoding = jest.fn(() => 'claude'); + + beforeEach(() => getEncoding.mockClear()); + + it('persists only compact tier and calibration metadata', () => { + const meta = resolveRunContextMeta({ + calibrationRatio: 1.25, + fadingTier: { ...fading, latched: true, messages: ['not persisted'] }, + fadingTiers: { + 'agent-a': { ...fading, latched: true, projection: { truncated: true } }, + 'agent-b': { v: 1, budgetTokens: 5_000, masked: false, latched: true }, + }, + getEncoding, + }); + + expect(meta).toEqual({ + calibrationRatio: 1.25, + encoding: 'claude', + fading, + fadingTiers: [ + { agentId: 'agent-a', v: 1, budgetTokens: 20_000, masked: true }, + { agentId: 'agent-b', v: 1, budgetTokens: 5_000, masked: false }, + ], + }); + expect(Object.keys(meta ?? {})).toEqual([ + 'calibrationRatio', + 'encoding', + 'fading', + 'fadingTiers', + ]); + }); + + it('persists per-agent tiers alone when only a non-default agent latched one', () => { + expect( + resolveRunContextMeta({ + calibrationRatio: 1, + fadingTier: undefined, + fadingTiers: { 'agent-b': { v: 1, budgetTokens: 5_000, masked: true } }, + getEncoding, + }), + ).toEqual({ + calibrationRatio: 1, + encoding: 'claude', + fadingTiers: [{ agentId: 'agent-b', v: 1, budgetTokens: 5_000, masked: true }], + }); + }); + + it('persists a latched fading tier even at a neutral calibration ratio', () => { + expect(resolveRunContextMeta({ calibrationRatio: 1, fadingTier: fading, getEncoding })).toEqual( + { calibrationRatio: 1, encoding: 'claude', fading }, + ); + }); + + it('persists calibration alone when the run exposes no tier', () => { + expect( + resolveRunContextMeta({ calibrationRatio: 1.2345, fadingTier: undefined, getEncoding }), + ).toEqual({ calibrationRatio: 1.235, encoding: 'claude' }); + }); + + it('persists nothing, without resolving the encoding, when there is nothing to keep', () => { + expect( + resolveRunContextMeta({ calibrationRatio: 0, fadingTier: { v: 1 }, getEncoding }), + ).toBeUndefined(); + expect(getEncoding).not.toHaveBeenCalled(); + }); +}); + +describe('persisted tier round trip through the SDK pruner', () => { + const window = 20_000; + const tokenCounter = (message: BaseMessage): number => + Math.ceil(JSON.stringify(message.content).length / 4); + const countMap = (messages: BaseMessage[]): Record => + Object.fromEntries(messages.map((message, index) => [String(index), tokenCounter(message)])); + const serializeToolExchange = (context: BaseMessage[]): string => + JSON.stringify( + context + .filter((message) => message._getType() === 'tool' || message._getType() === 'ai') + .map((message) => [message._getType(), message.content, message.additional_kwargs]), + ); + /** Fresh instances per run: the pruner rewrites the messages it is handed, as the Graph's per-Run projection clone allows. */ + const makeTurnOne = (): BaseMessage[] => [ + new HumanMessage('Fetch the alpha report and summarize it.'), + new AIMessage({ + content: '', + tool_calls: [{ id: 'call-1', name: 'fetch_report', args: { name: 'alpha' } }], + }), + new ToolMessage({ + tool_call_id: 'call-1', + name: 'fetch_report', + content: Array.from( + { length: 4_000 }, + (_, i) => `row ${i}: value ${(i * 7919) % 10007}`, + ).join('\n'), + }), + new AIMessage('The alpha report lists 4000 rows.'), + ]; + const prune = ( + messages: BaseMessage[], + overrides: { + calibrationRatio: number; + instructionTokens: number; + fadingTier?: IAgentFadingTier; + }, + ) => + createPruneMessages({ + startIndex: messages.length, + tokenCounter, + maxTokens: window, + indexTokenCountMap: countMap(messages), + summarizationEnabled: true, + calibrationRatio: overrides.calibrationRatio, + getInstructionTokens: () => overrides.instructionTokens, + ...(overrides.fadingTier == null ? {} : { fadingTier: overrides.fadingTier }), + })({ messages }); + + it("reproduces the first run's projection bytes when the next run is seeded from contextMeta", () => { + const turnOne = makeTurnOne(); + const first = prune(turnOne, { calibrationRatio: 1, instructionTokens: 12_000 }); + expect(first.fadingTier.budgetTokens).toBeLessThan(window); + + const meta = resolveRunContextMeta({ + calibrationRatio: first.calibrationRatio ?? 1, + fadingTier: first.fadingTier, + getEncoding: () => 'claude', + }); + expect(meta?.fading).toEqual({ + v: 1, + budgetTokens: first.fadingTier.budgetTokens, + masked: first.fadingTier.masked, + }); + expect(resolveFadingCaps(seedFadingTier(window, meta?.fading))).toEqual( + resolveFadingCaps(first.fadingTier), + ); + + const makeTurnTwo = (): BaseMessage[] => [ + ...makeTurnOne(), + new HumanMessage('Thanks. Anything else?'), + new AIMessage('No.'), + new HumanMessage('Say ok.'), + ]; + const firstBytes = serializeToolExchange(first.context); + const drifted = { calibrationRatio: 1, instructionTokens: 9_000 }; + const seeded = prune(makeTurnTwo(), { ...drifted, fadingTier: meta?.fading }); + expect(serializeToolExchange(seeded.context.slice(0, turnOne.length))).toBe(firstBytes); + expect(seeded.fadingTier.budgetTokens).toBe(first.fadingTier.budgetTokens); + + const recalibrated = prune(makeTurnTwo(), { + calibrationRatio: 1.2, + instructionTokens: 9_000, + fadingTier: meta?.fading, + }); + expect(serializeToolExchange(recalibrated.context.slice(0, turnOne.length))).toBe(firstBytes); + + const unseeded = prune(makeTurnTwo(), drifted); + expect(unseeded.fadingTier.budgetTokens).toBeGreaterThan(first.fadingTier.budgetTokens); + expect(serializeToolExchange(unseeded.context.slice(0, turnOne.length))).not.toBe(firstBytes); + }); +}); diff --git a/packages/api/src/agents/fading.ts b/packages/api/src/agents/fading.ts new file mode 100644 index 0000000000..aa83607e4d --- /dev/null +++ b/packages/api/src/agents/fading.ts @@ -0,0 +1,100 @@ +import { + isAgentFadingTier, + isAgentFadingTierEntry, + isAgentFadingTierEntries, + AGENT_FADING_TIER_VERSION, +} from '@librechat/data-schemas'; +import type { + IAgentEventActorContextMeta, + IAgentFadingTierEntry, + IAgentFadingTier, +} from '@librechat/data-schemas'; + +export { isAgentFadingTier, isAgentFadingTierEntries }; + +/** Latched tiers keyed by agent ID, the shape `RunConfig.fadingTiers` takes. */ +export type RunFadingTiers = Record; + +/** + * Normalizes the tier a run exposes for persistence, or undefined when there + * is none. `Run.getFadingTier()` already returns only tiers that carry + * information (masking active, a budget below the pruner's window, or a tier + * restored from host state), so the host validates the shape and strips + * anything else the SDK may attach, such as its `latched` provenance flag. + */ +export function resolvePersistableFadingTier(tier: unknown): IAgentFadingTier | undefined { + if (!isAgentFadingTier(tier)) { + return undefined; + } + return { v: AGENT_FADING_TIER_VERSION, budgetTokens: tier.budgetTokens, masked: tier.masked }; +} + +/** + * Normalizes the per-agent tiers a run exposes (`Run.getFadingTiers()`) into + * persisted entries. Only own enumerable keys are read and every tier is + * validated and stripped, so an agent ID never reaches storage as a field name + * and nothing beyond the compact tier is persisted. + */ +export function resolvePersistableFadingTiers(tiers: unknown): IAgentFadingTierEntry[] | undefined { + if (typeof tiers !== 'object' || tiers === null) { + return undefined; + } + const candidates: Array<[string, unknown]> = Object.entries(tiers); + const entries = candidates.flatMap(([agentId, candidate]) => { + const tier = resolvePersistableFadingTier(candidate); + const entry = tier == null ? undefined : { agentId, ...tier }; + return isAgentFadingTierEntry(entry) ? [entry] : []; + }); + return entries.length > 0 ? entries : undefined; +} + +/** + * Rebuilds `RunConfig.fadingTiers` from persisted entries on a null-prototype + * record, so an agent ID such as `__proto__` stays an own key and can never + * touch the prototype chain. + */ +export function resolveRunFadingTiers(entries: unknown): RunFadingTiers | undefined { + if (!isAgentFadingTierEntries(entries) || entries.length === 0) { + return undefined; + } + const tiers: RunFadingTiers = Object.create(null); + for (const { agentId, v, budgetTokens, masked } of entries) { + tiers[agentId] = { v, budgetTokens, masked }; + } + return tiers; +} + +export type RunContextMetaParams = { + calibrationRatio: number; + fadingTier: unknown; + /** Per-agent tiers from `Run.getFadingTiers()` or the live graph. */ + fadingTiers?: unknown; + /** Resolved lazily: only consulted when there is something to persist. */ + getEncoding: () => string; +}; + +/** + * Builds the compact `contextMeta` a response message persists from a run, or + * undefined when neither calibration nor fading carries information. Only the + * calibration ratio, its encoding and the latched tiers are kept; message + * content, canonical tool results and projection state never are. A latched + * tier is persisted even at a neutral calibration ratio, with the ratio + * recorded as 1 so the stored shape stays valid. + */ +export function resolveRunContextMeta( + params: RunContextMetaParams, +): IAgentEventActorContextMeta | undefined { + const { calibrationRatio } = params; + const fading = resolvePersistableFadingTier(params.fadingTier); + const fadingTiers = resolvePersistableFadingTiers(params.fadingTiers); + const calibrated = calibrationRatio > 0 && calibrationRatio !== 1; + if (!calibrated && fading == null && fadingTiers == null) { + return undefined; + } + return { + calibrationRatio: calibrationRatio > 0 ? Math.round(calibrationRatio * 1000) / 1000 : 1, + encoding: params.getEncoding(), + ...(fading == null ? {} : { fading }), + ...(fadingTiers == null ? {} : { fadingTiers }), + }; +} diff --git a/packages/api/src/agents/index.ts b/packages/api/src/agents/index.ts index 81326fb658..dd414940c1 100644 --- a/packages/api/src/agents/index.ts +++ b/packages/api/src/agents/index.ts @@ -52,6 +52,8 @@ export * from './skillConfigurable'; export * from './skillFiles'; export * from './codeFilesSession'; export * from './run'; +export * from './fading'; +export * from './publication'; export * from './runtime'; export * from './testHook'; export * from './tools'; diff --git a/packages/api/src/agents/publication.spec.ts b/packages/api/src/agents/publication.spec.ts new file mode 100644 index 0000000000..1b990071d6 --- /dev/null +++ b/packages/api/src/agents/publication.spec.ts @@ -0,0 +1,201 @@ +import type { IAgentEventActorContextMeta } from '@librechat/data-schemas'; +import { createContextMetaPublisher, selectRunContextMetaToPublish } from './publication'; + +const tier = (budgetTokens: number): IAgentEventActorContextMeta => ({ + calibrationRatio: 1.2, + encoding: 'claude', + fading: { v: 1, budgetTokens, masked: true }, +}); + +type Deferred = { promise: Promise; resolve: () => void; reject: (error: Error) => void }; +const deferred = (): Deferred => { + let resolve!: () => void; + let reject!: (error: Error) => void; + const promise = new Promise((res, rej) => { + resolve = res; + reject = rej; + }); + return { promise, resolve, reject }; +}; + +const noDelay = async (): Promise => undefined; + +describe('createContextMetaPublisher', () => { + it('shares one write between equal callers and skips a repeat of the latest record', async () => { + const write = jest.fn(async () => undefined); + const publisher = createContextMetaPublisher({ write, delay: noDelay }); + + const first = publisher.publish(tier(50_000)); + const second = publisher.publish(tier(50_000)); + expect(second).toBe(first); + await Promise.all([first, second]); + await publisher.publish(tier(50_000)); + + expect(write).toHaveBeenCalledTimes(1); + expect(publisher.hasPublished).toBe(true); + }); + + it('issues distinct records in order so the newest settles last', async () => { + const gate = deferred(); + const writes: number[] = []; + const write = jest.fn(async (contextMeta: IAgentEventActorContextMeta) => { + writes.push(contextMeta.fading?.budgetTokens ?? 0); + if (writes.length === 1) { + await gate.promise; + } + }); + const publisher = createContextMetaPublisher({ write, delay: noDelay }); + + const older = publisher.publish(tier(50_000)); + const newer = publisher.publish(tier(25_000)); + await Promise.resolve(); + expect(writes).toEqual([50_000]); + + gate.resolve(); + await Promise.all([older, newer]); + expect(writes).toEqual([50_000, 25_000]); + }); + + it('retries a failed write before reporting it', async () => { + const write = jest + .fn, [IAgentEventActorContextMeta]>() + .mockRejectedValueOnce(new Error('transient')) + .mockResolvedValueOnce(undefined); + const onFailure = jest.fn(); + const delays: number[] = []; + const publisher = createContextMetaPublisher({ + write, + onFailure, + retryDelayMs: 10, + delay: async (ms) => { + delays.push(ms); + }, + }); + + await publisher.publish(tier(50_000)); + + expect(write).toHaveBeenCalledTimes(2); + expect(delays).toEqual([10]); + expect(onFailure).not.toHaveBeenCalled(); + expect(publisher.hasPublished).toBe(true); + }); + + it('reports an exhausted publication once, never rejects, and writes again next time', async () => { + const write = jest + .fn, [IAgentEventActorContextMeta]>() + .mockRejectedValue(new Error('down')); + const onFailure = jest.fn(); + const publisher = createContextMetaPublisher({ + write, + onFailure, + attempts: 2, + delay: noDelay, + }); + + await expect(publisher.publish(tier(50_000))).resolves.toBeUndefined(); + expect(write).toHaveBeenCalledTimes(2); + expect(onFailure).toHaveBeenCalledTimes(1); + expect(publisher.hasPublished).toBe(false); + + write.mockResolvedValue(undefined); + await publisher.publish(tier(50_000)); + expect(write).toHaveBeenCalledTimes(3); + expect(publisher.hasPublished).toBe(true); + }); + + it('still counts an earlier committed record after a later publication is exhausted', async () => { + const write = jest + .fn, [IAgentEventActorContextMeta]>() + .mockResolvedValueOnce(undefined) + .mockRejectedValue(new Error('down')); + const onFailure = jest.fn(); + const publisher = createContextMetaPublisher({ + write, + onFailure, + attempts: 2, + delay: noDelay, + }); + const neutral = { calibrationRatio: 1, encoding: 'claude' }; + + await publisher.publish(tier(50_000)); + await publisher.publish(tier(25_000)); + expect(onFailure).toHaveBeenCalledTimes(1); + expect(publisher.hasPublished).toBe(true); + expect( + selectRunContextMetaToPublish({ + live: true, + captured: undefined, + inherited: undefined, + hasPublished: publisher.hasPublished, + getEncoding: () => 'claude', + }), + ).toEqual(neutral); + + write.mockResolvedValue(undefined); + await publisher.publish(neutral); + expect(write).toHaveBeenLastCalledWith(neutral); + }); + + it('keeps a newer publication when an older one fails after it was superseded', async () => { + const olderGate = deferred(); + const write = jest.fn(async (contextMeta: IAgentEventActorContextMeta) => { + if (contextMeta.fading?.budgetTokens === 50_000) { + await olderGate.promise; + } + }); + const publisher = createContextMetaPublisher({ write, attempts: 1, delay: noDelay }); + + const older = publisher.publish(tier(50_000)); + const newer = publisher.publish(tier(25_000)); + olderGate.reject(new Error('lost')); + await Promise.all([older, newer]); + + expect(publisher.hasPublished).toBe(true); + expect(write).toHaveBeenLastCalledWith(tier(25_000)); + }); +}); + +describe('selectRunContextMetaToPublish', () => { + const getEncoding = () => 'claude'; + const inherited = tier(30_000); + const captured = tier(20_000); + + it('prefers the captured state and falls back to the inherited seed before the run', () => { + expect( + selectRunContextMetaToPublish({ + live: false, + captured, + inherited, + hasPublished: false, + getEncoding, + }), + ).toBe(captured); + expect( + selectRunContextMetaToPublish({ + live: false, + captured: undefined, + inherited, + hasPublished: false, + getEncoding, + }), + ).toBe(inherited); + }); + + it('publishes a neutral record only for a live snapshot after an earlier publication', () => { + const neutral = { live: true, captured: undefined, inherited, getEncoding }; + expect(selectRunContextMetaToPublish({ ...neutral, hasPublished: false })).toBeUndefined(); + expect(selectRunContextMetaToPublish({ ...neutral, hasPublished: true })).toEqual({ + calibrationRatio: 1, + encoding: 'claude', + }); + expect( + selectRunContextMetaToPublish({ + live: false, + captured: undefined, + inherited: undefined, + hasPublished: true, + getEncoding, + }), + ).toBeUndefined(); + }); +}); diff --git a/packages/api/src/agents/publication.ts b/packages/api/src/agents/publication.ts new file mode 100644 index 0000000000..d8600a0d0d --- /dev/null +++ b/packages/api/src/agents/publication.ts @@ -0,0 +1,141 @@ +import type { IAgentEventActorContextMeta } from '@librechat/data-schemas'; + +export type ContextMetaWriter = (contextMeta: IAgentEventActorContextMeta) => Promise; + +export interface ContextMetaPublisherOptions { + /** Durable write of one record, e.g. the job metadata writer. */ + write: ContextMetaWriter; + /** Called once per publication whose every attempt failed. */ + onFailure?: (error: unknown, contextMeta: IAgentEventActorContextMeta) => void; + /** Write attempts per publication, including the first. */ + attempts?: number; + /** Backoff before the second attempt, doubled for each later one. */ + retryDelayMs?: number; + delay?: (ms: number) => Promise; +} + +export interface ContextMetaPublisher { + /** + * Publishes a record unless it equals the latest publication, in which case + * the caller shares that publication's promise. Distinct records are written + * in call order, each after the previous publication settles, so the store's + * last-writer-wins keeps the newest snapshot. Never rejects. + */ + publish(contextMeta: IAgentEventActorContextMeta): Promise; + /** Whether the job may carry a record: a write has committed, or one is in flight. */ + readonly hasPublished: boolean; +} + +type Publication = { + serialized: string; + promise: Promise; +}; + +const DEFAULT_ATTEMPTS = 3; +const DEFAULT_RETRY_DELAY_MS = 50; + +const sleep = (ms: number): Promise => new Promise((resolve) => setTimeout(resolve, ms)); + +async function writeWithRetry( + contextMeta: IAgentEventActorContextMeta, + options: Required< + Pick + >, +): Promise { + let lastError: unknown; + for (let attempt = 0; attempt < options.attempts; attempt += 1) { + try { + await options.write(contextMeta); + return; + } catch (error) { + lastError = error; + if (attempt + 1 < options.attempts) { + await options.delay(options.retryDelayMs * 2 ** attempt); + } + } + } + throw lastError; +} + +/** + * Coordinates the durable publication of a run's compact context meta onto its + * job: one write per distinct record, shared by equal concurrent callers, + * ordered so the newest record wins, retried on transient failure. A caller + * that awaits `publish` before its model call therefore knows the job carries + * the record that describes that call. An exhausted publication is reported + * through `onFailure` and forgotten, so the next snapshot writes again, while + * a record committed earlier still counts as published so that a later neutral + * snapshot overwrites it instead of leaving it on the job. + */ +export function createContextMetaPublisher( + options: ContextMetaPublisherOptions, +): ContextMetaPublisher { + const { + write, + onFailure, + attempts = DEFAULT_ATTEMPTS, + retryDelayMs = DEFAULT_RETRY_DELAY_MS, + delay = sleep, + } = options; + let latest: Publication | undefined; + let committed = false; + return { + get hasPublished() { + return committed || latest != null; + }, + publish(contextMeta) { + const serialized = JSON.stringify(contextMeta); + if (latest?.serialized === serialized) { + return latest.promise; + } + const previous = latest?.promise ?? Promise.resolve(); + const publication: Publication = { serialized, promise: Promise.resolve() }; + publication.promise = previous + .then(() => writeWithRetry(contextMeta, { write, attempts, retryDelayMs, delay })) + .then(() => { + committed = true; + }) + .catch((error: unknown) => { + if (latest === publication) { + latest = undefined; + } + onFailure?.(error, contextMeta); + }); + latest = publication; + return publication.promise; + }, + }; +} + +export type RunContextMetaSelection = { + /** True for a snapshot of the running graph; false for the pre-run seed publish. */ + live: boolean; + /** The run's own compact state, when it has any. */ + captured: IAgentEventActorContextMeta | undefined; + /** The state inherited from the parent response. */ + inherited: IAgentEventActorContextMeta | undefined; + hasPublished: boolean; + getEncoding: () => string; +}; + +/** + * Chooses what a publication should carry: the run's own state when it has + * any, the inherited seed before the run exists, and otherwise a neutral + * record (ratio 1, no tier) once something has been published, since a + * running job's fields cannot be deleted through the metadata writer and a + * neutral record seeds the next turn exactly as no record would. A fresh + * conversation's first neutral snapshot publishes nothing. + */ +export function selectRunContextMetaToPublish( + selection: RunContextMetaSelection, +): IAgentEventActorContextMeta | undefined { + const { live, captured, inherited, hasPublished } = selection; + const contextMeta = captured ?? (live ? undefined : inherited); + if (contextMeta != null) { + return contextMeta; + } + if (!live || !hasPublished) { + return undefined; + } + return { calibrationRatio: 1, encoding: selection.getEncoding() }; +} diff --git a/packages/api/src/agents/run.ts b/packages/api/src/agents/run.ts index 40fa88ed18..f015da0573 100644 --- a/packages/api/src/agents/run.ts +++ b/packages/api/src/agents/run.ts @@ -41,10 +41,10 @@ import type { ReasoningResponseKey, SummarizationConfig, } from 'librechat-data-provider'; +import type { AppConfig, IAgentFadingTier, IUser } from '@librechat/data-schemas'; import type { CallbackHandlerMethods } from '@langchain/core/callbacks/base'; import type { BaseMessage } from '@librechat/agents/langchain/messages'; import type { Callbacks } from '@langchain/core/callbacks/manager'; -import type { AppConfig, IUser } from '@librechat/data-schemas'; import type { ModelBoundChatModelCallback } from '~/middleware/modelBoundContent'; import type { ToolInputValidationError } from '~/agents/toolValidation'; import type { ResolvedToolApprovalHook } from '~/agents/hitl/hooks'; @@ -53,6 +53,7 @@ import type { ResolvedAlwaysApplySkill } from '~/agents/skills'; import type { CodeExecutionContext } from '~/agents/execution'; import type { MCPToolAlias } from '~/tools/classification'; import type { SubagentUsageEvent } from '~/agents/usage'; +import type { RunFadingTiers } from './fading'; import type * as t from '~/types'; import { assertAttachedCodeEnvironmentApprovalSupported, @@ -1416,6 +1417,8 @@ export async function createRun({ initialSummary, modelCallbacks, calibrationRatio, + fadingTier, + fadingTiers, appConfig, subagentUsageSink, subagentTasks, @@ -1467,6 +1470,20 @@ export async function createRun({ modelCallbacks?: readonly ModelBoundChatModelCallback[]; /** Calibration ratio from previous run's contextMeta, seeds the pruner EMA */ calibrationRatio?: number; + /** + * Default agent's latched context-fading tier from the previous run's + * contextMeta. It seeds the pruner so the provider-only projection of + * historical tool results keeps the same bytes across runs; graph messages + * stay canonical. Ships in `@librechat/agents` after 3.7.13; older SDK + * versions ignore it. + */ + fadingTier?: IAgentFadingTier | null; + /** + * Latched tiers keyed by agent ID from the previous run's contextMeta, so + * every agent of a multi-agent run restores its own tier. Same SDK + * availability as `fadingTier`. + */ + fadingTiers?: RunFadingTiers | null; /** * Resolved app config. Used to translate custom-endpoint provider names * (e.g. "Ollama") in the summarization config to SDK-recognized providers. @@ -2108,6 +2125,8 @@ export async function createRun({ customHandlers, initialSessions, calibrationRatio, + fadingTier, + fadingTiers, indexTokenCountMap, subagentUsageSink, subagentTasks, diff --git a/packages/api/src/stream/ApprovalLifecycle.ts b/packages/api/src/stream/ApprovalLifecycle.ts index c44d381aa6..e95aefde01 100644 --- a/packages/api/src/stream/ApprovalLifecycle.ts +++ b/packages/api/src/stream/ApprovalLifecycle.ts @@ -1,5 +1,8 @@ import { logger } from '@librechat/data-schemas'; -import type { ICompactionSemanticIndexProjection } from '@librechat/data-schemas'; +import type { + IAgentEventActorContextMeta, + ICompactionSemanticIndexProjection, +} from '@librechat/data-schemas'; import type { Agents } from 'librechat-data-provider'; import type { IJobStoreV2, @@ -32,7 +35,10 @@ export interface ApprovalPauseOptions { discoveredTools?: string[]; activityPhaseSnapshot?: ActivityPhaseSnapshot; compactionSemanticIndex?: ICompactionSemanticIndexProjection; + /** Calibration and fading state at the pause; seeds the resumed run's pruner. */ + contextMeta?: IAgentEventActorContextMeta; /** Generation identity observed by the interrupted run. */ + expectedCreatedAt?: number; /** Hold Stop/resume until the paused assistant row is durably unfinished. */ persistencePending?: boolean; @@ -116,6 +122,7 @@ export class ApprovalLifecycle { const discoveredTools = options.discoveredTools; const activityPhaseSnapshot = options.activityPhaseSnapshot; const compactionSemanticIndex = options.compactionSemanticIndex; + const contextMeta = options.contextMeta; /** The normal receipt TTL matches a running job, but a review pause can * live for 24h or an explicit later expiry. The store extends every * receipt in the SAME CAS that closes running enqueues: a separate pass @@ -148,10 +155,14 @@ export class ApprovalLifecycle { : {}), ...(activityPhaseSnapshot != null ? { activityPhaseSnapshot } : {}), ...(compactionSemanticIndex != null ? { compactionSemanticIndex } : {}), + ...(contextMeta != null ? { contextMeta } : {}), ...(options.agentEventSuspension != null ? { agentEventSuspension: options.agentEventSuspension } : {}), }, + /** A re-pause with nothing persistable must not leave the previous + * segment's calibration and tier on the job for the next resume. */ + ...(contextMeta == null ? { clear: ['contextMeta' as const] } : {}), expectCreatedAt: expectedCreatedAt, notAfterMs: pendingAction.expiresAt, steerReceiptTtlSeconds: pauseReceiptTtl, diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index aec4f26a9d..a2d7c91ca3 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -2791,6 +2791,7 @@ class GenerationJobManagerClass { discoveredTools: jobData.discoveredTools, activityPhaseSnapshot: jobData.activityPhaseSnapshot, compactionSemanticIndex: jobData.compactionSemanticIndex, + contextMeta: jobData.contextMeta, // Surface the owning replica's seal capability so the steer route can // honour it instead of probing its own (possibly older) SDK. preemptCapable: jobData.preemptCapable, @@ -4388,6 +4389,22 @@ class GenerationJobManagerClass { contentError, ); } + /** The owner publishes run state (context meta) ahead of each model call + * and awaits that write, so a publish that landed after the initial read + * describes the call whose partial output the snapshot above carries. + * Read it back from the same epoch so the stopped response persists the + * tier that produced its bytes, not the one seen before the claim. */ + try { + const refreshed = await this.jobStore.getJob(streamId); + if (refreshed?.createdAt === jobData.createdAt && refreshed.contextMeta != null) { + jobData = { ...jobData, contextMeta: refreshed.contextMeta }; + } + } catch (metadataError) { + logger.warn( + `[GenerationJobManager] Failed to refresh committed abort metadata for ${streamId}:`, + metadataError, + ); + } if (options?.transformAbortContent) { content = options.transformAbortContent( content as TMessageContentParts[], diff --git a/packages/api/src/stream/__tests__/RedisJobStore.spec.ts b/packages/api/src/stream/__tests__/RedisJobStore.spec.ts index 0c3f039026..6b3eab814f 100644 --- a/packages/api/src/stream/__tests__/RedisJobStore.spec.ts +++ b/packages/api/src/stream/__tests__/RedisJobStore.spec.ts @@ -352,6 +352,12 @@ describe('RedisJobStore', () => { preserveForScheduleReconcile: true, promptTokens: 0, discoveredTools: [], + contextMeta: { + calibrationRatio: 1.25, + encoding: 'claude', + fading: { v: 1, budgetTokens: 50_000, masked: true }, + fadingTiers: [{ agentId: 'agent-123', v: 1, budgetTokens: 50_000, masked: true }], + }, compactionSemanticIndex: { version: 1, entries: [ @@ -438,6 +444,12 @@ describe('RedisJobStore', () => { preserveForScheduleReconcile: true, promptTokens: 0, discoveredTools: [], + contextMeta: { + calibrationRatio: 1.25, + encoding: 'claude', + fading: { v: 1, budgetTokens: 50_000, masked: true }, + fadingTiers: [{ agentId: 'agent-123', v: 1, budgetTokens: 50_000, masked: true }], + }, compactionSemanticIndex: { version: 1, entries: [ diff --git a/packages/api/src/stream/__tests__/abortContextMeta.spec.ts b/packages/api/src/stream/__tests__/abortContextMeta.spec.ts new file mode 100644 index 0000000000..4ee7c0bc6b --- /dev/null +++ b/packages/api/src/stream/__tests__/abortContextMeta.spec.ts @@ -0,0 +1,79 @@ +/** + * A Stop persists the response from job data alone. The run owner publishes its + * context meta ahead of each model call and awaits that write, so a publish that + * lands between the abort's initial read and its terminal claim describes the call + * whose partial output the abort snapshot carries. The abort must read it back. + */ +import type { IAgentEventActorContextMeta } from '@librechat/data-schemas'; + +/** Suppress winston Console transport output (survives jest.resetModules) */ +jest.spyOn(console, 'log').mockImplementation(); + +async function configureManager() { + const { GenerationJobManager } = await import('../GenerationJobManager'); + const { InMemoryJobStore } = await import('../implementations/InMemoryJobStore'); + const { InMemoryEventTransport } = await import('../implementations/InMemoryEventTransport'); + + const jobStore = new InMemoryJobStore(); + GenerationJobManager.configure({ + jobStore, + eventTransport: new InMemoryEventTransport(), + isRedis: false, + cleanupOnComplete: false, + }); + GenerationJobManager.initialize(); + return { manager: GenerationJobManager, jobStore }; +} + +const earlier: IAgentEventActorContextMeta = { + calibrationRatio: 1.1, + encoding: 'claude', + fading: { v: 1, budgetTokens: 50_000, masked: false }, +}; +const later: IAgentEventActorContextMeta = { + calibrationRatio: 1.3, + encoding: 'claude', + fading: { v: 1, budgetTokens: 25_000, masked: true }, +}; + +describe('abortJob context meta', () => { + beforeEach(() => { + jest.resetModules(); + }); + + it('carries the context meta published while the abort was claiming the job', async () => { + const { manager, jobStore } = await configureManager(); + const streamId = 'abort-context-meta-race'; + const job = await manager.createJob(streamId, 'user-1'); + await manager.updateMetadata(streamId, { contextMeta: earlier }, job.createdAt); + + const original = jobStore.transitionStatusAndDrainSteers.bind(jobStore); + jest + .spyOn(jobStore, 'transitionStatusAndDrainSteers') + .mockImplementationOnce(async (...args) => { + await jobStore.updateJob(streamId, { contextMeta: later }, job.createdAt); + return original(...args); + }); + + const result = await manager.abortJob(streamId); + + expect(result.success).toBe(true); + expect(result.jobData?.contextMeta).toEqual(later); + + await manager.destroy(); + }); + + it('keeps the context meta it read when nothing newer was published', async () => { + const { manager } = await configureManager(); + const streamId = 'abort-context-meta-stable'; + const job = await manager.createJob(streamId, 'user-1'); + await manager.updateMetadata(streamId, { contextMeta: earlier }, job.createdAt); + + const result = await manager.abortJob(streamId); + + expect(result.success).toBe(true); + expect(result.jobData?.contextMeta).toEqual(earlier); + + await manager.destroy(); + }); +}); diff --git a/packages/api/src/stream/__tests__/pendingAction.spec.ts b/packages/api/src/stream/__tests__/pendingAction.spec.ts index 2787db2664..11270fa97b 100644 --- a/packages/api/src/stream/__tests__/pendingAction.spec.ts +++ b/packages/api/src/stream/__tests__/pendingAction.spec.ts @@ -279,6 +279,44 @@ describe('ApprovalLifecycle via GenerationJobManager.approvals (in-memory)', () expect(paused?.metadata.compactionSemanticIndex).toEqual(compactionSemanticIndex); }); + test('persists context meta in the same transition as the pause', async () => { + const streamId = 'stream-pause-context-meta'; + await manager.createJob(streamId, 'user-1'); + const contextMeta = { + calibrationRatio: 1.25, + encoding: 'claude', + fading: { v: 1 as const, budgetTokens: 50_000, masked: true }, + fadingTiers: [{ agentId: 'agent-123', v: 1 as const, budgetTokens: 50_000, masked: true }], + }; + + expect(await manager.approvals.pause(streamId, buildAction(streamId), { contextMeta })).toBe( + true, + ); + + const paused = await manager.getJob(streamId); + expect(paused?.metadata.contextMeta).toEqual(contextMeta); + }); + + test('clears the previous pause context meta when a re-pause has none', async () => { + const streamId = 'stream-pause-context-meta-cleared'; + await manager.createJob(streamId, 'user-1'); + const firstAction = buildAction(streamId); + const contextMeta = { + calibrationRatio: 1.25, + encoding: 'claude', + fading: { v: 1 as const, budgetTokens: 50_000, masked: true }, + fadingTiers: [{ agentId: 'agent-123', v: 1 as const, budgetTokens: 50_000, masked: true }], + }; + + expect(await manager.approvals.pause(streamId, firstAction, { contextMeta })).toBe(true); + expect(await manager.approvals.resolve(streamId, firstAction.actionId)).toBe(true); + expect(await manager.approvals.pause(streamId, buildAction(streamId))).toBe(true); + + const repaused = await manager.getJob(streamId); + expect(repaused?.status).toBe('requires_action'); + expect(repaused?.metadata.contextMeta).toBeUndefined(); + }); + test('does not write a stale pause or discoveries onto a replacement job', async () => { const streamId = 'stream-pause-replaced'; const original = await manager.createJob(streamId, 'user-1'); diff --git a/packages/api/src/stream/implementations/RedisJobStore.ts b/packages/api/src/stream/implementations/RedisJobStore.ts index 8670ddb953..b582e40e72 100644 --- a/packages/api/src/stream/implementations/RedisJobStore.ts +++ b/packages/api/src/stream/implementations/RedisJobStore.ts @@ -4999,6 +4999,9 @@ export class RedisJobStore implements IJobStoreV2 { compactionSemanticIndex: data.compactionSemanticIndex ? JSON.parse(data.compactionSemanticIndex) : undefined, + /** Calibration and fading state captured at a HITL pause; without this line + * every Redis deployment would resume with no tier to seed. */ + contextMeta: data.contextMeta ? JSON.parse(data.contextMeta) : undefined, /** The owning replica's seal capability. `serializeJob` writes every * boolean generically, but this mapper is explicit โ€” omitting it here * drops the flag on every read, so the steer route would compute diff --git a/packages/api/src/stream/interfaces/IJobStore.ts b/packages/api/src/stream/interfaces/IJobStore.ts index 0f93d54668..18064da049 100644 --- a/packages/api/src/stream/interfaces/IJobStore.ts +++ b/packages/api/src/stream/interfaces/IJobStore.ts @@ -1,10 +1,13 @@ +import type { + IAgentEventActorContextMeta, + ICompactionSemanticIndexProjection, +} from '@librechat/data-schemas'; import type { Agents, TFile, TPendingSteer, UserSubmittedMessageFieldPath, } from 'librechat-data-provider'; -import type { ICompactionSemanticIndexProjection } from '@librechat/data-schemas'; import type { RunStep, StandardGraph } from '@librechat/agents'; import type { AgentEventDetachedTerminalEvidence } from '~/agents/triggers/types'; import type { ActivityPhaseSnapshot } from '~/agents/activityPhases/runtime'; @@ -177,6 +180,8 @@ export interface SerializableJobData { activityPhaseSnapshot?: ActivityPhaseSnapshot; /** Exact bounded compaction guidance captured atomically with a HITL pause. */ compactionSemanticIndex?: ICompactionSemanticIndexProjection; + /** Calibration and fading state captured atomically with a HITL pause, so a resume seeds its rebuilt pruner from the same tier. */ + contextMeta?: IAgentEventActorContextMeta; /** * Whether the replica that OWNS this generation can seal mid-stream * (`PreemptBoundary` wiring). Recorded at createJob because the steer route @@ -456,6 +461,7 @@ export type JobMetadataPatch = Partial< | 'discoveredTools' | 'activityPhaseSnapshot' | 'compactionSemanticIndex' + | 'contextMeta' | 'preemptCapable' | 'steerQuotesCapable' | 'steerQuotesExecutionId' diff --git a/packages/api/src/stream/metadata.ts b/packages/api/src/stream/metadata.ts index 4ddc19252b..4720788c1f 100644 --- a/packages/api/src/stream/metadata.ts +++ b/packages/api/src/stream/metadata.ts @@ -107,5 +107,8 @@ export function sanitizeJobMetadata(metadata: Partial): J if (metadata.compactionSemanticIndex) { patch.compactionSemanticIndex = metadata.compactionSemanticIndex; } + if (metadata.contextMeta) { + patch.contextMeta = metadata.contextMeta; + } return patch; } diff --git a/packages/api/src/types/stream.ts b/packages/api/src/types/stream.ts index 952d62231d..d6ec148c87 100644 --- a/packages/api/src/types/stream.ts +++ b/packages/api/src/types/stream.ts @@ -1,5 +1,8 @@ +import type { + IAgentEventActorContextMeta, + ICompactionSemanticIndexProjection, +} from '@librechat/data-schemas'; import type { Agents, UserSubmittedMessageFieldPath } from 'librechat-data-provider'; -import type { ICompactionSemanticIndexProjection } from '@librechat/data-schemas'; import type { EventEmitter } from 'events'; import type { AgentEventDetachedTerminalEvidence, @@ -89,6 +92,8 @@ export interface GenerationJobMetadata { activityPhaseSnapshot?: ActivityPhaseSnapshot; /** Exact bounded compaction guidance captured atomically with a HITL pause. */ compactionSemanticIndex?: ICompactionSemanticIndexProjection; + /** Calibration and fading state captured atomically with a HITL pause, so a resume seeds its rebuilt pruner from the same tier. */ + contextMeta?: IAgentEventActorContextMeta; /** See `SerializableJobData.preemptCapable`. */ preemptCapable?: boolean; /** See `SerializableJobData.steerQuotesCapable`. */ diff --git a/packages/data-provider/src/schemas.ts b/packages/data-provider/src/schemas.ts index ef4c7e3c37..7916284858 100644 --- a/packages/data-provider/src/schemas.ts +++ b/packages/data-provider/src/schemas.ts @@ -835,6 +835,13 @@ export const tExampleSchema = z.object({ export type TExample = z.infer; +/** Compact context-fading tier persisted beside a message's calibration ratio. */ +const agentFadingTierSchema = z.object({ + v: z.literal(1), + budgetTokens: z.number().positive(), + masked: z.boolean(), +}); + export const tMessageSchema = z.object({ messageId: z.string(), endpoint: z.string().optional(), @@ -896,6 +903,15 @@ export const tMessageSchema = z.object({ .describe( 'Tokenizer encoding used when this ratio was computed (e.g. "claude", "o200k_base")', ), + fading: agentFadingTierSchema + .optional() + .describe( + 'Latched context-fading tier of the default agent; seeds the next run so the provider projection of history keeps the same bytes', + ), + fadingTiers: z + .array(agentFadingTierSchema.extend({ agentId: z.string().min(1) })) + .optional() + .describe('Latched context-fading tiers keyed by agent ID, stored as entries'), }) .optional(), /** diff --git a/packages/data-schemas/src/methods/conversation.spec.ts b/packages/data-schemas/src/methods/conversation.spec.ts index 9ef5807a2f..2a3cf7b9b4 100644 --- a/packages/data-schemas/src/methods/conversation.spec.ts +++ b/packages/data-schemas/src/methods/conversation.spec.ts @@ -4874,6 +4874,39 @@ describe('Conversation Operations', () => { }), ).rejects.toThrow('Event actor context calibration is invalid'); + await expect( + methods.commitAgentEventActorState({ + user: 'actor-context-user', + conversationId, + invocationId: 'invalid-fading', + expectedEpoch: 0, + action: { toolName: 'submit_move' }, + checkpoint: actorCheckpoint, + contextMeta: { + calibrationRatio: 1.25, + fading: { v: 1, budgetTokens: -5, masked: true }, + }, + }), + ).rejects.toThrow('Event actor context fading tier is invalid'); + + await expect( + methods.commitAgentEventActorState({ + user: 'actor-context-user', + conversationId, + invocationId: 'invalid-fading-tiers', + expectedEpoch: 0, + action: { toolName: 'submit_move' }, + checkpoint: actorCheckpoint, + contextMeta: { + calibrationRatio: 1.25, + fadingTiers: [ + { agentId: 'agent-a', v: 1, budgetTokens: 20_000, masked: true }, + { agentId: 'agent-a', v: 1, budgetTokens: 10_000, masked: true }, + ], + }, + }), + ).rejects.toThrow('Event actor context fading tiers are invalid'); + await expect( methods.commitAgentEventActorState({ user: 'actor-context-user', diff --git a/packages/data-schemas/src/methods/conversation.ts b/packages/data-schemas/src/methods/conversation.ts index 38b29f5828..036cf31d46 100644 --- a/packages/data-schemas/src/methods/conversation.ts +++ b/packages/data-schemas/src/methods/conversation.ts @@ -42,6 +42,7 @@ import { refreshChatProjectStatsForUser, updateChatProjectLastConversationForUser, } from './chatProject'; +import { isAgentFadingTier, isAgentFadingTierEntries } from '~/utils/fading'; import { isCompactionSemanticIndexProjection } from '~/types/compaction'; import { createTempChatExpirationDate } from '~/utils/tempChatRetention'; import { tenantSafeBulkWrite } from '~/utils/tenantBulkWrite'; @@ -1010,6 +1011,15 @@ export function createConversationMethods( ) { throw new RangeError('Event actor context calibration is invalid'); } + if ( + input.contextMeta?.fadingTiers != null && + !isAgentFadingTierEntries(input.contextMeta.fadingTiers) + ) { + throw new RangeError('Event actor context fading tiers are invalid'); + } + if (input.contextMeta?.fading != null && !isAgentFadingTier(input.contextMeta.fading)) { + throw new RangeError('Event actor context fading tier is invalid'); + } if ( input.compactionSemanticIndex != null && !isCompactionSemanticIndexProjection(input.compactionSemanticIndex) diff --git a/packages/data-schemas/src/methods/message.spec.ts b/packages/data-schemas/src/methods/message.spec.ts index 7dd63065f1..f9b3350087 100644 --- a/packages/data-schemas/src/methods/message.spec.ts +++ b/packages/data-schemas/src/methods/message.spec.ts @@ -158,6 +158,77 @@ describe('Message Operations', () => { expect(savedMessage?.isUserSubmitted).toBeUndefined(); }); + it('unsets a previously stored context meta in the same update as the terminal save', async () => { + await saveMessage(mockCtx, { + ...mockMessageData, + contextMeta: { + calibrationRatio: 1.2, + encoding: 'claude', + fading: { v: 1, budgetTokens: 50_000, masked: true }, + }, + }); + const partial = await Message.findOne({ messageId: 'msg123', user: 'user123' }); + expect(partial?.contextMeta?.fading?.budgetTokens).toBe(50_000); + + const findOneAndUpdate = jest.spyOn(Message, 'findOneAndUpdate'); + const updateOne = jest.spyOn(Message, 'updateOne'); + try { + const result = await saveMessage(mockCtx, { + ...mockMessageData, + text: 'Completed', + contextMeta: null, + }); + + expect(result?.contextMeta).toBeUndefined(); + expect(result?.text).toBe('Completed'); + expect(findOneAndUpdate).toHaveBeenCalledTimes(1); + expect(findOneAndUpdate).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ $unset: { contextMeta: 1 } }), + expect.anything(), + ); + expect(updateOne).not.toHaveBeenCalled(); + } finally { + findOneAndUpdate.mockRestore(); + updateOne.mockRestore(); + } + const completed = await Message.findOne({ messageId: 'msg123', user: 'user123' }).lean(); + expect(completed).not.toHaveProperty('contextMeta'); + expect(completed?.text).toBe('Completed'); + }); + + it('unsets a stored context meta atomically through the provenance merge as well', async () => { + await saveMessage(mockCtx, { + ...mockMessageData, + contextMeta: { calibrationRatio: 1.2, encoding: 'claude' }, + }); + + const findOneAndUpdate = jest.spyOn(Message, 'findOneAndUpdate'); + try { + const result = await saveMessage(mockCtx, { + ...mockMessageData, + isCreatedByUser: false, + isUserSubmitted: true, + userSubmittedPaths: ['/content/0/text'], + contextMeta: null, + }); + + expect(result?.contextMeta).toBeUndefined(); + expect(result?.userSubmittedPaths).toEqual(['/content/0/text']); + expect(findOneAndUpdate).toHaveBeenCalledTimes(1); + expect(findOneAndUpdate).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ $unset: { contextMeta: 1 } }), + expect.anything(), + ); + } finally { + findOneAndUpdate.mockRestore(); + } + const completed = await Message.findOne({ messageId: 'msg123', user: 'user123' }).lean(); + expect(completed).not.toHaveProperty('contextMeta'); + expect(completed?.isUserSubmitted).toBe(true); + }); + it('should persist optional user-submitted provenance independently of message role', async () => { const result = await saveMessage(mockCtx, { ...mockMessageData, diff --git a/packages/data-schemas/src/methods/message.ts b/packages/data-schemas/src/methods/message.ts index 26ea1194b8..5cba4f4a50 100644 --- a/packages/data-schemas/src/methods/message.ts +++ b/packages/data-schemas/src/methods/message.ts @@ -1,6 +1,6 @@ import { HITL_MESSAGE_FILTER_FIELDS, RetentionMode } from 'librechat-data-provider'; +import type { DeleteResult, FilterQuery, Model, Types, UpdateQuery } from 'mongoose'; import type { UserSubmittedMessageFieldPath } from 'librechat-data-provider'; -import type { DeleteResult, FilterQuery, Model, Types } from 'mongoose'; import type { SearchParams } from 'meilisearch'; import type { SchemaWithMeiliMethods } from '~/models/plugins/mongoMeili'; import type { AppConfig, IConversation, IMessage } from '~/types'; @@ -265,13 +265,32 @@ function getSteerUserSubmittedPaths(content: unknown): string[] { return paths; } +/** + * A terminal save that must drop a stored `contextMeta` unsets it in the same + * update that persists the response, so no failure between two writes can + * leave a completed row carrying a disconnect snapshot's state. + */ +function buildMessageSaveUpdate( + update: Record, + options: { stampModelOutputOnInsert: boolean; unsetContextMeta: boolean }, +): UpdateQuery { + if (!options.stampModelOutputOnInsert && !options.unsetContextMeta) { + return update; + } + return { + $set: update, + ...(options.stampModelOutputOnInsert && { $setOnInsert: { isUserSubmitted: false } }), + ...(options.unsetContextMeta && { $unset: { contextMeta: 1 } }), + }; +} + async function findOneAndMergeMessageProvenance( Message: Model, identity: FilterQuery, update: Record, userSubmittedPaths: readonly string[], userSubmittedMessageFieldPaths: readonly UserSubmittedMessageFieldPath[], - options: { upsert: boolean; stampModelOutputOnInsert?: boolean }, + options: { upsert: boolean; stampModelOutputOnInsert?: boolean; unsetContextMeta?: boolean }, ) { const safeUpdate = { ...update }; delete safeUpdate._id; @@ -311,7 +330,10 @@ async function findOneAndMergeMessageProvenance( try { const message = await Message.findOneAndUpdate( filter, - { $set: { ...safeUpdate, ...provenance } }, + { + $set: { ...safeUpdate, ...provenance }, + ...(options.unsetContextMeta && { $unset: { contextMeta: 1 } }), + }, { upsert: options.upsert && current == null, new: true }, ); if (message != null) { @@ -533,7 +555,10 @@ export interface MessageMethods { expiredAt?: Date; interfaceConfig?: AppConfig['interfaceConfig']; }, - params: Partial & { newMessageId?: string }, + params: Omit, 'contextMeta'> & { + newMessageId?: string; + contextMeta?: IMessage['contextMeta'] | null; + }, metadata?: { context?: string }, ): Promise; recordSubagentTaskControlReceipt(input: { @@ -717,7 +742,11 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa expiredAt?: Date; interfaceConfig?: AppConfig['interfaceConfig']; }, - params: Partial & { newMessageId?: string }, + params: Omit, 'contextMeta'> & { + newMessageId?: string; + /** `null` unsets a previously stored value; omission leaves it in place. */ + contextMeta?: IMessage['contextMeta'] | null; + }, metadata?: { context?: string }, ) { if (!userId) { @@ -770,6 +799,14 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa update.expiredAt = null; } + /** A response that ends with nothing to carry must drop what an earlier + * partial save (a disconnect snapshot) stored, or the next turn would seed + * from stale state; omission never unsets, and the unset rides the same + * update as the response. */ + const unsetContextMeta = update.contextMeta === null; + if (unsetContextMeta) { + delete update.contextMeta; + } if (update.tokenCount != null && isNaN(update.tokenCount as number)) { logger.warn( `Resetting invalid \`tokenCount\` for message \`${params.messageId}\`: ${update.tokenCount}`, @@ -797,13 +834,11 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa update, userSubmittedPaths, userSubmittedMessageFieldPaths, - { upsert: true, stampModelOutputOnInsert }, + { upsert: true, stampModelOutputOnInsert, unsetContextMeta }, ) : await Message.findOneAndUpdate( { messageId: params.messageId, user: userId }, - stampModelOutputOnInsert - ? { $set: update, $setOnInsert: { isUserSubmitted: false } } - : update, + buildMessageSaveUpdate(update, { stampModelOutputOnInsert, unsetContextMeta }), { upsert: true, new: true }, ); @@ -3117,15 +3152,21 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa sortOrder?: 1 | -1; limit?: number; cursor?: string | null; + /** Projection for the page, e.g. `CLIENT_MESSAGE_SELECT` for client-facing reads. */ + select?: string; } = {}, ) { const Message = mongoose.models.Message as Model; - const { sortField = 'createdAt', sortOrder = -1, limit = 25, cursor } = options; + const { sortField = 'createdAt', sortOrder = -1, limit = 25, cursor, select } = options; const queryFilter = { ...filter }; if (cursor) { queryFilter[sortField] = sortOrder === 1 ? { $gt: cursor } : { $lt: cursor }; } - const messages = await Message.find(queryFilter) + const query = Message.find(queryFilter); + if (select) { + query.select(select); + } + const messages = await query .sort({ [sortField]: sortOrder }) .limit(limit + 1) .lean(); diff --git a/packages/data-schemas/src/schema/convo.ts b/packages/data-schemas/src/schema/convo.ts index de7ac5a1d6..c853b1b160 100644 --- a/packages/data-schemas/src/schema/convo.ts +++ b/packages/data-schemas/src/schema/convo.ts @@ -13,6 +13,7 @@ import { MAX_AGENT_EVENT_ACTOR_SUMMARY_LENGTH, MAX_AGENT_EVENT_ACTOR_TOOL_NAME_LENGTH, } from '~/types/convo'; +import { agentFadingContextDefinition } from './fading'; import { conversationPreset } from './defaults'; import { IConversation } from '~/types'; @@ -147,6 +148,7 @@ const convoSchema: Schema = new Schema( maxlength: MAX_AGENT_EVENT_ACTOR_ENCODING_LENGTH, default: undefined, }, + ...agentFadingContextDefinition, }, _id: false, default: undefined, diff --git a/packages/data-schemas/src/schema/fading.ts b/packages/data-schemas/src/schema/fading.ts new file mode 100644 index 0000000000..3b303cd2f0 --- /dev/null +++ b/packages/data-schemas/src/schema/fading.ts @@ -0,0 +1,35 @@ +import type { SchemaDefinitionProperty } from 'mongoose'; +interface AgentFadingContextDefinition { + fading: SchemaDefinitionProperty; + fadingTiers: SchemaDefinitionProperty; +} + +/** Mongoose definition of one persisted context-fading tier (`IAgentFadingTier`). */ +const agentFadingTierDefinition = { + v: { type: Number, enum: [1], required: true }, + budgetTokens: { type: Number, min: 1, required: true }, + masked: { type: Boolean, required: true }, +}; + +/** + * Mongoose definition of the fading fields inside a persisted `contextMeta`: the + * default agent's tier and the per-agent entries. Shared by the message and + * conversation schemas so both persist exactly the same compact shape. + */ +export const agentFadingContextDefinition: AgentFadingContextDefinition = { + fading: { + type: agentFadingTierDefinition, + _id: false, + default: undefined, + }, + fadingTiers: { + type: [ + { + agentId: { type: String, required: true }, + ...agentFadingTierDefinition, + _id: false, + }, + ], + default: undefined, + }, +}; diff --git a/packages/data-schemas/src/schema/message.ts b/packages/data-schemas/src/schema/message.ts index b1ad4458eb..4418fff967 100644 --- a/packages/data-schemas/src/schema/message.ts +++ b/packages/data-schemas/src/schema/message.ts @@ -1,5 +1,6 @@ import mongoose, { Schema } from 'mongoose'; import type { IMessage } from '~/types/message'; +import { agentFadingContextDefinition } from './fading'; const messageSchema: Schema = new Schema( { @@ -235,6 +236,7 @@ const messageSchema: Schema = new Schema( type: { calibrationRatio: { type: Number }, encoding: { type: String }, + ...agentFadingContextDefinition, }, _id: false, default: undefined, diff --git a/packages/data-schemas/src/types/convo.ts b/packages/data-schemas/src/types/convo.ts index 42b706690a..0362977926 100644 --- a/packages/data-schemas/src/types/convo.ts +++ b/packages/data-schemas/src/types/convo.ts @@ -45,9 +45,37 @@ export interface IAgentEventActorSummary { tokenCount: number; } +/** + * Latched context-fading tier from `@librechat/agents`. Every cap the SDK + * applies to its provider-only projection of historical tool results derives + * from it alone, so persisting it keeps that projection byte-stable across runs + * for prefix-based provider prompt caches. Graph messages stay canonical. + */ +export interface IAgentFadingTier { + v: 1; + /** Token budget the caps derive from; never grows within a conversation. */ + budgetTokens: number; + /** Whether observation masking has activated. */ + masked: boolean; +} + +/** One agent's latched tier inside the persisted per-agent map. */ +export interface IAgentFadingTierEntry extends IAgentFadingTier { + agentId: string; +} + +/** + * Compact context state a run hands to its successor: calibration and the + * latched fading tiers. Messages themselves are never part of it; the SDK keeps + * graph history canonical and derives a provider-only projection per run. + */ export interface IAgentEventActorContextMeta { calibrationRatio: number; encoding?: string; + /** Default agent's tier, kept for single-agent seeding. */ + fading?: IAgentFadingTier; + /** Tiers keyed by agent ID, stored as entries so agent IDs never become field names. */ + fadingTiers?: IAgentFadingTierEntry[]; } /** Private committed checkpoint state for one event-bound child actor. */ diff --git a/packages/data-schemas/src/types/message.ts b/packages/data-schemas/src/types/message.ts index 69fd580630..18911f39d3 100644 --- a/packages/data-schemas/src/types/message.ts +++ b/packages/data-schemas/src/types/message.ts @@ -4,6 +4,7 @@ import type { UserSubmittedMessageFieldPath, } from 'librechat-data-provider'; import type { Document } from 'mongoose'; +import type { IAgentEventActorContextMeta } from './convo'; export type SubagentTaskControlAction = | 'steer' @@ -117,10 +118,7 @@ export interface IMessage extends Document { controlReceipts?: ISubagentTaskControlReceipt[]; }; subagentTriggerProjection?: SubagentTriggerProjection; - contextMeta?: { - calibrationRatio?: number; - encoding?: string; - }; + contextMeta?: Partial; attachments?: unknown[]; /** Skills the user invoked manually via the `$` popover on this turn. UI-only metadata for `SkillPills`. */ manualSkills?: string[]; diff --git a/packages/data-schemas/src/utils/fading.spec.ts b/packages/data-schemas/src/utils/fading.spec.ts new file mode 100644 index 0000000000..7de1327f95 --- /dev/null +++ b/packages/data-schemas/src/utils/fading.spec.ts @@ -0,0 +1,44 @@ +import { isAgentFadingTier, isAgentFadingTierEntries, isAgentFadingTierEntry } from './fading'; + +const tier = { v: 1, budgetTokens: 20_000, masked: true }; + +describe('isAgentFadingTier', () => { + it('accepts the compact shape and rejects anything else', () => { + expect(isAgentFadingTier(tier)).toBe(true); + expect(isAgentFadingTier({ ...tier, v: 2 })).toBe(false); + expect(isAgentFadingTier({ ...tier, budgetTokens: 0 })).toBe(false); + expect(isAgentFadingTier({ ...tier, masked: 'yes' })).toBe(false); + expect(isAgentFadingTier(null)).toBe(false); + }); +}); + +describe('isAgentFadingTierEntry', () => { + it('requires a non-empty agent ID on a valid tier, however long', () => { + expect(isAgentFadingTierEntry({ agentId: 'agent-a', ...tier })).toBe(true); + expect(isAgentFadingTierEntry({ agentId: '', ...tier })).toBe(false); + /** Ephemeral agent IDs encode endpoint, model and sender without a bound. */ + expect(isAgentFadingTierEntry({ agentId: 'ephemeral|'.repeat(200), ...tier })).toBe(true); + expect(isAgentFadingTierEntry({ agentId: 'agent-a', ...tier, budgetTokens: -1 })).toBe(false); + expect(isAgentFadingTierEntry(tier)).toBe(false); + }); +}); + +describe('isAgentFadingTierEntries', () => { + it('accepts unique valid entries and rejects duplicates or malformed members', () => { + expect( + isAgentFadingTierEntries([ + { agentId: 'agent-a', ...tier }, + { agentId: 'agent-b', ...tier, masked: false }, + ]), + ).toBe(true); + expect(isAgentFadingTierEntries([])).toBe(true); + expect( + isAgentFadingTierEntries([ + { agentId: 'agent-a', ...tier }, + { agentId: 'agent-a', ...tier }, + ]), + ).toBe(false); + expect(isAgentFadingTierEntries([{ agentId: 'agent-a', ...tier }, tier])).toBe(false); + expect(isAgentFadingTierEntries({ 'agent-a': tier })).toBe(false); + }); +}); diff --git a/packages/data-schemas/src/utils/fading.ts b/packages/data-schemas/src/utils/fading.ts new file mode 100644 index 0000000000..f7df158d25 --- /dev/null +++ b/packages/data-schemas/src/utils/fading.ts @@ -0,0 +1,47 @@ +import type { IAgentFadingTier, IAgentFadingTierEntry } from '~/types/convo'; + +/** Version of the persisted context-fading tier shape; must match `@librechat/agents`. */ +export const AGENT_FADING_TIER_VERSION = 1; + +/** Whether a persisted value is a well-formed context-fading tier. */ +export function isAgentFadingTier(value: unknown): value is IAgentFadingTier { + if (typeof value !== 'object' || value === null) { + return false; + } + const { v, budgetTokens, masked } = value as Partial>; + return ( + v === AGENT_FADING_TIER_VERSION && + typeof budgetTokens === 'number' && + Number.isFinite(budgetTokens) && + budgetTokens > 0 && + typeof masked === 'boolean' + ); +} + +/** + * Whether a persisted value is a well-formed per-agent tier entry. Agent IDs + * are server-generated and unbounded (an ephemeral agent's ID encodes its + * endpoint, model and sender), so only emptiness is rejected. + */ +export function isAgentFadingTierEntry(value: unknown): value is IAgentFadingTierEntry { + if (!isAgentFadingTier(value)) { + return false; + } + const { agentId } = value as Partial>; + return typeof agentId === 'string' && agentId.length > 0; +} + +/** Whether a persisted value is a list of per-agent tier entries with unique agent IDs. */ +export function isAgentFadingTierEntries(value: unknown): value is IAgentFadingTierEntry[] { + if (!Array.isArray(value)) { + return false; + } + const seen = new Set(); + for (const entry of value) { + if (!isAgentFadingTierEntry(entry) || seen.has(entry.agentId)) { + return false; + } + seen.add(entry.agentId); + } + return true; +} diff --git a/packages/data-schemas/src/utils/index.ts b/packages/data-schemas/src/utils/index.ts index 5ef6d70bb1..4adfa3488b 100644 --- a/packages/data-schemas/src/utils/index.ts +++ b/packages/data-schemas/src/utils/index.ts @@ -7,3 +7,4 @@ export * from './transactions'; export * from './objectId'; export * from './yaml'; export * from './stripUIResourceMarkers'; +export * from './fading';