From c20a5203ca49a7e2f8b04af9a7dc904ae220e95b Mon Sep 17 00:00:00 2001
From: Levi Neely <141506390+lneely@users.noreply.github.com>
Date: Fri, 15 May 2026 19:23:48 +0000
Subject: [PATCH] agent: add MaxSteps soft guardrail to agentConfig and loop

---
 pkg/agent/loop.go               |  17 ++++
 pkg/agent/loop_maxsteps_test.go | 143 ++++++++++++++++++++++++++++++++
 2 files changed, 160 insertions(+)
 create mode 100644 pkg/agent/loop_maxsteps_test.go

diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go
index 2723dd7..332dba6 100644
--- a/pkg/agent/loop.go
+++ b/pkg/agent/loop.go
@@ -45,6 +45,10 @@ type agentConfig struct {
 	PreTool              func(ctx context.Context, name string, args json.RawMessage) HookResult              // called before each tool; exit 2 blocks execution
 	PostTool             func(ctx context.Context, name string, args json.RawMessage, result string) HookResult // called after each tool; exit 0 appends, exit 2 replaces result
 	TurnError            func(ctx context.Context, errType, errMsg string) HookResult                          // called on first backend error; if ran, skips retries
+	// MaxSteps is the maximum number of tool-call rounds per turn.
+	// When reached, a soft nudge is injected and the loop exits cleanly.
+	// 0 means unlimited.
+	MaxSteps int
 }
 
 func run(ctx context.Context, cfg agentConfig, state state) error {
@@ -438,6 +442,19 @@ func run(ctx context.Context, cfg agentConfig, state state) error {
 			}, nil)
 		}
 
+		// Soft step-budget guardrail: when MaxSteps is set and the budget is
+		// exhausted, nudge the model to wrap up and exit cleanly. This is not
+		// a hard abort — the model gets one final turn without tools to emit
+		// a summary or hand-off message.
+		if cfg.MaxSteps > 0 && step >= cfg.MaxSteps-1 {
+			emit(cfg, Event{Role: "maxsteps", Content: fmt.Sprintf("%d", step+1)})
+			state.update(backend.Message{
+				Role:    "user",
+				Content: fmt.Sprintf("[system: step budget exhausted (%d/%d steps used). Stop calling tools. Summarize what you have done and what remains, then stop.]", step+1, cfg.MaxSteps),
+			}, nil)
+			break
+		}
+
 		if cfg.AutoCompact != nil {
 			cfg.AutoCompact(ctx)
 		}
