🩹 fix: Honor Static Token Config in Billing; Tighten Usage Freshness

- initializeCustom now uses a static endpoint tokenConfig as the agent's
  endpointTokenConfig (billing + balance checks), not just the advertised
  UI config — previously the gauge showed admin rates while the agent
  billed against built-in tables
- invalidate the token-config query alongside models on user-key add/
  revoke so context windows and pricing refresh without a reload
- include maxContextTokens in ChatForm's stabilized conversation memo so
  the gauge reflects a changed context-window setting immediately
- feed the live output estimate from the legacy content path (direct and
  assistants streams), setting from cumulative part text rather than
  accumulating deltas

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,
     ],
   );

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,

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;

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,

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);
+  });
 });

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] &&

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<unknow
     onSuccess: () => {
       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<unknown> => {
   return useMutation(() => dataService.revokeAllUserKeys(), {
     onSuccess: () => {
       queryClient.invalidateQueries([QueryKeys.name]);
+      queryClient.invalidateQueries([QueryKeys.tokenConfig]);
       queryClient.invalidateQueries([
         QueryKeys.assistants,
         s.EModelEndpoint.assistants,