diff options
Diffstat (limited to 'internal/executor/executor.go')
| -rw-r--r-- | internal/executor/executor.go | 46 |
1 files changed, 46 insertions, 0 deletions
diff --git a/internal/executor/executor.go b/internal/executor/executor.go index 4333f51..8614481 100644 --- a/internal/executor/executor.go +++ b/internal/executor/executor.go @@ -93,6 +93,14 @@ type Pool struct { dispatchDone chan struct{} // closed when the dispatch goroutine exits Classifier *Classifier LLM *llm.Client + // Budget gates paid providers against a rolling window; nil disables gating. + Budget BudgetGate +} + +// BudgetGate reports whether a task estimated at estCost may run on a provider +// without breaching its rolling spend window. Satisfied by *budget.Accountant. +type BudgetGate interface { + Allow(provider string, estCost float64) (bool, error) } // Result is emitted when a task execution completes. @@ -937,6 +945,44 @@ func (p *Pool) execute(ctx context.Context, t *task.Task) { agentType = "claude" } + // Budget gating: if running on this provider could breach its rolling window + // (estimated by the task's max_budget_usd), reroute to the free local runner + // when one is registered, otherwise stop and mark the task BUDGET_EXCEEDED. + if p.Budget != nil { + if ok, berr := p.Budget.Allow(agentType, t.Agent.MaxBudgetUSD); berr != nil { + p.logger.Warn("budget check failed; proceeding", "error", berr, "taskID", t.ID) + } else if !ok { + if _, hasLocal := p.runners["local"]; hasLocal && agentType != "local" { + p.logger.Info("budget window would be breached; rerouting to local", "taskID", t.ID, "from", agentType) + t.Agent.Type = "local" + t.Agent.Model = "" + agentType = "local" + if uerr := p.store.UpdateTaskAgent(t.ID, t.Agent); uerr != nil { + p.logger.Error("failed to persist rerouted agent", "error", uerr, "taskID", t.ID) + } + } else { + now := time.Now().UTC() + exec := &storage.Execution{ + ID: uuid.New().String(), + TaskID: t.ID, + StartTime: now, + EndTime: now, + Status: "BUDGET_EXCEEDED", + Agent: agentType, + ErrorMsg: fmt.Sprintf("rolling budget window cap reached for provider %q", agentType), + } + if createErr := p.store.CreateExecution(exec); createErr != nil { + p.logger.Error("failed to create execution record", "error", createErr) + } + if err := p.store.UpdateTaskState(t.ID, task.StateBudgetExceeded); err != nil { + p.logger.Error("failed to update task state", "taskID", t.ID, "state", task.StateBudgetExceeded, "error", err) + } + p.resultCh <- &Result{TaskID: t.ID, Execution: exec, Err: fmt.Errorf("budget cap reached for %s", agentType)} + return + } + } + } + // Check dependencies before taking the per-agent slot to avoid deadlock: // if a dependent task holds the slot while waiting for its dependency to run, // the dependency can never start (maxPerAgent=1). |