fix: persist accumulated cost in session DB to survive reload (#9191)

Signed-off-by: Douwe Osinga <douwe@squareup.com> Co-authored-by: Douwe Osinga <douwe@squareup.com>
2026-05-14 12:57:06 -04:00
parent c938a46417
commit 4671053985
13 changed files with 100 additions and 192 deletions
@@ -1272,6 +1272,7 @@ export type ScheduledJob = {
 };

 export type Session = {
+    accumulated_cost?: number | null;
    accumulated_input_tokens?: number | null;
    accumulated_output_tokens?: number | null;
    accumulated_total_tokens?: number | null;
@@ -1480,6 +1481,7 @@ export type ThinkingContent = {
 };

 export type TokenState = {
+    accumulatedCost?: number | null;
    accumulatedInputTokens: number;
    accumulatedOutputTokens: number;
    accumulatedTotalTokens: number;
@@ -28,7 +28,6 @@ import { RecipeHeader } from './RecipeHeader';
 import { RecipeWarningModal } from './ui/RecipeWarningModal';
 import { scanRecipe } from '../recipe';
 import { UserInput } from '../types/message';
-import { useCostTracking } from '../hooks/useCostTracking';
 import RecipeActivities from './recipes/RecipeActivities';
 import { useToolCount } from './alerts/useToolCount';
 import { getThinkingMessage, getTextAndImageContent } from '../types/message';
@@ -196,14 +195,6 @@ export default function BaseChat({
    handleSubmit(input);
  };

-  const { sessionCosts } = useCostTracking({
-    sessionInputTokens: session?.accumulated_input_tokens || 0,
-    sessionOutputTokens: session?.accumulated_output_tokens || 0,
-    localInputTokens: 0,
-    localOutputTokens: 0,
-    session,
-  });
-
  const sessionModel = session?.model_config?.model_name ?? null;
  const sessionProvider = session?.provider_name ?? null;
  const sessionLoaded = session !== undefined;
@@ -511,11 +502,14 @@ export default function BaseChat({
            accumulatedOutputTokens={
              tokenState?.accumulatedOutputTokens ?? session?.accumulated_output_tokens ?? undefined
            }
+            accumulatedCost={
+              tokenState?.accumulatedCost ?? session?.accumulated_cost ?? undefined
+            }
            droppedFiles={droppedFiles}
            onFilesProcessed={() => setDroppedFiles([])} // Clear dropped files after processing
            messages={messages}
            disableAnimation={disableAnimation}
-            sessionCosts={sessionCosts}
+
            recipe={recipe}
            recipeAccepted={!hasNotAcceptedRecipe}
            initialPrompt={initialPrompt}
@@ -167,14 +167,8 @@ interface ChatInputProps {
  totalTokens?: number;
  accumulatedInputTokens?: number;
  accumulatedOutputTokens?: number;
+  accumulatedCost?: number | null;
  messages?: Message[];
-  sessionCosts?: {
-    [key: string]: {
-      inputTokens: number;
-      outputTokens: number;
-      totalCost: number;
-    };
-  };
  disableAnimation?: boolean;
  recipe?: Recipe | null;
  recipeId?: string | null;
@@ -203,9 +197,9 @@ export default function ChatInput({
  totalTokens,
  accumulatedInputTokens,
  accumulatedOutputTokens,
+  accumulatedCost,
  messages = [],
  disableAnimation = false,
-  sessionCosts,
  recipe,
  recipeId,
  recipeAccepted,
@@ -1690,7 +1684,7 @@ export default function ChatInput({
                <CostTracker
                  inputTokens={accumulatedInputTokens}
                  outputTokens={accumulatedOutputTokens}
-                  sessionCosts={sessionCosts}
+                  accumulatedCost={accumulatedCost}
                  model={effectiveModel}
                  provider={effectiveProvider}
                />
@@ -107,7 +107,6 @@ export default function Hub({
          onFilesProcessed={() => {}}
          messages={[]}
          disableAnimation={false}
-          sessionCosts={undefined}
          toolCount={0}
          onWorkingDirChange={setWorkingDir}
          inputRef={inputRef}
@@ -14,10 +14,6 @@ const i18n = defineMessages({
    id: 'costTracker.costUnavailable',
    defaultMessage: 'Cost data not available for {model} ({inputTokens} input, {outputTokens} output tokens)',
  },
-  sessionCostBreakdown: {
-    id: 'costTracker.sessionCostBreakdown',
-    defaultMessage: 'Session cost breakdown:',
-  },
  totalSessionCost: {
    id: 'costTracker.totalSessionCost',
    defaultMessage: 'Total session cost: {cost}',
@@ -31,13 +27,7 @@ const i18n = defineMessages({
 interface CostTrackerProps {
  inputTokens?: number;
  outputTokens?: number;
-  sessionCosts?: {
-    [key: string]: {
-      inputTokens: number;
-      outputTokens: number;
-      totalCost: number;
-    };
-  };
+  accumulatedCost?: number | null;
  model: string | null;
  provider: string | null;
 }
@@ -45,7 +35,7 @@ interface CostTrackerProps {
 export function CostTracker({
  inputTokens = 0,
  outputTokens = 0,
-  sessionCosts,
+  accumulatedCost,
  model: currentModel,
  provider: currentProvider,
 }: CostTrackerProps) {
@@ -106,41 +96,7 @@ export function CostTracker({
  }

  const calculateCost = (): number => {
-    // If we have session costs, calculate the total across all models
-    if (sessionCosts) {
-      let totalCost = 0;
-
-      // Add up all historical costs from different models
-      Object.values(sessionCosts).forEach((modelCost) => {
-        totalCost += modelCost.totalCost;
-      });
-
-      // Add current model cost if we have pricing info
-      if (
-        costInfo &&
-        (costInfo.input_token_cost !== undefined || costInfo.output_token_cost !== undefined)
-      ) {
-        const currentInputCost = (inputTokens * (costInfo.input_token_cost || 0)) / 1_000_000;
-        const currentOutputCost = (outputTokens * (costInfo.output_token_cost || 0)) / 1_000_000;
-        totalCost += currentInputCost + currentOutputCost;
-      }
-
-      return totalCost;
-    }
-
-    // Fallback to simple calculation for current model only
-    if (
-      !costInfo ||
-      (costInfo.input_token_cost === undefined && costInfo.output_token_cost === undefined)
-    ) {
-      return 0;
-    }
-
-    const inputCost = (inputTokens * (costInfo.input_token_cost || 0)) / 1_000_000;
-    const outputCost = (outputTokens * (costInfo.output_token_cost || 0)) / 1_000_000;
-    const total = inputCost + outputCost;
-
-    return total;
+    return accumulatedCost ?? 0;
  };

  const formatCost = (cost: number): string => {
@@ -165,10 +121,10 @@ export function CostTracker({
    );
  }

-  // If no cost info found, try to return a default
  if (
-    !costInfo ||
-    (costInfo.input_token_cost === undefined && costInfo.output_token_cost === undefined)
+    accumulatedCost == null &&
+    (!costInfo ||
+      (costInfo.input_token_cost === undefined && costInfo.output_token_cost === undefined))
  ) {
    const freeProviders = ['ollama', 'local', 'localhost'];
    if (freeProviders.includes(currentProvider.toLowerCase())) {
@@ -216,38 +172,22 @@ export function CostTracker({

  // Build tooltip content
  const getTooltipContent = (): string => {
-    // Handle error states first
    if (pricingFailed) {
      return intl.formatMessage(i18n.pricingUnavailable, { model: `${currentProvider}/${currentModel}` });
    }

-    // Handle session costs
-    if (sessionCosts && Object.keys(sessionCosts).length > 0) {
-      // Show session breakdown
-      let tooltip = intl.formatMessage(i18n.sessionCostBreakdown) + '\n';
+    const currency = costInfo?.currency || '$';

-      Object.entries(sessionCosts).forEach(([modelKey, cost]) => {
-        const costStr = `${costInfo?.currency || '$'}${cost.totalCost.toFixed(6)}`;
-        tooltip += `${modelKey}: ${costStr} (${cost.inputTokens.toLocaleString()} in, ${cost.outputTokens.toLocaleString()} out)\n`;
-      });
-
-      // Add current model if it has costs
-      if (costInfo && (inputTokens > 0 || outputTokens > 0)) {
-        const currentCost =
-          (inputTokens * (costInfo.input_token_cost || 0) +
-            outputTokens * (costInfo.output_token_cost || 0)) /
-          1_000_000;
-        if (currentCost > 0) {
-          tooltip += `${currentProvider}/${currentModel} (current): ${costInfo.currency || '$'}${currentCost.toFixed(6)} (${inputTokens.toLocaleString()} in, ${outputTokens.toLocaleString()} out)\n`;
-        }
-      }
-
-      tooltip += '\n' + intl.formatMessage(i18n.totalSessionCost, { cost: `${costInfo?.currency || '$'}${totalCost.toFixed(6)}` });
-      return tooltip;
+    if (accumulatedCost != null) {
+      return intl.formatMessage(i18n.totalSessionCost, { cost: `${currency}${totalCost.toFixed(4)}` })
+        + `\n` + intl.formatMessage(i18n.inputOutputTooltip, {
+          inputTokens: inputTokens.toLocaleString(),
+          inputCost: `${currency}${((inputTokens * (costInfo?.input_token_cost || 0)) / 1_000_000).toFixed(6)}`,
+          outputTokens: outputTokens.toLocaleString(),
+          outputCost: `${currency}${((outputTokens * (costInfo?.output_token_cost || 0)) / 1_000_000).toFixed(6)}`,
+        });
    }

-    // Default tooltip for single model
-    const currency = costInfo?.currency || '$';
    const inputCostStr = `${currency}${((inputTokens * (costInfo?.input_token_cost || 0)) / 1_000_000).toFixed(6)}`;
    const outputCostStr = `${currency}${((outputTokens * (costInfo?.output_token_cost || 0)) / 1_000_000).toFixed(6)}`;
    return intl.formatMessage(i18n.inputOutputTooltip, {
@@ -1,92 +0,0 @@
-import { useEffect, useRef, useState } from 'react';
-import { fetchCanonicalModelInfo } from '../utils/canonical';
-import { Session } from '../api';
-
-interface UseCostTrackingProps {
-  sessionInputTokens: number;
-  sessionOutputTokens: number;
-  localInputTokens: number;
-  localOutputTokens: number;
-  session?: Session | null;
-}
-
-export const useCostTracking = ({
-  sessionInputTokens,
-  sessionOutputTokens,
-  localInputTokens,
-  localOutputTokens,
-  session,
-}: UseCostTrackingProps) => {
-  const [sessionCosts, setSessionCosts] = useState<{
-    [key: string]: {
-      inputTokens: number;
-      outputTokens: number;
-      totalCost: number;
-    };
-  }>({});
-
-  const currentModel = session?.model_config?.model_name ?? undefined;
-  const currentProvider = session?.provider_name ?? undefined;
-  const prevModelRef = useRef<string | undefined>(undefined);
-  const prevProviderRef = useRef<string | undefined>(undefined);
-
-  // Handle model changes and accumulate costs
-  useEffect(() => {
-    if (!currentModel || !currentProvider) return;
-
-    const handleModelChange = async () => {
-      if (
-        prevModelRef.current !== undefined &&
-        prevProviderRef.current !== undefined &&
-        (prevModelRef.current !== currentModel || prevProviderRef.current !== currentProvider)
-      ) {
-        // Model/provider has changed, save the costs for the previous model
-        const prevKey = `${prevProviderRef.current}/${prevModelRef.current}`;
-
-        // Get pricing info for the previous model
-        const prevCostInfo = await fetchCanonicalModelInfo(
-          prevProviderRef.current,
-          prevModelRef.current
-        );
-
-        if (prevCostInfo) {
-          const prevInputCost =
-            ((sessionInputTokens || localInputTokens) * (prevCostInfo.input_token_cost || 0)) /
-            1_000_000;
-          const prevOutputCost =
-            ((sessionOutputTokens || localOutputTokens) * (prevCostInfo.output_token_cost || 0)) /
-            1_000_000;
-          const prevTotalCost = prevInputCost + prevOutputCost;
-
-          // Save the accumulated costs for this model
-          setSessionCosts((prev) => ({
-            ...prev,
-            [prevKey]: {
-              inputTokens: sessionInputTokens || localInputTokens,
-              outputTokens: sessionOutputTokens || localOutputTokens,
-              totalCost: prevTotalCost,
-            },
-          }));
-        }
-
-      }
-
-      prevModelRef.current = currentModel || undefined;
-      prevProviderRef.current = currentProvider || undefined;
-    };
-
-    handleModelChange();
-  }, [
-    currentModel,
-    currentProvider,
-    sessionInputTokens,
-    sessionOutputTokens,
-    localInputTokens,
-    localOutputTokens,
-    session,
-  ]);
-
-  return {
-    sessionCosts,
-  };
-};
@@ -314,9 +314,7 @@
  "costTracker.pricingUnavailable": {
    "defaultMessage": "Pricing data unavailable for {model}"
  },
-  "costTracker.sessionCostBreakdown": {
-    "defaultMessage": "Session cost breakdown:"
-  },
+
  "costTracker.totalSessionCost": {
    "defaultMessage": "Total session cost: {cost}"
  },