rony-llm-agent/pkg/agent/loop_test.go

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_ForwardsTrailingUsageOnlyChunk(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)
				// No Delta/ReasoningDelta, as providers report usage in a
				// separate trailing event; it must still be forwarded.
				yield(llm.StreamChunk{Usage: llm.TokenUsage{InputTokens: 10, OutputTokens: 3, TotalTokens: 13}}, 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.Fatalf("expected 2 chunks (content + usage-only), got %d", len(chunks))
	}
	usage := chunks[len(chunks)-1].Usage
	if usage.InputTokens != 10 || usage.OutputTokens != 3 || usage.TotalTokens != 13 {
		t.Errorf("expected the trailing usage-only chunk to be forwarded, got %+v", usage)
	}
}

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")
	}
}