diff --git a/api/app/clients/BaseClient.js b/api/app/clients/BaseClient.js index dfa8d2201b..03653333b4 100644 --- a/api/app/clients/BaseClient.js +++ b/api/app/clients/BaseClient.js @@ -700,7 +700,6 @@ class BaseClient { user, ); this.savedMessageIds.add(responseMessage.messageId); - delete responseMessage.tokenCount; return responseMessage; } diff --git a/api/server/controllers/TokenConfigController.js b/api/server/controllers/TokenConfigController.js new file mode 100644 index 0000000000..6826ad304d --- /dev/null +++ b/api/server/controllers/TokenConfigController.js @@ -0,0 +1,50 @@ +const { logger } = require('@librechat/data-schemas'); +const { EModelEndpoint } = require('librechat-data-provider'); +const { buildTokenConfigMap, getTokenConfigKey, tokenConfigCache } = require('@librechat/api'); +const { getModelsConfig } = require('~/server/controllers/ModelController'); +const { getValueKey, getMultiplier, getCacheMultiplier } = require('~/models'); + +/** + * Returns server-resolved context windows (and pricing when + * `interface.contextCost` is enabled) for every configured model. + * @param {ServerRequest} req + * @param {ServerResponse} res + */ +async function tokenConfigController(req, res) { + try { + const appConfig = req.config; + const includePricing = appConfig?.interfaceConfig?.contextCost === true; + const modelsConfig = await getModelsConfig(req); + + /** @type {Record} */ + const endpointTokenConfigs = {}; + const customEndpoints = appConfig?.endpoints?.[EModelEndpoint.custom] ?? []; + const cache = tokenConfigCache(); + for (const endpointConfig of customEndpoints) { + const name = endpointConfig?.name; + if (!name) { + continue; + } + if (endpointConfig.tokenConfig != null) { + endpointTokenConfigs[name] = endpointConfig.tokenConfig; + continue; + } + const tokenKey = getTokenConfigKey(endpointConfig, name, req.user.id); + const cached = await cache.get(tokenKey); + if (cached) { + endpointTokenConfigs[name] = cached; + } + } + + const tokenConfigMap = buildTokenConfigMap( + { modelsConfig, endpointTokenConfigs, includePricing }, + { getValueKey, getMultiplier, getCacheMultiplier }, + ); + res.json(tokenConfigMap); + } catch (error) { + logger.error('[tokenConfigController]', error); + res.status(500).json({ error: 'Failed to resolve token config' }); + } +} + +module.exports = tokenConfigController; diff --git a/api/server/controllers/agents/__tests__/usageEvents.integration.spec.js b/api/server/controllers/agents/__tests__/usageEvents.integration.spec.js new file mode 100644 index 0000000000..9e20ab5ae7 --- /dev/null +++ b/api/server/controllers/agents/__tests__/usageEvents.integration.spec.js @@ -0,0 +1,252 @@ +const { z } = require('zod'); +const { tool } = require('@langchain/core/tools'); +const { ChatGenerationChunk } = require('@langchain/core/outputs'); +const { HumanMessage, AIMessageChunk } = require('@langchain/core/messages'); +const { + Run, + Providers, + GraphEvents, + FakeChatModel, + createContentAggregator, +} = require('@librechat/agents'); +const { GenerationJobManager } = require('@librechat/api'); +const { getDefaultHandlers } = require('~/server/controllers/agents/callbacks'); + +jest.mock('nanoid', () => ({ + nanoid: jest.fn(() => 'mock-nanoid'), +})); + +jest.mock('~/server/services/Files/Citations', () => ({ + processFileCitations: jest.fn(), +})); + +jest.mock('~/server/services/Files/Code/process', () => ({ + processCodeOutput: jest.fn(), + runPreviewFinalize: jest.fn(), +})); + +jest.mock('~/server/services/Files/process', () => ({ + saveBase64Image: jest.fn(), +})); + +/** Real pipeline guard: published lib versions without the event skip its assertions */ +const hasContextUsageEvent = GraphEvents.ON_CONTEXT_USAGE != null; + +/** + * FakeChatModel that attaches provider-style usage_metadata on a final + * empty chunk (the OpenAI streaming pattern), so CHAT_MODEL_END carries + * aggregated usage through the real @librechat/agents pipeline. + */ +class UsageFakeModel extends FakeChatModel { + constructor(options, usagePerCall) { + super(options); + this.usagePerCall = usagePerCall; + this.usageCallIndex = 0; + } + + async *_streamResponseChunks(messages, options, runManager) { + yield* super._streamResponseChunks(messages, options, runManager); + const index = Math.min(this.usageCallIndex, this.usagePerCall.length - 1); + this.usageCallIndex += 1; + yield new ChatGenerationChunk({ + text: '', + message: new AIMessageChunk({ content: '', usage_metadata: this.usagePerCall[index] }), + }); + } +} + +const addTool = tool(async ({ a, b }) => String(a + b), { + name: 'add', + description: 'Add two numbers', + schema: z.object({ a: z.number(), b: z.number() }), +}); + +const charCounter = (msg) => { + const content = msg.content; + if (typeof content === 'string') { + return content.length + 3; + } + if (Array.isArray(content)) { + let length = 3; + for (const part of content) { + if (typeof part === 'string') { + length += part.length; + } else if (typeof part?.text === 'string') { + length += part.text.length; + } + } + return length; + } + return 3; +}; + +function createMockRes() { + const events = []; + return { + events, + headersSent: true, + writableEnded: false, + write(payload) { + for (const line of String(payload).split('\n')) { + if (line.startsWith('data: ')) { + events.push(JSON.parse(line.slice(6))); + } + } + return true; + }, + }; +} + +const FIRST_CALL_USAGE = { + input_tokens: 100, + output_tokens: 20, + total_tokens: 120, +}; + +const SECOND_CALL_USAGE = { + input_tokens: 150, + output_tokens: 10, + total_tokens: 160, + input_token_details: { cache_creation: 30, cache_read: 50 }, +}; + +const MAX_CONTEXT_TOKENS = 8000; + +async function runToolLoop({ res, streamId = null, collectedUsage }) { + const { contentParts, aggregateContent } = createContentAggregator(); + const handlers = getDefaultHandlers({ + res, + aggregateContent, + toolEndCallback: () => {}, + collectedUsage, + streamId, + }); + + const run = await Run.create({ + runId: 'usage-e2e-response', + graphConfig: { + type: 'standard', + llmConfig: { + provider: Providers.OPENAI, + model: 'gpt-4o-mini', + streaming: true, + streamUsage: false, + }, + instructions: 'You are a helpful assistant.', + maxContextTokens: MAX_CONTEXT_TOKENS, + tools: [addTool], + }, + returnContent: true, + customHandlers: handlers, + tokenCounter: charCounter, + indexTokenCountMap: {}, + }); + + run.Graph.overrideModel = new UsageFakeModel( + { + responses: ['Let me calculate that.', 'The answer is 4.'], + toolCalls: [{ name: 'add', args: { a: 2, b: 2 }, id: 'tc_1', type: 'tool_call' }], + }, + [FIRST_CALL_USAGE, SECOND_CALL_USAGE], + ); + + await run.processStream( + { messages: [new HumanMessage('What is 2+2?')] }, + { + configurable: { thread_id: 'usage-e2e-thread', user_id: 'user-1' }, + streamMode: 'values', + version: 'v2', + }, + ); + + return { run, contentParts }; +} + +describe('usage events through the real agents pipeline', () => { + jest.setTimeout(30000); + + afterAll(async () => { + await GenerationJobManager.destroy(); + }); + + test('emits on_token_usage per model call with collectedUsage parity', async () => { + const res = createMockRes(); + const collectedUsage = []; + const { contentParts } = await runToolLoop({ res, collectedUsage }); + + const usageEvents = res.events.filter((e) => e.event === 'on_token_usage'); + expect(usageEvents).toHaveLength(2); + + expect(usageEvents[0].data).toMatchObject(FIRST_CALL_USAGE); + expect(usageEvents[1].data).toMatchObject(SECOND_CALL_USAGE); + expect(usageEvents[0].data.provider).toBe(Providers.OPENAI); + expect(usageEvents[0].data.model).toBeTruthy(); + expect(usageEvents[0].data.usage_type).toBeUndefined(); + + expect(collectedUsage).toHaveLength(2); + expect(collectedUsage[0]).toMatchObject(FIRST_CALL_USAGE); + expect(collectedUsage[1]).toMatchObject(SECOND_CALL_USAGE); + + const text = contentParts + .filter((part) => part?.type === 'text') + .map((part) => part.text) + .join(''); + expect(text).toContain('The answer is 4.'); + }); + + test('emits a context snapshot before each model call', async () => { + if (!hasContextUsageEvent) { + console.warn('Skipping: installed @librechat/agents predates ON_CONTEXT_USAGE'); + return; + } + const res = createMockRes(); + const { run } = await runToolLoop({ res, collectedUsage: [] }); + expect(run).toBeDefined(); + + const contextEvents = res.events.filter((e) => e.event === 'on_context_usage'); + expect(contextEvents).toHaveLength(2); + + for (const event of contextEvents) { + const { breakdown, contextBudget, remainingContextTokens, effectiveInstructionTokens } = + event.data; + expect(breakdown.maxContextTokens).toBe(MAX_CONTEXT_TOKENS); + expect(contextBudget).toBeGreaterThan(0); + expect(contextBudget).toBeLessThanOrEqual(MAX_CONTEXT_TOKENS); + expect(effectiveInstructionTokens).toBeGreaterThan(0); + expect(remainingContextTokens).toBeGreaterThan(0); + expect(remainingContextTokens).toBeLessThan(contextBudget); + } + + /** Tool loop grows the context between calls */ + expect(contextEvents[1].data.prePruneContextTokens).toBeGreaterThan( + contextEvents[0].data.prePruneContextTokens, + ); + + /** Snapshot precedes the call's usage event */ + const firstContextIndex = res.events.findIndex((e) => e.event === 'on_context_usage'); + const firstUsageIndex = res.events.findIndex((e) => e.event === 'on_token_usage'); + expect(firstContextIndex).toBeGreaterThanOrEqual(0); + expect(firstContextIndex).toBeLessThan(firstUsageIndex); + }); + + test('persists usage and context snapshot for resume via GenerationJobManager', async () => { + const streamId = `usage-e2e-stream-${Date.now()}`; + await GenerationJobManager.createJob(streamId, 'user-1', 'convo-1'); + + const res = createMockRes(); + await runToolLoop({ res, streamId, collectedUsage: [] }); + + const resumeState = await GenerationJobManager.getResumeState(streamId); + expect(resumeState).not.toBeNull(); + + expect(resumeState.collectedUsage).toHaveLength(2); + expect(resumeState.collectedUsage[0]).toMatchObject(FIRST_CALL_USAGE); + expect(resumeState.collectedUsage[1]).toMatchObject(SECOND_CALL_USAGE); + + if (hasContextUsageEvent) { + expect(resumeState.contextUsage.breakdown.maxContextTokens).toBe(MAX_CONTEXT_TOKENS); + /** Latest-wins: the persisted snapshot is the second call's */ + expect(resumeState.contextUsage.prePruneContextTokens).toBeGreaterThan(0); + } + }); +}); diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js index 0a8f8854f7..0ddc8d6939 100644 --- a/api/server/controllers/agents/callbacks.js +++ b/api/server/controllers/agents/callbacks.js @@ -1,6 +1,12 @@ const { nanoid } = require('nanoid'); const { logger } = require('@librechat/data-schemas'); -const { Tools, StepTypes, FileContext, ErrorTypes } = require('librechat-data-provider'); +const { + Tools, + StepTypes, + FileContext, + ErrorTypes, + UsageEvents, +} = require('librechat-data-provider'); const { GraphEvents, GraphNodeKeys, @@ -41,13 +47,16 @@ class ModelEndHandler { * Optional; when `null`, the handler is a no-op for signatures. Non-Vertex * providers don't emit `additional_kwargs.signatures`, so capture is also * a no-op for them even when the map is provided. + * @param {(data: Record) => Promise | void} [emitUsage] Optional + * callback to stream per-call token usage to the client. */ - constructor(collectedUsage, collectedThoughtSignatures = null) { + constructor(collectedUsage, collectedThoughtSignatures = null, emitUsage = null) { if (!Array.isArray(collectedUsage)) { throw new Error('collectedUsage must be an array'); } this.collectedUsage = collectedUsage; this.collectedThoughtSignatures = collectedThoughtSignatures; + this.emitUsage = emitUsage; } finalize(errorMessage) { @@ -104,6 +113,19 @@ class ModelEndHandler { this.collectedUsage.push(taggedUsage); + if (this.emitUsage) { + await this.emitUsage({ + input_tokens: taggedUsage.input_tokens, + output_tokens: taggedUsage.output_tokens, + total_tokens: taggedUsage.total_tokens, + input_token_details: taggedUsage.input_token_details, + model: taggedUsage.model, + provider: taggedUsage.provider, + usage_type: taggedUsage.usage_type, + runId: metadata?.run_id, + }); + } + /** * `additional_kwargs.signatures` is a flat array indexed by response * part position (text + functionCall interleaved). `tool_calls` is @@ -240,7 +262,11 @@ function getDefaultHandlers({ ); } const handlers = { - [GraphEvents.CHAT_MODEL_END]: new ModelEndHandler(collectedUsage, collectedThoughtSignatures), + [GraphEvents.CHAT_MODEL_END]: new ModelEndHandler( + collectedUsage, + collectedThoughtSignatures, + (data) => emitEvent(res, streamId, { event: UsageEvents.ON_TOKEN_USAGE, data }), + ), [GraphEvents.TOOL_END]: new ToolEndHandler(toolEndCallback, logger), [GraphEvents.ON_RUN_STEP]: { /** @@ -425,6 +451,20 @@ function getDefaultHandlers({ handlers[GraphEvents.ON_AGENT_LOG] = { handle: agentLogHandler }; + /** Guarded: no-op when the installed @librechat/agents predates the event */ + if (GraphEvents.ON_CONTEXT_USAGE) { + handlers[GraphEvents.ON_CONTEXT_USAGE] = { + /** + * Forward per-model-call context usage snapshots to the client. + * @param {string} event - The event name. + * @param {StreamEventData} data - The event data. + */ + handle: async (event, data) => { + await emitEvent(res, streamId, { event, data }); + }, + }; + } + return handlers; } diff --git a/api/server/routes/endpoints.js b/api/server/routes/endpoints.js index e7ff1c7000..03f225f560 100644 --- a/api/server/routes/endpoints.js +++ b/api/server/routes/endpoints.js @@ -1,9 +1,11 @@ const express = require('express'); const requireJwtAuth = require('~/server/middleware/requireJwtAuth'); const endpointController = require('~/server/controllers/EndpointController'); +const tokenConfigController = require('~/server/controllers/TokenConfigController'); const router = express.Router(); /** Auth required for role/tenant-scoped endpoint config resolution. */ router.get('/', requireJwtAuth, endpointController); +router.get('/token-config', requireJwtAuth, tokenConfigController); module.exports = router; diff --git a/client/src/components/Chat/Input/ChatForm.tsx b/client/src/components/Chat/Input/ChatForm.tsx index f8c058cae0..a0a2822187 100644 --- a/client/src/components/Chat/Input/ChatForm.tsx +++ b/client/src/components/Chat/Input/ChatForm.tsx @@ -5,12 +5,6 @@ import { useRecoilState, useRecoilValue } from 'recoil'; import { Constants, isAssistantsEndpoint, isAgentsEndpoint } from 'librechat-data-provider'; import type { TConversation } from 'librechat-data-provider'; import type { ExtendedFile, FileSetter, ConvoGenerator } from '~/common'; -import { - useChatContext, - useChatFormContext, - useAddedChatContext, - useAssistantsMapContext, -} from '~/Providers'; import { useTextarea, useAutoSave, @@ -21,6 +15,12 @@ import { useSubmitMessage, useFocusChatEffect, } from '~/hooks'; +import { + useChatContext, + useChatFormContext, + useAddedChatContext, + useAssistantsMapContext, +} from '~/Providers'; import PendingManualSkillsChips from './PendingManualSkillsChips'; import { cn, getModelSpec, removeFocusRings } from '~/utils'; import { useGetStartupConfig } from '~/data-provider'; @@ -28,11 +28,12 @@ import { mainTextareaId, BadgeItem } from '~/common'; import AttachFileChat from './Files/AttachFileChat'; import FileFormChat from './Files/FileFormChat'; import TextareaHeader from './TextareaHeader'; -import SkillsCommand from './SkillsCommand'; import PromptsCommand from './PromptsCommand'; +import SkillsCommand from './SkillsCommand'; import AudioRecorder from './AudioRecorder'; import CollapseChat from './CollapseChat'; import StreamAudio from './StreamAudio'; +import TokenUsage from './TokenUsage'; import StopButton from './StopButton'; import SendButton from './SendButton'; import EditBadges from './EditBadges'; @@ -379,6 +380,7 @@ const ChatForm = memo(function ChatForm({ } />
+ {SpeechToText && ( 0 ? Math.min((value / max) * 100, 100) : null; + return ( +
+ {label} + + {formatTokens(value)} + {percent != null && ( + + )} + +
+ ); +} + +interface BreakdownProps { + view: TokenUsageView; + showCost: boolean; +} + +export default function Breakdown({ view, showCost }: BreakdownProps) { + const localize = useLocalize(); + const { usedTokens, maxTokens, percent, snapshot, snapshotActive, usageTotals } = view; + + const breakdown = snapshotActive ? snapshot?.breakdown : undefined; + const instructionTokens = + snapshot?.effectiveInstructionTokens ?? breakdown?.instructionTokens ?? 0; + const systemTokens = + (breakdown?.systemMessageTokens ?? 0) + (breakdown?.dynamicInstructionTokens ?? 0); + const messageTokens = Math.max(0, usedTokens - instructionTokens); + const freeTokens = maxTokens != null ? Math.max(0, maxTokens - usedTokens) : null; + + return ( +
+
+ + {localize('com_ui_context_window')} + + + {maxTokens != null + ? `${formatTokens(usedTokens)} / ${formatTokens(maxTokens)} (${Math.round(percent)}%)` + : formatTokens(usedTokens)} + +
+ +
+
+
+ +
+ {breakdown ? ( + <> + + + + {breakdown.summaryTokens > 0 && ( + + )} + {freeTokens != null && ( + + )} + + ) : ( + <> + + + {maxTokens == null && ( +

{localize('com_ui_context_unknown')}

+ )} +

{localize('com_ui_estimated')}

+ + )} +
+ + {usageTotals.eventCount > 0 && ( + <> +
+
+ + + {usageTotals.cacheRead > 0 && ( + + )} + {usageTotals.cacheWrite > 0 && ( + + )} +
+ + )} + + {showCost && view.costUSD != null && ( + <> +
+
+ {localize('com_ui_session_cost')} + {formatCost(view.costUSD)} +
+ + )} +
+ ); +} diff --git a/client/src/components/Chat/Input/TokenUsage/Gauge.tsx b/client/src/components/Chat/Input/TokenUsage/Gauge.tsx new file mode 100644 index 0000000000..ba7e420302 --- /dev/null +++ b/client/src/components/Chat/Input/TokenUsage/Gauge.tsx @@ -0,0 +1,62 @@ +import { cn } from '~/utils'; + +const SIZE = 28; +const STROKE_WIDTH = 3.5; +const RADIUS = (SIZE - STROKE_WIDTH) / 2; +const CIRCUMFERENCE = 2 * Math.PI * RADIUS; + +interface GaugeProps { + /** 0–100, clamped by the caller */ + percent: number; + /** Max context unknown — render an empty track only */ + indeterminate: boolean; +} + +function getStrokeClass(percent: number, indeterminate: boolean): string { + if (indeterminate) { + return 'stroke-text-secondary'; + } + if (percent > 90) { + return 'stroke-red-500'; + } + if (percent > 75) { + return 'stroke-yellow-500'; + } + return 'stroke-text-secondary'; +} + +export default function Gauge({ percent, indeterminate }: GaugeProps) { + const offset = CIRCUMFERENCE - (percent / 100) * CIRCUMFERENCE; + + return ( + + ); +} diff --git a/client/src/components/Chat/Input/TokenUsage/index.tsx b/client/src/components/Chat/Input/TokenUsage/index.tsx new file mode 100644 index 0000000000..e0140f6f28 --- /dev/null +++ b/client/src/components/Chat/Input/TokenUsage/index.tsx @@ -0,0 +1,83 @@ +import { memo } from 'react'; +import { HoverCard, HoverCardTrigger, HoverCardContent, HoverCardPortal } from '@librechat/client'; +import type { TConversation } from 'librechat-data-provider'; +import useTokenUsage from '~/hooks/Chat/useTokenUsage'; +import { useGetStartupConfig } from '~/data-provider'; +import { formatTokens } from '~/utils'; +import { useLocalize } from '~/hooks'; +import Breakdown from './Breakdown'; +import Gauge from './Gauge'; + +interface TokenUsageProps { + index: number; + conversation: TConversation | null; + isSubmitting: boolean; +} + +function TokenUsageIndicator({ + index, + conversation, + isSubmitting, + showCost, +}: TokenUsageProps & { + showCost: boolean; +}) { + const localize = useLocalize(); + const view = useTokenUsage({ index, conversation, isSubmitting }); + + if (view.usedTokens <= 0 && view.maxTokens == null) { + return null; + } + + const hasMax = view.maxTokens != null && view.maxTokens > 0; + const ariaLabel = hasMax + ? localize('com_ui_context_usage_label', { + 0: formatTokens(view.usedTokens), + 1: formatTokens(view.maxTokens ?? 0), + 2: String(Math.round(view.percent)), + }) + : localize('com_ui_context_usage_label_unknown', { 0: formatTokens(view.usedTokens) }); + + return ( + + + + + + + + + + + ); +} + +/** Config gate kept outside the indicator so disabled deployments mount nothing */ +const TokenUsage = memo(function TokenUsage(props: TokenUsageProps) { + const { data: startupConfig } = useGetStartupConfig(); + if (startupConfig?.interface?.contextUsage === false) { + return null; + } + return ( + + ); +}); + +export default TokenUsage; diff --git a/client/src/data-provider/Endpoints/queries.ts b/client/src/data-provider/Endpoints/queries.ts index 17527c6dc1..a23a7c7fc5 100644 --- a/client/src/data-provider/Endpoints/queries.ts +++ b/client/src/data-provider/Endpoints/queries.ts @@ -1,6 +1,6 @@ import { useRecoilValue } from 'recoil'; -import { QueryKeys, dataService } from 'librechat-data-provider'; import { useQuery } from '@tanstack/react-query'; +import { QueryKeys, dataService } from 'librechat-data-provider'; import type { QueryObserverResult, UseQueryOptions } from '@tanstack/react-query'; import type t from 'librechat-data-provider'; import store from '~/store'; @@ -23,6 +23,20 @@ export const useGetEndpointsQuery = ( ); }; +export const useTokenConfigQuery = ( + config?: UseQueryOptions, +): QueryObserverResult => { + const queriesEnabled = useRecoilValue(store.queriesEnabled); + return useQuery([QueryKeys.tokenConfig], () => dataService.getTokenConfig(), { + staleTime: Infinity, + refetchOnWindowFocus: false, + refetchOnReconnect: false, + refetchOnMount: false, + ...config, + enabled: (config?.enabled ?? true) === true && queriesEnabled, + }); +}; + /** * Auth-aware query key so unauthenticated (login page) and authenticated * (chat page) configs are cached independently, preventing stale diff --git a/client/src/hooks/Chat/index.ts b/client/src/hooks/Chat/index.ts index 11754c6c5f..9d9af790ce 100644 --- a/client/src/hooks/Chat/index.ts +++ b/client/src/hooks/Chat/index.ts @@ -1,4 +1,6 @@ export { default as useChatHelpers } from './useChatHelpers'; +export { default as useTokenLimits } from './useTokenLimits'; +export { default as useTokenUsage } from './useTokenUsage'; export { default as useAddedResponse } from './useAddedResponse'; export { default as useChatFunctions } from './useChatFunctions'; export { default as useGetAddedConvo } from './useGetAddedConvo'; diff --git a/client/src/hooks/Chat/useTokenLimits.ts b/client/src/hooks/Chat/useTokenLimits.ts new file mode 100644 index 0000000000..053d96c4ec --- /dev/null +++ b/client/src/hooks/Chat/useTokenLimits.ts @@ -0,0 +1,59 @@ +import { useMemo } from 'react'; +import { isAgentsEndpoint } from 'librechat-data-provider'; +import type { TConversation, TModelTokenomics } from 'librechat-data-provider'; +import { useGetStartupConfig, useTokenConfigQuery, useGetAgentByIdQuery } from '~/data-provider'; +import { getModelSpec } from '~/utils'; + +export interface TokenLimits { + /** Statically resolved max context; live snapshots override this at run time */ + maxContextTokens?: number; + rates?: TModelTokenomics; + endpoint?: string; + model?: string; +} + +function toNumber(value: unknown): number | undefined { + const num = typeof value === 'string' ? parseFloat(value) : value; + return typeof num === 'number' && Number.isFinite(num) && num > 0 ? num : undefined; +} + +/** + * Mirrors the backend resolution chain (packages/api agents/initialize.ts): + * explicit conversation setting → agent params → model spec preset → + * server-resolved token config lookup. + */ +export default function useTokenLimits(conversation: TConversation | null): TokenLimits { + const { data: startupConfig } = useGetStartupConfig(); + const { data: tokenConfig } = useTokenConfigQuery(); + + const endpoint = conversation?.endpoint ?? ''; + const agentId = isAgentsEndpoint(endpoint) ? conversation?.agent_id : null; + const { data: agent } = useGetAgentByIdQuery(agentId); + + const spec = conversation?.spec; + const model = conversation?.model; + const maxContextSetting = conversation?.maxContextTokens; + + return useMemo(() => { + const specPreset = getModelSpec({ specName: spec, startupConfig })?.preset; + + let lookupEndpoint = endpoint; + let lookupModel = model ?? ''; + if (agent) { + lookupEndpoint = agent.provider ?? endpoint; + lookupModel = agent.model ?? lookupModel; + } else if (specPreset) { + lookupEndpoint = specPreset.endpoint ?? lookupEndpoint; + lookupModel = lookupModel || (specPreset.model ?? ''); + } + + const rates = tokenConfig?.[lookupEndpoint]?.[lookupModel]; + const maxContextTokens = + toNumber(maxContextSetting) ?? + toNumber(agent?.model_parameters?.maxContextTokens) ?? + toNumber(specPreset?.maxContextTokens) ?? + rates?.context; + + return { maxContextTokens, rates, endpoint: lookupEndpoint, model: lookupModel }; + }, [endpoint, model, spec, maxContextSetting, agent, startupConfig, tokenConfig]); +} diff --git a/client/src/hooks/Chat/useTokenUsage.ts b/client/src/hooks/Chat/useTokenUsage.ts new file mode 100644 index 0000000000..6ca1a4782f --- /dev/null +++ b/client/src/hooks/Chat/useTokenUsage.ts @@ -0,0 +1,151 @@ +import { useEffect, useMemo, useRef } from 'react'; +import { useAtomValue, useSetAtom } from 'jotai'; +import { useQueryClient } from '@tanstack/react-query'; +import { Constants, QueryKeys } from 'librechat-data-provider'; +import type { TMessage, TConversation, TModelTokenomics } from 'librechat-data-provider'; +import type { ContextSnapshot, UsageTotals } from '~/store/usage'; +import type { BranchTotals } from '~/utils/tokens'; +import { + liveTokensFamily, + usageTotalsFamily, + branchTotalsFamily, + contextSnapshotFamily, +} from '~/store/usage'; +import { useLatestMessageId } from '~/hooks/Messages/useLatestMessage'; +import { buildIndex, sumBranch } from '~/utils'; +import useTokenLimits from './useTokenLimits'; + +export interface TokenUsageParams { + index: number; + conversation: TConversation | null; + isSubmitting: boolean; +} + +export interface TokenUsageView { + usedTokens: number; + maxTokens?: number; + /** 0–100, clamped; 0 when max is unknown */ + percent: number; + /** True when derived from per-message counts instead of a backend snapshot */ + isEstimate: boolean; + snapshot: ContextSnapshot | null; + snapshotActive: boolean; + branchTotals: BranchTotals; + usageTotals: UsageTotals; + liveTokens: number; + rates?: TModelTokenomics; + /** Session cost from provider-reported usage; undefined until usage events arrive */ + costUSD?: number; +} + +/** + * View-model for the context usage indicator. Mount only inside the + * indicator so its subscriptions never re-render the chat tree. + */ +export default function useTokenUsage({ + index, + conversation, + isSubmitting, +}: TokenUsageParams): TokenUsageView { + const queryClient = useQueryClient(); + const conversationKey = conversation?.conversationId ?? Constants.NEW_CONVO; + + const tailId = useLatestMessageId(index); + const snapshot = useAtomValue(contextSnapshotFamily(conversationKey)); + const usageTotals = useAtomValue(usageTotalsFamily(conversationKey)); + const branchTotals = useAtomValue(branchTotalsFamily(conversationKey)); + const liveTokens = useAtomValue(liveTokensFamily(conversationKey)); + const setBranchTotals = useSetAtom(branchTotalsFamily(conversationKey)); + const limits = useTokenLimits(conversation); + + const isSubmittingRef = useRef(isSubmitting); + isSubmittingRef.current = isSubmitting; + const tailIdRef = useRef(tailId); + tailIdRef.current = tailId; + const anchorId = snapshot?.anchorMessageId ?? null; + const anchorIdRef = useRef(anchorId); + anchorIdRef.current = anchorId; + + useEffect(() => { + /** Cache `updated` events fire on every state transition — rebuild the + * O(n) index only when the data snapshot reference actually changed */ + let lastIndexed: TMessage[] | undefined; + const rebuild = (messages?: TMessage[]) => { + if (messages === lastIndexed && messages !== undefined) { + return; + } + lastIndexed = messages; + buildIndex(conversationKey, messages); + setBranchTotals(sumBranch(conversationKey, tailIdRef.current, anchorIdRef.current)); + }; + + rebuild(queryClient.getQueryData([QueryKeys.messages, conversationKey])); + + const unsubscribe = queryClient.getQueryCache().subscribe((event) => { + if (isSubmittingRef.current || event.type !== 'updated') { + return; + } + const queryKey = event.query.queryKey; + if ( + !Array.isArray(queryKey) || + queryKey[0] !== QueryKeys.messages || + queryKey[1] !== conversationKey + ) { + return; + } + rebuild(event.query.state.data as TMessage[] | undefined); + }); + return unsubscribe; + }, [conversationKey, queryClient, setBranchTotals]); + + useEffect(() => { + setBranchTotals(sumBranch(conversationKey, tailId, anchorId)); + }, [conversationKey, tailId, anchorId, setBranchTotals]); + + return useMemo(() => { + const snapshotActive = + snapshot != null && + (isSubmitting || snapshot.anchorMessageId == null || branchTotals.containsAnchor); + + if (snapshotActive && snapshot) { + const breakdown = snapshot.breakdown; + const maxTokens = snapshot.contextBudget ?? breakdown.maxContextTokens; + const instructionTokens = snapshot.effectiveInstructionTokens ?? breakdown.instructionTokens; + const baseUsed = + snapshot.remainingContextTokens != null + ? maxTokens - snapshot.remainingContextTokens + : instructionTokens + breakdown.messageTokens; + const usedTokens = Math.max(0, baseUsed) + liveTokens; + return { + usedTokens, + maxTokens, + percent: maxTokens > 0 ? Math.min((usedTokens / maxTokens) * 100, 100) : 0, + isEstimate: false, + snapshot, + snapshotActive, + branchTotals, + usageTotals, + liveTokens, + rates: limits.rates, + costUSD: usageTotals.eventCount > 0 ? usageTotals.costUSD : undefined, + }; + } + + const usedTokens = branchTotals.input + branchTotals.output + liveTokens; + const maxTokens = limits.maxContextTokens; + return { + usedTokens, + maxTokens, + percent: + maxTokens != null && maxTokens > 0 ? Math.min((usedTokens / maxTokens) * 100, 100) : 0, + isEstimate: true, + snapshot, + snapshotActive: false, + branchTotals, + usageTotals, + liveTokens, + rates: limits.rates, + costUSD: usageTotals.eventCount > 0 ? usageTotals.costUSD : undefined, + }; + }, [snapshot, isSubmitting, branchTotals, usageTotals, liveTokens, limits]); +} diff --git a/client/src/hooks/SSE/index.ts b/client/src/hooks/SSE/index.ts index 2829db76f6..5de4866d49 100644 --- a/client/src/hooks/SSE/index.ts +++ b/client/src/hooks/SSE/index.ts @@ -3,5 +3,6 @@ export { default as useResumableSSE } from './useResumableSSE'; export { default as useAdaptiveSSE } from './useAdaptiveSSE'; export { default as useResumeOnLoad } from './useResumeOnLoad'; export { default as useStepHandler } from './useStepHandler'; +export { default as useUsageHandler } from './useUsageHandler'; export { default as useContentHandler } from './useContentHandler'; export { default as useAttachmentHandler } from './useAttachmentHandler'; diff --git a/client/src/hooks/SSE/useResumableSSE.ts b/client/src/hooks/SSE/useResumableSSE.ts index 43ae4d2336..fd8cb7f2d4 100644 --- a/client/src/hooks/SSE/useResumableSSE.ts +++ b/client/src/hooks/SSE/useResumableSSE.ts @@ -10,6 +10,7 @@ import { ErrorTypes, StepEvents, apiBaseUrl, + UsageEvents, createPayload, ViolationTypes, removeNullishValues, @@ -29,6 +30,7 @@ import { clearAllDrafts, removeConvoFromAllQueries, upsertConvoInAllQueries, + countTrailingOutputChars, markStreamStartFailedMetadata, } from '~/utils'; import { @@ -39,6 +41,7 @@ import { } from '~/data-provider'; import useEventHandlers, { buildCreatedInitialResponse } from './useEventHandlers'; import { useAuthContext } from '~/hooks/AuthContext'; +import useUsageHandler from './useUsageHandler'; import store from '~/store'; type ChatHelpers = Pick< @@ -473,6 +476,15 @@ export default function useResumableSSE( const balanceQuery = useGetUserBalance({ enabled: !!isAuthenticated && startupConfig?.balance?.enabled, }); + const { + contextHandler, + usageHandler, + tapStream, + finalizeUsage, + backfillUsage, + resetLive, + seedLive, + } = useUsageHandler(); /** * Subscribe to stream via SSE library (supports custom headers) @@ -530,6 +542,7 @@ export default function useResumableSSE( } try { finalHandler(data, currentSubmission as EventSubmission); + finalizeUsage(data, { ...currentSubmission, userMessage }); } catch (error) { console.error('[ResumableSSE] Error in finalHandler:', error); setIsSubmitting(false); @@ -589,7 +602,23 @@ export default function useResumableSSE( return; } + if (data.event === UsageEvents.ON_CONTEXT_USAGE) { + contextHandler(data.data, { ...currentSubmission, userMessage }); + return; + } + + if (data.event === UsageEvents.ON_TOKEN_USAGE) { + usageHandler(data.data, { ...currentSubmission, userMessage }); + return; + } + if (data.event != null) { + if ( + data.event === StepEvents.ON_MESSAGE_DELTA || + data.event === StepEvents.ON_REASONING_DELTA + ) { + tapStream(data.data, { ...currentSubmission, userMessage }); + } if (!isResume && !createdStreamIdsRef.current.has(currentStreamId)) { if (isOAuthStepEvent(data)) { preCreatedStepEvents.push(data); @@ -619,6 +648,24 @@ export default function useResumableSSE( currentSubmission = resumeSubmission; submissionRef.current = resumeSubmission; userMessage = resumeSubmission.userMessage; + /** + * The run's collected usage is the source of truth at sync time: + * pre-snapshot events are never replayed (only represented via + * aggregated content), so totals are rebuilt from the backfill + * and token-usage replay events are skipped when it's present. + */ + const backfilledUsage = data.resumeState?.collectedUsage; + const hasUsageBackfill = (backfilledUsage?.length ?? 0) > 0; + backfillUsage(backfilledUsage ?? [], resumeSubmission); + if (data.resumeState?.contextUsage) { + contextHandler(data.resumeState.contextUsage, resumeSubmission); + } + /** Output streamed before this resume is not re-delivered as + * deltas — estimate it from the trailing aggregated content */ + seedLive( + countTrailingOutputChars(data.resumeState?.aggregatedContent), + resumeSubmission, + ); if (data.resumeState?.runSteps) { for (const runStep of data.resumeState.runSteps) { @@ -700,7 +747,13 @@ export default function useResumableSSE( `[ResumableSSE] Replaying ${data.resumeState.replayEvents.length} resume events`, ); for (const replayEvent of data.resumeState.replayEvents) { - if (replayEvent.event != null) { + if (replayEvent.event === UsageEvents.ON_CONTEXT_USAGE) { + contextHandler(replayEvent.data, resumeSubmission); + } else if (replayEvent.event === UsageEvents.ON_TOKEN_USAGE) { + if (!hasUsageBackfill) { + usageHandler(replayEvent.data, resumeSubmission); + } + } else if (replayEvent.event != null) { stepHandler(replayEvent, resumeSubmission); } } @@ -711,6 +764,12 @@ export default function useResumableSSE( for (const pendingEvent of data.pendingEvents) { if (pendingEvent.event === 'title') { titleHandler(pendingEvent); + } else if (pendingEvent.event === UsageEvents.ON_CONTEXT_USAGE) { + contextHandler(pendingEvent.data, resumeSubmission); + } else if (pendingEvent.event === UsageEvents.ON_TOKEN_USAGE) { + if (!hasUsageBackfill) { + usageHandler(pendingEvent.data, resumeSubmission); + } } else if (pendingEvent.event != null) { stepHandler(pendingEvent, resumeSubmission); } else if (pendingEvent.type != null) { @@ -821,6 +880,7 @@ export default function useResumableSSE( console.log('[ResumableSSE] Server-sent error event received:', e.data); sse.close(); removeActiveJob(currentStreamId); + resetLive({ ...currentSubmission, userMessage }); if ( !createdStreamIdsRef.current.has(currentStreamId) && optimisticStreamIdsRef.current.has(currentStreamId) @@ -994,6 +1054,13 @@ export default function useResumableSSE( balanceQuery, removeActiveJob, queryClient, + contextHandler, + usageHandler, + tapStream, + finalizeUsage, + backfillUsage, + resetLive, + seedLive, ], ); diff --git a/client/src/hooks/SSE/useSSE.ts b/client/src/hooks/SSE/useSSE.ts index 509b7c213e..9a2ba55870 100644 --- a/client/src/hooks/SSE/useSSE.ts +++ b/client/src/hooks/SSE/useSSE.ts @@ -2,13 +2,20 @@ import { useEffect, useState } from 'react'; import { v4 } from 'uuid'; import { SSE } from 'sse.js'; import { useSetRecoilState } from 'recoil'; -import { request, createPayload, removeNullishValues } from 'librechat-data-provider'; +import { + request, + UsageEvents, + StepEvents, + createPayload, + removeNullishValues, +} from 'librechat-data-provider'; import type { TMessage, TPayload, TSubmission, EventSubmission } from 'librechat-data-provider'; import type { EventHandlerParams } from './useEventHandlers'; import type { TResData } from '~/common'; import { useGetStartupConfig, useGetUserBalance } from '~/data-provider'; import { useAuthContext } from '~/hooks/AuthContext'; import useEventHandlers from './useEventHandlers'; +import useUsageHandler from './useUsageHandler'; import { clearAllDrafts } from '~/utils'; import store from '~/store'; @@ -60,6 +67,7 @@ export default function useSSE( const balanceQuery = useGetUserBalance({ enabled: !!isAuthenticated && startupConfig?.balance?.enabled, }); + const { contextHandler, usageHandler, tapStream, finalizeUsage, resetLive } = useUsageHandler(); useEffect(() => { if (submission == null || Object.keys(submission).length === 0) { @@ -96,6 +104,7 @@ export default function useSSE( clearAllDrafts(submission.conversation?.conversationId); try { finalHandler(data, submission as EventSubmission); + finalizeUsage(data, { ...submission, userMessage }); } catch (error) { console.error('Error in finalHandler:', error); setIsSubmitting(false); @@ -116,7 +125,17 @@ export default function useSSE( createdHandler(data, { ...submission, userMessage } as EventSubmission); } else if (data.event === 'title') { titleHandler(data); + } else if (data.event === UsageEvents.ON_CONTEXT_USAGE) { + contextHandler(data.data, { ...submission, userMessage }); + } else if (data.event === UsageEvents.ON_TOKEN_USAGE) { + usageHandler(data.data, { ...submission, userMessage }); } else if (data.event != null) { + if ( + data.event === StepEvents.ON_MESSAGE_DELTA || + data.event === StepEvents.ON_REASONING_DELTA + ) { + tapStream(data.data, { ...submission, userMessage }); + } stepHandler(data, { ...submission, userMessage } as EventSubmission); } else if (data.sync != null) { const runId = v4(); @@ -162,6 +181,7 @@ export default function useSSE( } setCompleted((prev) => new Set(prev.add(streamKey))); + resetLive({ ...submission, userMessage }); const latestMessages = getMessages(); const conversationId = latestMessages?.[latestMessages.length - 1]?.conversationId; try { @@ -206,6 +226,7 @@ export default function useSSE( console.log('error in server stream.'); (startupConfig?.balance?.enabled ?? false) && balanceQuery.refetch(); + resetLive({ ...submission, userMessage }); let data: TResData | undefined = undefined; try { diff --git a/client/src/hooks/SSE/useUsageHandler.ts b/client/src/hooks/SSE/useUsageHandler.ts new file mode 100644 index 0000000000..635c882d0b --- /dev/null +++ b/client/src/hooks/SSE/useUsageHandler.ts @@ -0,0 +1,212 @@ +import { useRef, useMemo } from 'react'; +import { getDefaultStore } from 'jotai'; +import { useQueryClient } from '@tanstack/react-query'; +import { Constants, QueryKeys } from 'librechat-data-provider'; +import type { + TMessage, + TConversation, + TTokenConfigMap, + TTokenUsageEvent, + TContextUsageEvent, +} from 'librechat-data-provider'; +import { + EMPTY_USAGE_TOTALS, + liveTokensFamily, + calibrationFamily, + usageTotalsFamily, + branchTotalsFamily, + removeUsageAtoms, + contextSnapshotFamily, +} from '~/store/usage'; +import { sumBranch, upsertEntries, migrateIndex, calcUsageCost, estimateTokens } from '~/utils'; + +const FLUSH_INTERVAL_MS = 250; + +interface UsageSubmissionLike { + userMessage?: Pick | null; + conversation?: Partial> | null; +} + +interface FinalDataLike { + requestMessage?: Partial | null; + responseMessage?: Partial | null; + conversation?: Partial | null; +} + +export interface UsageHandlers { + contextHandler: (data: TContextUsageEvent, submission: UsageSubmissionLike) => void; + usageHandler: (data: TTokenUsageEvent, submission: UsageSubmissionLike) => void; + tapStream: (data: { delta?: { content?: unknown } }, submission: UsageSubmissionLike) => void; + finalizeUsage: (data: FinalDataLike, submission: UsageSubmissionLike) => void; + resetLive: (submission: UsageSubmissionLike) => void; + /** Replaces accumulated totals with the run's collected usage on resume */ + backfillUsage: (entries: TTokenUsageEvent[], submission: UsageSubmissionLike) => void; + /** Seeds the live estimate from already-streamed output chars on resume */ + seedLive: (chars: number, submission: UsageSubmissionLike) => void; +} + +function getConvoKey(submission: UsageSubmissionLike): string { + return submission.conversation?.conversationId ?? Constants.NEW_CONVO; +} + +function countDeltaChars(content: unknown): number { + const parts = Array.isArray(content) ? content : [content]; + let chars = 0; + for (const part of parts) { + if (part == null || typeof part !== 'object') { + continue; + } + const { text, think } = part as { text?: unknown; think?: unknown }; + if (typeof text === 'string') { + chars += text.length; + } else if (typeof think === 'string') { + chars += think.length; + } + } + return chars; +} + +/** + * Imperative writers for the per-conversation token usage atoms, driven by + * SSE events. All state lives in refs/atoms — the returned handlers are + * stable and never cause re-renders themselves. + */ +export default function useUsageHandler(): UsageHandlers { + const queryClient = useQueryClient(); + /** Streamed chars since the last snapshot or model end (current call only) */ + const streamCharsRef = useRef(0); + /** Provider-confirmed output tokens since the last snapshot (current run) */ + const confirmedRef = useRef(0); + const lastFlushRef = useRef(0); + + return useMemo(() => { + const jotai = getDefaultStore(); + + const setLive = (convoKey: string, value: number) => { + jotai.set(liveTokensFamily(convoKey), value); + }; + + const contextHandler: UsageHandlers['contextHandler'] = (data, submission) => { + const convoKey = getConvoKey(submission); + jotai.set(contextSnapshotFamily(convoKey), { + ...data, + anchorMessageId: submission.userMessage?.messageId ?? null, + }); + if (data.calibrationRatio != null && data.calibrationRatio > 0) { + jotai.set(calibrationFamily(convoKey), data.calibrationRatio); + } + streamCharsRef.current = 0; + confirmedRef.current = 0; + setLive(convoKey, 0); + }; + + const foldUsage = (data: TTokenUsageEvent, submission: UsageSubmissionLike) => { + const convoKey = getConvoKey(submission); + const tokenConfig = queryClient.getQueryData([QueryKeys.tokenConfig]); + const endpoint = submission.conversation?.endpoint ?? ''; + const model = data.model ?? submission.conversation?.model ?? ''; + const rates = tokenConfig?.[endpoint]?.[model]; + + const totalsAtom = usageTotalsFamily(convoKey); + const prev = jotai.get(totalsAtom); + jotai.set(totalsAtom, { + input: prev.input + (data.input_tokens ?? 0), + output: prev.output + (data.output_tokens ?? 0), + cacheWrite: prev.cacheWrite + (data.input_token_details?.cache_creation ?? 0), + cacheRead: prev.cacheRead + (data.input_token_details?.cache_read ?? 0), + costUSD: prev.costUSD + calcUsageCost(data, rates), + eventCount: prev.eventCount + 1, + }); + }; + + const usageHandler: UsageHandlers['usageHandler'] = (data, submission) => { + foldUsage(data, submission); + + if (data.usage_type === 'summarization') { + return; + } + confirmedRef.current += data.output_tokens ?? 0; + streamCharsRef.current = 0; + setLive(getConvoKey(submission), confirmedRef.current); + }; + + const tapStream: UsageHandlers['tapStream'] = (data, submission) => { + const chars = countDeltaChars(data?.delta?.content); + if (chars <= 0) { + return; + } + streamCharsRef.current += chars; + const now = Date.now(); + if (now - lastFlushRef.current < FLUSH_INTERVAL_MS) { + return; + } + lastFlushRef.current = now; + const convoKey = getConvoKey(submission); + const ratio = jotai.get(calibrationFamily(convoKey)); + setLive(convoKey, confirmedRef.current + estimateTokens(streamCharsRef.current, ratio)); + }; + + const resetLive: UsageHandlers['resetLive'] = (submission) => { + streamCharsRef.current = 0; + confirmedRef.current = 0; + setLive(getConvoKey(submission), 0); + }; + + const backfillUsage: UsageHandlers['backfillUsage'] = (entries, submission) => { + jotai.set(usageTotalsFamily(getConvoKey(submission)), EMPTY_USAGE_TOTALS); + for (const entry of entries) { + foldUsage(entry, submission); + } + }; + + const seedLive: UsageHandlers['seedLive'] = (chars, submission) => { + if (chars <= 0) { + return; + } + const convoKey = getConvoKey(submission); + streamCharsRef.current = chars; + confirmedRef.current = 0; + setLive(convoKey, estimateTokens(chars, jotai.get(calibrationFamily(convoKey)))); + }; + + const finalizeUsage: UsageHandlers['finalizeUsage'] = (data, submission) => { + const fromKey = getConvoKey(submission); + const realId = data.conversation?.conversationId ?? fromKey; + + upsertEntries(fromKey, [data.requestMessage, data.responseMessage]); + + if (realId !== fromKey) { + migrateIndex(fromKey, realId); + jotai.set(contextSnapshotFamily(realId), jotai.get(contextSnapshotFamily(fromKey))); + jotai.set(usageTotalsFamily(realId), jotai.get(usageTotalsFamily(fromKey))); + jotai.set(calibrationFamily(realId), jotai.get(calibrationFamily(fromKey))); + removeUsageAtoms(fromKey); + } + + const responseMeta = data.responseMessage?.contextMeta; + if (responseMeta?.calibrationRatio != null && responseMeta.calibrationRatio > 0) { + jotai.set(calibrationFamily(realId), responseMeta.calibrationRatio); + } + + const tailId = data.responseMessage?.messageId ?? data.requestMessage?.messageId ?? null; + if (tailId) { + const anchorId = submission.userMessage?.messageId ?? null; + jotai.set(branchTotalsFamily(realId), sumBranch(realId, tailId, anchorId)); + } + + streamCharsRef.current = 0; + confirmedRef.current = 0; + setLive(realId, 0); + }; + + return { + contextHandler, + usageHandler, + tapStream, + finalizeUsage, + resetLive, + backfillUsage, + seedLive, + }; + }, [queryClient]); +} diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index b7d2e4ae96..d3fe0bc369 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -837,6 +837,8 @@ "com_ui_branch_error": "Failed to create branch", "com_ui_branch_message": "Create branch from this response", "com_ui_by_author": "by {{0}}", + "com_ui_cache_read": "Cache read", + "com_ui_cache_write": "Cache write", "com_ui_callback_url": "Callback URL", "com_ui_cancel": "Cancel", "com_ui_cancelled": "Cancelled", @@ -883,6 +885,16 @@ "com_ui_contact_admin_if_issue_persists": "Contact the Admin if the issue persists", "com_ui_context": "Context", "com_ui_context_filter_sort": "Filter and Sort by Context", + "com_ui_context_free": "Free space", + "com_ui_context_messages": "Messages", + "com_ui_context_summary": "Summary", + "com_ui_context_system": "System prompt", + "com_ui_context_tools": "Tool schemas", + "com_ui_context_unknown": "Context size unknown", + "com_ui_context_usage": "Context usage", + "com_ui_context_usage_label": "Context window: {{0}} of {{1}} tokens used ({{2}}%)", + "com_ui_context_usage_label_unknown": "Context usage: {{0}} tokens used", + "com_ui_context_window": "Context window", "com_ui_continue": "Continue", "com_ui_continue_oauth": "Continue with OAuth", "com_ui_control_bar": "Control bar", @@ -1029,6 +1041,7 @@ "com_ui_error_try_following_prefix": "Please try one of the following", "com_ui_error_unexpected": "Oops! Something Unexpected Occurred", "com_ui_error_updating_preferences": "Error updating preferences", + "com_ui_estimated": "Estimated from message history", "com_ui_everyone_permission_level": "Everyone's permission level", "com_ui_examples": "Examples", "com_ui_expand": "Expand", @@ -1508,6 +1521,7 @@ "com_ui_select_search_model": "Search model by name", "com_ui_select_search_provider": "Search provider by name", "com_ui_select_search_region": "Search region by name", + "com_ui_session_cost": "Session cost", "com_ui_set": "Set", "com_ui_share": "Share", "com_ui_share_create_message": "Your name and any messages you add after sharing stay private.", diff --git a/client/src/store/index.ts b/client/src/store/index.ts index 0584e9d805..7d7d9cfe61 100644 --- a/client/src/store/index.ts +++ b/client/src/store/index.ts @@ -1,21 +1,22 @@ import * as artifacts from './artifacts'; -import families from './families'; -import endpoints from './endpoints'; -import user from './user'; -import text from './text'; -import toast from './toast'; import submission from './submission'; +import isTemporary from './temporary'; +import endpoints from './endpoints'; +import families from './families'; +import settings from './settings'; +import prompts from './prompts'; import search from './search'; import preset from './preset'; -import prompts from './prompts'; import lang from './language'; -import settings from './settings'; +import toast from './toast'; +import user from './user'; +import text from './text'; import misc from './misc'; -import isTemporary from './temporary'; export * from './agents'; export * from './mcp'; export * from './favorites'; export * from './subagents'; +export * from './usage'; export default { ...artifacts, diff --git a/client/src/store/usage.ts b/client/src/store/usage.ts new file mode 100644 index 0000000000..c122efa20d --- /dev/null +++ b/client/src/store/usage.ts @@ -0,0 +1,60 @@ +import { atom } from 'jotai'; +import { atomFamily } from 'jotai/utils'; +import type { TContextUsageEvent } from 'librechat-data-provider'; +import type { BranchTotals } from '~/utils/tokens'; +import { EMPTY_BRANCH } from '~/utils/tokens'; + +/** Latest backend context snapshot, anchored to the run's user message for staleness checks */ +export interface ContextSnapshot extends TContextUsageEvent { + anchorMessageId: string | null; +} + +/** Cumulative provider-reported usage for the conversation's current session */ +export interface UsageTotals { + input: number; + output: number; + cacheWrite: number; + cacheRead: number; + costUSD: number; + eventCount: number; +} + +export const EMPTY_USAGE_TOTALS: UsageTotals = { + input: 0, + output: 0, + cacheWrite: 0, + cacheRead: 0, + costUSD: 0, + eventCount: 0, +}; + +/** + * Split into small per-conversation atoms so high-frequency writes (live + * stream estimates) re-render only the subscribers that draw them. + */ +export const branchTotalsFamily = atomFamily((_conversationId: string) => + atom(EMPTY_BRANCH), +); + +export const contextSnapshotFamily = atomFamily((_conversationId: string) => + atom(null), +); + +export const usageTotalsFamily = atomFamily((_conversationId: string) => + atom(EMPTY_USAGE_TOTALS), +); + +/** Throttled in-flight output token estimate for the current model call */ +export const liveTokensFamily = atomFamily((_conversationId: string) => atom(0)); + +/** Last known provider-vs-estimate calibration ratio for the conversation */ +export const calibrationFamily = atomFamily((_conversationId: string) => atom(1)); + +/** Jotai atomFamily entries are never GC'd — call on conversation switch/cleanup */ +export function removeUsageAtoms(conversationId: string): void { + branchTotalsFamily.remove(conversationId); + contextSnapshotFamily.remove(conversationId); + usageTotalsFamily.remove(conversationId); + liveTokensFamily.remove(conversationId); + calibrationFamily.remove(conversationId); +} diff --git a/client/src/utils/index.ts b/client/src/utils/index.ts index 14d76e7b73..0beb890c0b 100644 --- a/client/src/utils/index.ts +++ b/client/src/utils/index.ts @@ -21,6 +21,7 @@ export * from './presets'; export * from './prompts'; export * from './textarea'; export * from './messages'; +export * from './tokens'; export * from './redirect'; export * from './languages'; export * from './conversation'; diff --git a/client/src/utils/messages.ts b/client/src/utils/messages.ts index 91312567ec..4a9499530b 100644 --- a/client/src/utils/messages.ts +++ b/client/src/utils/messages.ts @@ -26,12 +26,29 @@ export type BranchSiblingIndex = { siblingIdx: number; }; +/** Keyed by array identity — cache selectors re-derive the tree from the same + * immutable snapshot once per subscriber, so repeat builds are pure waste. */ +const treeMemo = new WeakMap>(); + +const buildTreeMemoized = (messages: TMessage[] | null | undefined) => { + if (messages == null) { + return null; + } + const cached = treeMemo.get(messages); + if (cached !== undefined) { + return cached; + } + const tree = buildTree({ messages }); + treeMemo.set(messages, tree); + return tree; +}; + export const selectActiveBranchTail = ( messages: TMessage[] | null | undefined, rootSiblingKey: string | null | undefined, getSiblingIndex: SiblingIndexLookup = () => 0, ): TMessage | null => { - const messagesTree = buildTree({ messages: messages ?? null }); + const messagesTree = buildTreeMemoized(messages); if (!messagesTree?.length) { return null; } @@ -61,7 +78,7 @@ export const getMessageBranchSiblingParentIds = ( messages: TMessage[] | null | undefined, rootSiblingKey: string | null | undefined, ): (string | null)[] => { - const messagesTree = buildTree({ messages: messages ?? null }); + const messagesTree = buildTreeMemoized(messages); if (!messagesTree?.length) { return []; } diff --git a/client/src/utils/tokens.spec.ts b/client/src/utils/tokens.spec.ts new file mode 100644 index 0000000000..562ab5e0be --- /dev/null +++ b/client/src/utils/tokens.spec.ts @@ -0,0 +1,179 @@ +import { Constants } from 'librechat-data-provider'; +import type { TMessage } from 'librechat-data-provider'; +import { + buildIndex, + upsertEntries, + migrateIndex, + clearIndex, + hasIndex, + sumBranch, + estimateTokens, + calcUsageCost, + formatCost, + countTrailingOutputChars, + EMPTY_BRANCH, +} from './tokens'; + +const CONVO = 'convo-1'; + +function msg( + messageId: string, + parentMessageId: string | null, + isCreatedByUser: boolean, + tokenCount?: number, +): TMessage { + return { + messageId, + parentMessageId, + isCreatedByUser, + tokenCount, + conversationId: CONVO, + text: '', + } as TMessage; +} + +describe('token index', () => { + afterEach(() => { + clearIndex(CONVO); + clearIndex('convo-2'); + clearIndex(Constants.NEW_CONVO); + }); + + it('sums only the active branch via the parent chain', () => { + buildIndex(CONVO, [ + msg('u1', Constants.NO_PARENT, true, 10), + msg('a1', 'u1', false, 20), + msg('u2', 'a1', true, 30), + msg('a2', 'u2', false, 40), + /** Sibling branch that must not be counted */ + msg('a2-alt', 'u2', false, 999), + ]); + + const totals = sumBranch(CONVO, 'a2'); + expect(totals.input).toBe(40); + expect(totals.output).toBe(60); + expect(totals.total).toBe(4); + expect(totals.counted).toBe(4); + + const altTotals = sumBranch(CONVO, 'a2-alt'); + expect(altTotals.output).toBe(1019); + }); + + it('flags whether the anchor message is on the branch', () => { + buildIndex(CONVO, [ + msg('u1', Constants.NO_PARENT, true, 10), + msg('a1', 'u1', false, 20), + msg('a1-alt', 'u1', false, 25), + ]); + + expect(sumBranch(CONVO, 'a1', 'a1').containsAnchor).toBe(true); + expect(sumBranch(CONVO, 'a1-alt', 'a1').containsAnchor).toBe(false); + }); + + it('tracks uncounted messages and tolerates missing parents', () => { + buildIndex(CONVO, [msg('u2', 'missing-parent', true, undefined), msg('a2', 'u2', false, 15)]); + + const totals = sumBranch(CONVO, 'a2'); + expect(totals.total).toBe(2); + expect(totals.counted).toBe(1); + expect(totals.output).toBe(15); + }); + + it('returns EMPTY_BRANCH without an index or tail', () => { + expect(sumBranch('unknown', 'x')).toBe(EMPTY_BRANCH); + buildIndex(CONVO, []); + expect(sumBranch(CONVO, null)).toBe(EMPTY_BRANCH); + }); + + it('upserts incrementally and overwrites by messageId', () => { + buildIndex(CONVO, [msg('u1', Constants.NO_PARENT, true, 10)]); + upsertEntries(CONVO, [msg('a1', 'u1', false, 20), null, undefined]); + expect(sumBranch(CONVO, 'a1').output).toBe(20); + + upsertEntries(CONVO, [msg('a1', 'u1', false, 35)]); + expect(sumBranch(CONVO, 'a1').output).toBe(35); + }); + + it('migrates the index to a new conversation id', () => { + buildIndex(Constants.NEW_CONVO, [msg('u1', Constants.NO_PARENT, true, 10)]); + migrateIndex(Constants.NEW_CONVO, 'convo-2'); + + expect(hasIndex(Constants.NEW_CONVO)).toBe(false); + expect(sumBranch('convo-2', 'u1').input).toBe(10); + }); +}); + +describe('estimateTokens', () => { + it('estimates chars/4 scaled by calibration ratio', () => { + expect(estimateTokens(400)).toBe(100); + expect(estimateTokens(400, 1.1)).toBe(110); + expect(estimateTokens(0)).toBe(0); + expect(estimateTokens(100, 0)).toBe(25); + }); +}); + +describe('calcUsageCost', () => { + const rates = { prompt: 3, completion: 15, cacheWrite: 3.75, cacheRead: 0.3 }; + + it('prices additive cache usage (Anthropic pattern)', () => { + const cost = calcUsageCost( + { + input_tokens: 1000, + output_tokens: 500, + input_token_details: { cache_creation: 2000, cache_read: 10000 }, + }, + rates, + ); + expect(cost).toBeCloseTo((1000 * 3 + 2000 * 3.75 + 10000 * 0.3 + 500 * 15) / 1e6); + }); + + it('prices inclusive cache usage (OpenAI pattern)', () => { + const cost = calcUsageCost( + { + input_tokens: 10000, + output_tokens: 500, + input_token_details: { cache_read: 4000 }, + }, + rates, + ); + expect(cost).toBeCloseTo((6000 * 3 + 4000 * 0.3 + 500 * 15) / 1e6); + }); + + it('returns 0 without rates', () => { + expect(calcUsageCost({ input_tokens: 100, output_tokens: 100 })).toBe(0); + expect(calcUsageCost({ input_tokens: 100 }, { context: 1000 })).toBe(0); + }); +}); + +describe('formatCost', () => { + it('formats across magnitude bands', () => { + expect(formatCost(0)).toBe('$0.00'); + expect(formatCost(0.004)).toBe('<$0.01'); + expect(formatCost(0.0523)).toBe('$0.0523'); + expect(formatCost(1.234)).toBe('$1.23'); + }); +}); + +describe('countTrailingOutputChars', () => { + const text = (value: string) => ({ type: 'text', text: value }); + const think = (value: string) => ({ type: 'think', think: value }); + const tool = () => ({ type: 'tool_call', tool_call: { id: 'tc' } }); + + it('counts the trailing run of text and think parts', () => { + expect(countTrailingOutputChars([text('aaaa'), tool(), think('bb'), text('ccc')])).toBe(5); + }); + + it('skips trailing in-progress tool parts before collecting', () => { + expect(countTrailingOutputChars([text('aaaa'), tool(), text('ccc'), tool()])).toBe(3); + }); + + it('stops at the previous tool boundary', () => { + expect(countTrailingOutputChars([text('aaaa'), tool(), text('ccc')])).toBe(3); + }); + + it('handles empty and non-output content', () => { + expect(countTrailingOutputChars(undefined)).toBe(0); + expect(countTrailingOutputChars([])).toBe(0); + expect(countTrailingOutputChars([tool()])).toBe(0); + }); +}); diff --git a/client/src/utils/tokens.ts b/client/src/utils/tokens.ts new file mode 100644 index 0000000000..9e8dddc157 --- /dev/null +++ b/client/src/utils/tokens.ts @@ -0,0 +1,228 @@ +import { Constants } from 'librechat-data-provider'; +import type { TMessage, TTokenUsageEvent, TModelTokenomics } from 'librechat-data-provider'; + +export interface TokenEntry { + tokenCount: number; + isCreatedByUser: boolean; + parentMessageId: string | null; +} + +export interface BranchTotals { + /** Sum of user-message token counts on the active branch */ + input: number; + /** Sum of assistant-message token counts on the active branch */ + output: number; + /** Messages on the branch with a known tokenCount */ + counted: number; + /** Total messages on the branch */ + total: number; + tailId: string | null; + /** Whether the latest run's anchor message is on this branch */ + containsAnchor: boolean; +} + +export const EMPTY_BRANCH: BranchTotals = { + input: 0, + output: 0, + counted: 0, + total: 0, + tailId: null, + containsAnchor: false, +}; + +/** Module-level token index: conversationId → messageId → entry. Not render state. */ +const registry = new Map>(); + +function toEntry(message: Partial): TokenEntry { + return { + tokenCount: typeof message.tokenCount === 'number' ? message.tokenCount : 0, + isCreatedByUser: message.isCreatedByUser === true, + parentMessageId: message.parentMessageId ?? null, + }; +} + +/** Full O(n) rebuild — only on discrete cache replacements (load, refetch, edits) */ +export function buildIndex(conversationId: string, messages?: TMessage[] | null): void { + const index = new Map(); + if (messages != null) { + for (const message of messages) { + if (message?.messageId) { + index.set(message.messageId, toEntry(message)); + } + } + } + registry.set(conversationId, index); +} + +/** Incremental upsert — final SSE events touch at most two entries */ +export function upsertEntries( + conversationId: string, + messages: (Partial | null | undefined)[], +): void { + let index = registry.get(conversationId); + if (!index) { + index = new Map(); + registry.set(conversationId, index); + } + for (const message of messages) { + if (message?.messageId) { + index.set(message.messageId, toEntry(message)); + } + } +} + +/** Re-keys the index when `new` resolves to a real conversation id */ +export function migrateIndex(fromId: string, toId: string): void { + if (fromId === toId) { + return; + } + const index = registry.get(fromId); + if (!index) { + return; + } + registry.delete(fromId); + registry.set(toId, index); +} + +export function clearIndex(conversationId: string): void { + registry.delete(conversationId); +} + +export function hasIndex(conversationId: string): boolean { + return registry.has(conversationId); +} + +/** O(depth) walk from the branch tail up the parent chain */ +export function sumBranch( + conversationId: string, + tailId: string | null | undefined, + anchorId?: string | null, +): BranchTotals { + const index = registry.get(conversationId); + if (!index || !tailId) { + return EMPTY_BRANCH; + } + + const totals = { input: 0, output: 0, counted: 0, total: 0, containsAnchor: false }; + let currentId: string | null = tailId; + let guard = index.size; + + while (currentId && currentId !== Constants.NO_PARENT && guard-- > 0) { + const entry: TokenEntry | undefined = index.get(currentId); + if (!entry) { + break; + } + totals.total += 1; + if (anchorId != null && currentId === anchorId) { + totals.containsAnchor = true; + } + if (entry.tokenCount > 0) { + totals.counted += 1; + if (entry.isCreatedByUser) { + totals.input += entry.tokenCount; + } else { + totals.output += entry.tokenCount; + } + } + currentId = entry.parentMessageId; + } + + return { ...totals, tailId }; +} + +function getOutputChars(part: unknown): number | null { + if (part == null || typeof part !== 'object') { + return null; + } + const { text, think } = part as { text?: unknown; think?: unknown }; + if (typeof text === 'string') { + return text.length; + } + if (typeof think === 'string') { + return think.length; + } + return null; +} + +/** + * Chars of the trailing contiguous text/think parts in aggregated content — + * the model output produced since the last tool boundary, which the latest + * context snapshot does not yet account for. Trailing non-text parts + * (in-progress tool calls) are skipped before collecting. + */ +export function countTrailingOutputChars(content?: unknown[] | null): number { + if (!Array.isArray(content) || content.length === 0) { + return 0; + } + let i = content.length - 1; + while (i >= 0 && getOutputChars(content[i]) == null) { + i--; + } + let chars = 0; + for (; i >= 0; i--) { + const partChars = getOutputChars(content[i]); + if (partChars == null) { + break; + } + chars += partChars; + } + return chars; +} + +/** Rough live estimate for streaming text, calibrated by the last known provider ratio */ +export function estimateTokens(charCount: number, calibrationRatio = 1): number { + if (charCount <= 0) { + return 0; + } + const ratio = calibrationRatio > 0 ? calibrationRatio : 1; + return Math.round((charCount / 4) * ratio); +} + +/** + * USD cost of one model call. Mirrors the cache-detection heuristic used in + * token accounting: cache counts are additive when they exceed base input + * (Anthropic), otherwise cache reads are included in input tokens (OpenAI). + */ +export function calcUsageCost(usage: TTokenUsageEvent, rates?: TModelTokenomics): number { + if (!rates || rates.prompt == null || rates.completion == null) { + return 0; + } + const input = usage.input_tokens ?? 0; + const output = usage.output_tokens ?? 0; + const cacheWrite = usage.input_token_details?.cache_creation ?? 0; + const cacheRead = usage.input_token_details?.cache_read ?? 0; + const writeRate = rates.cacheWrite ?? rates.prompt; + const readRate = rates.cacheRead ?? rates.prompt; + + const cacheIsAdditive = cacheWrite + cacheRead > input; + const baseInput = cacheIsAdditive ? input : Math.max(0, input - cacheRead - cacheWrite); + + return ( + (baseInput * rates.prompt + + cacheWrite * writeRate + + cacheRead * readRate + + output * rates.completion) / + 1e6 + ); +} + +export function formatTokens(count: number): string { + const formatted = new Intl.NumberFormat(undefined, { + notation: 'compact', + maximumFractionDigits: 1, + }).format(count); + return formatted.replace(/\.0(?=[A-Za-z]|$)/, ''); +} + +export function formatCost(usd: number): string { + if (usd <= 0) { + return '$0.00'; + } + if (usd < 0.01) { + return '<$0.01'; + } + if (usd < 1) { + return `$${usd.toFixed(4)}`; + } + return `$${usd.toFixed(2)}`; +} diff --git a/e2e/setup/fake-model.js b/e2e/setup/fake-model.js index 9702f5282a..0e690a3679 100644 --- a/e2e/setup/fake-model.js +++ b/e2e/setup/fake-model.js @@ -10,6 +10,8 @@ * from the conversation and the agents' advertised tools. */ const { FakeChatModel } = require('@librechat/agents'); +const { ChatGenerationChunk } = require('@langchain/core/outputs'); +const { AIMessageChunk } = require('@langchain/core/messages'); const MOCK_REPLY = process.env.MOCK_LLM_REPLY || 'E2E mock reply: pong'; const CHUNK_DELAY_MS = Number(process.env.MOCK_LLM_CHUNK_DELAY_MS) || 10; @@ -298,9 +300,41 @@ function replyResponses(text) { return null; } +/** + * Attaches synthetic usage_metadata on a final empty chunk (the OpenAI + * streaming pattern) so token-usage SSE events flow end to end in mock runs. + */ +class UsageEmittingFakeChatModel extends FakeChatModel { + async *_streamResponseChunks(messages, options, runManager) { + let outputChars = 0; + for await (const chunk of super._streamResponseChunks(messages, options, runManager)) { + outputChars += typeof chunk.text === 'string' ? chunk.text.length : 0; + yield chunk; + } + const inputChars = (messages ?? []).reduce( + (sum, message) => sum + getContentText(message?.content).length, + 0, + ); + const input_tokens = Math.max(1, Math.ceil(inputChars / 4)); + const output_tokens = Math.max(1, Math.ceil(outputChars / 4)); + yield new ChatGenerationChunk({ + text: '', + message: new AIMessageChunk({ + content: '', + usage_metadata: { input_tokens, output_tokens, total_tokens: input_tokens + output_tokens }, + }), + }); + } +} + function overrideModel({ graph, responses, sleep, toolCalls, thrownError }) { if (!thrownError) { - graph.overrideTestModel(responses, sleep ?? CHUNK_DELAY_MS, toolCalls); + graph.overrideModel = new UsageEmittingFakeChatModel({ + responses, + sleep: sleep ?? CHUNK_DELAY_MS, + emitCustomEvent: true, + toolCalls, + }); return; } diff --git a/e2e/specs/mock/usage.spec.ts b/e2e/specs/mock/usage.spec.ts new file mode 100644 index 0000000000..332873f0e7 --- /dev/null +++ b/e2e/specs/mock/usage.spec.ts @@ -0,0 +1,50 @@ +import { expect, test } from '@playwright/test'; +import type { Page } from '@playwright/test'; +import { + MOCK_ENDPOINTS, + NEW_CHAT_PATH, + mockReply, + sendMessage, + selectMockEndpoint, +} from './helpers'; + +const gauge = (page: Page) => page.getByTestId('token-usage'); +const gaugeMeter = (page: Page) => gauge(page).getByRole('meter'); + +async function expectGaugeAboveZero(page: Page) { + await expect(gauge(page)).toBeVisible({ timeout: 20000 }); + await expect(gaugeMeter(page)).toHaveAttribute('aria-valuenow', /[1-9]/, { timeout: 20000 }); +} + +test.describe('context usage gauge', () => { + test('tracks usage from live SSE events and survives reload', async ({ page }) => { + test.setTimeout(120000); + await page.goto(NEW_CHAT_PATH, { timeout: 10000 }); + + // REQUIRED so the message streams without a real key. + await selectMockEndpoint(page, MOCK_ENDPOINTS[0]); + + const response = await sendMessage(page, 'hello'); + expect(response.ok()).toBeTruthy(); + await expect(mockReply(page)).toBeVisible({ timeout: 20000 }); + await expect(page).toHaveURL(/\/c\/(?!new)/, { timeout: 15000 }); + + /** Live path: the agents pipeline's context snapshot + usage events fill the gauge */ + await expectGaugeAboveZero(page); + + /** Breakdown popover: context section always, usage section from on_token_usage events */ + await gauge(page).hover(); + const popover = page.getByRole('region', { name: 'Context usage' }); + await expect(popover).toBeVisible({ timeout: 10000 }); + await expect(popover.getByText('Context window')).toBeVisible(); + await expect(popover.getByText('Input', { exact: true })).toBeVisible(); + await expect(popover.getByText('Output', { exact: true })).toBeVisible(); + await page.keyboard.press('Escape'); + + /** Fallback path: after reload there is no snapshot — the gauge rebuilds + * from per-message tokenCount history returned by the messages query */ + await page.reload({ timeout: 15000 }); + await expect(mockReply(page)).toBeVisible({ timeout: 20000 }); + await expectGaugeAboveZero(page); + }); +}); diff --git a/packages/api/src/endpoints/custom/initialize.ts b/packages/api/src/endpoints/custom/initialize.ts index 19d4e2c35e..80d0de430d 100644 --- a/packages/api/src/endpoints/custom/initialize.ts +++ b/packages/api/src/endpoints/custom/initialize.ts @@ -16,6 +16,26 @@ import { tokenConfigCache } from '~/cache'; const { PROXY } = process.env; +/** + * Cache key for an endpoint's fetched token config. User-scoped when the + * model fetch can resolve per-user: user-provided key/URL, or header + * templates forwarded against an admin-trusted base URL — making the + * response, and therefore the derived token config, user-specific. + */ +export function getTokenConfigKey( + endpointConfig: Partial, + endpoint: string, + userId: string, +): string { + const hasTokenConfig = (endpointConfig as Record).tokenConfig != null; + const userProvidesKey = isUserProvided(extractEnvVariable(endpointConfig.apiKey ?? '')); + const userProvidesURL = isUserProvided(extractEnvVariable(endpointConfig.baseURL ?? '')); + const willForwardUserScopedHeaders = !!endpointConfig?.headers && !userProvidesURL; + return !hasTokenConfig && (userProvidesKey || userProvidesURL || willForwardUserScopedHeaders) + ? `${endpoint}:${userId}` + : endpoint; +} + /** * Builds custom options from endpoint configuration */ @@ -138,17 +158,7 @@ export async function initializeCustom({ const cache = tokenConfigCache(); /** tokenConfig is an optional extended property on custom endpoints */ const hasTokenConfig = (endpointConfig as Record).tokenConfig != null; - // When `endpointConfig.headers` will be forwarded to the model fetch (i.e. - // base URL is admin-trusted, so the security guard below leaves them in - // place), header templates may resolve against the current user — making - // the response, and therefore the derived token config, user-specific. - // User-scope the token-config cache key in that case so a cached entry - // for one user can't be served to another. - const willForwardUserScopedHeaders = !!endpointConfig?.headers && !userProvidesURL; - const tokenKey = - !hasTokenConfig && (userProvidesKey || userProvidesURL || willForwardUserScopedHeaders) - ? `${endpoint}:${userId}` - : endpoint; + const tokenKey = getTokenConfigKey(endpointConfig, endpoint, userId); const cachedConfig = !hasTokenConfig && diff --git a/packages/api/src/endpoints/index.ts b/packages/api/src/endpoints/index.ts index 208181168e..98e0af0298 100644 --- a/packages/api/src/endpoints/index.ts +++ b/packages/api/src/endpoints/index.ts @@ -5,3 +5,4 @@ export * from './custom'; export * from './google'; export * from './models'; export * from './openai'; +export * from './pricing'; diff --git a/packages/api/src/endpoints/pricing.spec.ts b/packages/api/src/endpoints/pricing.spec.ts new file mode 100644 index 0000000000..487adbc39f --- /dev/null +++ b/packages/api/src/endpoints/pricing.spec.ts @@ -0,0 +1,89 @@ +import mongoose from 'mongoose'; +import { EModelEndpoint } from 'librechat-data-provider'; +import { createTxMethods } from '@librechat/data-schemas'; +import type { EndpointTokenConfig } from '~/types'; +import { matchModelName, findMatchingPattern } from '~/utils'; +import { buildTokenConfigMap } from './pricing'; + +const { getValueKey, getMultiplier, getCacheMultiplier } = createTxMethods(mongoose, { + matchModelName, + findMatchingPattern, +}); + +const deps = { getValueKey, getMultiplier, getCacheMultiplier }; + +describe('buildTokenConfigMap', () => { + it('resolves context windows without pricing by default', () => { + const map = buildTokenConfigMap( + { + modelsConfig: { + [EModelEndpoint.openAI]: ['gpt-4o'], + [EModelEndpoint.anthropic]: ['claude-3-5-sonnet-20241022'], + }, + }, + deps, + ); + + expect(map[EModelEndpoint.openAI]['gpt-4o'].context).toBeGreaterThan(100000); + expect(map[EModelEndpoint.anthropic]['claude-3-5-sonnet-20241022'].context).toBe(200000); + expect(map[EModelEndpoint.openAI]['gpt-4o'].prompt).toBeUndefined(); + expect(map[EModelEndpoint.openAI]['gpt-4o'].completion).toBeUndefined(); + }); + + it('includes pattern-matched pricing when enabled', () => { + const map = buildTokenConfigMap( + { + modelsConfig: { + [EModelEndpoint.anthropic]: ['claude-3-5-sonnet-20241022'], + }, + includePricing: true, + }, + deps, + ); + + const tokenomics = map[EModelEndpoint.anthropic]['claude-3-5-sonnet-20241022']; + expect(tokenomics.prompt).toBe(3); + expect(tokenomics.completion).toBe(15); + expect(tokenomics.cacheWrite).toBe(3.75); + expect(tokenomics.cacheRead).toBe(0.3); + }); + + it('treats unknown endpoints as custom for context lookups', () => { + const map = buildTokenConfigMap( + { + modelsConfig: { MyProxy: ['gpt-4o-mini'] }, + }, + deps, + ); + + expect(map.MyProxy['gpt-4o-mini'].context).toBeGreaterThan(0); + }); + + it('prefers endpoint token config overrides for context and rates', () => { + const override: EndpointTokenConfig = { + 'custom-model': { prompt: 1.5, completion: 4.5, context: 32000 }, + }; + const map = buildTokenConfigMap( + { + modelsConfig: { MyProxy: ['custom-model', 'gpt-4o-mini'] }, + endpointTokenConfigs: { MyProxy: override }, + includePricing: true, + }, + deps, + ); + + expect(map.MyProxy['custom-model']).toEqual({ + context: 32000, + prompt: 1.5, + completion: 4.5, + }); + /** Models absent from the override fall back to static tables */ + expect(map.MyProxy['gpt-4o-mini'].context).toBeGreaterThan(0); + expect(map.MyProxy['gpt-4o-mini'].prompt).toBeGreaterThan(0); + }); + + it('skips endpoints with empty model lists', () => { + const map = buildTokenConfigMap({ modelsConfig: { [EModelEndpoint.agents]: [] } }, deps); + expect(map[EModelEndpoint.agents]).toBeUndefined(); + }); +}); diff --git a/packages/api/src/endpoints/pricing.ts b/packages/api/src/endpoints/pricing.ts new file mode 100644 index 0000000000..48bb5f624e --- /dev/null +++ b/packages/api/src/endpoints/pricing.ts @@ -0,0 +1,98 @@ +import { EModelEndpoint } from 'librechat-data-provider'; +import type { TModelsConfig, TTokenConfigMap, TModelTokenomics } from 'librechat-data-provider'; +import type { TxMethods } from '@librechat/data-schemas'; +import type { EndpointTokenConfig } from '~/types'; +import { getModelMaxTokens, maxTokensMap } from '~/utils'; + +export interface TokenomicsDeps { + getValueKey: TxMethods['getValueKey']; + getMultiplier: TxMethods['getMultiplier']; + getCacheMultiplier: TxMethods['getCacheMultiplier']; +} + +export interface TokenConfigParams { + /** endpoint → model list, from the resolved models config */ + modelsConfig: TModelsConfig; + /** Per-endpoint overrides: fetched (e.g. OpenRouter) or yaml `tokenConfig` */ + endpointTokenConfigs?: Record; + /** Include USD-per-1M rates; gated by `interface.contextCost` */ + includePricing?: boolean; +} + +/** + * Resolves context windows (and optionally pricing) for every configured + * model, server-side, so the client never reimplements pattern matching. + */ +export function buildTokenConfigMap( + params: TokenConfigParams, + deps: TokenomicsDeps, +): TTokenConfigMap { + const { modelsConfig, endpointTokenConfigs, includePricing = false } = params; + const map: TTokenConfigMap = {}; + + for (const [endpoint, models] of Object.entries(modelsConfig)) { + if (!Array.isArray(models) || models.length === 0) { + continue; + } + const override = endpointTokenConfigs?.[endpoint]; + const endpointKey = ( + maxTokensMap[endpoint] != null ? endpoint : EModelEndpoint.custom + ) as EModelEndpoint; + + const entry: Record = {}; + for (const model of models) { + const tokenomics: TModelTokenomics = {}; + const context = override + ? (getModelMaxTokens(model, endpointKey, override) ?? getModelMaxTokens(model, endpointKey)) + : getModelMaxTokens(model, endpointKey); + if (context != null) { + tokenomics.context = context; + } + + if (includePricing) { + const overrideRates = override?.[model]; + if (overrideRates?.prompt != null || overrideRates?.completion != null) { + tokenomics.prompt = overrideRates.prompt; + tokenomics.completion = overrideRates.completion; + } else { + const valueKey = deps.getValueKey(model, endpoint); + tokenomics.prompt = deps.getMultiplier({ + valueKey, + model, + endpoint, + tokenType: 'prompt', + }); + tokenomics.completion = deps.getMultiplier({ + valueKey, + model, + endpoint, + tokenType: 'completion', + }); + const cacheWrite = deps.getCacheMultiplier({ + valueKey, + model, + endpoint, + cacheType: 'write', + }); + const cacheRead = deps.getCacheMultiplier({ + valueKey, + model, + endpoint, + cacheType: 'read', + }); + if (cacheWrite != null) { + tokenomics.cacheWrite = cacheWrite; + } + if (cacheRead != null) { + tokenomics.cacheRead = cacheRead; + } + } + } + + entry[model] = tokenomics; + } + map[endpoint] = entry; + } + + return map; +} diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index 238db4e6cb..9b1980a6b8 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -1,5 +1,5 @@ -import { Constants, parseTextParts } from 'librechat-data-provider'; import { logger, getTenantId, SYSTEM_TENANT_ID } from '@librechat/data-schemas'; +import { Constants, UsageEvents, parseTextParts } from 'librechat-data-provider'; import type { Agents, TMessageContentParts } from 'librechat-data-provider'; import type { StandardGraph } from '@librechat/agents'; import type { @@ -177,6 +177,9 @@ class GenerationJobManagerClass { /** Serializes replay-event read/modify/write updates per stream. */ private replayEventWriteQueues = new Map>(); + /** Serializes token-usage read/modify/write updates per stream. */ + private tokenUsageWriteQueues = new Map>(); + private cleanupInterval: NodeJS.Timeout | null = null; /** Whether we're using Redis stores */ @@ -649,6 +652,7 @@ class GenerationJobManagerClass { this.jobStore.clearContentState(streamId); this.runStepBuffers?.delete(streamId); this.replayEventWriteQueues.delete(streamId); + this.tokenUsageWriteQueues.delete(streamId); // For error jobs, DON'T delete immediately - keep around so late-connecting // clients can receive the error. This handles the race condition where error @@ -794,6 +798,7 @@ class GenerationJobManagerClass { this.jobStore.clearContentState(streamId); this.runStepBuffers?.delete(streamId); this.replayEventWriteQueues.delete(streamId); + this.tokenUsageWriteQueues.delete(streamId); // Immediate cleanup if configured (default: true) if (this._cleanupOnComplete) { @@ -1053,6 +1058,8 @@ class GenerationJobManagerClass { await this.trackUserMessage(streamId, event); await this.trackTitleEvent(streamId, event); await this.trackReplayEvent(streamId, event); + await this.trackContextUsage(streamId, event); + await this.trackTokenUsage(streamId, event); // For Redis mode, persist chunk for later reconstruction (fire-and-forget for resumability) if (this._isRedis) { @@ -1150,6 +1157,48 @@ class GenerationJobManagerClass { }); } + /** + * Persist the latest context usage snapshot (one per model call) so a + * resuming client can restore the context gauge without waiting for the + * next model call. + */ + private async trackContextUsage(streamId: string, event: t.ServerSentEvent): Promise { + if (!('event' in event) || event.event !== UsageEvents.ON_CONTEXT_USAGE) { + return; + } + + await this.jobStore.updateJob(streamId, { + contextUsage: JSON.stringify((event as { data?: unknown }).data ?? null), + }); + } + + /** + * Chains a read/modify/write job update onto the stream's queue so + * concurrent writers can't clobber each other's merged state. + */ + private async queueJobWrite( + queues: Map>, + streamId: string, + write: () => Promise, + ): Promise { + const previousWrite = queues.get(streamId) ?? Promise.resolve(); + const nextWrite = previousWrite + .catch(() => { + // Keep the queue moving even if a prior metadata write failed. + }) + .then(write); + + queues.set(streamId, nextWrite); + + try { + await nextWrite; + } finally { + if (queues.get(streamId) === nextWrite) { + queues.delete(streamId); + } + } + } + /** * Persist replay-only stream events that are needed to reconstruct active * UI state on resume but are not represented by aggregated message content. @@ -1159,22 +1208,45 @@ class GenerationJobManagerClass { return; } - const previousWrite = this.replayEventWriteQueues.get(streamId) ?? Promise.resolve(); - const nextWrite = previousWrite - .catch(() => { - // Keep the queue moving even if a prior replay metadata write failed. - }) - .then(() => this.persistReplayEvent(streamId, event)); + await this.queueJobWrite(this.replayEventWriteQueues, streamId, () => + this.persistReplayEvent(streamId, event), + ); + } - this.replayEventWriteQueues.set(streamId, nextWrite); + /** + * Persist per-model-call token usage so resuming clients can rebuild + * usage totals on any replica (the live collectedUsage array only exists + * on the generating instance). + */ + private async trackTokenUsage(streamId: string, event: t.ServerSentEvent): Promise { + if (!('event' in event) || event.event !== UsageEvents.ON_TOKEN_USAGE) { + return; + } - try { - await nextWrite; - } finally { - if (this.replayEventWriteQueues.get(streamId) === nextWrite) { - this.replayEventWriteQueues.delete(streamId); + await this.queueJobWrite(this.tokenUsageWriteQueues, streamId, () => + this.persistTokenUsage(streamId, event as { data?: unknown }), + ); + } + + private async persistTokenUsage(streamId: string, event: { data?: unknown }): Promise { + const jobData = await this.jobStore.getJob(streamId); + if (!jobData || event.data == null) { + return; + } + + let tokenUsage: unknown[] = []; + if (jobData.tokenUsage) { + try { + tokenUsage = JSON.parse(jobData.tokenUsage) as unknown[]; + } catch { + tokenUsage = []; } } + tokenUsage.push(event.data); + + await this.jobStore.updateJob(streamId, { + tokenUsage: JSON.stringify(tokenUsage), + }); } private async persistReplayEvent(streamId: string, event: t.ServerSentEvent): Promise { @@ -1342,10 +1414,32 @@ class GenerationJobManagerClass { } } + let contextUsage: t.ResumeState['contextUsage']; + if (jobData.contextUsage) { + try { + contextUsage = JSON.parse(jobData.contextUsage) as t.ResumeState['contextUsage']; + } catch { + // Ignore malformed persisted context usage. + } + } + + /** Persisted per model call by trackTokenUsage — unlike the live + * collectedUsage reference, this survives cross-replica resumes. */ + let collectedUsage: t.ResumeState['collectedUsage']; + if (jobData.tokenUsage) { + try { + const parsed = JSON.parse(jobData.tokenUsage) as t.ResumeState['collectedUsage']; + collectedUsage = parsed && parsed.length > 0 ? parsed : undefined; + } catch { + // Ignore malformed persisted token usage. + } + } + logger.debug(`[GenerationJobManager] getResumeState:`, { streamId, runStepsLength: runSteps.length, aggregatedContentLength: aggregatedContent.length, + collectedUsageLength: collectedUsage?.length ?? 0, }); return { @@ -1359,6 +1453,8 @@ class GenerationJobManagerClass { model: jobData.model, titleEvent, replayEvents, + collectedUsage, + contextUsage, }; } @@ -1576,6 +1672,7 @@ class GenerationJobManagerClass { this.syncRunningJobMetrics(); this.runStepBuffers?.clear(); this.replayEventWriteQueues.clear(); + this.tokenUsageWriteQueues.clear(); logger.debug('[GenerationJobManager] Destroyed'); } diff --git a/packages/api/src/stream/__tests__/usageResume.spec.ts b/packages/api/src/stream/__tests__/usageResume.spec.ts new file mode 100644 index 0000000000..59564730f9 --- /dev/null +++ b/packages/api/src/stream/__tests__/usageResume.spec.ts @@ -0,0 +1,96 @@ +import { InMemoryEventTransport } from '~/stream/implementations/InMemoryEventTransport'; +import { InMemoryJobStore } from '~/stream/implementations/InMemoryJobStore'; +import { GenerationJobManagerClass } from '~/stream/GenerationJobManager'; + +jest.spyOn(console, 'log').mockImplementation(); + +function createInMemoryManager(): GenerationJobManagerClass { + const manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore: new InMemoryJobStore({ ttlAfterComplete: 60000 }), + eventTransport: new InMemoryEventTransport(), + isRedis: false, + }); + manager.initialize(); + return manager; +} + +describe('GenerationJobManager usage resume state', () => { + let manager: GenerationJobManagerClass | undefined; + + afterEach(async () => { + await manager?.destroy(); + manager = undefined; + }); + + test('accumulates persisted token usage events in resume state', async () => { + manager = createInMemoryManager(); + const streamId = `usage-resume-${Date.now()}`; + await manager.createJob(streamId, 'user-1', streamId); + + const firstUsage = { + input_tokens: 1200, + output_tokens: 300, + input_token_details: { cache_creation: 100, cache_read: 800 }, + model: 'claude-3-5-sonnet', + provider: 'anthropic', + }; + const secondUsage = { + input_tokens: 1600, + output_tokens: 120, + model: 'claude-3-5-sonnet', + provider: 'anthropic', + usage_type: 'summarization', + }; + + await manager.emitChunk(streamId, { event: 'on_token_usage', data: firstUsage }); + await manager.emitChunk(streamId, { event: 'on_token_usage', data: secondUsage }); + + const resumeState = await manager.getResumeState(streamId); + expect(resumeState?.collectedUsage).toEqual([firstUsage, secondUsage]); + }); + + test('omits collected usage when none was recorded', async () => { + manager = createInMemoryManager(); + const streamId = `usage-resume-empty-${Date.now()}`; + await manager.createJob(streamId, 'user-1', streamId); + + const resumeState = await manager.getResumeState(streamId); + expect(resumeState?.collectedUsage).toBeUndefined(); + expect(resumeState?.contextUsage).toBeUndefined(); + }); + + test('persists the latest context usage snapshot for resume', async () => { + manager = createInMemoryManager(); + const streamId = `context-resume-${Date.now()}`; + await manager.createJob(streamId, 'user-1', streamId); + + const makeSnapshot = (messageTokens: number) => ({ + runId: 'run-1', + agentId: 'agent-1', + breakdown: { + maxContextTokens: 200000, + instructionTokens: 1500, + systemMessageTokens: 1000, + dynamicInstructionTokens: 0, + toolSchemaTokens: 500, + summaryTokens: 0, + toolCount: 3, + messageCount: 4, + messageTokens, + availableForMessages: 188500, + }, + contextBudget: 190000, + effectiveInstructionTokens: 1500, + prePruneContextTokens: messageTokens, + remainingContextTokens: 190000 - 1500 - messageTokens, + calibrationRatio: 1.05, + }); + + await manager.emitChunk(streamId, { event: 'on_context_usage', data: makeSnapshot(4000) }); + await manager.emitChunk(streamId, { event: 'on_context_usage', data: makeSnapshot(9000) }); + + const resumeState = await manager.getResumeState(streamId); + expect(resumeState?.contextUsage).toEqual(makeSnapshot(9000)); + }); +}); diff --git a/packages/api/src/stream/implementations/RedisJobStore.ts b/packages/api/src/stream/implementations/RedisJobStore.ts index d49affdc28..a00eeba6b4 100644 --- a/packages/api/src/stream/implementations/RedisJobStore.ts +++ b/packages/api/src/stream/implementations/RedisJobStore.ts @@ -923,6 +923,8 @@ export class RedisJobStore implements IJobStore { promptTokens: data.promptTokens ? parseInt(data.promptTokens, 10) : undefined, titleEvent: data.titleEvent || undefined, replayEvents: data.replayEvents || undefined, + contextUsage: data.contextUsage || undefined, + tokenUsage: data.tokenUsage || undefined, }; } } diff --git a/packages/api/src/stream/interfaces/IJobStore.ts b/packages/api/src/stream/interfaces/IJobStore.ts index 3e6786a49b..aab3c74d22 100644 --- a/packages/api/src/stream/interfaces/IJobStore.ts +++ b/packages/api/src/stream/interfaces/IJobStore.ts @@ -48,6 +48,12 @@ export interface SerializableJobData { /** Serialized replay-only stream events for active-stream resume */ replayEvents?: string; + /** Serialized latest context usage snapshot for active-stream resume */ + contextUsage?: string; + + /** Serialized token usage events for active-stream resume (cross-replica safe) */ + tokenUsage?: string; + /** Endpoint metadata for abort handling - avoids storing functions */ endpoint?: string; iconURL?: string; diff --git a/packages/data-provider/src/api-endpoints.ts b/packages/data-provider/src/api-endpoints.ts index bc3107d7ca..09864bea43 100644 --- a/packages/data-provider/src/api-endpoints.ts +++ b/packages/data-provider/src/api-endpoints.ts @@ -148,6 +148,8 @@ export const deletePreset = () => `${BASE_URL}/api/presets/delete`; export const aiEndpoints = () => `${BASE_URL}/api/endpoints`; +export const tokenConfig = () => `${BASE_URL}/api/endpoints/token-config`; + export const models = () => `${BASE_URL}/api/models`; export const tokenizer = () => `${BASE_URL}/api/tokenizer`; diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index 6309cfc279..5b50388dfd 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -1167,6 +1167,8 @@ export const interfaceSchema = z retainAgentFiles: z.boolean().optional(), runCode: z.boolean().optional(), webSearch: z.boolean().optional(), + contextUsage: z.boolean().optional(), + contextCost: z.boolean().optional(), peoplePicker: z .object({ users: z.boolean().optional(), @@ -1236,6 +1238,8 @@ export const interfaceSchema = z autoSubmitFromUrl: true, runCode: true, webSearch: true, + contextUsage: true, + contextCost: false, peoplePicker: { users: true, groups: true, diff --git a/packages/data-provider/src/data-service.ts b/packages/data-provider/src/data-service.ts index 95d362ac34..e008e8c514 100644 --- a/packages/data-provider/src/data-service.ts +++ b/packages/data-provider/src/data-service.ts @@ -245,6 +245,10 @@ export const getAIEndpoints = (): Promise => { return request.get(endpoints.aiEndpoints()); }; +export const getTokenConfig = (): Promise => { + return request.get(endpoints.tokenConfig()); +}; + export const getModels = async (): Promise => { return request.get(endpoints.models()); }; diff --git a/packages/data-provider/src/keys.ts b/packages/data-provider/src/keys.ts index 7a66cfd80d..04a65beb51 100644 --- a/packages/data-provider/src/keys.ts +++ b/packages/data-provider/src/keys.ts @@ -12,6 +12,7 @@ export enum QueryKeys { models = 'models', balance = 'balance', endpoints = 'endpoints', + tokenConfig = 'tokenConfig', presets = 'presets', searchResults = 'searchResults', tokenCount = 'tokenCount', diff --git a/packages/data-provider/src/schemas.ts b/packages/data-provider/src/schemas.ts index 4503d2af05..42dd6927dc 100644 --- a/packages/data-provider/src/schemas.ts +++ b/packages/data-provider/src/schemas.ts @@ -748,6 +748,8 @@ export const tMessageSchema = z.object({ feedback: feedbackSchema.optional(), /** metadata */ metadata: z.record(z.unknown()).optional(), + /** Output tokens for assistant messages, calibrated prompt-side estimate for user messages */ + tokenCount: z.number().optional(), contextMeta: z .object({ calibrationRatio: z diff --git a/packages/data-provider/src/types.ts b/packages/data-provider/src/types.ts index ac069b210d..68a0bdde92 100644 --- a/packages/data-provider/src/types.ts +++ b/packages/data-provider/src/types.ts @@ -454,6 +454,18 @@ export type TEndpointsConfig = export type TModelsConfig = Record; +/** Server-resolved context window and pricing for one model. Rates are USD per 1M tokens. */ +export type TModelTokenomics = { + context?: number; + prompt?: number; + completion?: number; + cacheWrite?: number; + cacheRead?: number; +}; + +/** endpoint → model → resolved tokenomics, from GET /api/endpoints/token-config */ +export type TTokenConfigMap = Record>; + export type TUpdateTokenCountResponse = { count: number; }; diff --git a/packages/data-provider/src/types/agents.ts b/packages/data-provider/src/types/agents.ts index 0eab753bd3..0dc96fb6cc 100644 --- a/packages/data-provider/src/types/agents.ts +++ b/packages/data-provider/src/types/agents.ts @@ -1,5 +1,6 @@ /* eslint-disable @typescript-eslint/no-namespace */ import type { FunctionToolCall, SummaryContentPart } from './assistants'; +import type { TTokenUsageEvent, TContextUsageEvent } from './runs'; import type { TAttachment, TPlugin } from 'src/schemas'; import { StepTypes, ContentTypes, ToolCallTypes } from './runs'; @@ -231,6 +232,10 @@ export namespace Agents { data?: unknown; [key: string]: unknown; }>; + /** Cumulative provider-reported usage for the run; backfills usage totals on resume */ + collectedUsage?: TTokenUsageEvent[]; + /** Latest context window snapshot; restores the usage gauge on resume */ + contextUsage?: TContextUsageEvent; } /** * Represents a run step delta i.e. any changed fields on a run step during diff --git a/packages/data-provider/src/types/runs.ts b/packages/data-provider/src/types/runs.ts index 706c380c6f..b086081b62 100644 --- a/packages/data-provider/src/types/runs.ts +++ b/packages/data-provider/src/types/runs.ts @@ -40,6 +40,57 @@ export enum StepEvents { ON_SUBAGENT_UPDATE = 'on_subagent_update', } +/** Token-tracking event names streamed to the client (separate from StepEvents dispatch). */ +export enum UsageEvents { + ON_CONTEXT_USAGE = 'on_context_usage', + ON_TOKEN_USAGE = 'on_token_usage', +} + +/** Mirrors TokenBudgetBreakdown from @librechat/agents (data-provider cannot import it). */ +export type TTokenBudgetBreakdown = { + maxContextTokens: number; + instructionTokens: number; + systemMessageTokens: number; + dynamicInstructionTokens: number; + toolSchemaTokens: number; + summaryTokens: number; + toolCount: number; + messageCount: number; + messageTokens: number; + availableForMessages: number; +}; + +/** Per-model-call context snapshot, dispatched after pruning and before the LLM call. */ +export type TContextUsageEvent = { + runId?: string; + agentId?: string; + breakdown: TTokenBudgetBreakdown; + /** Usable budget this call: maxContextTokens minus output reserve */ + contextBudget?: number; + /** Calibrated instruction overhead actually applied this call */ + effectiveInstructionTokens?: number; + /** Calibrated message tokens before pruning (excluding instructions) */ + prePruneContextTokens?: number; + /** Tokens still free after instructions + pruned messages */ + remainingContextTokens?: number; + calibrationRatio?: number; +}; + +/** Provider-reported usage for a single completed model call. */ +export type TTokenUsageEvent = { + input_tokens?: number; + output_tokens?: number; + total_tokens?: number; + input_token_details?: { + cache_creation?: number; + cache_read?: number; + }; + model?: string; + provider?: string; + usage_type?: 'summarization'; + runId?: string; +}; + /** Lifecycle phase carried on subagent-progress envelopes (mirrors SDK SubagentUpdatePhase). */ export type SubagentUpdatePhase = | 'start' diff --git a/packages/data-schemas/src/index.ts b/packages/data-schemas/src/index.ts index 32fb4aa85b..9753ca6e6d 100644 --- a/packages/data-schemas/src/index.ts +++ b/packages/data-schemas/src/index.ts @@ -14,6 +14,7 @@ export { cacheTokenValues, premiumTokenValues, defaultRate, + createTxMethods, permissionBitSupersets, partitionIssues, validateSkillName, diff --git a/packages/data-schemas/src/methods/index.ts b/packages/data-schemas/src/methods/index.ts index 6bab02a4f4..9ea9e5f1e5 100644 --- a/packages/data-schemas/src/methods/index.ts +++ b/packages/data-schemas/src/methods/index.ts @@ -89,7 +89,7 @@ import { createAgentMethods, type AgentMethods, type AgentDeps } from './agent'; import { createConfigMethods, type ConfigMethods } from './config'; export { RoleConflictError, DEFAULT_REFRESH_TOKEN_EXPIRY, DEFAULT_SESSION_EXPIRY }; -export { tokenValues, cacheTokenValues, premiumTokenValues, defaultRate }; +export { tokenValues, cacheTokenValues, premiumTokenValues, defaultRate, createTxMethods }; export { permissionBitSupersets }; export { partitionIssues,