diff --git a/pkg/agent/loop_maxsteps_test.go b/pkg/agent/loop_maxsteps_test.go
new file mode 100644
index 0000000..1c0ef13
--- /dev/null
+++ b/pkg/agent/loop_maxsteps_test.go
@@ -0,0 +1,143 @@
+package agent
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+
+	"ollie/pkg/backend"
+)
+
+// TestMaxStepsZeroUnlimited verifies that MaxSteps=0 does not trigger the
+// guardrail — the loop runs until the model stops calling tools.
+func TestMaxStepsZeroUnlimited(t *testing.T) {
+	var steps int
+	mb := &mockBackend{
+		responses: []mockResponse{
+			{toolCalls: []backend.ToolCall{{ID: "1", Name: "tool", Arguments: json.RawMessage(`{}`)}}, stopReason: "tool_calls"},
+			{toolCalls: []backend.ToolCall{{ID: "2", Name: "tool", Arguments: json.RawMessage(`{}`)}}, stopReason: "tool_calls"},
+			{content: "done", stopReason: "stop"},
+		},
+	}
+	cfg := agentConfig{
+		Backend:  mb,
+		Tools:    []backend.Tool{{Name: "tool"}},
+		MaxSteps: 0,
+		Exec: func(_ context.Context, _ string, _ json.RawMessage) (string, error) {
+			steps++
+			return "ok", nil
+		},
+	}
+	if err := run(context.Background(), cfg, newState()); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if steps != 2 {
+		t.Errorf("expected 2 tool executions, got %d", steps)
+	}
+}
+
+// TestMaxStepsSoftNudge verifies that when MaxSteps is reached the loop injects
+// the budget-exhausted nudge message, emits a maxsteps event, and exits cleanly
+// (no error returned). The model is given one final tool-free turn.
+func TestMaxStepsSoftNudge(t *testing.T) {
+	var nudgeSeen bool
+	var maxstepsEventSeen bool
+
+	// Backend: two tool-calling rounds, then a final text turn.
+	mb := &mockBackend{
+		responses: []mockResponse{
+			{toolCalls: []backend.ToolCall{{ID: "1", Name: "tool", Arguments: json.RawMessage(`{}`)}}, stopReason: "tool_calls"},
+			{toolCalls: []backend.ToolCall{{ID: "2", Name: "tool", Arguments: json.RawMessage(`{}`)}}, stopReason: "tool_calls"},
+			{content: "wrapping up", stopReason: "stop"},
+		},
+	}
+
+	// MaxSteps=1 means the guardrail fires after completing step 0 (the first
+	// tool round), before step 1 would begin.
+	cfg := agentConfig{
+		Backend:  mb,
+		Tools:    []backend.Tool{{Name: "tool"}},
+		MaxSteps: 1,
+		Exec: func(_ context.Context, _ string, _ json.RawMessage) (string, error) {
+			return "ok", nil
+		},
+		Output: func(ev Event) {
+			if ev.Role == "maxsteps" {
+				maxstepsEventSeen = true
+			}
+		},
+	}
+
+	// Intercept state updates to detect the nudge message.
+	s := newState()
+	origUpdate := s.update
+	_ = origUpdate // state.update is not a field; we'll check history post-run instead.
+
+	if err := run(context.Background(), cfg, s); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	if !maxstepsEventSeen {
+		t.Error("expected maxsteps event to be emitted")
+	}
+
+	// Confirm the nudge message is present in conversation history.
+	for _, m := range s.history() {
+		if m.Role == "user" && len(m.Content) > 0 {
+			if contains(m.Content, "step budget exhausted") {
+				nudgeSeen = true
+				break
+			}
+		}
+	}
+	if !nudgeSeen {
+		t.Error("expected step-budget nudge message in conversation history")
+	}
+}
+
+// TestMaxStepsExactBoundary checks that with MaxSteps=N the loop completes
+// exactly N tool rounds before nudging.
+func TestMaxStepsExactBoundary(t *testing.T) {
+	var toolRounds int
+
+	const limit = 3
+
+	// Build limit+1 tool responses so the model would run forever without the cap.
+	var responses []mockResponse
+	for i := range limit + 1 {
+		responses = append(responses, mockResponse{
+			toolCalls:  []backend.ToolCall{{ID: fmt.Sprintf("%d", i+1), Name: "tool", Arguments: json.RawMessage(`{}`)}},
+			stopReason: "tool_calls",
+		})
+	}
+	responses = append(responses, mockResponse{content: "done", stopReason: "stop"})
+
+	mb := &mockBackend{responses: responses}
+	cfg := agentConfig{
+		Backend:  mb,
+		Tools:    []backend.Tool{{Name: "tool"}},
+		MaxSteps: limit,
+		Exec: func(_ context.Context, _ string, _ json.RawMessage) (string, error) {
+			toolRounds++
+			return "ok", nil
+		},
+	}
+	if err := run(context.Background(), cfg, newState()); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if toolRounds != limit {
+		t.Errorf("expected %d tool rounds, got %d", limit, toolRounds)
+	}
+}
+
+func contains(s, substr string) bool {
+	return len(s) >= len(substr) && (s == substr || len(substr) == 0 ||
+		(func() bool {
+			for i := 0; i <= len(s)-len(substr); i++ {
+				if s[i:i+len(substr)] == substr {
+					return true
+				}
+			}
+			return false
+		})())
+}