feat: reflexion loop token cost tracking (#87)

bronz3beard · web-flow · commit 33313a1843d4 · 2026-07-16T15:19:04.000-03:00
* feat: track actual token costs in the reflexion loop

- Adds a centralized token pricing module (`PRICE_PER_MTOK`) with safe placeholder rates.
- Updates `buildRunner` (used across both environments) to accurately aggregate tokens into USD costs according to the defined model pricing.
- Gracefully logs warnings and ignores cost tracking for models absent from the pricing configuration.
- Adapts and expands tests to explicitly simulate cost consumption and proper budget triggering via `getUsage()`.

* feat: track actual token costs in the reflexion loop

- Adds a centralized token pricing module (`PRICE_PER_MTOK`) with safe placeholder rates.
- Updates `buildRunner` (used across both environments) to accurately aggregate tokens into USD costs according to the defined model pricing.
- Gracefully logs warnings and ignores cost tracking for models absent from the pricing configuration.
- Adapts and expands tests to explicitly simulate cost consumption and proper budget triggering via `getUsage()`.
- Fixes evaluator calibration in scripts/reflexion-eval.ts by updating the critic model.

* feat: track actual token costs in the reflexion loop

- Adds a centralized token pricing module (`PRICE_PER_MTOK`) with safe placeholder rates.
- Updates `buildRunner` (used across both environments) to accurately aggregate tokens into USD costs according to the defined model pricing.
- Gracefully logs warnings and ignores cost tracking for models absent from the pricing configuration.
- Adapts and expands tests to explicitly simulate cost consumption and proper budget triggering via `getUsage()`.
- Fixes evaluator calibration in scripts/reflexion-eval.ts by updating the critic model.
diff --git a/.ai/cursor-skills.manifest b/.ai/cursor-skills.manifest
@@ -39,6 +39,7 @@ workflow-code-review|.agents/workflows/code-review.md
 workflow-competitive-analysis|.agents/workflows/competitive-analysis.md
 workflow-design-requirements-to-architecture|.agents/workflows/design-requirements-to-architecture.md
 workflow-design-system-review|.agents/workflows/design-system-review.md
+workflow-dev-team-orchestrator|.agents/workflows/dev-team-orchestrator.md
 workflow-dev-team|.agents/workflows/dev-team.md
 workflow-feature-orchestrator|.agents/workflows/feature-orchestrator.md
 workflow-init|.agents/workflows/init.md
@@ -57,4 +58,5 @@ workflow-style-logic-exporter|.agents/workflows/style-logic-exporter.md
 workflow-ui-spec-generator|.agents/workflows/ui-spec-generator.md
 workflow-verify-changes|.agents/workflows/verify-changes.md
 workflow-vertical-slice|.agents/workflows/vertical-slice.md
+workflow-visual-verifier|.agents/workflows/visual-verifier.md
 workflow-weekly-leadership-report|.agents/workflows/weekly-leadership-report.md
diff --git a/src/lib/ai/reflexion/__tests__/engine.test.ts b/src/lib/ai/reflexion/__tests__/engine.test.ts
@@ -4,7 +4,7 @@ import {
   resumeReflexion,
   runReflexion,
 } from '../engine';
-import { Answers, ReflexionStateV2 } from '../schema';
+import { Answers, ReflexionStateV2, Interview } from '../schema';
 
 function createMockRunner(
   generateResponse = 'plan',
@@ -18,7 +18,7 @@ function createMockRunner(
     modernWeb: 9,
   },
   adjudicateResponse = 'verdict',
-  interviewResponse: any = {
+  interviewResponse: Interview = {
     runId: '123',
     revision: 0,
     recommendation: 'approve',
@@ -69,6 +69,21 @@ describe('engine v2', () => {
     expect(runner.generate).not.toHaveBeenCalled();
   });
 
+  it('budget gate does not trip if cost is below maxCostUsd', async () => {
+    const runner = createMockRunner('plan', undefined, undefined, undefined, {
+      tokens: 100,
+      costUsd: 0.5,
+    });
+    const cfg: ReflexionConfig = {
+      brief: 'test',
+      mode: 'auto',
+      budget: { maxCostUsd: 1 },
+    };
+    const res = await runReflexion(runner, cfg);
+    expect(res.stopReason).not.toBe('budget-exceeded');
+    expect(runner.generate).toHaveBeenCalled();
+  });
+
   it('zero-question auto-approve', async () => {
     const runner = createMockRunner(
       'plan',
diff --git a/src/lib/ai/reflexion/__tests__/providers.test.ts b/src/lib/ai/reflexion/__tests__/providers.test.ts
@@ -1,5 +1,15 @@
 import * as aiModule from 'ai';
 import { buildRunner } from '../providers-env';
+import { PRICE_PER_MTOK } from '../pricing';
+import { MODELS } from '@/app/api/chat/constants';
+
+jest.mock('../pricing', () => ({
+  PRICE_PER_MTOK: {
+    'gemini-3.5-flash': { inputUsdPerMTok: 1.0, outputUsdPerMTok: 2.0 },
+    'claude-sonnet-4-6': { inputUsdPerMTok: 3.0, outputUsdPerMTok: 4.0 },
+    'gemini-3.1-pro-preview': { inputUsdPerMTok: 5.0, outputUsdPerMTok: 6.0 },
+  },
+}));
 
 jest.mock('ai', () => ({
   generateText: jest.fn(),
@@ -13,48 +23,83 @@ describe('providers v2', () => {
     jest.clearAllMocks();
   });
 
-  it('buildRunner accumulates usage and provides interview()', async () => {
+  it('buildRunner accumulates usage and provides interview() with missing model pricing', async () => {
     (aiModule.generateText as jest.Mock)
-      .mockResolvedValueOnce({ text: 'text', usage: { totalTokens: 10 } })
+      .mockResolvedValueOnce({ text: 'text', usage: { promptTokens: 5, completionTokens: 5, totalTokens: 10 } })
       .mockResolvedValueOnce({
         output: { runId: '1', questions: [] },
-        usage: { totalTokens: 20 },
+        usage: { promptTokens: 10, completionTokens: 10, totalTokens: 20 },
       });
 
     const spy = jest.spyOn(console, 'warn').mockImplementation(() => {});
 
-    const mockModel: any = {};
+    const mockModel = {} as aiModule.LanguageModel;
     const runner = buildRunner(mockModel, mockModel, mockModel, {
-      creator: 'a',
-      critic: 'b',
-      adjudicator: 'c',
+      creator: 'unknown-a',
+      critic: 'unknown-b',
+      adjudicator: 'unknown-c',
     });
 
     expect(runner.getUsage().tokens).toBe(0);
+    expect(runner.getUsage().costUsd).toBe(0);
 
     await runner.generate('prompt', 'system');
     expect(runner.getUsage().tokens).toBe(10);
+    expect(runner.getUsage().costUsd).toBe(0); // Cost remains 0 for missing model
 
     const interviewRes = await runner.interview('prompt', 'system');
     expect(interviewRes.runId).toBe('1');
     expect(runner.getUsage().tokens).toBe(30);
+    expect(runner.getUsage().costUsd).toBe(0); // Cost remains 0
+
+    expect(spy).toHaveBeenCalledWith(
+      expect.stringContaining('cost tracking not available for provider unknown-a')
+    );
+    expect(spy).toHaveBeenCalledWith(
+      expect.stringContaining('cost tracking not available for provider unknown-b')
+    );
 
     spy.mockRestore();
   });
 
+  it('buildRunner computes correct costUsd using deterministic fixture rates', async () => {
+    // mock generate with gemini-3.5-flash (creator): 1,000 prompt (1.0/M), 2,000 completion (2.0/M) = $0.001 + $0.004 = $0.005
+    (aiModule.generateText as jest.Mock)
+      .mockResolvedValueOnce({ text: 'text', usage: { promptTokens: 1000, completionTokens: 2000, totalTokens: 3000 } })
+      // mock critique with claude-sonnet-4-6 (critic): 2,000 prompt (3.0/M), 1,000 completion (4.0/M) = $0.006 + $0.004 = $0.010
+      .mockResolvedValueOnce({
+        output: { score: 9 },
+        usage: { promptTokens: 2000, completionTokens: 1000, totalTokens: 3000 },
+      });
+
+    const mockModel = {} as aiModule.LanguageModel;
+    const runner = buildRunner(mockModel, mockModel, mockModel, {
+      creator: 'gemini-3.5-flash',
+      critic: 'claude-sonnet-4-6',
+      adjudicator: 'claude-sonnet-4-6',
+    });
+
+    await runner.generate('prompt', 'system');
+    expect(runner.getUsage().costUsd).toBeCloseTo(0.005);
+
+    await runner.critique('prompt', 'system');
+    expect(runner.getUsage().costUsd).toBeCloseTo(0.015); // Total cost = $0.005 + $0.010 = $0.015
+    expect(runner.getUsage().tokens).toBe(6000);
+  });
+
   it('401 on Claude falls back to gemini-3.1-pro-preview', async () => {
-    const mockCreator: any = { id: 'gemini-3.5-flash' };
-    const mockCritic: any = { id: 'claude-sonnet-4-6' };
-    const mockFallbackCritic: any = { id: 'gemini-3.1-pro-preview' };
+    const mockCreator = { id: 'gemini-3.5-flash' } as unknown as aiModule.LanguageModel;
+    const mockCritic = { id: 'claude-sonnet-4-6' } as unknown as aiModule.LanguageModel;
+    const mockFallbackCritic = { id: 'gemini-3.1-pro-preview' } as unknown as aiModule.LanguageModel;
 
-    const error401 = new Error('unauthorized');
-    (error401 as any).status = 401;
+    const error401 = new Error('unauthorized') as Error & { status?: number };
+    error401.status = 401;
 
     (aiModule.generateText as jest.Mock)
       .mockRejectedValueOnce(error401) // First call to Claude fails
       .mockResolvedValueOnce({
         output: { score: 7, passed: false, actionableFix: 'fix' },
-        usage: { totalTokens: 15 },
+        usage: { promptTokens: 1_000_000, completionTokens: 1_000_000, totalTokens: 2_000_000 },
       }); // Fallback succeeds
 
     const runner = buildRunner(
@@ -86,12 +131,15 @@ describe('providers v2', () => {
         model: mockFallbackCritic,
       })
     );
+
+    // Pro preview mock rates: $5.0 input, $6.0 output. With 1M tokens each, total is $11.0.
+    expect(runner.getUsage().costUsd).toBeCloseTo(11.0);
   });
 
   it('successful passed=false critique does not trigger fallback and wasDegraded remains false', async () => {
-    const mockCreator: any = { id: 'gemini-3.5-flash' };
-    const mockCritic: any = { id: 'claude-sonnet-4-6' };
-    const mockFallbackCritic: any = { id: 'gemini-3.1-pro-preview' };
+    const mockCreator = { id: 'gemini-3.5-flash' } as unknown as aiModule.LanguageModel;
+    const mockCritic = { id: 'claude-sonnet-4-6' } as unknown as aiModule.LanguageModel;
+    const mockFallbackCritic = { id: 'gemini-3.1-pro-preview' } as unknown as aiModule.LanguageModel;
 
     (aiModule.generateText as jest.Mock).mockResolvedValueOnce({
       output: { score: 5, passed: false, actionableFix: 'fix' },
diff --git a/src/lib/ai/reflexion/pricing.ts b/src/lib/ai/reflexion/pricing.ts
@@ -0,0 +1,23 @@
+import { MODELS } from '@/app/api/chat/constants';
+
+/**
+ * Operator-maintained rates per 1M tokens.
+ * These rates may drift from actual provider pricing over time.
+ * Do not hardcode prices outside of this single source of truth.
+ *
+ * TODO: Operator must fill these in from the providers' official pricing pages.
+ */
+export const PRICE_PER_MTOK: Record<string, { inputUsdPerMTok: number; outputUsdPerMTok: number }> = {
+  [MODELS.GEMINI]: {
+    inputUsdPerMTok: 0.0, // TODO: operator must fill
+    outputUsdPerMTok: 0.0, // TODO: operator must fill
+  },
+  [MODELS.CLAUDE]: {
+    inputUsdPerMTok: 0.0, // TODO: operator must fill
+    outputUsdPerMTok: 0.0, // TODO: operator must fill
+  },
+  [MODELS.GEMINI_FALLBACK_CRITIC]: {
+    inputUsdPerMTok: 0.0, // TODO: operator must fill
+    outputUsdPerMTok: 0.0, // TODO: operator must fill
+  },
+};
diff --git a/src/lib/ai/reflexion/providers-env.ts b/src/lib/ai/reflexion/providers-env.ts
@@ -8,6 +8,7 @@ import {
 } from 'ai';
 import { MODELS } from '../../../app/api/chat/constants';
 import type { ReflexionRunner } from './engine';
+import { PRICE_PER_MTOK } from './pricing';
 import { CritiqueSchema, InterviewSchema } from './schema';
 
 /**
@@ -57,17 +58,33 @@ export function buildRunner(
   fallbackCritic?: LanguageModel
 ): ReflexionRunner {
   let accumulatedTokens = 0;
+  let totalCostUsd = 0;
   let degraded = false;
 
-  function addUsage(usage?: LanguageModelUsage) {
-    if (usage) {
-      accumulatedTokens += usage.totalTokens || 0;
+  const warnedAboutCost = new Set<string>();
+
+  function addUsage(usage: LanguageModelUsage | undefined, modelId: string) {
+    if (!usage) return;
+    accumulatedTokens += usage.totalTokens || 0;
+
+    const pricing = PRICE_PER_MTOK[modelId];
+    if (!pricing) {
+      if (!warnedAboutCost.has(modelId)) {
+        console.warn(
+          `reflexion: cost tracking not available for provider ${modelId}; cost contribution will be 0.`
+        );
+        warnedAboutCost.add(modelId);
+      }
+      return;
     }
-  }
 
-  // Cost tracking is not readily available via standard AI SDK metadata out-of-the-box
-  // for all providers here without mapping model IDs to pricing tables.
-  let warnedAboutCost = false;
+    const usageFallback = usage as { promptTokens?: number; completionTokens?: number } & typeof usage;
+    const inputTokens = usageFallback.promptTokens ?? usage.inputTokens ?? 0;
+    const outputTokens = usageFallback.completionTokens ?? usage.outputTokens ?? 0;
+    const inputCost = (inputTokens / 1_000_000) * pricing.inputUsdPerMTok;
+    const outputCost = (outputTokens / 1_000_000) * pricing.outputUsdPerMTok;
+    totalCostUsd += inputCost + outputCost;
+  }
 
   return {
     models,
@@ -80,7 +97,7 @@ export function buildRunner(
         system,
         prompt,
       });
-      addUsage(usage);
+      addUsage(usage, models.creator);
       return text.trim();
     },
     async critique(prompt, system) {
@@ -91,7 +108,7 @@ export function buildRunner(
           system,
           prompt,
         });
-        addUsage(usage);
+        addUsage(usage, models.critic);
         return output;
       } catch (error) {
         if (isHardApiFailure(error) && fallbackCritic) {
@@ -102,7 +119,7 @@ export function buildRunner(
             system,
             prompt,
           });
-          addUsage(usage);
+          addUsage(usage, MODELS.GEMINI_FALLBACK_CRITIC);
           return output;
         }
         throw error;
@@ -115,7 +132,7 @@ export function buildRunner(
           system,
           prompt,
         });
-        addUsage(usage);
+        addUsage(usage, models.adjudicator);
         return text.trim();
       } catch (error) {
         if (isHardApiFailure(error) && fallbackCritic) {
@@ -125,7 +142,7 @@ export function buildRunner(
             system,
             prompt,
           });
-          addUsage(usage);
+          addUsage(usage, MODELS.GEMINI_FALLBACK_CRITIC);
           return text.trim();
         }
         throw error;
@@ -139,7 +156,7 @@ export function buildRunner(
           system,
           prompt,
         });
-        addUsage(usage);
+        addUsage(usage, models.critic);
         return { ...output, questions: output.questions.slice(0, 5) };
       } catch (error) {
         if (isHardApiFailure(error) && fallbackCritic) {
@@ -150,22 +167,16 @@ export function buildRunner(
             system,
             prompt,
           });
-          addUsage(usage);
+          addUsage(usage, MODELS.GEMINI_FALLBACK_CRITIC);
           return { ...output, questions: output.questions.slice(0, 5) };
         }
         throw error;
       }
     },
     getUsage() {
-      if (!warnedAboutCost) {
-        console.warn(
-          'reflexion: cost tracking not available for this provider; maxCostUsd cap will not fire.'
-        );
-        warnedAboutCost = true;
-      }
       return {
         tokens: accumulatedTokens,
-        costUsd: 0, // Fallback as per requirements
+        costUsd: Number(totalCostUsd.toFixed(6)),
       };
     },
   };
