From d6e46f1bc77b9a736df0aa0ab263b381492cfc54 Mon Sep 17 00:00:00 2001 From: Lakshman Patel Date: Tue, 3 Nov 2026 23:24:13 +0530 Subject: [PATCH] fix(cli): reset operational limits for each request --- internal/engine/lifecycle/limits.go | 17 ++++- internal/engine/lifecycle/limits_test.go | 44 +++++++++++++ internal/engine/stream.go | 6 ++ internal/engine/stream_guards.go | 10 ++- internal/engine/stream_request_limits_test.go | 65 +++++++++++++++++++ 5 files changed, 139 insertions(+), 3 deletions(-) create mode 100644 internal/engine/stream_request_limits_test.go diff --git a/internal/engine/lifecycle/limits.go b/internal/engine/lifecycle/limits.go index f19357c91..587d1212b 100644 --- a/internal/engine/lifecycle/limits.go +++ b/internal/engine/lifecycle/limits.go @@ -6,7 +6,8 @@ import ( "sync" ) -// SafetyLimits caps what the agent can do in a single session. +// SafetyLimits caps operations within one user request. Spending limits apply +// to the whole session so starting a follow-up cannot bypass its budget. type SafetyLimits struct { MaxToolCalls int // max total tool invocations (default: 200) MaxFileWrites int // max files created/modified (default: 50) @@ -34,6 +35,18 @@ func NewLimitTracker(limits SafetyLimits) *LimitTracker { return &LimitTracker{limits: limits} } +// BeginRequest starts fresh operational counters without changing configured +// limits or the session's accumulated spending. +func (lt *LimitTracker) BeginRequest() { + lt.mu.Lock() + defer lt.mu.Unlock() + lt.toolCalls = 0 + lt.fileWrites = 0 + lt.bashCmds = 0 + lt.turns = 0 + lt.outTokens = 0 +} + // RecordToolCall records a tool invocation. Tools named "Bash" or "bash" also // increment the bash command counter; "Write" or "Edit" increment file writes. func (lt *LimitTracker) RecordToolCall(toolName string) { @@ -55,7 +68,7 @@ func (lt *LimitTracker) RecordCost(usd float64) { lt.costUSD += usd } -// RecordTokens adds output tokens to the running total. +// RecordTokens adds output tokens to the current request's running total. func (lt *LimitTracker) RecordTokens(n int) { lt.mu.Lock() defer lt.mu.Unlock() diff --git a/internal/engine/lifecycle/limits_test.go b/internal/engine/lifecycle/limits_test.go index 7362341b0..a1f4a3486 100644 --- a/internal/engine/lifecycle/limits_test.go +++ b/internal/engine/lifecycle/limits_test.go @@ -143,6 +143,50 @@ func TestLimitTracker_BashAndFileTracking(t *testing.T) { } } +func TestLimitTrackerBeginRequestResetsOperationsAndPreservesLimits(t *testing.T) { + for _, tc := range []struct { + name string + limits SafetyLimits + record func(*LimitTracker) + }{ + {name: "tools", limits: SafetyLimits{MaxToolCalls: 1}, record: func(lt *LimitTracker) { lt.RecordToolCall("Read") }}, + {name: "bash", limits: SafetyLimits{MaxBashCommands: 1}, record: func(lt *LimitTracker) { lt.RecordToolCall("Bash") }}, + {name: "writes", limits: SafetyLimits{MaxFileWrites: 1}, record: func(lt *LimitTracker) { lt.RecordToolCall("Write") }}, + {name: "turns", limits: SafetyLimits{MaxTurns: 1}, record: func(lt *LimitTracker) { lt.RecordTurn() }}, + {name: "output", limits: SafetyLimits{MaxOutputTokens: 1}, record: func(lt *LimitTracker) { lt.RecordTokens(1) }}, + } { + t.Run(tc.name, func(t *testing.T) { + lt := NewLimitTracker(tc.limits) + tc.record(lt) + if exceeded, _ := lt.IsExceeded(); !exceeded { + t.Fatal("fixture did not exhaust limit") + } + lt.BeginRequest() + if exceeded, reason := lt.IsExceeded(); exceeded { + t.Fatalf("previous request blocked follow-up: %s", reason) + } + tc.record(lt) + if exceeded, _ := lt.IsExceeded(); !exceeded { + t.Fatal("configured limit was lost after request reset") + } + }) + } +} + +func TestLimitTrackerBeginRequestPreservesSessionSpending(t *testing.T) { + lt := NewLimitTracker(SafetyLimits{MaxBudgetUSD: 1}) + lt.RecordCost(0.6) + lt.BeginRequest() + lt.RecordCost(0.5) + if exceeded, reason := lt.IsExceeded(); !exceeded || !strings.Contains(reason, "cost limit") { + t.Fatalf("follow-up bypassed session spending limit: exceeded=%v reason=%q", exceeded, reason) + } + lt.BeginRequest() + if exceeded, _ := lt.IsExceeded(); !exceeded { + t.Fatal("new request cleared an exhausted spending limit") + } +} + func TestLimitTracker_ConcurrentAccessors(t *testing.T) { lt := NewLimitTracker(DefaultLimits()) diff --git a/internal/engine/stream.go b/internal/engine/stream.go index c427344b3..9a00ab0e0 100644 --- a/internal/engine/stream.go +++ b/internal/engine/stream.go @@ -156,6 +156,12 @@ func (s *Session) agentLoop(ctx context.Context, ch chan<- StreamEvent) { } } sessionStart := time.Now() + // Each Stream handles a new user request, including all of its model/tool + // iterations. Earlier requests must not consume this request's allowance. + // The tracker retains session spending and its configured limits. + if limits := s.LifecycleSvc().Limits(); limits != nil && ctx.Err() == nil { + limits.BeginRequest() + } // Start session-level trace span var sessionSpan *oteltrace.Span diff --git a/internal/engine/stream_guards.go b/internal/engine/stream_guards.go index cd578fc2a..e49a0e15d 100644 --- a/internal/engine/stream_guards.go +++ b/internal/engine/stream_guards.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "os" + "strings" "github.com/GrayCodeAI/reflect/internal/engine/control" @@ -82,7 +83,14 @@ func (s *Session) emitExhaustion(ctx context.Context, ch chan<- StreamEvent, rea return } } - ch <- StreamEvent{Type: "content", Content: fmt.Sprintf("\n\nLimit reached: %s", reason)} + message := fmt.Sprintf("\n\nLimit reached: %s", reason) + for _, prefix := range []string{"tool call limit", "file write limit", "bash command limit", "turn limit", "output token limit"} { + if strings.HasPrefix(reason, prefix) { + message += ". Send a follow-up to continue; this limit resets with each new request." + break + } + } + ch <- StreamEvent{Type: "content", Content: message} ch <- StreamEvent{Type: "done"} } diff --git a/internal/engine/stream_request_limits_test.go b/internal/engine/stream_request_limits_test.go new file mode 100644 index 000000000..febae9021 --- /dev/null +++ b/internal/engine/stream_request_limits_test.go @@ -0,0 +1,65 @@ +package engine + +import ( + "context" + "strings" + "testing" + "time" + + "github.com/GrayCodeAI/reflect/internal/engine/lifecycle" + "github.com/GrayCodeAI/reflect/internal/engine/safety" + "github.com/GrayCodeAI/reflect/internal/tool" + "github.com/GrayCodeAI/reflect/internal/types" +) + +func requestLimitReply(t *testing.T, s *Session, prompt string) string { + t.Helper() + s.AddUser(prompt) + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + events, err := s.Stream(ctx) + if err != nil { + t.Fatal(err) + } + var out strings.Builder + for event := range events { + if event.Type == "error" { + t.Errorf("request failed: %s", event.Content) + } + if event.Type == "content" { + out.WriteString(event.Content) + } + } + return out.String() +} + +func TestNewRequestDoesNotInheritBashCommandCount(t *testing.T) { + client := &conversationalStreamClient{resilienceBoundaryClient: &resilienceBoundaryClient{streams: [][]types.RelayStreamEvent{ + {{Type: "content", Content: "The repository has cmd and internal directories."}, {Type: "done", StopReason: "end_turn"}}, + }}} + s := NewSessionWithClient(client, "test", "test-model", "", nil, false) + for range 100 { + s.LifecycleSvc().Limits().RecordToolCall("Bash") + } + if out := requestLimitReply(t, s, "How is this repository structured?"); !strings.Contains(out, "cmd and internal") || len(client.options) != 1 { + t.Fatalf("previous requests blocked repository question: requests=%d reply=%q", len(client.options), out) + } +} + +func TestRequestToolLimitStopsCurrentRunAndAllowsFollowUp(t *testing.T) { + client := &conversationalStreamClient{resilienceBoundaryClient: &resilienceBoundaryClient{streams: [][]types.RelayStreamEvent{ + {{Type: "tool_call", ToolCall: &types.ToolCall{ID: "read", Name: "Read", Arguments: map[string]interface{}{"id": 1}}}, {Type: "done", StopReason: "tool_use"}}, + {{Type: "content", Content: "Here is the repository structure."}, {Type: "done", StopReason: "end_turn"}}, + }}} + s := NewSessionWithClient(client, "test", "test-model", "", tool.NewRegistry(&countedReadTool{}), false) + s.LifecycleSvc().limits = lifecycle.NewLimitTracker(lifecycle.SafetyLimits{MaxToolCalls: 1}) + s.PermSvc().SetAutonomy(safety.AutonomyYOLO) + first := requestLimitReply(t, s, "Inspect a source file") + if !strings.Contains(first, "tool call limit reached (1/1)") || !strings.Contains(first, "follow-up") || len(client.options) != 1 { + t.Fatalf("current request failed to enforce or explain its limit: requests=%d reply=%q", len(client.options), first) + } + second := requestLimitReply(t, s, "Explain the repository structure") + if !strings.Contains(second, "Here is the repository structure.") || len(client.options) != 2 { + t.Fatalf("follow-up remained blocked: requests=%d reply=%q", len(client.options), second) + } +}