palaryn 0.1.0 → 0.3.2

package/src/budget/manager.ts

@@ -0,0 +1,679 @@
+import { BudgetConfig, BudgetState, CostEstimate } from '../types/budget';
+import { BudgetReport, UsageData } from '../types/tool-result';
+import { BudgetStore, WorkspaceBudgetConfig } from '../storage/interfaces';
+import { ToolCall } from '../types/tool-call';
+
+/** Small epsilon for floating-point comparison tolerance in budget checks */
+const BUDGET_EPSILON = 1e-10;
+
+// Default cost table for different tool types (used when no config override provided)
+const DEFAULT_COST_TABLE: Record<string, number> = {
+  'http.request': 0.001,
+  'http.request.write': 0.002,
+  'slack.post': 0.003,
+  'git.operation': 0.002,
+  'db.query': 0.005,
+  'browser.navigate': 0.01,
+  'default': 0.001,
+};
+
+export interface CostRecord {
+  estimated_cost_usd: number;
+  actual_cost_usd?: number;
+  usage?: UsageData;
+}
+
+export class BudgetManager {
+  private config: BudgetConfig;
+  private costTable: Record<string, number>;
+  private externalStore?: BudgetStore;
+  // In-memory stores (fallback when no external store)
+  private taskStates: Map<string, BudgetState>;
+  private userDailySpend: Map<string, number>;     // key: "actor_id:date"
+  private userMonthlySpend: Map<string, number>;    // key: "actor_id:month"
+  private workspaceDailySpend: Map<string, number>; // key: "workspace_id:date"
+  private workspaceMonthlySpend: Map<string, number>; // key: "workspace_id:month"
+  private callRetryCounts: Map<string, number>;     // key: tool_call_id
+  private costRecords: Map<string, CostRecord>;     // key: tool_call_id
+  private reservations: Map<string, { amount: number; taskId: string; workspaceId: string; actorId: string; createdAt: number }>;
+  /** Promise-based mutex per budget key to serialize concurrent reservations */
+  private budgetLocks: Map<string, Promise<void>>;
+
+  constructor(config: BudgetConfig, store?: BudgetStore) {
+    this.config = {
+      task_budget_usd: config.task_budget_usd ?? 2.0,
+      user_daily_budget_usd: config.user_daily_budget_usd,
+      user_monthly_budget_usd: config.user_monthly_budget_usd,
+      workspace_daily_budget_usd: config.workspace_daily_budget_usd,
+      workspace_monthly_budget_usd: config.workspace_monthly_budget_usd,
+      max_steps_per_task: config.max_steps_per_task ?? 50,
+      max_retries_per_call: config.max_retries_per_call ?? 3,
+      max_wall_clock_ms: config.max_wall_clock_ms ?? 300000,
+      token_pricing: config.token_pricing,
+    };
+
+    // Merge config.cost_table over defaults
+    this.costTable = { ...DEFAULT_COST_TABLE, ...(config.cost_table || {}) };
+
+    this.externalStore = store;
+    this.taskStates = new Map();
+    this.userDailySpend = new Map();
+    this.userMonthlySpend = new Map();
+    this.workspaceDailySpend = new Map();
+    this.workspaceMonthlySpend = new Map();
+    this.callRetryCounts = new Map();
+    this.costRecords = new Map();
+    this.reservations = new Map();
+    this.budgetLocks = new Map();
+  }
+
+  /**
+   * Populate in-memory maps from the external store.
+   * Call once at startup after the store itself has been hydrated.
+   */
+  hydrateFromStore(): void {
+    if (!this.externalStore) return;
+
+    if (this.externalStore.getAllTaskStates) {
+      for (const [id, state] of this.externalStore.getAllTaskStates()) {
+        if (!this.taskStates.has(id)) this.taskStates.set(id, state);
+      }
+    }
+
+    if (this.externalStore.getAllCounters) {
+      for (const [key, value] of this.externalStore.getAllCounters()) {
+        if (!this.userDailySpend.has(key)) this.userDailySpend.set(key, value);
+        if (!this.userMonthlySpend.has(key)) this.userMonthlySpend.set(key, value);
+        if (!this.workspaceDailySpend.has(key)) this.workspaceDailySpend.set(key, value);
+        if (!this.workspaceMonthlySpend.has(key)) this.workspaceMonthlySpend.set(key, value);
+      }
+    }
+
+    if (this.externalStore.getAllRetryCounts) {
+      for (const [id, count] of this.externalStore.getAllRetryCounts()) {
+        if (!this.callRetryCounts.has(id)) this.callRetryCounts.set(id, count);
+      }
+    }
+  }
+
+  // Estimate cost for a tool call
+  estimateCost(toolCall: ToolCall): CostEstimate {
+    const toolName = toolCall.tool.name;
+    const capability = toolCall.tool.capability;
+
+    // First try an exact match with capability suffix for write/delete operations
+    let cost: number | undefined;
+
+    if (capability === 'write' || capability === 'delete' || capability === 'admin') {
+      cost = this.costTable[`${toolName}.${capability}`] ?? this.costTable[`${toolName}.write`];
+    }
+
+    // Fall back to base tool name, then default
+    if (cost === undefined) {
+      cost = this.costTable[toolName] ?? this.costTable['default'];
+    }
+
+    // Apply a multiplier for destructive or elevated operations
+    if (capability === 'delete') {
+      cost *= 1.5;
+    } else if (capability === 'admin') {
+      cost *= 2.0;
+    }
+
+    // Respect per-call max_cost_usd constraint as an upper bound on the estimate
+    if (toolCall.constraints?.max_cost_usd !== undefined && toolCall.constraints.max_cost_usd > 0) {
+      cost = Math.min(cost, toolCall.constraints.max_cost_usd);
+    }
+
+    return {
+      tool_name: toolName,
+      capability,
+      estimated_cost_usd: cost,
+    };
+  }
+
+  // Check if a tool call is within budget (returns allowed: false with reason if not)
+  // Optional overrides allow per-workspace budget configuration.
+  check(toolCall: ToolCall, overrides?: WorkspaceBudgetConfig): { allowed: boolean; reason?: string; report: BudgetReport } {
+    const cfg = overrides ? { ...this.config, ...overrides } : this.config;
+    const state = this.getTaskState(toolCall);
+    const estimate = this.estimateCost(toolCall);
+    const estimatedCost = estimate.estimated_cost_usd;
+
+    // Check task budget
+    const taskBudget = cfg.task_budget_usd!;
+    if (state.spent_usd + estimatedCost > taskBudget + BUDGET_EPSILON) {
+      return {
+        allowed: false,
+        reason: `Task budget exceeded: spent $${state.spent_usd.toFixed(4)} + estimated $${estimatedCost.toFixed(4)} > limit $${taskBudget.toFixed(4)}`,
+        report: this.getReport(toolCall, estimatedCost, cfg),
+      };
+    }
+
+    // Check step limit
+    const maxSteps = cfg.max_steps_per_task!;
+    if (state.steps + 1 > maxSteps) {
+      return {
+        allowed: false,
+        reason: `Step limit exceeded: ${state.steps} steps taken, max is ${maxSteps}`,
+        report: this.getReport(toolCall, estimatedCost, cfg),
+      };
+    }
+
+    // Check wall clock time
+    const maxWallClock = cfg.max_wall_clock_ms!;
+    const elapsed = Date.now() - new Date(state.started_at).getTime();
+    if (elapsed > maxWallClock) {
+      return {
+        allowed: false,
+        reason: `Wall clock time exceeded: ${elapsed}ms elapsed, max is ${maxWallClock}ms`,
+        report: this.getReport(toolCall, estimatedCost, cfg),
+      };
+    }
+
+    // Check user daily budget
+    const userDailyLimit = cfg.user_daily_budget_usd;
+    if (userDailyLimit !== undefined) {
+      const dailyKey = `${state.actor_id}:${this.getDateKey()}`;
+      const dailySpend = this.userDailySpend.get(dailyKey) ?? 0;
+      if (dailySpend + estimatedCost > userDailyLimit + BUDGET_EPSILON) {
+        return {
+          allowed: false,
+          reason: `User daily budget exceeded: spent $${dailySpend.toFixed(4)} + estimated $${estimatedCost.toFixed(4)} > limit $${userDailyLimit.toFixed(4)}`,
+          report: this.getReport(toolCall, estimatedCost, cfg),
+        };
+      }
+    }
+
+    // Check user monthly budget
+    const userMonthlyLimit = cfg.user_monthly_budget_usd;
+    if (userMonthlyLimit !== undefined) {
+      const monthlyKey = `${state.actor_id}:${this.getMonthKey()}`;
+      const monthlySpend = this.userMonthlySpend.get(monthlyKey) ?? 0;
+      if (monthlySpend + estimatedCost > userMonthlyLimit + BUDGET_EPSILON) {
+        return {
+          allowed: false,
+          reason: `User monthly budget exceeded: spent $${monthlySpend.toFixed(4)} + estimated $${estimatedCost.toFixed(4)} > limit $${userMonthlyLimit.toFixed(4)}`,
+          report: this.getReport(toolCall, estimatedCost, cfg),
+        };
+      }
+    }
+
+    // Check workspace daily budget
+    const wsDailyLimit = cfg.workspace_daily_budget_usd;
+    if (wsDailyLimit !== undefined) {
+      const dailyKey = `${state.workspace_id}:${this.getDateKey()}`;
+      const dailySpend = this.workspaceDailySpend.get(dailyKey) ?? 0;
+      if (dailySpend + estimatedCost > wsDailyLimit + BUDGET_EPSILON) {
+        return {
+          allowed: false,
+          reason: `Workspace daily budget exceeded: spent $${dailySpend.toFixed(4)} + estimated $${estimatedCost.toFixed(4)} > limit $${wsDailyLimit.toFixed(4)}`,
+          report: this.getReport(toolCall, estimatedCost, cfg),
+        };
+      }
+    }
+
+    // Check workspace monthly budget
+    const wsMonthlyLimit = cfg.workspace_monthly_budget_usd;
+    if (wsMonthlyLimit !== undefined) {
+      const monthlyKey = `${state.workspace_id}:${this.getMonthKey()}`;
+      const monthlySpend = this.workspaceMonthlySpend.get(monthlyKey) ?? 0;
+      if (monthlySpend + estimatedCost > wsMonthlyLimit + BUDGET_EPSILON) {
+        return {
+          allowed: false,
+          reason: `Workspace monthly budget exceeded: spent $${monthlySpend.toFixed(4)} + estimated $${estimatedCost.toFixed(4)} > limit $${wsMonthlyLimit.toFixed(4)}`,
+          report: this.getReport(toolCall, estimatedCost, cfg),
+        };
+      }
+    }
+
+    return {
+      allowed: true,
+      report: this.getReport(toolCall, estimatedCost, cfg),
+    };
+  }
+
+  /**
+   * Atomically check budget and reserve the estimated cost.
+   * The reservation temporarily reduces available budget so concurrent
+   * requests cannot both pass the check.
+   * Uses a per-task lock to serialize concurrent reservations for the same budget key.
+   */
+  async reserveAndCheck(toolCall: ToolCall, overrides?: WorkspaceBudgetConfig): Promise<{ allowed: boolean; reason?: string; report: BudgetReport; reservationKey?: string }> {
+    // Serialize concurrent reservations for the same task to prevent interleaving
+    const lockKey = toolCall.task_id;
+    const existingLock = this.budgetLocks.get(lockKey) ?? Promise.resolve();
+
+    let releaseLock!: () => void;
+    const newLock = new Promise<void>(resolve => { releaseLock = resolve; });
+    this.budgetLocks.set(lockKey, existingLock.then(() => newLock));
+
+    // Wait for prior reservation with timeout to prevent deadlocks
+    const LOCK_TIMEOUT_MS = 5000;
+    let timeoutHandle: ReturnType<typeof setTimeout> | undefined;
+    try {
+      await Promise.race([
+        existingLock,
+        new Promise<void>((_, reject) => {
+          timeoutHandle = setTimeout(() => reject(new Error('Budget lock timeout')), LOCK_TIMEOUT_MS);
+        }),
+      ]);
+    } catch {
+      // On timeout, proceed without lock (accept brief race risk vs permanent hang)
+      console.warn(`[budget] Lock timeout for task ${lockKey}, proceeding without lock`);
+    } finally {
+      if (timeoutHandle !== undefined) clearTimeout(timeoutHandle);
+    }
+
+    try {
+      return this._reserveAndCheckSync(toolCall, overrides);
+    } finally {
+      releaseLock();
+      // Clean up lock if no one else is waiting
+      if (this.budgetLocks.get(lockKey) === existingLock.then(() => newLock)) {
+        this.budgetLocks.delete(lockKey);
+      }
+    }
+  }
+
+  /**
+   * Synchronous check-and-reserve. Must be called while holding the per-task lock.
+   */
+  private _reserveAndCheckSync(toolCall: ToolCall, overrides?: WorkspaceBudgetConfig): { allowed: boolean; reason?: string; report: BudgetReport; reservationKey?: string } {
+    const result = this.check(toolCall, overrides);
+    if (!result.allowed) {
+      return result;
+    }
+
+    // Reserve the estimated cost
+    const estimate = this.estimateCost(toolCall);
+    const reservationKey = `reservation:${toolCall.tool_call_id}`;
+
+    // Temporarily record the estimated cost so concurrent requests see it
+    const state = this.getTaskState(toolCall);
+    state.spent_usd += estimate.estimated_cost_usd;
+    this.taskStates.set(toolCall.task_id, state);
+
+    // Also reserve against user/workspace counters
+    const dateKey = this.getDateKey();
+    const monthKey = this.getMonthKey();
+
+    const userDailyKey = `${state.actor_id}:${dateKey}`;
+    this.userDailySpend.set(userDailyKey, (this.userDailySpend.get(userDailyKey) ?? 0) + estimate.estimated_cost_usd);
+
+    const userMonthlyKey = `${state.actor_id}:${monthKey}`;
+    this.userMonthlySpend.set(userMonthlyKey, (this.userMonthlySpend.get(userMonthlyKey) ?? 0) + estimate.estimated_cost_usd);
+
+    const workspaceDailyKey = `${state.workspace_id}:${dateKey}`;
+    this.workspaceDailySpend.set(workspaceDailyKey, (this.workspaceDailySpend.get(workspaceDailyKey) ?? 0) + estimate.estimated_cost_usd);
+
+    const workspaceMonthlyKey = `${state.workspace_id}:${monthKey}`;
+    this.workspaceMonthlySpend.set(workspaceMonthlyKey, (this.workspaceMonthlySpend.get(workspaceMonthlyKey) ?? 0) + estimate.estimated_cost_usd);
+
+    // Track the reservation so we can commit or release later
+    this.reservations.set(reservationKey, {
+      amount: estimate.estimated_cost_usd,
+      taskId: toolCall.task_id,
+      workspaceId: state.workspace_id,
+      actorId: state.actor_id,
+      createdAt: Date.now(),
+    });
+
+    return { ...result, reservationKey };
+  }
+
+  /**
+   * Commit a reservation with the actual cost. Releases the difference
+   * between the reserved (estimated) amount and the actual cost.
+   */
+  commitReservation(reservationKey: string, actualCost: number): void {
+    const reservation = this.reservations.get(reservationKey);
+    if (!reservation) return;
+
+    const diff = reservation.amount - actualCost;
+    if (diff !== 0) {
+      // Adjust task state: remove reserved amount, add actual
+      const state = this.taskStates.get(reservation.taskId);
+      if (state) {
+        state.spent_usd -= diff;
+        this.taskStates.set(reservation.taskId, state);
+      }
+
+      // Adjust user/workspace counters
+      const dateKey = this.getDateKey();
+      const monthKey = this.getMonthKey();
+
+      const userDailyKey = `${reservation.actorId}:${dateKey}`;
+      const curUD = this.userDailySpend.get(userDailyKey);
+      if (curUD !== undefined) this.userDailySpend.set(userDailyKey, curUD - diff);
+
+      const userMonthlyKey = `${reservation.actorId}:${monthKey}`;
+      const curUM = this.userMonthlySpend.get(userMonthlyKey);
+      if (curUM !== undefined) this.userMonthlySpend.set(userMonthlyKey, curUM - diff);
+
+      const workspaceDailyKey = `${reservation.workspaceId}:${dateKey}`;
+      const curWD = this.workspaceDailySpend.get(workspaceDailyKey);
+      if (curWD !== undefined) this.workspaceDailySpend.set(workspaceDailyKey, curWD - diff);
+
+      const workspaceMonthlyKey = `${reservation.workspaceId}:${monthKey}`;
+      const curWM = this.workspaceMonthlySpend.get(workspaceMonthlyKey);
+      if (curWM !== undefined) this.workspaceMonthlySpend.set(workspaceMonthlyKey, curWM - diff);
+    }
+
+    this.reservations.delete(reservationKey);
+  }
+
+  /**
+   * Release a reservation entirely (e.g. on error before execution completes).
+   * Removes the reserved amount from all counters.
+   */
+  releaseReservation(reservationKey: string): void {
+    const reservation = this.reservations.get(reservationKey);
+    if (!reservation) return;
+
+    // Remove the reserved amount from task state
+    const state = this.taskStates.get(reservation.taskId);
+    if (state) {
+      state.spent_usd -= reservation.amount;
+      state.steps = Math.max(0, state.steps - 0); // steps not incremented during reserve
+      this.taskStates.set(reservation.taskId, state);
+    }
+
+    // Remove from user/workspace counters
+    const dateKey = this.getDateKey();
+    const monthKey = this.getMonthKey();
+
+    const userDailyKey = `${reservation.actorId}:${dateKey}`;
+    const curUD = this.userDailySpend.get(userDailyKey);
+    if (curUD !== undefined) this.userDailySpend.set(userDailyKey, curUD - reservation.amount);
+
+    const userMonthlyKey = `${reservation.actorId}:${monthKey}`;
+    const curUM = this.userMonthlySpend.get(userMonthlyKey);
+    if (curUM !== undefined) this.userMonthlySpend.set(userMonthlyKey, curUM - reservation.amount);
+
+    const workspaceDailyKey = `${reservation.workspaceId}:${dateKey}`;
+    const curWD = this.workspaceDailySpend.get(workspaceDailyKey);
+    if (curWD !== undefined) this.workspaceDailySpend.set(workspaceDailyKey, curWD - reservation.amount);
+
+    const workspaceMonthlyKey = `${reservation.workspaceId}:${monthKey}`;
+    const curWM = this.workspaceMonthlySpend.get(workspaceMonthlyKey);
+    if (curWM !== undefined) this.workspaceMonthlySpend.set(workspaceMonthlyKey, curWM - reservation.amount);
+
+    this.reservations.delete(reservationKey);
+  }
+
+  /**
+   * Record a completed tool call's cost.
+   * When skipCostIncrement is true (e.g. because a reservation already accounts
+   * for the cost), only step counting and cost record storage are performed.
+   */
+  record(toolCall: ToolCall, costOrRecord: number | CostRecord, skipCostIncrement = false): void {
+    const actualCostUsd = typeof costOrRecord === 'number' ? costOrRecord : (costOrRecord.actual_cost_usd ?? costOrRecord.estimated_cost_usd);
+
+    // Store cost record for later retrieval
+    if (typeof costOrRecord !== 'number') {
+      this.costRecords.set(toolCall.tool_call_id, costOrRecord);
+    }
+
+    const state = this.getTaskState(toolCall);
+
+    if (!skipCostIncrement) {
+      // Update task state cost
+      state.spent_usd += actualCostUsd;
+    }
+    // Always increment steps
+    state.steps += 1;
+    this.taskStates.set(toolCall.task_id, state);
+
+    // Persist to external store if available
+    if (this.externalStore) {
+      this.externalStore.setTaskState(toolCall.task_id, state);
+    }
+
+    if (!skipCostIncrement) {
+      // Update user daily spend
+      const dateKey = this.getDateKey();
+      const monthKey = this.getMonthKey();
+
+      const userDailyKey = `${state.actor_id}:${dateKey}`;
+      const currentUserDaily = this.userDailySpend.get(userDailyKey) ?? 0;
+      this.userDailySpend.set(userDailyKey, currentUserDaily + actualCostUsd);
+
+      // Update user monthly spend
+      const userMonthlyKey = `${state.actor_id}:${monthKey}`;
+      const currentUserMonthly = this.userMonthlySpend.get(userMonthlyKey) ?? 0;
+      this.userMonthlySpend.set(userMonthlyKey, currentUserMonthly + actualCostUsd);
+
+      // Update workspace daily spend
+      const workspaceDailyKey = `${state.workspace_id}:${dateKey}`;
+      const currentWorkspaceDaily = this.workspaceDailySpend.get(workspaceDailyKey) ?? 0;
+      this.workspaceDailySpend.set(workspaceDailyKey, currentWorkspaceDaily + actualCostUsd);
+
+      // Update workspace monthly spend
+      const workspaceMonthlyKey = `${state.workspace_id}:${monthKey}`;
+      const currentWorkspaceMonthly = this.workspaceMonthlySpend.get(workspaceMonthlyKey) ?? 0;
+      this.workspaceMonthlySpend.set(workspaceMonthlyKey, currentWorkspaceMonthly + actualCostUsd);
+
+      // Persist counter updates to external store
+      if (this.externalStore) {
+        this.externalStore.incrementCounter(userDailyKey, actualCostUsd);
+        this.externalStore.incrementCounter(userMonthlyKey, actualCostUsd);
+        this.externalStore.incrementCounter(workspaceDailyKey, actualCostUsd);
+        this.externalStore.incrementCounter(workspaceMonthlyKey, actualCostUsd);
+      }
+
+      // Prune stale counter keys (from prior days/months) to prevent unbounded growth
+      this.pruneStaleCounters(dateKey, monthKey);
+    }
+  }
+
+  /** Remove counter entries whose date/month suffix doesn't match today */
+  private pruneStaleCounters(currentDateKey: string, currentMonthKey: string): void {
+    // Only prune periodically (when maps get large)
+    if (this.userDailySpend.size <= 100) return;
+
+    for (const key of this.userDailySpend.keys()) {
+      if (!key.endsWith(`:${currentDateKey}`)) {
+        this.userDailySpend.delete(key);
+      }
+    }
+    for (const key of this.workspaceDailySpend.keys()) {
+      if (!key.endsWith(`:${currentDateKey}`)) {
+        this.workspaceDailySpend.delete(key);
+      }
+    }
+    for (const key of this.userMonthlySpend.keys()) {
+      if (!key.endsWith(`:${currentMonthKey}`)) {
+        this.userMonthlySpend.delete(key);
+      }
+    }
+    for (const key of this.workspaceMonthlySpend.keys()) {
+      if (!key.endsWith(`:${currentMonthKey}`)) {
+        this.workspaceMonthlySpend.delete(key);
+      }
+    }
+  }
+
+  // Record a retry for a tool call
+  recordRetry(toolCallId: string): { allowed: boolean; count: number } {
+    const currentCount = this.callRetryCounts.get(toolCallId) ?? 0;
+    const newCount = currentCount + 1;
+    this.callRetryCounts.set(toolCallId, newCount);
+
+    if (this.externalStore) {
+      this.externalStore.incrementRetryCount(toolCallId);
+    }
+
+    const maxRetries = this.config.max_retries_per_call!;
+
+    return {
+      allowed: newCount <= maxRetries,
+      count: newCount,
+    };
+  }
+
+  // Get budget report for a tool call
+  getReport(toolCall: ToolCall, estimatedCost: number, effectiveConfig?: BudgetConfig): BudgetReport {
+    const cfg = effectiveConfig ?? this.config;
+    const state = this.getTaskState(toolCall);
+    return {
+      estimated_cost_usd: estimatedCost,
+      spent_cost_usd_task: state.spent_usd,
+      remaining_cost_usd_task: Math.max(0, (cfg.task_budget_usd ?? 2.0) - state.spent_usd),
+    };
+  }
+
+  // Get budget report with actual cost and usage data
+  getReportWithActual(toolCall: ToolCall, estimatedCost: number, actualCostUsd?: number, usage?: UsageData): BudgetReport {
+    const base = this.getReport(toolCall, estimatedCost);
+    return {
+      ...base,
+      actual_cost_usd: actualCostUsd,
+      usage,
+    };
+  }
+
+  // Get the active cost table (for introspection)
+  getCostTable(): Record<string, number> {
+    return { ...this.costTable };
+  }
+
+  // Get the cost record for a specific tool call
+  getCostRecord(toolCallId: string): CostRecord | undefined {
+    return this.costRecords.get(toolCallId);
+  }
+
+  // Get or create task state
+  private getTaskState(toolCall: ToolCall): BudgetState {
+    const existing = this.taskStates.get(toolCall.task_id);
+    if (existing) {
+      return existing;
+    }
+
+    const newState: BudgetState = {
+      task_id: toolCall.task_id,
+      workspace_id: toolCall.workspace_id,
+      actor_id: toolCall.actor.id,
+      spent_usd: 0,
+      steps: 0,
+      started_at: toolCall.timestamp ?? new Date().toISOString(),
+    };
+
+    this.taskStates.set(toolCall.task_id, newState);
+    return newState;
+  }
+
+  // Helper to get date key (YYYY-MM-DD)
+  private getDateKey(): string {
+    return new Date().toISOString().split('T')[0];
+  }
+
+  // Helper to get month key (YYYY-MM)
+  private getMonthKey(): string {
+    return new Date().toISOString().slice(0, 7);
+  }
+
+  // Reset task state (for testing)
+  async flush(): Promise<void> {
+    if (this.externalStore?.flush) await this.externalStore.flush();
+  }
+
+  reset(): void {
+    this.taskStates.clear();
+    this.userDailySpend.clear();
+    this.userMonthlySpend.clear();
+    this.workspaceDailySpend.clear();
+    this.workspaceMonthlySpend.clear();
+    this.callRetryCounts.clear();
+    this.costRecords.clear();
+    this.reservations.clear();
+    this.budgetLocks.clear();
+    this.externalStore?.reset();
+  }
+
+  /** Get spending data for a specific actor */
+  getActorSpending(actorId: string, days?: number): {
+    daily_spend: number;
+    monthly_spend: number;
+    task_count: number;
+    total_spend: number;
+  } {
+    const dateKey = this.getDateKey();
+    const monthKey = this.getMonthKey();
+
+    // Daily spend for today
+    const dailyKey = `${actorId}:${dateKey}`;
+    const dailySpend = this.userDailySpend.get(dailyKey) ?? 0;
+
+    // Monthly spend for current month
+    const monthlyKey = `${actorId}:${monthKey}`;
+    const monthlySpend = this.userMonthlySpend.get(monthlyKey) ?? 0;
+
+    // Count tasks where this actor is involved
+    let taskCount = 0;
+    let totalSpend = 0;
+    for (const state of this.taskStates.values()) {
+      if (state.actor_id === actorId) {
+        taskCount++;
+        totalSpend += state.spent_usd;
+      }
+    }
+
+    return {
+      daily_spend: dailySpend,
+      monthly_spend: monthlySpend,
+      task_count: taskCount,
+      total_spend: totalSpend,
+    };
+  }
+
+  // Get current config
+  getConfig(): BudgetConfig {
+    return { ...this.config };
+  }
+
+  /** Return aggregate spending across all tracked entities */
+  getSpendingSummary(): {
+    task_total: number;
+    user_daily_total: number;
+    user_monthly_total: number;
+    workspace_daily_total: number;
+    workspace_monthly_total: number;
+  } {
+    let taskTotal = 0;
+    for (const state of this.taskStates.values()) {
+      taskTotal += state.spent_usd;
+    }
+
+    const dateKey = this.getDateKey();
+    const monthKey = this.getMonthKey();
+
+    let userDailyTotal = 0;
+    for (const [key, val] of this.userDailySpend) {
+      if (key.endsWith(`:${dateKey}`)) userDailyTotal += val;
+    }
+
+    let userMonthlyTotal = 0;
+    for (const [key, val] of this.userMonthlySpend) {
+      if (key.endsWith(`:${monthKey}`)) userMonthlyTotal += val;
+    }
+
+    let workspaceDailyTotal = 0;
+    for (const [key, val] of this.workspaceDailySpend) {
+      if (key.endsWith(`:${dateKey}`)) workspaceDailyTotal += val;
+    }
+
+    let workspaceMonthlyTotal = 0;
+    for (const [key, val] of this.workspaceMonthlySpend) {
+      if (key.endsWith(`:${monthKey}`)) workspaceMonthlyTotal += val;
+    }
+
+    return {
+      task_total: taskTotal,
+      user_daily_total: userDailyTotal,
+      user_monthly_total: userMonthlyTotal,
+      workspace_daily_total: workspaceDailyTotal,
+      workspace_monthly_total: workspaceMonthlyTotal,
+    };
+  }
+}