server/runtime/usage.ts

/**
 * Usage tracking module — matches original rust/crates/runtime/src/usage.rs exactly.
 *
 * Provides:
 * - ModelPricing with per-model cost rates
 * - pricing_for_model() lookup
 * - TokenUsage with cache_creation/cache_read tokens
 * - UsageCostEstimate with total_cost_usd()
 * - UsageTracker with record(), cumulative_usage(), current_turn_usage(), turns()
 * - summary_lines_for_model() with pricing=estimated-default fallback
 * - format_usd() helper
 */

// ─── Constants (match original DEFAULT_*_COST_PER_MILLION) ──────────────────
const DEFAULT_INPUT_COST_PER_MILLION = 15.0;
const DEFAULT_OUTPUT_COST_PER_MILLION = 75.0;
const DEFAULT_CACHE_CREATION_COST_PER_MILLION = 18.75;
const DEFAULT_CACHE_READ_COST_PER_MILLION = 1.5;

// ─── ModelPricing ───────────────────────────────────────────────────────────
export interface ModelPricing {
  input_cost_per_million: number;
  output_cost_per_million: number;
  cache_creation_cost_per_million: number;
  cache_read_cost_per_million: number;
}

export function defaultSonnetTierPricing(): ModelPricing {
  return {
    input_cost_per_million: DEFAULT_INPUT_COST_PER_MILLION,
    output_cost_per_million: DEFAULT_OUTPUT_COST_PER_MILLION,
    cache_creation_cost_per_million: DEFAULT_CACHE_CREATION_COST_PER_MILLION,
    cache_read_cost_per_million: DEFAULT_CACHE_READ_COST_PER_MILLION,
  };
}

/**
 * Matches original pricing_for_model() — returns pricing for known models,
 * null for unknown models (caller should use default).
 */
export function pricingForModel(model: string): ModelPricing | null {
  const normalized = model.toLowerCase();
  if (normalized.includes("haiku")) {
    return {
      input_cost_per_million: 1.0,
      output_cost_per_million: 5.0,
      cache_creation_cost_per_million: 1.25,
      cache_read_cost_per_million: 0.1,
    };
  }
  if (normalized.includes("opus")) {
    return {
      input_cost_per_million: 15.0,
      output_cost_per_million: 75.0,
      cache_creation_cost_per_million: 18.75,
      cache_read_cost_per_million: 1.5,
    };
  }
  if (normalized.includes("sonnet")) {
    return defaultSonnetTierPricing();
  }
  return null;
}

// ─── TokenUsage ─────────────────────────────────────────────────────────────
export interface TokenUsage {
  input_tokens: number;
  output_tokens: number;
  cache_creation_input_tokens: number;
  cache_read_input_tokens: number;
}

export function emptyTokenUsage(): TokenUsage {
  return {
    input_tokens: 0,
    output_tokens: 0,
    cache_creation_input_tokens: 0,
    cache_read_input_tokens: 0,
  };
}

export function totalTokens(usage: TokenUsage): number {
  return (
    usage.input_tokens +
    usage.output_tokens +
    usage.cache_creation_input_tokens +
    usage.cache_read_input_tokens
  );
}

// ─── UsageCostEstimate ──────────────────────────────────────────────────────
export interface UsageCostEstimate {
  input_cost_usd: number;
  output_cost_usd: number;
  cache_creation_cost_usd: number;
  cache_read_cost_usd: number;
}

export function totalCostUsd(estimate: UsageCostEstimate): number {
  return (
    estimate.input_cost_usd +
    estimate.output_cost_usd +
    estimate.cache_creation_cost_usd +
    estimate.cache_read_cost_usd
  );
}

function costForTokens(tokens: number, costPerMillion: number): number {
  return (tokens / 1_000_000) * costPerMillion;
}

export function estimateCostUsd(usage: TokenUsage): UsageCostEstimate {
  return estimateCostUsdWithPricing(usage, defaultSonnetTierPricing());
}

export function estimateCostUsdWithPricing(
  usage: TokenUsage,
  pricing: ModelPricing
): UsageCostEstimate {
  return {
    input_cost_usd: costForTokens(usage.input_tokens, pricing.input_cost_per_million),
    output_cost_usd: costForTokens(usage.output_tokens, pricing.output_cost_per_million),
    cache_creation_cost_usd: costForTokens(
      usage.cache_creation_input_tokens,
      pricing.cache_creation_cost_per_million
    ),
    cache_read_cost_usd: costForTokens(
      usage.cache_read_input_tokens,
      pricing.cache_read_cost_per_million
    ),
  };
}

// ─── format_usd ─────────────────────────────────────────────────────────────
export function formatUsd(amount: number): string {
  return `$${amount.toFixed(4)}`;
}

// ─── summary_lines_for_model ────────────────────────────────────────────────
export function summaryLines(usage: TokenUsage, label: string): string[] {
  return summaryLinesForModel(usage, label, undefined);
}

export function summaryLinesForModel(
  usage: TokenUsage,
  label: string,
  model?: string
): string[] {
  const pricing = model ? pricingForModel(model) : null;
  const effectivePricing = pricing ?? defaultSonnetTierPricing();
  const cost = estimateCostUsdWithPricing(usage, effectivePricing);
  const total = totalCostUsd(cost);

  const pricingLabel = pricing ? `model=${model}` : "pricing=estimated-default";

  const line1 = [
    `${label}:`,
    `estimated_cost=${formatUsd(total)}`,
    pricingLabel,
    `input=${formatUsd(cost.input_cost_usd)}`,
    `output=${formatUsd(cost.output_cost_usd)}`,
  ].join(" ");

  const line2 = [
    `  tokens:`,
    `input=${usage.input_tokens}`,
    `output=${usage.output_tokens}`,
    `cache_creation=${usage.cache_creation_input_tokens}`,
    `cache_read=${usage.cache_read_input_tokens}`,
    `cache_creation=${formatUsd(cost.cache_creation_cost_usd)}`,
    `cache_read=${formatUsd(cost.cache_read_cost_usd)}`,
  ].join(" ");

  return [line1, line2];
}

// ─── UsageTracker ───────────────────────────────────────────────────────────
export interface SessionMessage {
  usage?: TokenUsage | null;
}

export class UsageTracker {
  private latestTurn: TokenUsage = emptyTokenUsage();
  private cumulative: TokenUsage = emptyTokenUsage();
  private _turns: number = 0;

  static new(): UsageTracker {
    return new UsageTracker();
  }

  /**
   * Matches original UsageTracker::from_session() — reconstructs tracker
   * from all messages that have usage data.
   */
  static fromSession(messages: SessionMessage[]): UsageTracker {
    const tracker = new UsageTracker();
    for (const msg of messages) {
      if (msg.usage) {
        tracker.record(msg.usage);
      }
    }
    return tracker;
  }

  record(usage: TokenUsage): void {
    this.latestTurn = usage;
    this.cumulative.input_tokens += usage.input_tokens;
    this.cumulative.output_tokens += usage.output_tokens;
    this.cumulative.cache_creation_input_tokens += usage.cache_creation_input_tokens;
    this.cumulative.cache_read_input_tokens += usage.cache_read_input_tokens;
    this._turns += 1;
  }

  currentTurnUsage(): TokenUsage {
    return { ...this.latestTurn };
  }

  cumulativeUsage(): TokenUsage {
    return { ...this.cumulative };
  }

  turns(): number {
    return this._turns;
  }
}