diff --git a/client/src/components/Chat/Input/ChatForm.tsx b/client/src/components/Chat/Input/ChatForm.tsx index a0a2822187..c1efba72fc 100644 --- a/client/src/components/Chat/Input/ChatForm.tsx +++ b/client/src/components/Chat/Input/ChatForm.tsx @@ -447,6 +447,7 @@ function ChatFormWrapper({ index = 0, placeholder }: { index?: number; placehold conversation?.spec, conversation?.useResponsesApi, conversation?.model, + conversation?.maxContextTokens, hasMessages, ], ); diff --git a/client/src/hooks/SSE/useResumableSSE.ts b/client/src/hooks/SSE/useResumableSSE.ts index 71b3a98789..3e8ade6718 100644 --- a/client/src/hooks/SSE/useResumableSSE.ts +++ b/client/src/hooks/SSE/useResumableSSE.ts @@ -481,6 +481,7 @@ export default function useResumableSSE( contextHandler, usageHandler, tapStream, + tapContent, finalizeUsage, backfillUsage, resetLive, @@ -798,6 +799,7 @@ export default function useResumableSSE( if (text != null && index !== textIndex) { textIndex = index; } + tapContent(text, { ...currentSubmission, userMessage }); contentHandler({ data, submission: currentSubmission as EventSubmission }); return; } @@ -1067,6 +1069,7 @@ export default function useResumableSSE( contextHandler, usageHandler, tapStream, + tapContent, finalizeUsage, backfillUsage, resetLive, diff --git a/client/src/hooks/SSE/useSSE.ts b/client/src/hooks/SSE/useSSE.ts index 9a2ba55870..2e4e4b60cd 100644 --- a/client/src/hooks/SSE/useSSE.ts +++ b/client/src/hooks/SSE/useSSE.ts @@ -67,7 +67,8 @@ export default function useSSE( const balanceQuery = useGetUserBalance({ enabled: !!isAuthenticated && startupConfig?.balance?.enabled, }); - const { contextHandler, usageHandler, tapStream, finalizeUsage, resetLive } = useUsageHandler(); + const { contextHandler, usageHandler, tapStream, tapContent, finalizeUsage, resetLive } = + useUsageHandler(); useEffect(() => { if (submission == null || Object.keys(submission).length === 0) { @@ -148,6 +149,7 @@ export default function useSSE( textIndex = index; } + tapContent(text, { ...submission, userMessage }); contentHandler({ data, submission: submission as EventSubmission }); } else { const text = data.text ?? data.response; diff --git a/client/src/hooks/SSE/useUsageHandler.ts b/client/src/hooks/SSE/useUsageHandler.ts index a49ba384d8..bdace63922 100644 --- a/client/src/hooks/SSE/useUsageHandler.ts +++ b/client/src/hooks/SSE/useUsageHandler.ts @@ -41,6 +41,9 @@ export interface UsageHandlers { contextHandler: (data: TContextUsageEvent, submission: UsageSubmissionLike) => void; usageHandler: (data: TTokenUsageEvent, submission: UsageSubmissionLike) => void; tapStream: (data: { delta?: { content?: unknown } }, submission: UsageSubmissionLike) => void; + /** Live estimate for the legacy content path, which streams cumulative + * (not incremental) text per part — sets rather than accumulates */ + tapContent: (text: 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 */ @@ -53,6 +56,15 @@ function getConvoKey(submission: UsageSubmissionLike): string { return submission.conversation?.conversationId ?? Constants.NEW_CONVO; } +/** Cumulative text of a content-path part: a raw string or a `{ value }` part */ +function extractContentText(text: unknown): string { + if (typeof text === 'string') { + return text; + } + const value = (text as { value?: unknown })?.value; + return typeof value === 'string' ? value : ''; +} + function countDeltaChars(content: unknown): number { const parts = Array.isArray(content) ? content : [content]; let chars = 0; @@ -167,6 +179,23 @@ export default function useUsageHandler(): UsageHandlers { setLive(convoKey, confirmedRef.current + estimateTokens(streamCharsRef.current, ratio)); }; + const tapContent: UsageHandlers['tapContent'] = (text, submission) => { + const value = extractContentText(text); + if (value.length === 0) { + return; + } + /** Cumulative per part — replace the running char count, don't add */ + streamCharsRef.current = value.length; + 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; @@ -236,6 +265,7 @@ export default function useUsageHandler(): UsageHandlers { contextHandler, usageHandler, tapStream, + tapContent, finalizeUsage, resetLive, backfillUsage, diff --git a/packages/api/src/endpoints/custom/initialize.spec.ts b/packages/api/src/endpoints/custom/initialize.spec.ts index cafb06ec19..9aa7f0112a 100644 --- a/packages/api/src/endpoints/custom/initialize.spec.ts +++ b/packages/api/src/endpoints/custom/initialize.spec.ts @@ -351,4 +351,34 @@ describe('initializeCustom – token-config fetch header forwarding', () => { }), ); }); + + it('uses a static tokenConfig for billing and skips the model fetch', async () => { + const tokenConfig = { + 'gpt-4': { prompt: 1.5, completion: 4.5, context: 32000, cacheRead: 0.3, cacheWrite: 1.8 }, + }; + mockGetCustomEndpointConfig.mockReturnValue({ + apiKey: 'sk-test-key', + baseURL: 'https://openrouter.ai/api/v1', + models: { fetch: true }, + tokenConfig, + }); + + const params: BaseInitializeParams = { + req: { + user: { id: 'user-1', email: 'user@example.com' }, + body: { key: '2099-01-01' }, + config: {}, + } as unknown as BaseInitializeParams['req'], + endpoint: 'openrouter', + model_parameters: { model: 'gpt-4' }, + db: { + getUserKeyValues: jest.fn(), + } as unknown as BaseInitializeParams['db'], + }; + + const result = (await initializeCustom(params)) as { endpointTokenConfig?: unknown }; + + expect(fetchModels).not.toHaveBeenCalled(); + expect(result.endpointTokenConfig).toEqual(tokenConfig); + }); }); diff --git a/packages/api/src/endpoints/custom/initialize.ts b/packages/api/src/endpoints/custom/initialize.ts index 9a983e5a0f..9a9e334dfe 100644 --- a/packages/api/src/endpoints/custom/initialize.ts +++ b/packages/api/src/endpoints/custom/initialize.ts @@ -159,12 +159,16 @@ export async function initializeCustom({ const hasTokenConfig = endpointConfig.tokenConfig != null; const tokenKey = getTokenConfigKey(endpointConfig, endpoint, userId); - const cachedConfig = - !hasTokenConfig && - FetchTokenConfig[endpoint.toLowerCase() as keyof typeof FetchTokenConfig] && - (await cache.get(tokenKey)); - - endpointTokenConfig = (cachedConfig as EndpointTokenConfig) || undefined; + if (hasTokenConfig) { + /** A static override is authoritative — use it for the agent's billing + * and balance checks, not just the advertised UI token config */ + endpointTokenConfig = endpointConfig.tokenConfig as EndpointTokenConfig; + } else { + const cachedConfig = + FetchTokenConfig[endpoint.toLowerCase() as keyof typeof FetchTokenConfig] && + (await cache.get(tokenKey)); + endpointTokenConfig = (cachedConfig as EndpointTokenConfig) || undefined; + } if ( FetchTokenConfig[endpoint.toLowerCase() as keyof typeof FetchTokenConfig] && diff --git a/packages/data-provider/src/react-query/react-query-service.ts b/packages/data-provider/src/react-query/react-query-service.ts index b865ab3841..59430c78d1 100644 --- a/packages/data-provider/src/react-query/react-query-service.ts +++ b/packages/data-provider/src/react-query/react-query-service.ts @@ -4,17 +4,17 @@ import type { UseMutationResult, QueryObserverResult, } from '@tanstack/react-query'; +import { MCPServerConnectionStatusResponse } from '../types/queries'; import { Constants, initialModelsConfig } from '../config'; import { defaultOrderQuery } from '../types/assistants'; -import { MCPServerConnectionStatusResponse } from '../types/queries'; +import * as permissions from '../accessPermissions'; +import { ResourceType } from '../accessPermissions'; import * as dataService from '../data-service'; import * as m from '../types/mutations'; import * as q from '../types/queries'; import { QueryKeys } from '../keys'; import * as s from '../schemas'; import * as t from '../types'; -import * as permissions from '../accessPermissions'; -import { ResourceType } from '../accessPermissions'; export { hasPermissions } from '../accessPermissions'; @@ -122,6 +122,8 @@ export const useUpdateUserKeysMutation = (): UseMutationResult< onSuccess: (data, variables) => { queryClient.invalidateQueries([QueryKeys.name, variables.name]); queryClient.invalidateQueries([QueryKeys.models]); + /** token-config is derived from the same per-user model fetch */ + queryClient.invalidateQueries([QueryKeys.tokenConfig]); }, }); }; @@ -141,6 +143,7 @@ export const useRevokeUserKeyMutation = (name: string): UseMutationResult { queryClient.invalidateQueries([QueryKeys.name, name]); queryClient.invalidateQueries([QueryKeys.models]); + queryClient.invalidateQueries([QueryKeys.tokenConfig]); if (s.isAssistantsEndpoint(name)) { queryClient.invalidateQueries([QueryKeys.assistants, name, defaultOrderQuery]); queryClient.invalidateQueries([QueryKeys.assistantDocs]); @@ -159,6 +162,7 @@ export const useRevokeAllUserKeysMutation = (): UseMutationResult => { return useMutation(() => dataService.revokeAllUserKeys(), { onSuccess: () => { queryClient.invalidateQueries([QueryKeys.name]); + queryClient.invalidateQueries([QueryKeys.tokenConfig]); queryClient.invalidateQueries([ QueryKeys.assistants, s.EModelEndpoint.assistants,