From 4b39f52081da98cbd8f64df4bf8ccac3b9d5a60c Mon Sep 17 00:00:00 2001 From: Victor Vargas Date: Tue, 30 Jun 2026 23:53:26 -0700 Subject: [PATCH] feat(agent): add Agent loop with iteration control, tool handling, and streaming support --- pkg/agent/integration_test.go | 157 ++++++++++ pkg/agent/loop.go | 239 +++++++++++++++ pkg/agent/loop_test.go | 527 ++++++++++++++++++++++++++++++++++ 3 files changed, 923 insertions(+) create mode 100644 pkg/agent/integration_test.go create mode 100644 pkg/agent/loop.go create mode 100644 pkg/agent/loop_test.go diff --git a/pkg/agent/integration_test.go b/pkg/agent/integration_test.go new file mode 100644 index 0000000..98cfd49 --- /dev/null +++ b/pkg/agent/integration_test.go @@ -0,0 +1,157 @@ +package agent_test + +import ( + "context" + "encoding/json" + "os" + "testing" + + "github.com/VictorVargas/rony-llm-agent/pkg/agent" + "github.com/VictorVargas/rony-llm-agent/pkg/llm" + "github.com/VictorVargas/rony-llm-agent/pkg/llm/providers/llamacpp" + "github.com/VictorVargas/rony-llm-agent/pkg/persona" + "github.com/VictorVargas/rony-llm-agent/pkg/tools" +) + +// TestIntegration_LlamaCPP_Generate is an integration test that requires llama.cpp running on localhost:8080. +// Run with: go test ./pkg/agent/ -run TestIntegration_LlamaCPP_Generate -tags=integration +func TestIntegration_LlamaCPP_Generate(t *testing.T) { + client, err := llamacpp.New(llamacpp.Config{ + BaseURL: "http://localhost:8080/v1", + }) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + resp, err := client.Generate(context.Background(), llm.CompletionRequest{ + Messages: []llm.Message{ + {Role: llm.RoleUser, Content: "What is 2+2? Answer with just the number."}, + }, + }) + if err != nil { + t.Fatalf("generate failed: %v", err) + } + + if resp.Content == "" { + t.Fatal("expected non-empty response") + } +} + +// TestIntegration_LlamaCPP_Stream is an integration test that requires llama.cpp running on localhost:8080. +func TestIntegration_LlamaCPP_Stream(t *testing.T) { + client, err := llamacpp.New(llamacpp.Config{ + BaseURL: "http://localhost:8080/v1", + }) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + stream := client.Stream(context.Background(), llm.CompletionRequest{ + Messages: []llm.Message{ + {Role: llm.RoleUser, Content: "Say hello in 5 words."}, + }, + }) + + var chunks []llm.StreamChunk + for chunk, err := range stream { + if err != nil { + t.Fatalf("stream error: %v", err) + } + chunks = append(chunks, chunk) + } + + if len(chunks) == 0 { + t.Fatal("expected at least one chunk") + } +} + +// TestIntegration_AgentLoop_Generate is an integration test for the agent loop with llama.cpp. +func TestIntegration_AgentLoop_Generate(t *testing.T) { + client, err := llamacpp.New(llamacpp.Config{ + BaseURL: "http://localhost:8080/v1", + }) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + registry := tools.NewRegistry() + registry.Register(tools.Tool{ + Name: "add_numbers", + Description: "Add two numbers together", + InputSchema: json.RawMessage(`{"type": "function", "function": {"name": "add_numbers", "description": "Add two numbers together", "parameters": {"type": "object", "properties": {"a": {"type": "number"}, "b": {"type": "number"}}}}}`), + Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) { + var input struct { + A float64 `json:"a"` + B float64 `json:"b"` + } + if err := json.Unmarshal(args, &input); err != nil { + return tools.ToolResult{IsError: true}, err + } + return tools.ToolResult{Content: "42"}, nil + }, + Permission: tools.Allow, + }) + + loop := agent.New(agent.Config{ + LLM: client, + Persona: persona.DefaultPersona(), + Tools: registry, + Sandbox: &mockSandbox{}, + MaxIters: 3, + }) + + resp, err := loop.Run(context.Background(), "What is 20+22? Use the add_numbers tool.") + if err != nil { + t.Fatalf("run failed: %v", err) + } + + if resp.Content == "" { + t.Fatal("expected non-empty response") + } +} + +// TestIntegration_AgentLoop_Stream is an integration test for the agent loop with streaming. +func TestIntegration_AgentLoop_Stream(t *testing.T) { + client, err := llamacpp.New(llamacpp.Config{ + BaseURL: "http://localhost:8080/v1", + }) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + loop := agent.New(agent.Config{ + LLM: client, + Persona: persona.DefaultPersona(), + Tools: tools.NewRegistry(), + MaxIters: 3, + }) + + stream := loop.RunStream(context.Background(), "Say something interesting.") + + var chunks []llm.StreamChunk + for chunk, err := range stream { + if err != nil { + t.Fatalf("stream error: %v", err) + } + chunks = append(chunks, chunk) + } + + if len(chunks) == 0 { + t.Fatal("expected at least one chunk") + } +} + +// mockSandbox is a simple sandbox that allows all calls. +type mockSandbox struct{} + +func (m *mockSandbox) ValidateToolCall(tool tools.Tool, call llm.ToolCall) error { + return nil +} + +func TestMain(m *testing.M) { + // Skip integration tests unless explicitly enabled + if os.Getenv("INTEGRATION_TESTS") != "1" { + return + } + os.Exit(m.Run()) +} diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go new file mode 100644 index 0000000..aa988b1 --- /dev/null +++ b/pkg/agent/loop.go @@ -0,0 +1,239 @@ +package agent + +import ( + "context" + "encoding/json" + "fmt" + "iter" + "time" + + "github.com/VictorVargas/rony-llm-agent/pkg/llm" + "github.com/VictorVargas/rony-llm-agent/pkg/persona" + "github.com/VictorVargas/rony-llm-agent/pkg/tools" +) + +// MaxIterations is the default maximum number of agent iterations. +const DefaultMaxIterations = 50 + +// MaxToolOutputBytes is the default maximum size for tool output in bytes. +const DefaultMaxToolOutputBytes = 50 * 1024 // 50KB + +// Sandbox defines the interface for filesystem sandbox operations. +type Sandbox interface { + ValidateToolCall(tool tools.Tool, call llm.ToolCall) error +} + +// Approver is called before executing tools with Ask permission. +// Return true to allow, false to deny. +type Approver func(tool tools.Tool, call llm.ToolCall) bool + +// OnIterationHook is called after each iteration. +type OnIterationHook func(Iteration) + +// Config holds the dependencies and settings for the agent loop. +type Config struct { + LLM llm.LLMClient + Persona persona.Persona + Tools tools.Registry + Sandbox Sandbox + MaxIters int + Approver Approver + OnIteration OnIterationHook + ToolTimeout time.Duration +} + +// Iteration represents a single cycle of the agent loop. +type Iteration struct { + Number int + ToolCalls []llm.ToolCall + ToolsUsed int + Duration time.Duration +} + +// Response is the final output of the agent loop. +type Response struct { + Content string + ToolCalls []llm.ToolCall + Iterations int + Duration time.Duration + TokenUsage llm.TokenUsage +} + +// Loop is the main agent loop that orchestrates LLM calls and tool execution. +type Loop struct { + cfg Config +} + +// New creates a new Loop with the given configuration. +func New(cfg Config) *Loop { + if cfg.MaxIters == 0 { + cfg.MaxIters = DefaultMaxIterations + } + return &Loop{cfg: cfg} +} + +// Run executes the agent loop and returns the final response. +func (l *Loop) Run(ctx context.Context, input string) (Response, error) { + start := time.Now() + + messages := l.buildInitialMessages(input) + var finalContent string + var allToolCalls []llm.ToolCall + var totalUsage llm.TokenUsage + iterations := 0 + + for iterations < l.cfg.MaxIters { + iterations++ + + resp, err := l.cfg.LLM.Generate(ctx, llm.CompletionRequest{ + Messages: messages, + Tools: l.getToolSchemas(), + }) + if err != nil { + return Response{}, fmt.Errorf("LLM generate failed: %w", err) + } + + totalUsage.InputTokens += resp.Usage.InputTokens + totalUsage.OutputTokens += resp.Usage.OutputTokens + totalUsage.TotalTokens += resp.Usage.TotalTokens + + if len(resp.ToolCalls) == 0 { + finalContent = resp.Content + break + } + + for _, call := range resp.ToolCalls { + result, err := l.executeTool(ctx, call) + if err != nil { + messages = append(messages, llm.Message{ + Role: llm.RoleTool, + Content: fmt.Sprintf("Error: %v", err), + }) + continue + } + messages = append(messages, llm.Message{ + Role: llm.RoleTool, + Content: result.Content, + }) + allToolCalls = append(allToolCalls, call) + } + } + + duration := time.Since(start) + if iterations >= l.cfg.MaxIters { + return Response{}, fmt.Errorf("max iterations (%d) reached", l.cfg.MaxIters) + } + + return Response{ + Content: finalContent, + ToolCalls: allToolCalls, + Iterations: iterations, + Duration: duration, + TokenUsage: totalUsage, + }, nil +} + +// RunStream executes the agent loop with streaming output. +func (l *Loop) RunStream(ctx context.Context, input string) iter.Seq2[llm.StreamChunk, error] { + return func(yield func(llm.StreamChunk, error) bool) { + messages := l.buildInitialMessages(input) + iterations := 0 + + for iterations < l.cfg.MaxIters { + iterations++ + + stream := l.cfg.LLM.Stream(ctx, llm.CompletionRequest{ + Messages: messages, + Tools: l.getToolSchemas(), + }) + + var hasToolCalls bool + for chunk, err := range stream { + if err != nil { + yield(llm.StreamChunk{}, err) + return + } + + if len(chunk.ToolCalls) > 0 { + hasToolCalls = true + for _, tc := range chunk.ToolCalls { + result, err := l.executeTool(ctx, tc) + if err != nil { + messages = append(messages, llm.Message{ + Role: llm.RoleTool, + Content: fmt.Sprintf("Error: %v", err), + }) + continue + } + messages = append(messages, llm.Message{ + Role: llm.RoleTool, + Content: result.Content, + }) + } + } + + if !hasToolCalls && chunk.Delta != "" { + if !yield(chunk, nil) { + return + } + } + } + + if !hasToolCalls { + return + } + } + + if iterations >= l.cfg.MaxIters { + yield(llm.StreamChunk{}, fmt.Errorf("max iterations (%d) reached", l.cfg.MaxIters)) + } + } +} + +func (l *Loop) buildInitialMessages(input string) []llm.Message { + systemPrompt := persona.AssembleSystemPrompt(l.cfg.Persona, "") + return []llm.Message{ + {Role: llm.RoleSystem, Content: systemPrompt}, + {Role: llm.RoleUser, Content: input}, + } +} + +func (l *Loop) getToolSchemas() []json.RawMessage { + var schemas []json.RawMessage + for _, tool := range l.cfg.Tools.List() { + schemas = append(schemas, tool.InputSchema) + } + return schemas +} + +func (l *Loop) executeTool(ctx context.Context, call llm.ToolCall) (tools.ToolResult, error) { + tool, found := l.cfg.Tools.Get(call.Name) + if !found { + return tools.ToolResult{IsError: true}, fmt.Errorf("tool %q not found", call.Name) + } + + if tool.Permission == tools.Ask && l.cfg.Approver != nil && !l.cfg.Approver(tool, call) { + return tools.ToolResult{Content: "Tool execution denied by user"}, nil + } + + if l.cfg.Sandbox != nil { + if err := l.cfg.Sandbox.ValidateToolCall(tool, call); err != nil { + return tools.ToolResult{IsError: true}, fmt.Errorf("sandbox violation: %w", err) + } + } + + if l.cfg.OnIteration != nil { + l.cfg.OnIteration(Iteration{ + Number: 1, + ToolCalls: []llm.ToolCall{call}, + ToolsUsed: 1, + }) + } + + result, err := tool.Handler(ctx, call.Arguments) + if err != nil { + return tools.ToolResult{IsError: true, Content: fmt.Sprintf("Tool error: %v", err)}, nil + } + + return result, nil +} diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go new file mode 100644 index 0000000..a4d5c86 --- /dev/null +++ b/pkg/agent/loop_test.go @@ -0,0 +1,527 @@ +package agent + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "iter" + "testing" + "time" + + llm "github.com/VictorVargas/rony-llm-agent/pkg/llm" + "github.com/VictorVargas/rony-llm-agent/pkg/persona" + "github.com/VictorVargas/rony-llm-agent/pkg/tools" +) + +type mockLLM struct { + generateFunc func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) + streamFunc func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] +} + +func (m *mockLLM) Generate(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + return m.generateFunc(ctx, req) +} + +func (m *mockLLM) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] { + return m.streamFunc(ctx, req) +} + +func (m *mockLLM) Name() string { return "mock" } +func (m *mockLLM) Capabilities() llm.ProviderCapabilities { return llm.ProviderCapabilities{} } + +type mockSandbox struct { + validateFunc func(tool tools.Tool, call llm.ToolCall) error +} + +func (m *mockSandbox) ValidateToolCall(tool tools.Tool, call llm.ToolCall) error { + if m.validateFunc != nil { + return m.validateFunc(tool, call) + } + return nil +} + +func TestNew(t *testing.T) { + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + return llm.CompletionResponse{Content: "done"}, nil + }, + } + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: tools.NewRegistry(), + }) + + if loop == nil { + t.Fatal("expected non-nil loop") + } + if loop.cfg.MaxIters != DefaultMaxIterations { + t.Errorf("expected default max iters %d, got %d", DefaultMaxIterations, loop.cfg.MaxIters) + } +} + +func TestNew_CustomMaxIters(t *testing.T) { + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + return llm.CompletionResponse{Content: "done"}, nil + }, + } + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: tools.NewRegistry(), + MaxIters: 10, + }) + + if loop.cfg.MaxIters != 10 { + t.Errorf("expected max iters 10, got %d", loop.cfg.MaxIters) + } +} + +func TestRun_NoToolCalls(t *testing.T) { + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + return llm.CompletionResponse{ + Content: "I understand.", + Usage: llm.TokenUsage{ + InputTokens: 10, + OutputTokens: 5, + TotalTokens: 15, + }, + }, nil + }, + } + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: tools.NewRegistry(), + }) + + resp, err := loop.Run(context.Background(), "Hello") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.Content != "I understand." { + t.Errorf("expected 'I understand.', got %q", resp.Content) + } + if resp.Iterations != 1 { + t.Errorf("expected 1 iteration, got %d", resp.Iterations) + } + if resp.TokenUsage.InputTokens != 10 { + t.Errorf("expected 10 input tokens, got %d", resp.TokenUsage.InputTokens) + } +} + +func TestRun_ToolCalls(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(tools.Tool{ + Name: "greet", + Description: "Greet someone", + InputSchema: json.RawMessage(`{"type":"object","properties":{"name":{"type":"string"}}}`), + Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) { + return tools.ToolResult{Content: "Hello!"}, nil + }, + Permission: tools.Allow, + }) + + callCount := 0 + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + callCount++ + if callCount == 1 { + return llm.CompletionResponse{ + ToolCalls: []llm.ToolCall{ + {ID: "call-1", Name: "greet", Arguments: json.RawMessage(`{"name":"World"}`)}, + }, + }, nil + } + return llm.CompletionResponse{ + Content: "Done!", + }, nil + }, + } + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: registry, + }) + + resp, err := loop.Run(context.Background(), "Say hi") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.Content != "Done!" { + t.Errorf("expected 'Done!', got %q", resp.Content) + } + if resp.Iterations != 2 { + t.Errorf("expected 2 iterations, got %d", resp.Iterations) + } + if len(resp.ToolCalls) != 1 { + t.Errorf("expected 1 tool call, got %d", len(resp.ToolCalls)) + } +} + +func TestRun_MaxIterations(t *testing.T) { + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + return llm.CompletionResponse{ + ToolCalls: []llm.ToolCall{{ID: "1", Name: "x", Arguments: json.RawMessage("{}")}}, + }, nil + }, + } + + registry := tools.NewRegistry() + registry.Register(tools.Tool{ + Name: "x", + Description: "x", + InputSchema: json.RawMessage(`{}`), + Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) { + return tools.ToolResult{Content: "ok"}, nil + }, + Permission: tools.Allow, + }) + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: registry, + MaxIters: 3, + }) + + _, err := loop.Run(context.Background(), "test") + if err == nil { + t.Fatal("expected error, got nil") + } + if err.Error() != "max iterations (3) reached" { + t.Errorf("expected 'max iterations (3) reached', got %q", err.Error()) + } +} + +func TestRun_ToolNotFound(t *testing.T) { + callCount := 0 + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + callCount++ + if callCount == 1 { + return llm.CompletionResponse{ + ToolCalls: []llm.ToolCall{{ID: "1", Name: "nonexistent", Arguments: json.RawMessage("{}")}}, + }, nil + } + return llm.CompletionResponse{Content: "done"}, nil + }, + } + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: tools.NewRegistry(), + }) + + resp, err := loop.Run(context.Background(), "test") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.Content != "done" { + t.Errorf("expected 'done', got %q", resp.Content) + } +} + +func TestRun_ApprovalDenied(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(tools.Tool{ + Name: "dangerous", + Description: "Do something dangerous", + InputSchema: json.RawMessage(`{}`), + Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) { + return tools.ToolResult{Content: "executed"}, nil + }, + Permission: tools.Ask, + }) + + callCount := 0 + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + callCount++ + if callCount == 1 { + return llm.CompletionResponse{ + ToolCalls: []llm.ToolCall{{ID: "1", Name: "dangerous", Arguments: json.RawMessage("{}")}}, + }, nil + } + return llm.CompletionResponse{Content: "done"}, nil + }, + } + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: registry, + Approver: func(tool tools.Tool, call llm.ToolCall) bool { + return false // deny all + }, + }) + + resp, err := loop.Run(context.Background(), "test") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(resp.ToolCalls) != 1 { + t.Errorf("expected 1 tool call, got %d", len(resp.ToolCalls)) + } +} + +func TestRun_SandboxViolation(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(tools.Tool{ + Name: "restricted", + Description: "Restricted tool", + InputSchema: json.RawMessage(`{}`), + Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) { + return tools.ToolResult{Content: "executed"}, nil + }, + Permission: tools.Allow, + }) + + callCount := 0 + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + callCount++ + if callCount == 1 { + return llm.CompletionResponse{ + ToolCalls: []llm.ToolCall{{ID: "1", Name: "restricted", Arguments: json.RawMessage("{}")}}, + }, nil + } + return llm.CompletionResponse{Content: "done"}, nil + }, + } + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: registry, + Sandbox: &mockSandbox{ + validateFunc: func(tool tools.Tool, call llm.ToolCall) error { + return fmt.Errorf("path traversal detected") + }, + }, + }) + + resp, err := loop.Run(context.Background(), "test") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(resp.ToolCalls) != 0 { + t.Errorf("expected 0 tool calls (sandbox rejected), got %d", len(resp.ToolCalls)) + } +} + +func TestRun_OnIterationHook(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(tools.Tool{ + Name: "greet", + Description: "Greet", + InputSchema: json.RawMessage(`{}`), + Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) { + return tools.ToolResult{Content: "hello"}, nil + }, + Permission: tools.Allow, + }) + + callCount := 0 + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + callCount++ + if callCount == 1 { + return llm.CompletionResponse{ + ToolCalls: []llm.ToolCall{{ID: "1", Name: "greet", Arguments: json.RawMessage("{}")}}, + }, nil + } + return llm.CompletionResponse{Content: "done"}, nil + }, + } + + var iterations []Iteration + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: registry, + OnIteration: func(iter Iteration) { + iterations = append(iterations, iter) + }, + }) + + _, err := loop.Run(context.Background(), "test") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(iterations) != 1 { + t.Errorf("expected 1 iteration hook call, got %d", len(iterations)) + } + if iterations[0].ToolsUsed != 1 { + t.Errorf("expected 1 tool used, got %d", iterations[0].ToolsUsed) + } +} + +func TestRun_Stream_NoToolCalls(t *testing.T) { + mockClient := &mockLLM{ + streamFunc: func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] { + return func(yield func(llm.StreamChunk, error) bool) { + yield(llm.StreamChunk{Delta: "Hello"}, nil) + yield(llm.StreamChunk{Delta: " world"}, nil) + yield(llm.StreamChunk{FinishReason: "stop"}, nil) + } + }, + } + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: tools.NewRegistry(), + }) + + var chunks []llm.StreamChunk + stream := loop.RunStream(context.Background(), "test") + for chunk, err := range stream { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + chunks = append(chunks, chunk) + } + + if len(chunks) != 2 { + t.Errorf("expected 2 chunks, got %d", len(chunks)) + } + if chunks[0].Delta != "Hello" { + t.Errorf("expected 'Hello', got %q", chunks[0].Delta) + } + if chunks[1].Delta != " world" { + t.Errorf("expected ' world', got %q", chunks[1].Delta) + } +} + +func TestRun_Stream_WithToolCalls(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(tools.Tool{ + Name: "greet", + Description: "Greet", + InputSchema: json.RawMessage(`{}`), + Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) { + return tools.ToolResult{Content: "greeted"}, nil + }, + Permission: tools.Allow, + }) + + callCount := 0 + mockClient := &mockLLM{ + streamFunc: func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] { + return func(yield func(llm.StreamChunk, error) bool) { + callCount++ + if callCount == 1 { + yield(llm.StreamChunk{ + ToolCalls: []llm.ToolCall{{ID: "1", Name: "greet", Arguments: json.RawMessage("{}")}}, + }, nil) + } else { + yield(llm.StreamChunk{Delta: "done"}, nil) + yield(llm.StreamChunk{FinishReason: "stop"}, nil) + } + } + }, + } + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: registry, + }) + + var chunks []llm.StreamChunk + stream := loop.RunStream(context.Background(), "test") + for chunk, err := range stream { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + chunks = append(chunks, chunk) + } + + if len(chunks) != 1 { + t.Errorf("expected 1 chunk (only the 'done' chunk), got %d", len(chunks)) + } + if chunks[0].Delta != "done" { + t.Errorf("expected 'done', got %q", chunks[0].Delta) + } +} + +func TestRun_Stream_MaxIterations(t *testing.T) { + mockClient := &mockLLM{ + streamFunc: func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] { + return func(yield func(llm.StreamChunk, error) bool) { + yield(llm.StreamChunk{ + ToolCalls: []llm.ToolCall{{ID: "1", Name: "x", Arguments: json.RawMessage("{}")}}, + }, nil) + } + }, + } + + registry := tools.NewRegistry() + registry.Register(tools.Tool{ + Name: "x", + Description: "x", + InputSchema: json.RawMessage(`{}`), + Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) { + return tools.ToolResult{Content: "ok"}, nil + }, + Permission: tools.Allow, + }) + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: registry, + MaxIters: 1, + }) + + var chunks []llm.StreamChunk + stream := loop.RunStream(context.Background(), "test") + for chunk, err := range stream { + if err != nil { + // expect max iterations error + if !errors.Is(err, context.DeadlineExceeded) && err.Error() != "max iterations (1) reached" { + t.Fatalf("expected max iterations error, got: %v", err) + } + continue + } + chunks = append(chunks, chunk) + } +} + +func TestRun_Timeout(t *testing.T) { + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + select { + case <-ctx.Done(): + return llm.CompletionResponse{}, ctx.Err() + default: + return llm.CompletionResponse{Content: "done"}, nil + } + }, + } + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) // ensure context is cancelled before calling + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: tools.NewRegistry(), + }) + + _, err := loop.Run(ctx, "test") + if err == nil { + t.Fatal("expected timeout error, got nil") + } +}