feat(agentloop): Token-/Kosten-Budget mit Hard-Cap (issue #151)

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)

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")
+	}
+}

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"

cmd/sin-code/internal/ledger/store.go

@@ -37,6 +37,9 @@
 	// 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.