diff --git a/agent/budget.go b/agent/budget.go new file mode 100644 index 0000000..7b47c4c --- /dev/null +++ b/agent/budget.go @@ -0,0 +1,271 @@ +package agent + +import ( + "context" + "fmt" + "math" + "os" + "strconv" + "strings" + "sync" + + "github.com/Autumn-27/norma/llm" +) + +// Hierarchical run budgets (off by default). Levels: task -> intent (one worker +// run) -> planner round (counted inside the task). A child's effective budget is +// min(its own budget, the parent's remaining above the wrap-up reserve). When the +// task can no longer afford a worker, dispatch stops and the run settles through +// one forced wrap-up planner round (same pattern as the task-timeout endgame). +// Wall-clock stays in taskclock/endgame. Counting matches the llm_usage ledger: +// one turn per metered call (one row), tokens = input + output; norma's Usage +// already folds the cache buckets into InputTokens, so they are not added again. + +// Unlimited is the cap of a dimension with no budget (env value 0 or unset). +const Unlimited int64 = math.MaxInt64 + +const ( + BudgetMarginPercent = 5 // wrap-up reserve = max(5% of the task budget, floor) + BudgetMarginMinTokens = 20_000 // token floor of the reserve + BudgetMarginMinTurns = 2 // turn floor of the reserve + DefaultWorkerEstimateTokens = 50_000 // one worker run, until a worker of this task finished + DefaultWorkerEstimateTurns = 10 +) + +// Budget caps tokens and turns; a dimension equal to Unlimited has no cap. +type Budget struct{ Tokens, Turns int64 } + +// Spend is what a level has consumed. +type Spend struct{ Tokens, Turns int64 } + +// NewBudget builds a budget from configured values: <=0 means Unlimited. +func NewBudget(tokens, turns int64) Budget { return Budget{budgetCap(tokens), budgetCap(turns)} } + +func budgetCap(n int64) int64 { + if n <= 0 { + return Unlimited + } + return n +} + +// EffectiveBudget is the hierarchical calculator: per dimension, a child never +// gets more than its parent has left. +func EffectiveBudget(own, parentRemaining Budget) Budget { + return Budget{min(own.Tokens, parentRemaining.Tokens), min(own.Turns, parentRemaining.Turns)} +} + +// Remaining is the budget left after s (never negative; Unlimited stays Unlimited). +func (b Budget) Remaining(s Spend) Budget { + return Budget{left(b.Tokens, s.Tokens), left(b.Turns, s.Turns)} +} + +func left(capacity, spent int64) int64 { + if capacity == Unlimited { + return Unlimited + } + return max(capacity-spent, 0) +} + +// WrapupMargin is the reserve kept for the forced wrap-up (0 when unlimited). +func WrapupMargin(task Budget) Budget { + m := func(capacity, floor int64) int64 { + if capacity == Unlimited { + return 0 + } + return max(capacity*BudgetMarginPercent/100, floor) + } + return Budget{m(task.Tokens, BudgetMarginMinTokens), m(task.Turns, BudgetMarginMinTurns)} +} + +func addUsage(s Spend, u llm.Usage) Spend { + return Spend{s.Tokens + int64(u.InputTokens+u.OutputTokens), s.Turns + 1} +} + +// BudgetsFromEnv reads ARTEX_TASK_BUDGET_{TOKENS,TURNS} and +// ARTEX_INTENT_BUDGET_{TOKENS,TURNS}. Unset, 0 or invalid = Unlimited. +func BudgetsFromEnv() (task, intent Budget) { + env := func(k string) int64 { + n, _ := strconv.ParseInt(strings.TrimSpace(os.Getenv(k)), 10, 64) + return n + } + return NewBudget(env("ARTEX_TASK_BUDGET_TOKENS"), env("ARTEX_TASK_BUDGET_TURNS")), + NewBudget(env("ARTEX_INTENT_BUDGET_TOKENS"), env("ARTEX_INTENT_BUDGET_TURNS")) +} + +// BudgetLedger meters one task, in memory (a process restart starts afresh). +type BudgetLedger struct { + task, intent Budget + mu sync.Mutex + spent Spend // planner rounds + worker runs + workerRuns int64 + workerSpent Spend + wrapup bool // the forced wrap-up has started +} + +// NewBudgetLedger returns nil when no budget is set: callers then change nothing. +func NewBudgetLedger(task, intent Budget) *BudgetLedger { + if task == NewBudget(0, 0) && intent == NewBudget(0, 0) { + return nil + } + return &BudgetLedger{task: task, intent: intent} +} + +// Record meters one task-level call (the planner's own rounds). +func (l *BudgetLedger) Record(u llm.Usage) { + l.mu.Lock() + defer l.mu.Unlock() + l.spent = addUsage(l.spent, u) +} + +// BudgetState is one decision snapshot of the task level. +type BudgetState struct { + Task, Remaining, Margin Budget + Spent, Estimate Spend + AtMargin bool // remaining is at or under the wrap-up reserve + Dispatch bool // a new worker may start +} + +// State reports the remaining budget and whether a worker may be dispatched: +// not once the wrap-up started or the reserve is reached, and not when the +// remaining budget is below one worker's estimated cost (the average of this +// task's finished worker runs, else the defaults; never above the intent cap). +func (l *BudgetLedger) State() BudgetState { + l.mu.Lock() + defer l.mu.Unlock() + rem, m := l.task.Remaining(l.spent), WrapupMargin(l.task) + est := Spend{DefaultWorkerEstimateTokens, DefaultWorkerEstimateTurns} + if l.workerRuns > 0 { + est = Spend{l.workerSpent.Tokens / l.workerRuns, l.workerSpent.Turns / l.workerRuns} + } + est = Spend{min(est.Tokens, l.intent.Tokens), min(est.Turns, l.intent.Turns)} + at := (rem.Tokens != Unlimited && rem.Tokens <= m.Tokens) || (rem.Turns != Unlimited && rem.Turns <= m.Turns) + return BudgetState{Task: l.task, Remaining: rem, Margin: m, Spent: l.spent, Estimate: est, AtMargin: at, + Dispatch: !l.wrapup && !at && rem.Tokens >= est.Tokens && rem.Turns >= est.Turns} +} + +// BeginWrapup marks the forced wrap-up; true only for the first caller, so the +// wrap-up runs exactly once. Running workers see an exhausted budget from here on. +func (l *BudgetLedger) BeginWrapup() bool { + l.mu.Lock() + defer l.mu.Unlock() + first := !l.wrapup + l.wrapup = true + return first +} + +// Line is the one-line remaining budget shown to the planner and logged. +func (s BudgetState) Line() string { + of := func(rem, capacity int64) string { + if capacity == Unlimited { + return "unlimited" + } + return fmt.Sprintf("%d of %d", rem, capacity) + } + return fmt.Sprintf("[Task budget] remaining: tokens %s, turns %s (wrap-up reserve: %d tokens, %d turns)", + of(s.Remaining.Tokens, s.Task.Tokens), of(s.Remaining.Turns, s.Task.Turns), s.Margin.Tokens, s.Margin.Turns) +} + +// IntentMeter meters one worker run against min(intent budget, task remaining +// above the reserve). +type IntentMeter struct { + l *BudgetLedger + spent Spend // guarded by l.mu +} + +// StartIntent opens the meter of one worker run. +func (l *BudgetLedger) StartIntent() *IntentMeter { return &IntentMeter{l: l} } + +// Record meters one worker call into the intent and the task. +func (m *IntentMeter) Record(u llm.Usage) { + m.l.mu.Lock() + defer m.l.mu.Unlock() + m.spent, m.l.spent = addUsage(m.spent, u), addUsage(m.l.spent, u) +} + +// Remaining is the run's effective remaining budget; zero once the wrap-up started. +func (m *IntentMeter) Remaining() Budget { + l := m.l + l.mu.Lock() + defer l.mu.Unlock() + parent := Budget{} + if !l.wrapup { + rem, res := l.task.Remaining(l.spent), WrapupMargin(l.task) + parent = Budget{left(rem.Tokens, res.Tokens), left(rem.Turns, res.Turns)} + } + return EffectiveBudget(l.intent.Remaining(m.spent), parent) +} + +// Exhausted reports whether the run has no budget left in some dimension. +func (m *IntentMeter) Exhausted() bool { + r := m.Remaining() + return r.Tokens == 0 || r.Turns == 0 +} + +// Finish folds the run into the task's per-worker average. +func (m *IntentMeter) Finish() { + m.l.mu.Lock() + defer m.l.mu.Unlock() + m.l.workerRuns++ + m.l.workerSpent = Spend{m.l.workerSpent.Tokens + m.spent.Tokens, m.l.workerSpent.Turns + m.spent.Turns} +} + +// BudgetWrapupBlocks reports whether a worker tool is refused once its budget is +// exhausted: all but the write-back tools, so the run records facts and stops. +func BudgetWrapupBlocks(tool string) bool { + return tool != "record_fact" && tool != "insert_assets" && tool != "report_finding" +} + +// WorkerBudgetWrapup is returned in place of a refused tool call. +const WorkerBudgetWrapup = "[Budget reached] This intent's token/turn budget is used up; the tool call was not executed. Do not run any more commands or probes. Write back what you identified but have not persisted yet (insert_assets, record_fact, report_finding), then, in a single standalone sentence of plain text, summarize what you did and your key conclusions." + +const plannerBudgetWrapup = "**The task's token/turn budget is exhausted and the task is about to end** (a normal end, not an error). Based on the current facts and findings, make the final goal judgment: call prove_goal for goals proven met by evidence. **Do not generate any new intents** (they will not be executed). Then, in a single standalone sentence of plain text, summarize where the task stands." + +// BudgetRound carries one planner round's budget on ctx (zero = no budget). +type BudgetRound struct { + Line string // one-line remaining budget for the planner context + WrapUp bool // the forced wrap-up round +} + +type budgetRoundKey struct{} +type intentMeterKey struct{} + +// WithBudgetRound attaches a planner round's budget to ctx. +func WithBudgetRound(ctx context.Context, r BudgetRound) context.Context { + return context.WithValue(ctx, budgetRoundKey{}, r) +} + +// WithIntentMeter attaches a worker run's meter to ctx (turn clamp). +func WithIntentMeter(ctx context.Context, m *IntentMeter) context.Context { + return context.WithValue(ctx, intentMeterKey{}, m) +} + +// plannerBudgetBlock is the code-owned budget text appended to the planner's +// situational block (same pattern as the endgame words); empty without a budget. +func plannerBudgetBlock(ctx context.Context) string { + r, _ := ctx.Value(budgetRoundKey{}).(BudgetRound) + s := "" + if r.Line != "" { + s += "\n\n" + r.Line + } + if r.WrapUp { + s += "\n\n[Task budget wrap-up (special instruction for this round, overriding the normal planning flow above)]: " + plannerBudgetWrapup + } + return s +} + +// turnCap is the worker's MaxTurns for this run: its own, clamped by the +// intent's effective remaining turns (never 0, which harness reads as unlimited). +func (w *Worker) turnCap(ctx context.Context) int { + m, _ := ctx.Value(intentMeterKey{}).(*IntentMeter) + if m == nil { + return w.maxTurns + } + rem := m.Remaining().Turns + if rem == Unlimited { + return w.maxTurns + } + if c := int(max(rem, 1)); w.maxTurns <= 0 || c < w.maxTurns { + return c + } + return w.maxTurns +} diff --git a/agent/budget_test.go b/agent/budget_test.go new file mode 100644 index 0000000..b35f358 --- /dev/null +++ b/agent/budget_test.go @@ -0,0 +1,146 @@ +package agent + +import ( + "context" + "strings" + "testing" + + "github.com/Autumn-27/norma/llm" +) + +// spend records n tokens as one task-level call; finishWorker records one +// finished worker run of n tokens (it feeds the per-worker estimate). +func spend(l *BudgetLedger, n int) { l.Record(llm.Usage{InputTokens: n}) } + +func finishWorker(l *BudgetLedger, n int) { + m := l.StartIntent() + m.Record(llm.Usage{OutputTokens: n}) + m.Finish() +} + +func TestEffectiveBudgetNeverExceedsParent(t *testing.T) { + for _, c := range []struct{ own, parent, want Budget }{ + {NewBudget(50, 5), NewBudget(100, 10), Budget{50, 5}}, + {NewBudget(200, 20), NewBudget(100, 10), Budget{100, 10}}, + {NewBudget(0, 0), NewBudget(100, 10), Budget{100, 10}}, // own unlimited + {NewBudget(70, 7), NewBudget(0, 0), Budget{70, 7}}, // parent unlimited + {NewBudget(0, 0), NewBudget(0, 0), Budget{Unlimited, Unlimited}}, + {NewBudget(70, 7), Budget{0, 0}, Budget{0, 0}}, // parent exhausted + } { + if got := EffectiveBudget(c.own, c.parent); got != c.want || got.Tokens > c.parent.Tokens || got.Turns > c.parent.Turns { + t.Fatalf("EffectiveBudget(%+v, %+v) = %+v, want %+v", c.own, c.parent, got, c.want) + } + } + // Through the meter: intent 100k is cut to the task's remaining above the + // reserve (task 100k, reserve 20k, 70k spent → 10k); intent turns clamp MaxTurns. + l := NewBudgetLedger(NewBudget(100_000, 0), NewBudget(100_000, 5)) + spend(l, 70_000) + m := l.StartIntent() + if got := m.Remaining(); got != (Budget{10_000, 5}) { + t.Fatalf("intent remaining = %+v, want {10000 5}", got) + } + if got := (&Worker{maxTurns: 40}).turnCap(WithIntentMeter(context.Background(), m)); got != 5 { + t.Fatalf("worker MaxTurns = %d, want the intent's 5", got) + } + if m.Record(llm.Usage{InputTokens: 10_000}); !m.Exhausted() { + t.Fatal("intent not exhausted once the task's remaining above the reserve is spent") + } +} + +func TestWrapupMarginTriggersOnceAndStopsDispatch(t *testing.T) { + if a, b := WrapupMargin(NewBudget(1_000_000, 100)), WrapupMargin(NewBudget(100_000, 10)); a != (Budget{50_000, 5}) || b != (Budget{20_000, 2}) { + t.Fatalf("margins %+v %+v, want 5%% then the floors", a, b) + } + l := NewBudgetLedger(NewBudget(100_000, 0), NewBudget(0, 0)) + finishWorker(l, 10) // tiny estimate: only the reserve can stop dispatch here + spend(l, 100_000-10-20_001) + if st := l.State(); st.AtMargin || !st.Dispatch { + t.Fatalf("one token above the reserve: %+v, want dispatch", st) + } + spend(l, 1) // remaining == reserve + st := l.State() + if !st.AtMargin || st.Dispatch { + t.Fatalf("at the reserve: %+v, want AtMargin and no dispatch", st) + } + running, wrapups := l.StartIntent(), 0 + for i := 0; i < 3; i++ { // several planner wake-ups at the reserve + if l.BeginWrapup() { + wrapups++ + } + } + if wrapups != 1 || l.State().Dispatch { + t.Fatalf("wrap-up started %d times (want 1), dispatch=%v", wrapups, l.State().Dispatch) + } + if !running.Exhausted() || !BudgetWrapupBlocks("Bash") || BudgetWrapupBlocks("record_fact") { + t.Fatal("after the wrap-up a running worker must be held to its write-back tools") + } + block := plannerBudgetBlock(WithBudgetRound(context.Background(), BudgetRound{Line: st.Line(), WrapUp: true})) + if strings.Count(block, plannerBudgetWrapup) != 1 || !strings.Contains(block, "[Task budget] remaining: tokens 20000 of 100000") { + t.Fatalf("wrap-up block:\n%s", block) + } +} + +func TestDispatchStopsUnderOneWorkerEstimate(t *testing.T) { + // No worker finished yet: the default estimate applies. + l := NewBudgetLedger(NewBudget(200_000, 0), NewBudget(0, 0)) + spend(l, 200_000-DefaultWorkerEstimateTokens) + if st := l.State(); !st.Dispatch { + t.Fatalf("remaining == default estimate: %+v, want dispatch", st) + } + if spend(l, 1); l.State().Dispatch || l.State().AtMargin { + t.Fatalf("below the default estimate, above the reserve: %+v, want no dispatch", l.State()) + } + // Finished workers: their average spend is the estimate (90k, 30k → 60k). + l = NewBudgetLedger(NewBudget(300_000, 0), NewBudget(0, 0)) + finishWorker(l, 90_000) + finishWorker(l, 30_000) + spend(l, 120_000) // remaining 60k + if st := l.State(); st.Estimate.Tokens != 60_000 || !st.Dispatch { + t.Fatalf("average estimate: %+v, want 60000 and dispatch", st) + } + if spend(l, 1); l.State().Dispatch || l.State().AtMargin { + t.Fatalf("below the 60000 estimate, above the reserve: %+v, want no dispatch", l.State()) + } +} + +func TestBudgetUnsetIsPassthrough(t *testing.T) { + for _, v := range []string{"", "0", "abc"} { + for _, k := range []string{"ARTEX_TASK_BUDGET_TOKENS", "ARTEX_TASK_BUDGET_TURNS", "ARTEX_INTENT_BUDGET_TOKENS", "ARTEX_INTENT_BUDGET_TURNS"} { + t.Setenv(k, v) + } + if NewBudgetLedger(BudgetsFromEnv()) != nil { + t.Fatalf("env %q: got a ledger, want nil (no budget)", v) + } + } + // Without a budget on ctx the planner input and the worker turns are unchanged. + ctx := context.Background() + situational := plannerSituational(nil, triggerLookup{}, nil, newDataEnvelope()) + if plannerRoundInput("l", situational, "", "g", "") != plannerRoundInput("l", situational+plannerBudgetBlock(ctx), "", "g", "") { + t.Fatal("planner input changed without a budget") + } + if got := (&Worker{maxTurns: 40}).turnCap(ctx); got != 40 { + t.Fatalf("worker turns without budget = %d, want 40", got) + } +} + +func TestBudgetCountsMatchUsageFixture(t *testing.T) { + // One entry per llm_usage row; InputTokens already includes the cache buckets. + planner := []llm.Usage{ + {InputTokens: 12_000, OutputTokens: 800, CacheReadTokens: 9_000}, + {InputTokens: 15_500, OutputTokens: 1_200, CacheReadTokens: 12_000, CacheWriteTokens: 3_000}, + } + worker := []llm.Usage{{InputTokens: 4_000, OutputTokens: 300}, {InputTokens: 6_100, CacheWriteTokens: 6_000}, {}} // {} = an errored call's row + l := NewBudgetLedger(NewBudget(1_000_000, 1_000), NewBudget(0, 0)) + for _, u := range planner { + l.Record(u) + } + m := l.StartIntent() + for _, u := range worker { + m.Record(u) + } + m.Finish() + const want = 12_000 + 800 + 15_500 + 1_200 + 4_000 + 300 + 6_100 + if st := l.State(); st.Spent != (Spend{want, 5}) || st.Remaining != (Budget{1_000_000 - want, 995}) || st.Estimate != (Spend{10_400, 3}) { + t.Fatalf("spent %+v remaining %+v estimate %+v, want {%d 5}, remaining, {10400 3}", st.Spent, st.Remaining, st.Estimate, want) + } +} diff --git a/agent/planner.go b/agent/planner.go index 8139ffa..3b38967 100644 --- a/agent/planner.go +++ b/agent/planner.go @@ -500,6 +500,9 @@ func (p *Planner) Plan(ctx context.Context, taskID int64, as *db.AssetStore, ts if tc.Final { situational += "\n\n[Task endgame wrap-up (special instruction for this round, overriding the normal planning flow above)]: " + resolveTaskTimeoutWrapup("planner") } + // Task budget (off by default, see budget.go): the remaining-budget line, and on + // the forced wrap-up round the budget wrap-up instruction. Empty without a budget. + situational += plannerBudgetBlock(ctx) // This task's work dir /tasks/; create it first. taskDir := ensureRunDir(p.workDir, taskID, 0) ctx = intercept.WithReviewContext(ctx, taskDir, intercept.ReviewBackground{}) diff --git a/agent/worker.go b/agent/worker.go index 582911b..2d9338e 100644 --- a/agent/worker.go +++ b/agent/worker.go @@ -503,7 +503,7 @@ func (w *Worker) execute(ctx context.Context, name string, taskID int64, as *db. // default (tools need no -x/-k). BashEnv: proxyEnv(w.proxyAddr, w.proxyCACert), WorkingDir: runDir, - MaxTurns: w.maxTurns, // 0 = unlimited (configurable in agent management) + MaxTurns: w.turnCap(ctx), // 0 = unlimited (configurable in agent management); clamped by the intent budget // Wall-clock budget, judged at turn boundaries without interrupting mid-flight; // 0 = unlimited. When a task-level deadline exists, clamp to min(own budget, time // until the deadline) so this run naturally enters wrap-up when the task hits its diff --git a/llmrec/llmrec.go b/llmrec/llmrec.go index f3c87df..2c5f2fa 100644 --- a/llmrec/llmrec.go +++ b/llmrec/llmrec.go @@ -60,6 +60,15 @@ func workerFrom(ctx context.Context) string { return strings.TrimSpace(worker) } +type usageSinkContextKey struct{} + +// WithUsageSink attaches a callback that receives the usage of every metered +// call made under ctx — the same value, at the same point, as the llm_usage row. +// The task runtime uses it to meter run budgets. +func WithUsageSink(ctx context.Context, sink func(llm.Usage)) context.Context { + return context.WithValue(ctx, usageSinkContextKey{}, sink) +} + // Recorder wraps an llm.Provider and records every completion call. type Recorder struct { inner llm.Provider @@ -168,7 +177,7 @@ func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.S status = "error" } // Lightweight metering row — always written. - r.recordUsage(taskID, expID, worker, usage, int(time.Since(start).Milliseconds()), status) + r.recordUsage(ctx, taskID, expID, worker, usage, int(time.Since(start).Milliseconds()), status) // Heavy trace row — only when body recording is on. if recordBodies { r.record(req, session, taskID, worker, reqBody, capt, start, textBuf.String(), thinkingBuf.String(), usage, stopReason, err) @@ -246,7 +255,7 @@ func (r *Recorder) Complete(ctx context.Context, req llm.CompletionRequest) (llm if err != nil { status = "error" } - r.recordUsage(taskID, expID, worker, usage, int(time.Since(start).Milliseconds()), status) + r.recordUsage(ctx, taskID, expID, worker, usage, int(time.Since(start).Milliseconds()), status) if recordBodies { r.record(req, session, taskID, worker, reqBody, capt, start, msg.Text(), thinkingText(msg), usage, stopReason, err) } @@ -264,15 +273,19 @@ func thinkingText(msg llm.Message) string { } // recordUsage appends one lightweight metering row to llm_usage (no bodies). Skips -// zero-token calls with no model, which carry nothing worth metering. -func (r *Recorder) recordUsage(taskID string, expID int64, worker string, usage llm.Usage, latencyMs int, status string) { - if r.pg == nil { - return - } +// zero-token calls with no model, which carry nothing worth metering. The ctx +// usage sink (if any) sees exactly the calls that get a row. +func (r *Recorder) recordUsage(ctx context.Context, taskID string, expID int64, worker string, usage llm.Usage, latencyMs int, status string) { if r.model == "" && usage.InputTokens == 0 && usage.OutputTokens == 0 && usage.CacheReadTokens == 0 && usage.CacheWriteTokens == 0 { return } + if sink, _ := ctx.Value(usageSinkContextKey{}).(func(llm.Usage)); sink != nil { + sink(usage) + } + if r.pg == nil { + return + } err := r.pg.InsertLLMUsage(&db.LLMUsage{ TaskID: taskID, ExplorationID: expID, diff --git a/llmrec/llmrec_test.go b/llmrec/llmrec_test.go index 2fd0e0e..d5f643c 100644 --- a/llmrec/llmrec_test.go +++ b/llmrec/llmrec_test.go @@ -68,6 +68,32 @@ func TestCompleteForwardsAtomicResponse(t *testing.T) { } } +// The usage sink sees each metered call once, with the usage the llm_usage row gets. +func TestUsageSinkSeesMeteredCalls(t *testing.T) { + var got []llm.Usage + ctx := WithUsageSink(context.Background(), func(u llm.Usage) { got = append(got, u) }) + recorder := Wrap(meteredStreamProvider{}, nil, "model", "profile", "", "", func() bool { return false }) + for _, err := range recorder.Stream(ctx, llm.CompletionRequest{}) { + if err != nil { + t.Fatal(err) + } + } + if _, _, _, err := recorder.Complete(ctx, llm.CompletionRequest{}); err != nil { + t.Fatal(err) + } + want := []llm.Usage{{InputTokens: 23, OutputTokens: 5}, {InputTokens: 7, OutputTokens: 3}} + if len(got) != len(want) || got[0] != want[0] || got[1] != want[1] { + t.Fatalf("sink saw %+v, want %+v", got, want) + } + // A call that gets no llm_usage row (no model, zero usage) is not seen either. + got = nil + empty := Wrap(completeProvider{}, nil, "", "", "", "", nil) + empty.recordUsage(ctx, "", 0, "", llm.Usage{}, 0, "ok") + if len(got) != 0 { + t.Fatalf("sink saw an unmetered call: %+v", got) + } +} + type meteredStreamProvider struct{ completeProvider } func (meteredStreamProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] { diff --git a/server/engine.go b/server/engine.go index ab24682..80067cd 100644 --- a/server/engine.go +++ b/server/engine.go @@ -146,6 +146,7 @@ type Engine struct { stamped sync.Map // taskID -> bool, whether first_run_at has been stamped (once per process) inflight sync.Map // taskID -> *int64, count of in-flight planner.Plan + worker.Execute (for drain) coordStarted sync.Map // taskID -> bool, whether the deadline coordinator has started (deduplicated across Run/reload) + budgets sync.Map // taskID -> *agent.BudgetLedger, run budget meter (nil = no budget configured; see engine_budget.go) // resolve returns a task's dedicated planner/worker (wired by the server as the // authoritative task-router). nil,nil means this task is deliberately unavailable @@ -289,6 +290,7 @@ func (e *Engine) StopTask(taskID string) { e.stamped.Delete(taskID) e.inflight.Delete(taskID) e.coordStarted.Delete(taskID) + e.budgets.Delete(taskID) e.deleteMu.Lock() e.deleting.Delete(taskID) e.deleteMu.Unlock() @@ -532,6 +534,9 @@ type steerHooks struct { limit int // label looks like "worker-1 · #42"; used only for logging. label string + // budget meters this intent's run (nil = no budget). Once exhausted, every tool + // but the write-back ones is refused with the wrap-up instruction. + budget *agent.IntentMeter } // Default number of reruns for idle turns (thinking only, neither body text nor tool @@ -573,6 +578,9 @@ func (h steerHooks) PreToolUse(ctx context.Context, name string, input []byte) ( return true, "[Planner real-time steering] " + msg + "\n(This is the planner's immediate instruction for this intent; the tool call was not executed — adjust your next step accordingly. If it conflicts with your current plan, this takes precedence.)", nil } + if h.budget != nil && agent.BudgetWrapupBlocks(name) && h.budget.Exhausted() { + return true, agent.WorkerBudgetWrapup, nil + } if h.inner != nil { return h.inner.PreToolUse(ctx, name, input) } @@ -879,6 +887,11 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) { } return // the goalless branch never enters planner.Plan } + // Run budget (off by default): no worker affordable any more → the forced + // wrap-up (once) instead of a normal round. + if e.budgetStop(ctx, t) { + return + } // quiet-mode (#79): an edge round is redundant while work is in flight // and the open-goal set is unchanged since the last completed round — // there is no new decision for the plan to make. Drain the triggers so @@ -910,7 +923,7 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) { // coalesces a burst; empty for time/heartbeat wakes). triggers := t.drainTriggers() taskIDInt, _ := strconv.ParseInt(t.ID, 10, 64) - pctx := e.plannerGateContext(ectx, t, intercept.DefaultTargetWaitBound) + pctx := e.budgetRoundCtx(e.plannerGateContext(ectx, t, intercept.DefaultTargetWaitBound), t, false) e.BeginLLMCall(t.ID) met, reason, err := planner.Plan(pctx, taskIDInt, e.m.assets, t.Store, t.Goal, triggers, t.TargetGuard.Hooks(), emit) e.EndLLMCall(t.ID) @@ -936,6 +949,9 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) { if fp, err := t.Store.OpenGoalFingerprint(); err == nil { e.goalFP.Store(t.ID, fp) } + if e.budgetBlocksClaim(t.ID) { + t.Notify() // this round crossed the budget: settle on the next round, not at the next heartbeat + } e.touch(t.ID) } @@ -1041,6 +1057,12 @@ func (e *Engine) workerLoop(ctx context.Context, t *Task, name string) { } continue // task timeout settlement in progress: claim no new intents (running ones settle themselves; the coordinator waits for them to drain) } + if e.budgetBlocksClaim(t.ID) { + if sleepCtx(ctx, 1000*time.Millisecond) { + return + } + continue // run budget below one worker's estimate: dispatch no new worker (the planner loop settles the task) + } if isTerminalStatus(e.m.TaskStatus(t.ID)) { if sleepCtx(ctx, 1000*time.Millisecond) { return @@ -1116,6 +1138,9 @@ func (e *Engine) runIntent(ctx context.Context, t *Task, name string, worker *ag label := fmt.Sprintf("%s · #%d", name, iid) workCtx = intercept.WithTaskContext(workCtx, t.ID, label, taskEmit) workCtx = intercept.WithTargetScope(workCtx, t.targetScope(intercept.DefaultTargetWaitBound)) + // Run budget (off by default): meter this run's calls; its effective budget is + // min(intent budget, task remaining above the wrap-up reserve). + workCtx, meter := e.budgetIntent(workCtx, t.ID) // nudges deliberately sits outside the model_error rerun loop: the idle-turn rerun // cap is a per-intent total, and one rerun round must not reset it back to zero. hooks := steerHooks{ @@ -1124,6 +1149,7 @@ func (e *Engine) runIntent(ctx context.Context, t *Task, name string, worker *ag nudges: &atomic.Int64{}, limit: e.emptyTurnNudgeLimit(), label: label, + budget: meter, } wTaskID, _ := strconv.ParseInt(t.ID, 10, 64) e.BeginLLMCall(t.ID) @@ -1157,6 +1183,9 @@ func (e *Engine) runIntent(ctx context.Context, t *Task, name string, worker *ag e.EndLLMCall(t.ID) } stopHeartbeat() // no longer busy: the state writes below follow within the just-renewed claim + if meter != nil { + meter.Finish() // feeds the per-worker cost estimate + } // Capture kill state before detachWork cancels workCtx. kill = this work's // ctx was cancelled (planner kill_work) while the TASK ctx kept running; a // pause cancels the task ctx (ectx) instead. Checking workCtx.Err() AFTER diff --git a/server/engine_budget.go b/server/engine_budget.go new file mode 100644 index 0000000..b6021b2 --- /dev/null +++ b/server/engine_budget.go @@ -0,0 +1,73 @@ +package server + +import ( + "context" + "log" + + "github.com/Autumn-27/artex/agent" + "github.com/Autumn-27/artex/llmrec" +) + +// Hierarchical run budgets (agent/budget.go), off unless an ARTEX_*_BUDGET_* env +// var is set. Every helper here is a no-op when budgetFor returns nil. + +// budgetFor returns the task's budget ledger, created on first use from the env; +// nil when no budget is configured (passthrough). +func (e *Engine) budgetFor(taskID string) *agent.BudgetLedger { + if v, ok := e.budgets.Load(taskID); ok { + return v.(*agent.BudgetLedger) + } + v, _ := e.budgets.LoadOrStore(taskID, agent.NewBudgetLedger(agent.BudgetsFromEnv())) + return v.(*agent.BudgetLedger) +} + +// budgetStop is the enforcement point of a planner round: it logs the remaining +// budget and, once no worker can be dispatched any more, starts the forced +// wrap-up (exactly once: the settle sequence with the budget words) and reports +// that the normal round must not run. ctx is the planner loop's context. +func (e *Engine) budgetStop(ctx context.Context, t *Task) bool { + l := e.budgetFor(t.ID) + if l == nil { + return false + } + st := l.State() + log.Printf("[budget] task %s planner round: %s", t.ID, st.Line()) + if st.Dispatch { + return false + } + if l.BeginWrapup() { + log.Printf("[budget] task %s cannot afford another worker (estimate %d tokens / %d turns), forced wrap-up", + t.ID, st.Estimate.Tokens, st.Estimate.Turns) + e.settleTask(ctx, t, true) + } + return true +} + +// budgetRoundCtx attaches a planner round's budget (the one-line remaining +// budget, the wrap-up flag) and meters the round's own calls into the task. +func (e *Engine) budgetRoundCtx(ctx context.Context, t *Task, wrapup bool) context.Context { + l := e.budgetFor(t.ID) + if l == nil { + return ctx + } + ctx = agent.WithBudgetRound(ctx, agent.BudgetRound{Line: l.State().Line(), WrapUp: wrapup}) + return llmrec.WithUsageSink(ctx, l.Record) +} + +// budgetBlocksClaim reports whether a worker must not claim a new intent: the +// remaining task budget is below one worker's cost estimate, or the wrap-up started. +func (e *Engine) budgetBlocksClaim(taskID string) bool { + l := e.budgetFor(taskID) + return l != nil && !l.State().Dispatch +} + +// budgetIntent opens the meter of one worker run: ctx meters its calls and +// clamps its turns; the returned meter feeds the tool hook. nil = no budget. +func (e *Engine) budgetIntent(ctx context.Context, taskID string) (context.Context, *agent.IntentMeter) { + l := e.budgetFor(taskID) + if l == nil { + return ctx, nil + } + m := l.StartIntent() + return llmrec.WithUsageSink(agent.WithIntentMeter(ctx, m), m.Record), m +} diff --git a/server/engine_budget_test.go b/server/engine_budget_test.go new file mode 100644 index 0000000..6c51fe8 --- /dev/null +++ b/server/engine_budget_test.go @@ -0,0 +1,44 @@ +package server + +import ( + "context" + "testing" + + "github.com/Autumn-27/artex/agent" + "github.com/Autumn-27/norma/llm" +) + +// With the budget env vars unset the engine decides exactly as before (no round +// stopped, no claim blocked, contexts untouched, every tool call let through); +// a task budget flips the same decisions once the task cannot afford a worker. +func TestEngineBudgetPassthroughAndGate(t *testing.T) { + for _, k := range []string{"ARTEX_TASK_BUDGET_TOKENS", "ARTEX_TASK_BUDGET_TURNS", "ARTEX_INTENT_BUDGET_TOKENS", "ARTEX_INTENT_BUDGET_TURNS"} { + t.Setenv(k, "") + } + ctx, task := context.Background(), &Task{ID: "7"} + noSteer := func() (string, bool) { return "", false } + e := NewEngine(nil) + wctx, meter := e.budgetIntent(ctx, task.ID) + if e.budgetStop(ctx, task) || e.budgetBlocksClaim(task.ID) || e.budgetRoundCtx(ctx, task, false) != ctx || wctx != ctx || meter != nil { + t.Fatal("no budget: a decision or a context changed") + } + if block, _, _ := (steerHooks{drain: noSteer}).PreToolUse(ctx, "Bash", nil); block { + t.Fatal("no budget: a tool call was refused") + } + + t.Setenv("ARTEX_TASK_BUDGET_TOKENS", "100000") + e = NewEngine(nil) + _, meter = e.budgetIntent(ctx, task.ID) + meter.Record(llm.Usage{InputTokens: 60_000}) // remaining 40k < default worker estimate 50k + h := steerHooks{drain: noSteer, budget: meter} + if block, _, _ := h.PreToolUse(ctx, "Bash", nil); !e.budgetBlocksClaim(task.ID) || block { + t.Fatal("budget: claim must be blocked, the running worker (above the reserve) not") + } + meter.Record(llm.Usage{OutputTokens: 20_000}) // reaches the 20k reserve + if block, msg, _ := h.PreToolUse(ctx, "Bash", nil); !block || msg != agent.WorkerBudgetWrapup { + t.Fatalf("budget: exhausted run not sent to wrap-up (block=%v msg=%q)", block, msg) + } + if block, _, _ := h.PreToolUse(ctx, "record_fact", nil); block { + t.Fatal("budget: write-back tool refused during the wrap-up") + } +} diff --git a/server/engine_timeout.go b/server/engine_timeout.go index 1d3ab55..933b059 100644 --- a/server/engine_timeout.go +++ b/server/engine_timeout.go @@ -90,6 +90,7 @@ func (e *Engine) resetTimeoutRevival(taskID string) { e.deadline.Delete(taskID) e.stamped.Delete(taskID) e.coordStarted.Delete(taskID) + e.budgets.Delete(taskID) // a revived task also gets a fresh run budget } // stampFirstRun records first_run_at + deadline_at on the FIRST real run (LLM ready) @@ -176,17 +177,23 @@ func (e *Engine) deadlineCoordinator(ctx context.Context, t *Task) { } continue } - e.settleTask(ctx, t) + e.settleTask(ctx, t, false) return } } -// settleTask runs the ordered settle sequence once (§4 steps ①–⑥). -func (e *Engine) settleTask(ctx context.Context, t *Task) { +// settleTask runs the ordered settle sequence once (§4 steps ①–⑥). budget = the +// task run budget is exhausted (budget.go) instead of the deadline: same sequence, +// the final round carries the budget wrap-up words. +func (e *Engine) settleTask(ctx context.Context, t *Task, budget bool) { if !e.markSettling(t.ID) { return } - log.Printf("[deadline] task %s reached the timeout limit, entering the settle sequence", t.ID) + if budget { + log.Printf("[budget] task %s run budget exhausted, entering the settle sequence", t.ID) + } else { + log.Printf("[deadline] task %s reached the timeout limit, entering the settle sequence", t.ID) + } // ④ wait for running workers/planners to drain (running runs enter wrap-up on their own // via the clamped MaxDuration); if not drained within the grace period, hard-cancel the @@ -205,7 +212,7 @@ func (e *Engine) settleTask(ctx context.Context, t *Task) { } // ⑤ one final planner round (task-timeout words, final goal judgment, no new intents). - met := e.runFinalPlannerRound(ctx, t) + met := e.runFinalPlannerRound(ctx, t, budget) if !e.beginTaskOperation(t.ID) { return } @@ -213,11 +220,17 @@ func (e *Engine) settleTask(ctx context.Context, t *Task) { // ⑥ set the terminal state (guarded): met → done (completed); otherwise timeout. If the // normal path already settled done first, the guard (SetTaskStatusGuarded) rejects the - // overwrite, preserving the completed semantics. + // overwrite, preserving the completed semantics. A budget-exhausted run reuses the + // timeout state (a resource limit ended it; no new task status) and is marked by an + // activity record. status := "timeout" if met { status = "done" } + if budget { + e.emitActivity(t, db.Activity{Worker: "system", Kind: "text", + Summary: "Task run budget exhausted: forced wrap-up done, task stopped (final state " + status + ")"}) + } won, err := e.m.SetTaskStatusGuarded(t.ID, status) switch { case err != nil: @@ -232,7 +245,7 @@ func (e *Engine) settleTask(ctx context.Context, t *Task) { // runFinalPlannerRound drives exactly ONE terminal planner round with the // task-timeout planner words (final goal judgment; no new intents). Waits for the // LLM to be ready (bounded by ctx) so a completable task isn't mis-judged timeout. -func (e *Engine) runFinalPlannerRound(ctx context.Context, t *Task) (met bool) { +func (e *Engine) runFinalPlannerRound(ctx context.Context, t *Task, budget bool) (met bool) { if e.IsDeleting(t.ID) { return false } @@ -251,14 +264,19 @@ func (e *Engine) runFinalPlannerRound(ctx context.Context, t *Task) (met bool) { } // Independent ctx (not wired to execCancel, so a pause/hard cancel cannot interrupt this // final round), with Final injecting the task-timeout words. - fctx := e.clockCtx(ctx, t, true) + fctx := e.clockCtx(ctx, t, !budget) + label := "task timeout settle" + if budget { + fctx = e.budgetRoundCtx(fctx, t, true) + label = "task budget settle" + } if !e.beginTaskOperation(t.ID) { return false } defer e.decInflight(t.ID) emit := func(r db.Activity) { e.emitActivity(t, r) } e.emitActivity(t, db.Activity{Worker: "planner", Kind: "round", - Summary: fmt.Sprintf("task timeout settle · final judgment (round %d)", e.nextPlannerRound(t.ID))}) + Summary: fmt.Sprintf("%s · final judgment (round %d)", label, e.nextPlannerRound(t.ID))}) tTaskID, _ := strconv.ParseInt(t.ID, 10, 64) // The final round judges, it does not explore: a gated call never waits here. fctx = e.plannerGateContext(fctx, t, 0) diff --git a/server/target_gate_scope_test.go b/server/target_gate_scope_test.go index ae07bf0..2f4b1dc 100644 --- a/server/target_gate_scope_test.go +++ b/server/target_gate_scope_test.go @@ -229,7 +229,7 @@ func (f *targetScopeFixture) finalPlannerRound(t *testing.T, task *Task) scopePr t.Helper() probe := newScopeProbeProvider() f.usePlanner(probe) - f.e.runFinalPlannerRound(context.Background(), task) + f.e.runFinalPlannerRound(context.Background(), task, false) return probe.first(t) }