agent: add MaxSteps soft guardrail to agentConfig and loop
This commit is contained in:
parent
46940eef07
commit
c20a5203ca
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
})())
|
||||
}
|
||||
Reference in New Issue