diff options
| author | Claude <noreply@anthropic.com> | 2026-05-26 20:31:24 +0000 |
|---|---|---|
| committer | Claude <noreply@anthropic.com> | 2026-05-26 20:31:24 +0000 |
| commit | 32715355fe2eed321df4f7083dfe580d35f8a62a (patch) | |
| tree | 76c3fd6a45646a2f6bdca95e46f66a00d6f08e59 /internal/executor/executor_test.go | |
| parent | 39f74dbbc7adca0e2058112408bb99694deefcaf (diff) | |
feat(executor): budget-gate the dispatcher — reroute to local or block (Phase 6)
The pool now consults a BudgetGate before taking an agent slot. If running on
the chosen paid provider could breach its rolling window (estimated by the
task's max_budget_usd), it reroutes to the free local runner when one is
registered; otherwise it stops before running and marks the task
BUDGET_EXCEEDED. serve.go builds a budget.Accountant from [budget] config and
wires it into the pool (only when limits are configured).
Allows the QUEUED → BUDGET_EXCEEDED transition, since the gate blocks a queued
task before it ever reaches RUNNING. Covered by pool tests for both the block
and reroute paths against a fake gate.
https://claude.ai/code/session_01SESwn7kQ7oP62trWw6pc39
Diffstat (limited to 'internal/executor/executor_test.go')
| -rw-r--r-- | internal/executor/executor_test.go | 57 |
1 files changed, 57 insertions, 0 deletions
diff --git a/internal/executor/executor_test.go b/internal/executor/executor_test.go index d98efbf..b737e22 100644 --- a/internal/executor/executor_test.go +++ b/internal/executor/executor_test.go @@ -217,6 +217,63 @@ func TestPool_Submit_Subtask_GoesToCompleted(t *testing.T) { } } +type fakeGate struct{ deny map[string]bool } + +func (g fakeGate) Allow(provider string, _ float64) (bool, error) { return !g.deny[provider], nil } + +func TestPool_BudgetGate_BlocksWhenNoLocalFallback(t *testing.T) { + store := testStore(t) + runner := &mockRunner{} + pool := NewPool(2, map[string]Runner{"claude": runner}, store, slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))) + pool.Budget = fakeGate{deny: map[string]bool{"claude": true}} + + tk := makeTask("bg-1") + store.CreateTask(tk) + if err := pool.Submit(context.Background(), tk); err != nil { + t.Fatalf("submit: %v", err) + } + + result := <-pool.Results() + if result.Execution.Status != "BUDGET_EXCEEDED" { + t.Errorf("status: want BUDGET_EXCEEDED, got %q", result.Execution.Status) + } + if runner.callCount() != 0 { + t.Errorf("claude runner should not have run, got %d calls", runner.callCount()) + } + got, _ := store.GetTask("bg-1") + if got.State != task.StateBudgetExceeded { + t.Errorf("task state: want BUDGET_EXCEEDED, got %v", got.State) + } +} + +func TestPool_BudgetGate_ReroutesToLocal(t *testing.T) { + store := testStore(t) + claudeRunner := &mockRunner{} + localRunner := &mockRunner{} + pool := NewPool(2, map[string]Runner{"claude": claudeRunner, "local": localRunner}, store, slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))) + pool.Budget = fakeGate{deny: map[string]bool{"claude": true}} // local is allowed + + tk := makeTask("bg-2") + store.CreateTask(tk) + if err := pool.Submit(context.Background(), tk); err != nil { + t.Fatalf("submit: %v", err) + } + + result := <-pool.Results() + if result.Err != nil { + t.Errorf("expected success after reroute, got %v", result.Err) + } + if localRunner.callCount() != 1 { + t.Errorf("local runner should have run once, got %d", localRunner.callCount()) + } + if claudeRunner.callCount() != 0 { + t.Errorf("claude runner should not have run, got %d", claudeRunner.callCount()) + } + if result.Execution.Agent != "local" { + t.Errorf("execution agent: want local, got %q", result.Execution.Agent) + } +} + func TestPool_Submit_Failure(t *testing.T) { store := testStore(t) runner := &mockRunner{err: fmt.Errorf("boom"), exitCode: 1} |
