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';