From a5f9ae3351bcd3965b00ed7cd8a1c2aae8b2dbc7 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sat, 13 Jun 2026 13:52:30 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=A9=B9=20fix:=20Move=20Token-Config=20Res?= =?UTF-8?q?olution=20to=20TS;=20Key=20Live=20Usage=20by=20Created=20Convo?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - extract the token-config resolution (override gathering + cache lookup + buildTokenConfigMap) into resolveTokenConfigMap in packages/api, leaving the /api controller a thin request-scoped wrapper (CLAUDE.md TS rule) - getConvoKey prefers the user message's real conversationId once the `created` event stamps it, so a new chat's first-response live gauge and totals land under the id TokenUsage subscribes to instead of NEW_CONVO --- .../controllers/TokenConfigController.js | 37 ++------- client/src/hooks/SSE/useUsageHandler.ts | 9 ++- packages/api/src/endpoints/index.ts | 1 + .../api/src/endpoints/tokenConfig.spec.ts | 79 +++++++++++++++++++ packages/api/src/endpoints/tokenConfig.ts | 52 ++++++++++++ 5 files changed, 145 insertions(+), 33 deletions(-) create mode 100644 packages/api/src/endpoints/tokenConfig.spec.ts create mode 100644 packages/api/src/endpoints/tokenConfig.ts diff --git a/api/server/controllers/TokenConfigController.js b/api/server/controllers/TokenConfigController.js index f6cf7ac554..f139effd6b 100644 --- a/api/server/controllers/TokenConfigController.js +++ b/api/server/controllers/TokenConfigController.js @@ -1,47 +1,20 @@ const { logger } = require('@librechat/data-schemas'); -const { EModelEndpoint, normalizeEndpointName } = require('librechat-data-provider'); -const { buildTokenConfigMap, getTokenConfigKey, tokenConfigCache } = require('@librechat/api'); +const { resolveTokenConfigMap } = 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. + * `interface.contextCost` is enabled) for every configured model. Resolution + * lives in `@librechat/api`; this controller only supplies request-scoped deps. * @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) { - /** Models config and the token-config cache key by the normalized name */ - const name = normalizeEndpointName(endpointConfig?.name); - if (!name) { - continue; - } - if (endpointConfig.tokenConfig != null) { - endpointTokenConfigs[name] = endpointConfig.tokenConfig; - continue; - } - /** Model fetches and chat initialization both store under this key — - * user-scoped whenever the fetched config can be user-specific, so a - * plain-name fallback would risk serving another user's entry */ - 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 }, + const tokenConfigMap = await resolveTokenConfigMap( + { appConfig: req.config, modelsConfig, userId: req.user.id }, { getValueKey, getMultiplier, getCacheMultiplier }, ); res.json(tokenConfigMap); diff --git a/client/src/hooks/SSE/useUsageHandler.ts b/client/src/hooks/SSE/useUsageHandler.ts index eece21fe1c..e4aac5509e 100644 --- a/client/src/hooks/SSE/useUsageHandler.ts +++ b/client/src/hooks/SSE/useUsageHandler.ts @@ -28,7 +28,7 @@ import { const FLUSH_INTERVAL_MS = 250; interface UsageSubmissionLike { - userMessage?: Pick | null; + userMessage?: Pick | null; conversation?: Partial> | null; } @@ -54,6 +54,13 @@ export interface UsageHandlers { } function getConvoKey(submission: UsageSubmissionLike): string { + /** On a new chat's first turn the `created` event stamps the real id onto + * the user message while the submission's conversation is still `new`; + * prefer the real id so live writes land where TokenUsage is subscribed */ + const fromUserMessage = submission.userMessage?.conversationId; + if (fromUserMessage != null && fromUserMessage !== Constants.NEW_CONVO) { + return fromUserMessage; + } return submission.conversation?.conversationId ?? Constants.NEW_CONVO; } diff --git a/packages/api/src/endpoints/index.ts b/packages/api/src/endpoints/index.ts index 98e0af0298..9e6e9dbac0 100644 --- a/packages/api/src/endpoints/index.ts +++ b/packages/api/src/endpoints/index.ts @@ -6,3 +6,4 @@ export * from './google'; export * from './models'; export * from './openai'; export * from './pricing'; +export * from './tokenConfig'; diff --git a/packages/api/src/endpoints/tokenConfig.spec.ts b/packages/api/src/endpoints/tokenConfig.spec.ts new file mode 100644 index 0000000000..a6ad3d27d4 --- /dev/null +++ b/packages/api/src/endpoints/tokenConfig.spec.ts @@ -0,0 +1,79 @@ +import { EModelEndpoint } from 'librechat-data-provider'; +import type { AppConfig } from '@librechat/data-schemas'; +import type { TokenomicsDeps } from './pricing'; + +const mockCacheGet = jest.fn(); +jest.mock('~/cache', () => ({ + tokenConfigCache: jest.fn(() => ({ get: (key: string) => mockCacheGet(key) })), +})); + +import { resolveTokenConfigMap } from './tokenConfig'; + +const deps: TokenomicsDeps = { + getValueKey: () => 'value-key', + getMultiplier: ({ tokenType }) => (tokenType === 'prompt' ? 5 : 15), + getCacheMultiplier: () => null, +}; + +function appConfigWith(custom: unknown[], contextCost = true): AppConfig { + return { + interfaceConfig: { contextCost }, + endpoints: { [EModelEndpoint.custom]: custom }, + } as unknown as AppConfig; +} + +describe('resolveTokenConfigMap', () => { + beforeEach(() => mockCacheGet.mockReset().mockResolvedValue(null)); + + it('uses a static tokenConfig override without consulting the cache', async () => { + const appConfig = appConfigWith([ + { + name: 'MyProxy', + tokenConfig: { 'custom-model': { prompt: 1.5, completion: 4.5, context: 32000 } }, + }, + ]); + + const map = await resolveTokenConfigMap( + { appConfig, modelsConfig: { MyProxy: ['custom-model'] }, userId: 'user-1' }, + deps, + ); + + expect(map.MyProxy['custom-model']).toEqual({ context: 32000, prompt: 1.5, completion: 4.5 }); + expect(mockCacheGet).not.toHaveBeenCalled(); + }); + + it('falls back to the cached fetched config when no static override exists', async () => { + mockCacheGet.mockResolvedValue({ 'custom-model': { prompt: 2, completion: 6, context: 8000 } }); + const appConfig = appConfigWith([{ name: 'MyProxy', models: { fetch: true } }]); + + const map = await resolveTokenConfigMap( + { appConfig, modelsConfig: { MyProxy: ['custom-model'] }, userId: 'user-1' }, + deps, + ); + + expect(mockCacheGet).toHaveBeenCalled(); + expect(map.MyProxy['custom-model'].context).toBe(8000); + expect(map.MyProxy['custom-model'].prompt).toBe(2); + }); + + it('omits pricing when contextCost is disabled', async () => { + const appConfig = appConfigWith( + [ + { + name: 'MyProxy', + tokenConfig: { 'custom-model': { prompt: 1.5, completion: 4.5, context: 32000 } }, + }, + ], + false, + ); + + const map = await resolveTokenConfigMap( + { appConfig, modelsConfig: { MyProxy: ['custom-model'] }, userId: 'user-1' }, + deps, + ); + + expect(map.MyProxy['custom-model'].context).toBe(32000); + expect(map.MyProxy['custom-model'].prompt).toBeUndefined(); + expect(map.MyProxy['custom-model'].completion).toBeUndefined(); + }); +}); diff --git a/packages/api/src/endpoints/tokenConfig.ts b/packages/api/src/endpoints/tokenConfig.ts new file mode 100644 index 0000000000..1c29220e04 --- /dev/null +++ b/packages/api/src/endpoints/tokenConfig.ts @@ -0,0 +1,52 @@ +import { EModelEndpoint, normalizeEndpointName } from 'librechat-data-provider'; +import type { TModelsConfig, TTokenConfigMap, TEndpoint } from 'librechat-data-provider'; +import type { AppConfig } from '@librechat/data-schemas'; +import type { EndpointTokenConfig } from '~/types'; +import { buildTokenConfigMap, type TokenomicsDeps } from '~/endpoints/pricing'; +import { getTokenConfigKey } from '~/endpoints/custom/initialize'; +import { tokenConfigCache } from '~/cache'; + +export interface ResolveTokenConfigParams { + appConfig?: AppConfig; + /** endpoint → model list, from the resolved models config */ + modelsConfig: TModelsConfig; + userId: string; +} + +/** + * Server-side token-config resolution for the `/endpoints/token-config` + * route. Gathers each custom endpoint's override — a static yaml `tokenConfig` + * or the cached fetched config — then resolves context windows (and pricing, + * when `interface.contextCost` is enabled) via {@link buildTokenConfigMap}. + * Lives in TypeScript so the `/api` controller stays a thin wrapper. + */ +export async function resolveTokenConfigMap( + { appConfig, modelsConfig, userId }: ResolveTokenConfigParams, + deps: TokenomicsDeps, +): Promise { + const includePricing = appConfig?.interfaceConfig?.contextCost === true; + const customEndpoints = (appConfig?.endpoints?.[EModelEndpoint.custom] ?? []) as TEndpoint[]; + const cache = tokenConfigCache(); + const endpointTokenConfigs: Record = {}; + + for (const endpointConfig of customEndpoints) { + /** Models config and the token-config cache key by the normalized name */ + const name = normalizeEndpointName(endpointConfig?.name); + if (!name) { + continue; + } + if (endpointConfig.tokenConfig != null) { + endpointTokenConfigs[name] = endpointConfig.tokenConfig as EndpointTokenConfig; + continue; + } + /** Model fetches and chat initialization both store under this key — + * user-scoped whenever the fetched config can be user-specific */ + const tokenKey = getTokenConfigKey(endpointConfig, name, userId); + const cached = (await cache.get(tokenKey)) as EndpointTokenConfig | undefined; + if (cached) { + endpointTokenConfigs[name] = cached; + } + } + + return buildTokenConfigMap({ modelsConfig, endpointTokenConfigs, includePricing }, deps); +}