agent_step_limit_test.go RawHistoryFolder

1package core
2
3import (
4	"context"
5	"encoding/json"
6	"errors"
7	"fmt"
8	"sync/atomic"
9	"testing"
10
11	"github.com/bnema/zut/packages/provider"
12)
13
14type loopingToolClient struct {
15	calls int32
16}
17
18func (c *loopingToolClient) Name() string { return "looping-tool" }
19
20func (c *loopingToolClient) Stream(_ context.Context, req provider.Request) (<-chan provider.Event, error) {
21	call := atomic.AddInt32(&c.calls, 1)
22	callID := fmt.Sprintf("call-%d", call)
23	out := make(chan provider.Event, 4)
24	out <- provider.EventStart{Provider: c.Name(), Model: req.Model}
25	out <- provider.EventToolStart{ID: callID, Name: "echo"}
26	out <- provider.EventToolEnd{ID: callID}
27	out <- provider.EventDone{Stop: provider.StopToolUse, Message: provider.Message{
28		Role: provider.RoleAssistant,
29		Content: []provider.Content{provider.ToolCallBlock{
30			ID: callID, Name: "echo", Arguments: json.RawMessage(`{}`),
31		}},
32	}}
33	close(out)
34	return out, nil
35}
36
37func TestAgentReturnsTypedStepLimitAfterModelLoopBudget(t *testing.T) {
38	client := &loopingToolClient{}
39	agent := NewAgent(client, "model", "system", Registry{"echo": &recordingTool{}})
40	agent.MaxSteps = 2
41
42	err := agent.Prompt(context.Background(), "keep using tools", nil, nil)
43	if !errors.Is(err, ErrMaxSteps) {
44		t.Fatalf("prompt error = %v, want ErrMaxSteps", err)
45	}
46	if got := atomic.LoadInt32(&client.calls); got != 2 {
47		t.Fatalf("provider calls = %d, want 2", got)
48	}
49}