mirror of
https://github.com/BaaRai/ARTEX.git
synced 2026-11-13 11:40:53 +08:00
Per-level token/turn budgets: task -> intent (worker run) -> planner round (counted inside the task). A child's effective budget is min(own budget, parent's remaining above the wrap-up reserve). When the task can no longer afford one worker, workers stop claiming intents and the task settles through one forced wrap-up planner round (the existing settle sequence, budget words instead of the timeout words). Config: ARTEX_TASK_BUDGET_TOKENS, ARTEX_TASK_BUDGET_TURNS, ARTEX_INTENT_BUDGET_TOKENS, ARTEX_INTENT_BUDGET_TURNS. 0 or unset = unlimited; with all four unset nothing changes.
74 lines
2.5 KiB
Go
74 lines
2.5 KiB
Go
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
|
|
}
|