diff --git a/cmd/sin-code/internal/agentloop/loop.go b/cmd/sin-code/internal/agentloop/loop.go index 500cd40c..8e90c32a 100644 --- a/cmd/sin-code/internal/agentloop/loop.go +++ b/cmd/sin-code/internal/agentloop/loop.go @@ -27,6 +27,18 @@ type Completion struct { Text string ToolCalls []ToolCall Raw session.Message + // Usage carries token accounting returned by the model provider. All + // fields optional; zero values mean "unknown" and never trigger the + // budget guard (issue #151). + Usage Usage +} + +// Usage carries token accounting returned by the model provider. All fields +// optional; zero values mean "unknown" and never trigger the budget guard. +type Usage struct { + PromptTokens int + CompletionTokens int + TotalTokens int } type ToolSpec struct { @@ -97,6 +109,16 @@ type Loop struct { // disables stall detection. Recommended: 3. Independent of MaxStopRejects. StallThreshold int + // MaxTokens is a hard cap on cumulative tokens (prompt+completion) across + // the whole run. Zero means unlimited. When exceeded the run checkpoints + // (if AllowContinuation) or errors, rather than continuing to spend. + // Issue #151. + MaxTokens int + + // BudgetWarnRatio, if set, fires hooks.BudgetWarn once when token usage + // crosses this fraction of MaxTokens (e.g. 0.8). Useful for alerting. + BudgetWarnRatio float64 + // AllowContinuation switches the maxTurns outcome from a hard error to a // checkpointed, resumable Result (Continuation=true). Daemons set this so // a long task is re-enqueued and resumed rather than abandoned; one-shot @@ -259,6 +281,8 @@ func (l *Loop) Run(ctx context.Context, sess *session.Session, prompt string) (* stopRejects := 0 // tracks how many times the stop-gate rejected completion lastCritFingerprint := "" stallCount := 0 + totalTokens := 0 // issue #151: cumulative tokens across the run + warnedBudget := false // fires hooks.BudgetWarn once per run toolsSeen := map[string]bool{} var toolsUsed []string @@ -284,6 +308,40 @@ func (l *Loop) Run(ctx context.Context, sess *session.Session, prompt string) (* return nil, fmt.Errorf("turn %d: %w", turn, err) } msgs = append(msgs, resp.Raw) + // Token budget accounting (issue #151). Provider usage is optional; + // if zero we simply skip the guard for that turn. + if u := resp.Usage.TotalTokens; u > 0 { + totalTokens += u + } else { + totalTokens += resp.Usage.PromptTokens + resp.Usage.CompletionTokens + } + if l.MaxTokens > 0 { + if !warnedBudget && l.BudgetWarnRatio > 0 && + float64(totalTokens) >= l.BudgetWarnRatio*float64(l.MaxTokens) { + warnedBudget = true + l.fire(ctx, hooks.BudgetWarn, "", map[string]any{ + "total_tokens": totalTokens, "max_tokens": l.MaxTokens, + }) + } + if totalTokens >= l.MaxTokens { + if serr := sess.SaveHistory(msgs); serr != nil { + return nil, serr + } + l.fire(ctx, hooks.BudgetExhausted, "", map[string]any{ + "total_tokens": totalTokens, "max_tokens": l.MaxTokens, + }) + l.record(ctx, ledger.TypeTokenBudgetExhausted, + map[string]any{"total_tokens": totalTokens, "max_tokens": l.MaxTokens}, + fmt.Sprintf("token budget exhausted: %d/%d", totalTokens, l.MaxTokens)) + if l.AllowContinuation { + return &Result{ + SessionID: sess.ID, Summary: lastText, Verified: false, + Turns: turn + 1, Continuation: true, OpenCriteria: lastOpen, + }, nil + } + return nil, fmt.Errorf("token budget exhausted: %d/%d tokens used", totalTokens, l.MaxTokens) + } + } if len(resp.ToolCalls) == 0 { vpre := l.fire(ctx, hooks.VerifyPre, "", nil) diff --git a/cmd/sin-code/internal/agentloop/loop_token_test.go b/cmd/sin-code/internal/agentloop/loop_token_test.go new file mode 100644 index 00000000..f9f13edc --- /dev/null +++ b/cmd/sin-code/internal/agentloop/loop_token_test.go @@ -0,0 +1,94 @@ +// SPDX-License-Identifier: MIT +// Purpose: tests for issue #151 — token budget with hard cap. The +// loop accumulates prompt+completion tokens and bails out (or +// checkpoints) when MaxTokens is reached. +package agentloop + +import ( + "context" + "strings" + "testing" + + "github.com/OpenSIN-Code/SIN-Code/cmd/sin-code/internal/session" +) + +// MaxTokens=0 disables the budget guard (legacy behavior). +func TestTokenBudget_DisabledWhenZero(t *testing.T) { + s := setupSession(t) + calls := 0 + loop := &Loop{ + Gate: passGate(), + Workspace: "/tmp", + // MaxTokens=0 -> unlimited + StopGate: func(ctx context.Context, snap StopSnapshot) StopDecision { + calls++ + return StopDecision{Complete: true} + }, + Completion: func(ctx context.Context, msgs []session.Message, tools []ToolSpec) (*Completion, error) { + return &Completion{Text: "ok", Usage: Usage{TotalTokens: 999999}}, nil + }, + } + _, err := loop.Run(context.Background(), s, "x") + if err != nil { + t.Fatalf("expected success with MaxTokens=0, got %v", err) + } + if calls != 1 { + t.Fatalf("expected 1 stop-gate call, got %d", calls) + } +} + +// MaxTokens with TotalTokens exceeds the cap -> error (without AllowContinuation). +func TestTokenBudget_Exhausts(t *testing.T) { + s := setupSession(t) + calls := 0 + loop := &Loop{ + Gate: passGate(), + Workspace: "/tmp", + MaxTokens: 100, + MaxStopRejects: 100, // high — must NOT be the trigger + StopGate: func(ctx context.Context, snap StopSnapshot) StopDecision { + calls++ + return StopDecision{Complete: false, OpenCriteria: []string{"k"}} + }, + Completion: func(ctx context.Context, msgs []session.Message, tools []ToolSpec) (*Completion, error) { + calls++ + return &Completion{Text: "ok", Usage: Usage{TotalTokens: 50}}, nil + }, + } + // After turn 1, total=50. After turn 2, total=100 (cap). The budget + // check fires after appending msgs. + _, err := loop.Run(context.Background(), s, "x") + if err == nil { + t.Fatal("expected token budget error") + } + if !strings.Contains(err.Error(), "token budget exhausted") { + t.Errorf("expected 'token budget exhausted' in error, got: %v", err) + } +} + +// MaxTokens + AllowContinuation -> checkpoint, not error. +func TestTokenBudget_AllowContinuation(t *testing.T) { + s := setupSession(t) + loop := &Loop{ + Gate: passGate(), + Workspace: "/tmp", + MaxTokens: 50, + AllowContinuation: true, + StopGate: func(ctx context.Context, snap StopSnapshot) StopDecision { + return StopDecision{Complete: true} + }, + Completion: func(ctx context.Context, msgs []session.Message, tools []ToolSpec) (*Completion, error) { + return &Completion{Text: "done", Usage: Usage{TotalTokens: 50}}, nil + }, + } + res, err := loop.Run(context.Background(), s, "x") + if err != nil { + t.Fatalf("expected checkpoint, got error: %v", err) + } + if res == nil { + t.Fatal("expected non-nil result") + } + if !res.Continuation { + t.Errorf("expected Continuation=true, got false") + } +} diff --git a/cmd/sin-code/internal/hooks/hooks.go b/cmd/sin-code/internal/hooks/hooks.go index 5f176832..0c56b68d 100644 --- a/cmd/sin-code/internal/hooks/hooks.go +++ b/cmd/sin-code/internal/hooks/hooks.go @@ -53,6 +53,9 @@ const ( // StopStalled fires when the stop-gate returns identical open criteria // StallThreshold turns in a row (no-progress escalation). StopStalled = "stop.stalled" + // Token budget lifecycle (issue #151). + BudgetWarn = "budget.warn" + BudgetExhausted = "budget.exhausted" AgentSpawn = "agent.spawn" AgentComplete = "agent.complete" diff --git a/cmd/sin-code/internal/ledger/store.go b/cmd/sin-code/internal/ledger/store.go index e3faeda4..57972d0a 100644 --- a/cmd/sin-code/internal/ledger/store.go +++ b/cmd/sin-code/internal/ledger/store.go @@ -37,6 +37,9 @@ const ( // criteria StallThreshold times in a row — the worker is not making // progress and the run escalates early instead of burning the full budget. TypeStallDetected EntryType = "stall_detected" + // TypeTokenBudgetExhausted is recorded when cumulative token usage exceeds + // MaxTokens and the run stops spending (issue #151). + TypeTokenBudgetExhausted EntryType = "token_budget_exhausted" ) // Entry is one row in the ledger.