1 files changed, 62 insertions, 3 deletions

diff --git a/internal/executor/executor_test.go b/internal/executor/executor_test.go
index 9214872..b737e22 100644
--- a/internal/executor/executor_test.go
+++ b/internal/executor/executor_test.go
@@ -12,6 +12,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/thepeterstone/claudomator/internal/event"
 	"github.com/thepeterstone/claudomator/internal/storage"
 	"github.com/thepeterstone/claudomator/internal/task"
 )
@@ -63,7 +64,7 @@ type mockRunner struct {
 	onRun    func(*task.Task, *storage.Execution) error
 }
 
-func (m *mockRunner) Run(ctx context.Context, t *task.Task, e *storage.Execution) error {
+func (m *mockRunner) Run(ctx context.Context, t *task.Task, e *storage.Execution, _ AgentChannel) error {
 	m.mu.Lock()
 	m.calls++
 	cb := m.onRun
@@ -216,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}
@@ -391,11 +449,11 @@ func (m *logPatherMockRunner) ExecLogDir(execID string) string {
 	return filepath.Join(m.logDir, execID)
 }
 
-func (m *logPatherMockRunner) Run(ctx context.Context, t *task.Task, e *storage.Execution) error {
+func (m *logPatherMockRunner) Run(ctx context.Context, t *task.Task, e *storage.Execution, ch AgentChannel) error {
 	m.mu.Lock()
 	m.capturedPath = e.StdoutPath
 	m.mu.Unlock()
-	return m.mockRunner.Run(ctx, t, e)
+	return m.mockRunner.Run(ctx, t, e, ch)
 }
 
 // TestPool_Execute_LogPathsPreSetBeforeRun verifies that when the runner
@@ -1194,6 +1252,7 @@ func (m *minimalMockStore) GetStory(_ string) (*task.Story, error)
 func (m *minimalMockStore) ListTasksByStory(_ string) ([]*task.Task, error)     { return nil, nil }
 func (m *minimalMockStore) UpdateStoryStatus(_ string, _ task.StoryState) error { return nil }
 func (m *minimalMockStore) CreateTask(_ *task.Task) error                      { return nil }
+func (m *minimalMockStore) CreateEvent(_ *event.Event) error                   { return nil }
 func (m *minimalMockStore) UpdateTaskCheckerReport(_ string, _ string) error   { return nil }
 func (m *minimalMockStore) GetCheckerTask(_ string) (*task.Task, error)        { return nil, nil }