From 612327f781c6a2e03ffe570c84aa7f0f77842e89 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 28 Jul 2026 08:48:33 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=A7=B0=20fix:=20Harden=20Label=20Resoluti?= =?UTF-8?q?on,=20Output=20Bounds,=20and=20Cache=20Billing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-ten review (all P2, feature-scoped); the sixth finding is the documented edited+reconnect index-space limitation, answered on-thread as deliberately out of scope for this PR. - Rejected-LLM memoization (runtime.ts): the hook cached a rejected `resolveLLM()` promise permanently, failing every later batch and silently defeating the host resolver's own rejected-cache eviction. The memo now evicts on rejection so the next batch retries. - `current_model` precedence (host.ts): an explicit `activityModel: current_model` resolved to `undefined` and then lost to a configured `titleModel`. The sentinel now resolves straight to the run model; the title fallback applies only when `activityModel` is absent. - Output bounds (runtime.ts): label text was persisted verbatim; a model ignoring the 4–9-word instruction (or steered by injection in untrusted tool output) could emit thousands of tokens duplicated through SSE, the chunk log, persistence, and the UI. `normalizeLabelOutput` keeps the first non-empty line, collapses whitespace, and hard-caps at 200 chars on both generation paths. - Cache-token billing (host.ts, client.js): the usage mapper dropped cache fields, vanishing Anthropic cache tokens from billing and charging OpenAI cache reads at the full input rate. The mapper now normalizes Anthropic/OpenAI/LangChain cache shapes into `input_token_details`, and the emit + cost path carries them with the label endpoint's `provider` (additive-provider adjustment). - Usage-type union (runs.ts): `TTokenUsageEvent.usage_type` now includes the emitted `activity-label` literal; the lone consumer keys on `usage_type != null`, so this is type-level completion. Tests: sentinel/title/explicit model precedence and all three cache shapes (host.spec), transient-resolution retry and output normalization with truncation (runtime.spec), the new usage literal (runs.spec). --- api/server/controllers/agents/client.js | 16 ++- .../activityLabels/__tests__/host.spec.ts | 111 +++++++++++++++++- .../activityLabels/__tests__/runtime.spec.ts | 41 +++++++ .../api/src/agents/activityLabels/host.ts | 70 ++++++++--- .../api/src/agents/activityLabels/runtime.ts | 42 ++++++- packages/data-provider/src/types/runs.spec.ts | 13 ++ packages/data-provider/src/types/runs.ts | 5 +- 7 files changed, 273 insertions(+), 25 deletions(-) diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index e8b65ca520..8ec4de9909 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -399,6 +399,10 @@ class AgentClient extends BaseClient { * scope that closed while this was in flight still suppresses the write. * Defaults open for callers that own no scope. */ scopeOpen = () => true, + /** The LABEL endpoint's provider — cost math needs it to know whether + * cache tokens are folded into `input_tokens` (additive providers like + * Bedrock keep them separate). */ + provider = undefined, ) { const appConfig = this.options.req?.config; const collectedUsage = mapCollectedMetadataToUsage(collectedMetadata); @@ -432,6 +436,12 @@ class AgentClient extends BaseClient { const data = { input_tokens: usage.input_tokens, output_tokens: usage.output_tokens, + /** Cache tokens ride along (subagent-event shape) so display and + * aggregation price cached label calls at cache rates. */ + ...(usage.input_token_details != null && { + input_token_details: usage.input_token_details, + }), + ...(provider != null && { provider }), model, usage_type: 'activity-label', /** @@ -454,7 +464,7 @@ class AgentClient extends BaseClient { * `interface.contextCost` is on. */ cost: includeCost ? computeUsageCostUSD( - { ...usage, model }, + { ...usage, model, provider }, { getMultiplier: db.getMultiplier, getCacheMultiplier: db.getCacheMultiplier }, labelTokenConfig, ) @@ -560,6 +570,7 @@ class AgentClient extends BaseClient { endpointTokenConfig, sameEndpoint, scopeStillOpen, + provider, ); }; /** @@ -775,7 +786,7 @@ class AgentClient extends BaseClient { return { callbacks: [{ handleLLMEnd }], collect: async () => { - const { clientOptions, endpointTokenConfig, sameEndpoint } = + const { provider, clientOptions, endpointTokenConfig, sameEndpoint } = await this.resolveActivityLabelLLM(); await this.recordActivityLabelUsage( collected, @@ -786,6 +797,7 @@ class AgentClient extends BaseClient { * fallback label whose fill was dropped as out-of-scope * still billed and emitted after finalization. */ () => labelScope.closed !== true, + provider, ); }, }; diff --git a/packages/api/src/agents/activityLabels/__tests__/host.spec.ts b/packages/api/src/agents/activityLabels/__tests__/host.spec.ts index 889fe1a34e..d31f465d09 100644 --- a/packages/api/src/agents/activityLabels/__tests__/host.spec.ts +++ b/packages/api/src/agents/activityLabels/__tests__/host.spec.ts @@ -1,5 +1,22 @@ import type { AppConfig } from '@librechat/data-schemas'; -import { resolveActivityConfig } from '../host'; +import type { EndpointDbMethods, ServerRequest } from '~/types'; +import { + mapCollectedMetadataToUsage, + resolveActivityConfig, + resolveActivityLabelModel, +} from '../host'; + +const mockGetOptions = jest.fn(async (_params: unknown) => ({ + llmConfig: { model: 'resolved' }, +})); +jest.mock('~/endpoints/config/providers', () => ({ + getProviderConfig: jest.fn(() => ({ + getOptions: (params: unknown) => mockGetOptions(params), + customEndpointConfig: undefined, + })), +})); +jest.mock('~/utils/headers', () => ({ resolveConfigHeaders: jest.fn() })); +jest.mock('~/utils/env', () => ({ createSafeUser: jest.fn(() => undefined) })); const appConfig = (endpoints: Record): AppConfig => ({ endpoints }) as unknown as AppConfig; @@ -62,3 +79,95 @@ describe('resolveActivityConfig', () => { expect(config.model).toBe('gpt-4o-mini'); }); }); + +describe('resolveActivityLabelModel model precedence', () => { + const db = {} as EndpointDbMethods; + const resolve = (endpointConfig: Record) => + resolveActivityLabelModel({ + req: { config: appConfig({ openAI: endpointConfig }) } as unknown as ServerRequest, + agent: { endpoint: 'openAI', model_parameters: { model: 'run-model' } }, + ids: {}, + db, + }); + + beforeEach(() => { + mockGetOptions.mockClear(); + }); + + /** An EXPLICIT `activityModel: current_model` names the run model — a + * configured `titleModel` must not shadow it via the fallback chain. */ + it('resolves an explicit current_model sentinel to the run model over titleModel', async () => { + await resolve({ activityLabel: true, activityModel: 'current_model', titleModel: 'haiku' }); + expect(mockGetOptions).toHaveBeenCalledWith( + expect.objectContaining({ model_parameters: { model: 'run-model' } }), + ); + }); + + it('falls back to titleModel only when activityModel is absent', async () => { + await resolve({ activityLabel: true, titleModel: 'haiku' }); + expect(mockGetOptions).toHaveBeenCalledWith( + expect.objectContaining({ model_parameters: { model: 'haiku' } }), + ); + }); + + it('prefers an explicit activityModel over everything', async () => { + await resolve({ activityLabel: true, activityModel: 'label-model', titleModel: 'haiku' }); + expect(mockGetOptions).toHaveBeenCalledWith( + expect.objectContaining({ model_parameters: { model: 'label-model' } }), + ); + }); +}); + +describe('mapCollectedMetadataToUsage cache tokens', () => { + it('carries Anthropic raw cache fields as normalized details', () => { + const [usage] = mapCollectedMetadataToUsage([ + { + usage: { + input_tokens: 100, + output_tokens: 9, + cache_read_input_tokens: 80, + cache_creation_input_tokens: 10, + }, + }, + ]); + expect(usage).toEqual({ + input_tokens: 100, + output_tokens: 9, + input_token_details: { cache_read: 80, cache_creation: 10 }, + }); + }); + + it('maps OpenAI cached_tokens to cache_read', () => { + const [usage] = mapCollectedMetadataToUsage([ + { + usage: { + prompt_tokens: 50, + completion_tokens: 7, + prompt_tokens_details: { cached_tokens: 40 }, + }, + }, + ]); + expect(usage.input_token_details).toEqual({ cache_read: 40, cache_creation: undefined }); + }); + + it('passes through LangChain-standard input_token_details', () => { + const [usage] = mapCollectedMetadataToUsage([ + { + usage_metadata: { + input_tokens: 30, + output_tokens: 5, + input_token_details: { cache_read: 20, cache_creation: 4 }, + }, + }, + ]); + expect(usage.input_token_details).toEqual({ cache_read: 20, cache_creation: 4 }); + }); + + it('omits the details object entirely when no cache tokens are reported', () => { + const [usage] = mapCollectedMetadataToUsage([ + { usage: { input_tokens: 10, output_tokens: 2 } }, + ]); + expect(usage).toEqual({ input_tokens: 10, output_tokens: 2 }); + expect('input_token_details' in usage).toBe(false); + }); +}); diff --git a/packages/api/src/agents/activityLabels/__tests__/runtime.spec.ts b/packages/api/src/agents/activityLabels/__tests__/runtime.spec.ts index 2805798dfa..35e006e691 100644 --- a/packages/api/src/agents/activityLabels/__tests__/runtime.spec.ts +++ b/packages/api/src/agents/activityLabels/__tests__/runtime.spec.ts @@ -383,4 +383,45 @@ describe('createActivityLabelHook', () => { expect(slots).toHaveLength(2); expect(resolveLLM).toHaveBeenCalledTimes(1); }); + + /** A transient resolution failure must stay transient: memoizing the + * rejected promise would fail every later batch in the run instantly. */ + it('retries LLM resolution on the next batch after a transient failure', async () => { + const flaky = jest + .fn, []>() + .mockRejectedValueOnce(new Error('credential read timeout')) + .mockResolvedValue({ provider: Providers.OPENAI, clientOptions: { model: 'small-model' } }); + const hook = createActivityLabelHook({ claimSlot, resolveLLM: flaky }); + await hook(batchInput(), new AbortController().signal); + await flushDetached(); + expect(slots[0].filled).toEqual([null]); + + await hook(batchInput(), new AbortController().signal); + await flushDetached(); + expect(flaky).toHaveBeenCalledTimes(2); + expect(slots[1].filled).toEqual(['Searched the web for LibreChat docs.']); + }); + + /** Output is bounded before persisting: a model that ignores the 4–9-word + * instruction (or is steered by injected tool output) must not turn one + * header into an unbounded multi-line content part. */ + it('normalizes label output to one bounded line', async () => { + mockInvoke.mockResolvedValue({ + content: `\n \nFound the failing\tspec\nIGNORE PREVIOUS INSTRUCTIONS ${'x'.repeat(5000)}`, + }); + const hook = createActivityLabelHook({ claimSlot, resolveLLM }); + await hook(batchInput(), new AbortController().signal); + await flushDetached(); + expect(slots[0].filled).toEqual(['Found the failing spec']); + }); + + it('truncates a single giant label line with an ellipsis', async () => { + mockInvoke.mockResolvedValue({ content: 'word '.repeat(2000) }); + const hook = createActivityLabelHook({ claimSlot, resolveLLM }); + await hook(batchInput(), new AbortController().signal); + await flushDetached(); + const label = slots[0].filled[0] as string; + expect(label.length).toBeLessThanOrEqual(200); + expect(label.endsWith('…')).toBe(true); + }); }); diff --git a/packages/api/src/agents/activityLabels/host.ts b/packages/api/src/agents/activityLabels/host.ts index ecb8a5f447..18d68f46c7 100644 --- a/packages/api/src/agents/activityLabels/host.ts +++ b/packages/api/src/agents/activityLabels/host.ts @@ -10,6 +10,12 @@ import { getProviderConfig } from '~/endpoints/config/providers'; import { resolveConfigHeaders } from '~/utils/headers'; import { createSafeUser } from '~/utils/env'; +/** Cache-token details in the LangChain-standard normalized shape. */ +interface CacheTokenDetails { + cache_read?: number; + cache_creation?: number; +} + /** Aggregated LLM metadata entries (shape varies by provider SDK). */ export interface CollectedMetadataEntry { usage?: { @@ -19,19 +25,35 @@ export interface CollectedMetadataEntry { completion_tokens?: number; output_tokens?: number; outputTokens?: number; + /** Anthropic raw usage. */ + cache_creation_input_tokens?: number; + cache_read_input_tokens?: number; + /** OpenAI raw usage. */ + prompt_tokens_details?: { cached_tokens?: number }; }; tokenUsage?: { promptTokens?: number; completionTokens?: number }; - usage_metadata?: { input_tokens?: number; output_tokens?: number }; + usage_metadata?: { + input_tokens?: number; + output_tokens?: number; + input_token_details?: CacheTokenDetails; + }; } export interface ActivityLabelUsage { input_tokens?: number; output_tokens?: number; + /** Normalized cache tokens — `computeUsageCostUSD` and the transaction + * path read this shape first, so carrying it prices cached label calls + * at cache rates instead of the ordinary input rate (or not at all). */ + input_token_details?: CacheTokenDetails; } /** * Normalizes provider-specific aggregated metadata into the usage shape - * `recordCollectedUsage` expects. Mirrors the title path's inline mapping. + * `recordCollectedUsage` expects, cache-token details included — dropping + * them made Anthropic cache tokens vanish from billing and charged OpenAI + * cache reads at the full input rate. Mirrors the title path's inline + * mapping otherwise. */ export function mapCollectedMetadataToUsage( collected: CollectedMetadataEntry[], @@ -39,18 +61,31 @@ export function mapCollectedMetadataToUsage( return collected.map((item) => { let input_tokens: number | undefined; let output_tokens: number | undefined; + let cache_read: number | undefined; + let cache_creation: number | undefined; if (item.usage) { input_tokens = item.usage.prompt_tokens ?? item.usage.input_tokens ?? item.usage.inputTokens; output_tokens = item.usage.completion_tokens ?? item.usage.output_tokens ?? item.usage.outputTokens; + cache_read = + item.usage.cache_read_input_tokens ?? item.usage.prompt_tokens_details?.cached_tokens; + cache_creation = item.usage.cache_creation_input_tokens; } else if (item.tokenUsage) { input_tokens = item.tokenUsage.promptTokens; output_tokens = item.tokenUsage.completionTokens; } else if (item.usage_metadata) { input_tokens = item.usage_metadata.input_tokens; output_tokens = item.usage_metadata.output_tokens; + cache_read = item.usage_metadata.input_token_details?.cache_read; + cache_creation = item.usage_metadata.input_token_details?.cache_creation; } - return { input_tokens, output_tokens }; + return { + input_tokens, + output_tokens, + ...(cache_read != null || cache_creation != null + ? { input_token_details: { cache_read, cache_creation } } + : {}), + }; }); } @@ -192,24 +227,29 @@ export async function resolveActivityLabelModel({ * the credential target silently change the model and its cost. The * destination supplies credentials, never the model choice. */ const titleModel = originatingTitleModel; - /** `current_model` means "the agent's model" for BOTH overrides. The - * activity options are documented as title-shaped, so an `activityModel` - * set to the sentinel must resolve the same way `titleModel` does — passing - * the literal through would send `model: "current_model"` to the provider - * and fail every label. */ - const activityModel = - activity.model != null && activity.model !== Constants.CURRENT_MODEL - ? activity.model - : undefined; /** `model_parameters.model` FIRST: `initializeAgent` merges the request's * `endpointOption` override into it and the run itself gives it precedence, * so the saved `agent.model` can be a stale or entirely different model. * Reading it first is what makes "current model" mean the model the * conversation is actually running on. */ const runModel = agent.model_parameters?.model ?? agent.model; - const model = - activityModel ?? - (titleModel != null && titleModel !== Constants.CURRENT_MODEL ? titleModel : runModel); + /** `current_model` means "the agent's model" for BOTH overrides — passing + * the literal through would send `model: "current_model"` to the provider + * and fail every label. An EXPLICIT `activityModel: current_model` resolves + * straight to the run model: the admin asked for it by name, so letting a + * configured `titleModel` win instead would route labels to an unintended + * model with different behavior and cost. The title fallback applies only + * when `activityModel` is absent. */ + let model: string | undefined; + if (activity.model === Constants.CURRENT_MODEL) { + model = runModel; + } else if (activity.model != null) { + model = activity.model; + } else if (titleModel != null && titleModel !== Constants.CURRENT_MODEL) { + model = titleModel; + } else { + model = runModel; + } const options = await providerConfig.getOptions({ req, endpoint, diff --git a/packages/api/src/agents/activityLabels/runtime.ts b/packages/api/src/agents/activityLabels/runtime.ts index cee39df2f9..0bdd8cc744 100644 --- a/packages/api/src/agents/activityLabels/runtime.ts +++ b/packages/api/src/agents/activityLabels/runtime.ts @@ -169,6 +169,28 @@ const DEFAULT_MAX_PER_RUN = 20; const DEFAULT_CHAR_LIMIT = 600; const INPUT_CHAR_LIMIT = 200; const SUMMARY_TIMEOUT_MS = 12_000; +/** Hard bound on the PERSISTED label. The instruction asks for 4–9 words, but + * a model that ignores it — or is steered by injection through untrusted + * tool output — could otherwise turn one header into thousands of tokens + * duplicated through SSE, the durable chunk log, persistence, and the UI. */ +const LABEL_OUTPUT_CHAR_LIMIT = 200; + +/** + * Normalizes raw model output into a header: the first non-empty line, + * whitespace collapsed, hard-capped at {@link LABEL_OUTPUT_CHAR_LIMIT}. A + * header renders as one line, so everything past the first line break is + * noise at best and injected payload at worst. + */ +export function normalizeLabelOutput(text: string | null | undefined): string { + if (text == null) { + return ''; + } + const firstLine = text.split(/\r?\n/).find((line) => line.trim().length > 0) ?? ''; + const collapsed = firstLine.replace(/\s+/g, ' ').trim(); + return collapsed.length > LABEL_OUTPUT_CHAR_LIMIT + ? `${collapsed.slice(0, LABEL_OUTPUT_CHAR_LIMIT - 1)}…` + : collapsed; +} function truncate(value: string, limit: number): string { return value.length > limit ? `${value.slice(0, limit)}…` : value; @@ -375,7 +397,16 @@ export function createActivityLabelHook( let llmPromise: Promise | null = null; const getLLM = (): Promise => { - llmPromise = llmPromise ?? opts.resolveLLM(); + llmPromise = + llmPromise ?? + opts.resolveLLM().catch((error) => { + /** Never cache a rejection: memoizing it would fail every later + * batch in the run instantly — and silently defeat the host + * resolver's own rejected-cache eviction, which exists precisely so + * a transient credential read failure stays transient. */ + llmPromise = null; + throw error; + }); return llmPromise; }; @@ -468,10 +499,11 @@ export function createActivityLabelHook( } else { text = await generateDirect(); } - /** Trim centrally: a whitespace-only label from either path must - * fill null so the UI keeps the deterministic counts fallback. */ - const trimmed = text?.trim() ?? ''; - const committed = (await slot.fill(trimmed.length > 0 ? trimmed : null)) === true; + /** Normalize centrally — BOTH paths: single line, bounded length, + * whitespace-only becomes null so the UI keeps the deterministic + * counts fallback. */ + const normalized = normalizeLabelOutput(text); + const committed = (await slot.fill(normalized.length > 0 ? normalized : null)) === true; await collectDeferredUsage(committed); } catch (error) { logger.warn( diff --git a/packages/data-provider/src/types/runs.spec.ts b/packages/data-provider/src/types/runs.spec.ts index f4058d1c7b..d2090b9e94 100644 --- a/packages/data-provider/src/types/runs.spec.ts +++ b/packages/data-provider/src/types/runs.spec.ts @@ -33,6 +33,19 @@ describe('promptTokensFromUsage', () => { expect(promptTokensFromUsage({ provider: 'anthropic' })).toBe(0); }); + it('accepts the activity-label usage bucket emitted on the wire', () => { + /** Type-level pin: the backend emits this literal for fast-model header + * calls, so the union must be able to represent the actual payload. */ + const event: TTokenUsageEvent = { + input_tokens: 120, + output_tokens: 9, + usage_type: 'activity-label', + runId: 'msg-1:1700000000000', + seq: -1, + }; + expect(promptTokensFromUsage(event)).toBe(120); + }); + it('uses the magnitude heuristic when the provider is absent (cache ≤ input ⇒ included)', () => { /** OpenAI-compatible/custom payload with no provider: cache already folded * into input_tokens, so it must NOT be re-added. */ diff --git a/packages/data-provider/src/types/runs.ts b/packages/data-provider/src/types/runs.ts index c97d90b552..47362c94db 100644 --- a/packages/data-provider/src/types/runs.ts +++ b/packages/data-provider/src/types/runs.ts @@ -205,8 +205,9 @@ export type TTokenUsageEvent = { provider?: string; /** Non-primary buckets fold into session cost/totals but not the live * context gauge: hidden sequential-agent calls (`sequential`), summary - * passes (`summarization`), and isolated subagent runs (`subagent`) */ - usage_type?: 'summarization' | 'subagent' | 'sequential'; + * passes (`summarization`), isolated subagent runs (`subagent`), and + * fast-model activity headers (`activity-label`) */ + usage_type?: 'summarization' | 'subagent' | 'sequential' | 'activity-label'; runId?: string; /** Per-run emission sequence; keeps identical payloads from distinct model calls unique */ seq?: number;