fix: meter extract atoms haiku calls (#2371)

src/core/config.ts

@@ -928,6 +928,8 @@ export const KNOWN_CONFIG_KEYS: readonly string[] = [
   'models.tier.subagent',
   'models.aliases',
   'models.dream.synthesize',
+  'models.dream.extract_atoms',
+  'cycle.extract_atoms.budget_usd',
   'models.dream.patterns',
   'models.dream.synthesize_verdict',
   'models.drift',

src/core/cycle/extract-atoms.ts

@@ -51,13 +51,15 @@ import type { BrainEngine } from '../engine.ts';
 import type { PhaseResult } from '../cycle.ts';
 import type { GBrainConfig } from '../config.ts';
 import type { ProgressReporter } from '../progress.ts';
-import { chat as gatewayChat } from '../ai/gateway.ts';
+import { chat as gatewayChat, withBudgetTracker } from '../ai/gateway.ts';
+import { BudgetExhausted, BudgetTracker } from '../budget/budget-tracker.ts';
 import { writeReceipt } from '../extract/receipt-writer.ts';
 import { upsertExtractRollup } from '../extract/rollup-writer.ts';
 import { createHash } from 'crypto';
 import { slugifySegment } from '../sync.ts';
 
 const DEFAULT_BUDGET_USD = 0.3;
+const DEFAULT_EXTRACT_ATOMS_MODEL = 'anthropic:claude-haiku-4-5';
 
 // v0.42+ TODO: read atom_type enum from active pack manifest at runtime.
 const ATOM_TYPES = [
@@ -500,7 +502,24 @@ export async function runPhaseExtractAtoms(
   let pagesSkipped = 0;
   const failures: Array<{ source: string; error: string }> = [];
   let estimatedSpendUsd = 0;
-  const budgetCap = DEFAULT_BUDGET_USD;
+  let budgetExhausted = false;
+  let extractModel = DEFAULT_EXTRACT_ATOMS_MODEL;
+  let budgetCap = DEFAULT_BUDGET_USD;
+  try {
+    const configuredModel = await engine.getConfig('models.dream.extract_atoms');
+    if (configuredModel) extractModel = configuredModel;
+    const configuredBudget = await engine.getConfig('cycle.extract_atoms.budget_usd');
+    if (configuredBudget) {
+      const n = Number(configuredBudget);
+      if (Number.isFinite(n) && n > 0) budgetCap = n;
+    }
+  } catch {
+    // Keep safe defaults: Haiku + $0.30.
+  }
+  const budgetTracker = new BudgetTracker({
+    maxCostUsd: budgetCap,
+    label: 'cycle.extract_atoms',
+  });
 
   // v0.41.19.0 (T3): throttled yield helper. Fires `opts.yieldDuringPhase`
   // every 30s. Cycle.ts threads `buildYieldDuringPhase(lock, outer)` so
@@ -525,9 +544,10 @@ export async function runPhaseExtractAtoms(
     }
   }
 
+  await withBudgetTracker(budgetTracker, async () => {
   for (const item of work) {
     await maybeYield();
-    if (estimatedSpendUsd >= budgetCap) {
+    if (budgetExhausted || budgetTracker.totalSpent >= budgetCap) {
       if (item.kind === 'transcript') transcriptsSkipped++;
       else pagesSkipped++;
       continue;
@@ -536,6 +556,7 @@ export async function runPhaseExtractAtoms(
     const originLabel = item.kind === 'transcript' ? item.filePath : item.slug;
     try {
       const result = await chat({
+        model: extractModel,
         system: EXTRACT_PROMPT,
         messages: [
           {
@@ -550,9 +571,7 @@ export async function runPhaseExtractAtoms(
       // actual refresh rate so this is cheap when calls are fast.
       await maybeYield();
 
-      // Rough cost estimate — Haiku at ~$0.80/M input + $4/M output
-      estimatedSpendUsd +=
-        (result.usage.input_tokens * 0.8 + result.usage.output_tokens * 4.0) / 1_000_000;
+      estimatedSpendUsd = budgetTracker.totalSpent;
 
       const atoms = parseAtomsResponse(result.text);
       if (atoms.length === 0) {
@@ -605,12 +624,20 @@ export async function runPhaseExtractAtoms(
       // Reporter rate-limits to ~1 line/sec; safe to tick every iter.
       opts.progress?.tick(1, `${totalAtomsExtracted} atoms / ${duplicatesSkipped} skipped`);
     } catch (err) {
+      if (err instanceof BudgetExhausted) {
+        budgetExhausted = true;
+        if (item.kind === 'transcript') transcriptsSkipped++;
+        else pagesSkipped++;
+        continue;
+      }
       failures.push({
         source: originLabel,
         error: err instanceof Error ? err.message : String(err),
       });
     }
   }
+  });
+  estimatedSpendUsd = budgetTracker.totalSpent;
 
   // v0.42 Wave B2: write extract receipt + rollup row when the phase
   // actually extracted atoms. Both are best-effort per F-OUT-19 —
@@ -668,6 +695,8 @@ export async function runPhaseExtractAtoms(
       failures,
       estimated_spend_usd: estimatedSpendUsd,
       budget_usd: budgetCap,
+      model: extractModel,
+      budget_exhausted: budgetExhausted,
       source_id: sourceId,
       dry_run: opts.dryRun ?? false,
     },

test/cycle/extract-atoms-progress.test.ts

@@ -117,6 +117,43 @@ describe('extract_atoms progress wiring (T4)', () => {
     expect(ticks[0].note).toMatch(/atoms.*skipped/);
   });
 
+  test('passes an explicit Haiku model to chat calls', async () => {
+    const seenModels: Array<string | undefined> = [];
+    const validAtomJson = JSON.stringify([
+      { title: 'A', atom_type: 'insight', body: 'body a' },
+    ]);
+    await runPhaseExtractAtoms(engine, {
+      sourceId: 'default',
+      _transcripts: [
+        { filePath: '/tmp/t1.txt', content: 'transcript 1 body', contentHash: 'h1'.repeat(8) },
+      ],
+      _pages: [],
+      _chat: async (o: ChatOpts) => {
+        seenModels.push(o.model);
+        return stubChat(validAtomJson)(o);
+      },
+    });
+    expect(seenModels).toEqual(['anthropic:claude-haiku-4-5']);
+  });
+
+  test('DB config can override the extract_atoms budget and model', async () => {
+    await engine.setConfig('models.dream.extract_atoms', 'anthropic:claude-haiku-4-5-20251001');
+    await engine.setConfig('cycle.extract_atoms.budget_usd', '0.12');
+    const validAtomJson = JSON.stringify([
+      { title: 'A', atom_type: 'insight', body: 'body a' },
+    ]);
+    const result = await runPhaseExtractAtoms(engine, {
+      sourceId: 'default',
+      _transcripts: [
+        { filePath: '/tmp/t1.txt', content: 'transcript 1 body', contentHash: 'h1'.repeat(8) },
+      ],
+      _pages: [],
+      _chat: stubChat(validAtomJson),
+    });
+    expect(result.details.model).toBe('anthropic:claude-haiku-4-5-20251001');
+    expect(result.details.budget_usd).toBe(0.12);
+  });
+
   test('no progress wiring required — opts.progress is optional', async () => {
     // Sanity: phase works without a reporter.
     const result = await runPhaseExtractAtoms(engine, {