interface TokenUsage {
  inputTokens: number;
  outputTokens: number;
  cacheReadTokens?: number;   // Prompt caching hits
  cacheWriteTokens?: number;  // Prompt caching misses
}

interface CostCalculation {
  inputCostUsd: number;
  outputCostUsd: number;
  cacheSavingsUsd?: number;
  totalCostUsd: number;
}

function calculateCost(usage: TokenUsage, model: string): CostCalculation {
  const pricing = MODEL_PRICING[model];

  const inputCostUsd = (usage.inputTokens / 1_000_000) * pricing.inputPerMTok;
  const outputCostUsd = (usage.outputTokens / 1_000_000) * pricing.outputPerMTok;

  return {
    inputCostUsd,
    outputCostUsd,
    totalCostUsd: inputCostUsd + outputCostUsd,
  };
}

class CostAccrualTracker {
  private totalInputTokens = 0;
  private totalOutputTokens = 0;
  private accruedCostUsd = 0;
  private readonly model: string;

  constructor(model: string) {
    this.model = model;
  }

  /**
   * Called after each API response (streaming or complete)
   */
  recordUsage(usage: TokenUsage): void {
    this.totalInputTokens += usage.inputTokens;
    this.totalOutputTokens += usage.outputTokens;

    const cost = calculateCost(usage, this.model);
    this.accruedCostUsd += cost.totalCostUsd;
  }

  /**
   * Get current accrued cost (for real-time display)
   */
  getCurrentCost(): number {
    return this.accruedCostUsd;
  }

  /**
   * Finalize on completion or abort
   */
  finalize(reason: 'completed' | 'aborted' | 'failed'): CostReport {
    return {
      totalInputTokens: this.totalInputTokens,
      totalOutputTokens: this.totalOutputTokens,
      totalCostUsd: this.accruedCostUsd,
      completionReason: reason,
      finalizedAt: Date.now(),
    };
  }
}

// In execution handler
const tracker = new CostAccrualTracker(model);

try {
  for await (const chunk of executeStream(request)) {
    if (abortSignal.aborted) {
      // CRITICAL: Capture cost BEFORE throwing
      const partialCost = tracker.finalize('aborted');
      onCostUpdate(partialCost);
      throw new AbortError('Execution aborted');
    }

    tracker.recordUsage(chunk.usage);
    onCostUpdate(tracker.getCurrentCost());
  }

  return tracker.finalize('completed');
} catch (error) {
  if (error instanceof AbortError) {
    throw error; // Already handled
  }
  return tracker.finalize('failed');
}

interface BudgetConfig {
  maxCostUsd: number;
  warnAtPercentage: number;  // e.g., 0.8 for 80%
  onWarn?: (current: number, max: number) => void;
  onExceed?: (current: number, max: number) => void;
}

function createBudgetGuard(config: BudgetConfig) {
  return {
    check(currentCostUsd: number): 'ok' | 'warn' | 'exceed' {
      const percentage = currentCostUsd / config.maxCostUsd;

      if (percentage >= 1.0) {
        config.onExceed?.(currentCostUsd, config.maxCostUsd);
        return 'exceed';
      }

      if (percentage >= config.warnAtPercentage) {
        config.onWarn?.(currentCostUsd, config.maxCostUsd);
        return 'warn';
      }

      return 'ok';
    }
  };
}

                    ┌─────────────────────────────────────────┐
                    │           CostAccrualTracker            │
                    └─────────────────────────────────────────┘
                                        │
              ┌─────────────────────────┼─────────────────────────┐
              │                         │                         │
              ▼                         ▼                         ▼
    ┌─────────────────┐     ┌─────────────────┐     ┌─────────────────┐
    │  recordUsage()  │     │ getCurrentCost()│     │   finalize()    │
    │                 │     │                 │     │                 │
    │ After each API  │     │ For real-time   │     │ On completion,  │
    │ response        │     │ display         │     │ abort, or fail  │
    └─────────────────┘     └─────────────────┘     └─────────────────┘
              │                         │                         │
              │                         │                         │
              ▼                         ▼                         ▼
    ┌─────────────────────────────────────────────────────────────────┐
    │                        CostReport                               │
    │  { inputTokens, outputTokens, totalCostUsd, completionReason }  │
    └─────────────────────────────────────────────────────────────────┘

// Poll every 2 seconds while executing
useEffect(() => {
  if (status !== 'running') return;

  const interval = setInterval(async () => {
    const response = await fetch(`/api/execute/${executionId}`);
    const data = await response.json();
    setAccruedCost(data.cost.accruedUsd);
    setTokens({
      input: data.cost.inputTokens,
      output: data.cost.outputTokens,
    });
  }, 2000);

  return () => clearInterval(interval);
}, [executionId, status]);

// Display format
<div className="cost-display">
  <span className="cost-amount">${accruedCost.toFixed(4)}</span>
  <span className="token-count">
    {tokens.input.toLocaleString()} in / {tokens.output.toLocaleString()} out
  </span>
</div>