agentu/internal/agent/agent_test.go

package agent

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"testing"

	"agentu/internal/tools"
	"agentu/pkg/llm"
)

type fakeProvider struct {
	calls int
	t     *testing.T
}

func (p *fakeProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
	p.calls++
	switch p.calls {
	case 1:
		if len(req.Tools) == 0 {
			p.t.Fatal("expected tools")
		}
		if req.ToolChoice != "auto" {
			p.t.Fatalf("tool_choice = %q", req.ToolChoice)
		}
		return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{
			{Index: 0, ID: "call_1", Type: "function", Name: "test_echo", Arguments: `{"text":"he`},
			{Index: 0, Arguments: `llo"}`},
		}})
	case 2:
		last := req.Messages[len(req.Messages)-1]
		if last.Role != llm.RoleTool || last.Content != "hello" {
			p.t.Fatalf("last message = %#v", last)
		}
		return emit(llm.StreamEvent{Content: "done"})
	default:
		p.t.Fatalf("unexpected call count %d", p.calls)
		return nil
	}
}

type echoTool struct{}

func (echoTool) Definition() llm.Tool {
	return llm.Tool{
		Type: "function",
		Function: llm.ToolFunction{
			Name:        "test_echo",
			Description: "echo test",
			Parameters:  json.RawMessage(`{"type":"object"}`),
		},
	}
}

func (echoTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
	var args struct {
		Text string `json:"text"`
	}
	if err := json.Unmarshal(raw, &args); err != nil {
		return "", err
	}
	return args.Text, nil
}

func TestAgentRunsToolLoop(t *testing.T) {
	provider := &fakeProvider{t: t}
	a := New(Options{
		Provider:     provider,
		Model:        "test",
		SystemPrompt: "system",
		ToolRegistry: tools.NewRegistry(echoTool{}),
	})

	var out strings.Builder
	var logs strings.Builder
	if err := a.RunTurn(context.Background(), "use a tool", &out, &logs); err != nil {
		t.Fatal(err)
	}
	if out.String() != "done" {
		t.Fatalf("out = %q", out.String())
	}
	if !strings.Contains(logs.String(), "[tool] test_echo") {
		t.Fatalf("logs = %q", logs.String())
	}
}

type contentProvider struct {
	called bool
	text   string
}

func (p *contentProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
	p.called = true
	return emit(llm.StreamEvent{Content: p.text})
}

type compactProvider struct {
	summaryCalls int
	answerCalls  int
	lastSummary  string
}

func (p *compactProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
	if len(req.Messages) == 2 && req.Messages[0].Role == llm.RoleSystem && req.Messages[0].Content == summarizerSystemPrompt {
		p.summaryCalls++
		p.lastSummary = req.Messages[1].Content
		if err := emit(llm.StreamEvent{Content: "old work summarized"}); err != nil {
			return err
		}
		return emit(llm.StreamEvent{Usage: &llm.Usage{PromptTokens: 10, CompletionTokens: 2, TotalTokens: 12}})
	}

	p.answerCalls++
	return emit(llm.StreamEvent{Content: "done"})
}

func TestAgentCompact(t *testing.T) {
	provider := &compactProvider{}
	a := New(Options{Provider: provider, Model: "test", SystemPrompt: "system"})
	a.SetMessages(longConversation())

	if err := a.Compact(context.Background()); err != nil {
		t.Fatal(err)
	}

	messages := a.Messages()
	if len(messages) != 8 {
		t.Fatalf("messages = %d, want 8", len(messages))
	}
	if messages[0].Role != llm.RoleSystem || messages[0].Content != "system" {
		t.Fatalf("system message = %#v", messages[0])
	}
	if messages[1].Role != llm.RoleAssistant || !strings.HasPrefix(messages[1].Content, contextSummaryPrefix) {
		t.Fatalf("summary message = %#v", messages[1])
	}
	if !strings.Contains(messages[1].Content, "old work summarized") {
		t.Fatalf("summary content = %q", messages[1].Content)
	}
	if !strings.Contains(provider.lastSummary, "old-user-0") {
		t.Fatalf("summary input did not include old history: %q", provider.lastSummary)
	}
	if messages[len(messages)-1].Content != "recent-answer-2" {
		t.Fatalf("last recent message = %#v", messages[len(messages)-1])
	}
	if messages[2].Content != "recent-user-0" {
		t.Fatalf("first retained recent message = %#v", messages[2])
	}
	if usage := a.LastUsage(); usage == nil || usage.TotalTokens != 12 {
		t.Fatalf("usage = %#v", usage)
	}
}

func TestAgentAutoCompact(t *testing.T) {
	provider := &compactProvider{}
	a := New(Options{
		Provider:          provider,
		Model:             "test",
		SystemPrompt:      "system",
		MaxContextTokens:  1,
		ToolRegistry:      tools.NewRegistry(),
		MaxToolIterations: 1,
	})
	a.SetMessages(longConversation())

	var out strings.Builder
	if err := a.RunTurn(context.Background(), "new question", &out, nil); err != nil {
		t.Fatal(err)
	}
	if provider.summaryCalls != 1 {
		t.Fatalf("summaryCalls = %d, want 1", provider.summaryCalls)
	}
	if provider.answerCalls != 1 {
		t.Fatalf("answerCalls = %d, want 1", provider.answerCalls)
	}
	if out.String() != "done" {
		t.Fatalf("out = %q", out.String())
	}
	if !strings.Contains(joinMessageContents(a.Messages()), contextSummaryPrefix) {
		t.Fatalf("messages missing summary: %#v", a.Messages())
	}
}

func TestAgentCompactSkipsShortHistory(t *testing.T) {
	provider := &contentProvider{text: "unused"}
	a := New(Options{Provider: provider, Model: "test", SystemPrompt: "system"})
	a.SetMessages([]llm.Message{
		{Role: llm.RoleSystem, Content: "system"},
		{Role: llm.RoleUser, Content: "question"},
		{Role: llm.RoleAssistant, Content: "answer"},
	})

	if err := a.Compact(context.Background()); err != nil {
		t.Fatal(err)
	}
	if provider.called {
		t.Fatal("provider should not be called for short history")
	}
	if len(a.Messages()) != 3 {
		t.Fatalf("messages = %d, want 3", len(a.Messages()))
	}
}

func longConversation() []llm.Message {
	messages := []llm.Message{{Role: llm.RoleSystem, Content: "system"}}
	for i := 0; i < 3; i++ {
		messages = append(messages,
			llm.Message{Role: llm.RoleUser, Content: fmt.Sprintf("old-user-%d", i)},
			llm.Message{Role: llm.RoleAssistant, Content: fmt.Sprintf("old-answer-%d", i)},
		)
	}
	for i := 0; i < 3; i++ {
		messages = append(messages,
			llm.Message{Role: llm.RoleUser, Content: fmt.Sprintf("recent-user-%d", i)},
			llm.Message{Role: llm.RoleAssistant, Content: fmt.Sprintf("recent-answer-%d", i)},
		)
	}
	return messages
}

func manyTurnConversation(turns int) []llm.Message {
	messages := []llm.Message{{Role: llm.RoleSystem, Content: "system"}}
	for i := 0; i < turns; i++ {
		messages = append(messages,
			llm.Message{Role: llm.RoleUser, Content: fmt.Sprintf("user-%02d", i)},
			llm.Message{Role: llm.RoleAssistant, Content: fmt.Sprintf("assistant-%02d", i)},
		)
	}
	return messages
}

func TestAgentCompactKeepsTwentyPercentRecentContext(t *testing.T) {
	provider := &compactProvider{}
	a := New(Options{Provider: provider, Model: "test", SystemPrompt: "system"})
	a.SetMessages(manyTurnConversation(25))

	if err := a.Compact(context.Background()); err != nil {
		t.Fatal(err)
	}

	messages := a.Messages()
	// 1 system + 1 summary + 10 retained recent messages (5 user/assistant turns)
	if len(messages) != 12 {
		t.Fatalf("messages = %d, want 12; contents: %v", len(messages), joinMessageContents(messages))
	}
	if provider.summaryCalls != 1 {
		t.Fatalf("summaryCalls = %d, want 1", provider.summaryCalls)
	}
	if messages[2].Content != "user-20" {
		t.Fatalf("first retained message = %#v, want user-20", messages[2])
	}
	if messages[len(messages)-1].Content != "assistant-24" {
		t.Fatalf("last retained message = %#v, want assistant-24", messages[len(messages)-1])
	}
	if !strings.Contains(provider.lastSummary, "user-19") {
		t.Fatalf("summary should contain user-19 (part of summarized 80%%): %q", provider.lastSummary)
	}
	if strings.Contains(provider.lastSummary, "user-20") {
		t.Fatalf("summary should not contain user-20 (part of retained 20%%): %q", provider.lastSummary)
	}
}

func joinMessageContents(messages []llm.Message) string {
	var b strings.Builder
	for _, msg := range messages {
		b.WriteString(msg.Content)
		b.WriteByte('\n')
	}
	return b.String()
}
func TestAgentLetsProviderDecideShellRequestWhenShellDisabled(t *testing.T) {
	provider := &contentProvider{text: "shell is disabled"}
	a := New(Options{
		Provider:     provider,
		Model:        "test",
		SystemPrompt: "system",
		ToolRegistry: tools.ReadOnlyBuiltins(t.TempDir()),
	})

	var out strings.Builder
	if err := a.RunTurn(context.Background(), "能帮我看看 `ssh home-dev` 的运行情况么", &out, nil); err != nil {
		t.Fatal(err)
	}
	if !provider.called {
		t.Fatal("provider should be called so the model can decide whether to call a tool or answer")
	}
	if out.String() != "shell is disabled" {
		t.Fatalf("out = %q", out.String())
	}
}

func TestAgentDoesNotBlockFileRequest(t *testing.T) {
	provider := &contentProvider{text: "ok"}
	a := New(Options{
		Provider:     provider,
		Model:        "test",
		SystemPrompt: "system",
		ToolRegistry: tools.ReadOnlyBuiltins(t.TempDir()),
	})

	var out strings.Builder
	if err := a.RunTurn(context.Background(), "请读取 `README.md`", &out, nil); err != nil {
		t.Fatal(err)
	}
	if !provider.called {
		t.Fatal("provider should be called for file requests")
	}
	if out.String() != "ok" {
		t.Fatalf("out = %q", out.String())
	}
}

// alwaysToolProvider always returns a tool call with the same args on every request.
type alwaysToolProvider struct {
	calls int
}

func (p *alwaysToolProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
	p.calls++
	return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{
		{Index: 0, ID: fmt.Sprintf("call_%d", p.calls), Type: "function", Name: "test_echo", Arguments: `{"text":"loop"}`},
	}})
}

func TestAgentStopsOnDoomLoop(t *testing.T) {
	provider := &alwaysToolProvider{}
	a := New(Options{
		Provider:          provider,
		Model:             "test",
		SystemPrompt:      "system",
		ToolRegistry:      tools.NewRegistry(echoTool{}),
		MaxToolIterations: 50, // High hard limit — doom loop should trigger first
	})

	var out strings.Builder
	err := a.RunTurn(context.Background(), "loop forever", &out, nil)
	if err == nil {
		t.Fatal("expected error when doom loop is detected")
	}
	if !strings.Contains(err.Error(), "doom loop detected") {
		t.Fatalf("unexpected error: %v", err)
	}
	// Should have stopped at doomLoopThreshold (3) rounds, not reached hard limit.
	if provider.calls != doomLoopThreshold {
		t.Fatalf("calls = %d, want %d (doom loop threshold)", provider.calls, doomLoopThreshold)
	}
}

// varyingToolProvider alternates between different tool calls to avoid doom loop detection.
type varyingToolProvider struct {
	calls int
}

func (p *varyingToolProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
	p.calls++
	// Vary the arguments each round to avoid doom loop detection.
	return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{
		{Index: 0, ID: fmt.Sprintf("call_%d", p.calls), Type: "function", Name: "test_echo",
			Arguments: fmt.Sprintf(`{"text":"round-%d"}`, p.calls)},
	}})
}

func TestAgentHardLimitStopsVaryingCalls(t *testing.T) {
	provider := &varyingToolProvider{}
	const limit = 5
	a := New(Options{
		Provider:          provider,
		Model:             "test",
		SystemPrompt:      "system",
		ToolRegistry:      tools.NewRegistry(echoTool{}),
		MaxToolIterations: limit,
	})

	var out strings.Builder
	err := a.RunTurn(context.Background(), "vary forever", &out, nil)
	if err == nil {
		t.Fatal("expected error when hard limit is reached")
	}
	if !strings.Contains(err.Error(), "tool round limit reached") {
		t.Fatalf("unexpected error: %v", err)
	}
	if provider.calls != limit {
		t.Fatalf("calls = %d, want %d", provider.calls, limit)
	}
}

// slowTool blocks on a channel before returning.
type slowTool struct {
	name    string
	started chan<- string
	release <-chan struct{}
}

func (t *slowTool) Definition() llm.Tool {
	return llm.Tool{
		Type: "function",
		Function: llm.ToolFunction{
			Name:        t.name,
			Description: "slow test tool",
			Parameters:  json.RawMessage(`{"type":"object","properties":{"text":{"type":"string"}}}`),
		},
	}
}

func (t *slowTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
	var args struct {
		Text string `json:"text"`
	}
	if err := json.Unmarshal(raw, &args); err != nil {
		return "", err
	}
	t.started <- t.name
	<-t.release
	return args.Text, nil
}

// mutatingSlowTool is like slowTool but implements MutatingTool.
type mutatingSlowTool struct {
	slowTool
}

func (*mutatingSlowTool) Mutates() bool { return true }

func TestAgentExecutesIndependentToolCallsInParallel(t *testing.T) {
	started := make(chan string, 2)
	release := make(chan struct{})

	tool1 := &slowTool{name: "slow_one", started: started, release: release}
	tool2 := &slowTool{name: "slow_two", started: started, release: release}

	provider := &parallelProvider{
		responses: [][]llm.ToolCallDelta{
			{
				{Index: 0, ID: "call_1", Type: "function", Name: "slow_one", Arguments: `{"text":"one"}`},
				{Index: 1, ID: "call_2", Type: "function", Name: "slow_two", Arguments: `{"text":"two"}`},
			},
		},
		finalContent: "done",
	}

	a := New(Options{
		Provider:     provider,
		Model:        "test",
		SystemPrompt: "system",
		ToolRegistry: tools.NewRegistry(tool1, tool2),
	})

	var out strings.Builder
	done := make(chan error, 1)
	go func() {
		done <- a.RunTurn(context.Background(), "run both", &out, nil)
	}()

	// Both tools must start before we release them.
	names := make(map[string]bool)
	for i := 0; i < 2; i++ {
		select {
		case name := <-started:
			names[name] = true
		case <-make(chan struct{}):
			// Use a timer in real code, but for test simplicity we rely on deadlock detection.
			t.Fatal("timed out waiting for tool to start")
		}
	}
	if !names["slow_one"] || !names["slow_two"] {
		t.Fatalf("expected both tools to start, got %v", names)
	}
	close(release)

	if err := <-done; err != nil {
		t.Fatal(err)
	}
	if out.String() != "done" {
		t.Fatalf("out = %q", out.String())
	}
	// Verify tool results were in order.
	msgs := a.Messages()
	var toolMsgs []string
	for _, m := range msgs {
		if m.Role == llm.RoleTool {
			toolMsgs = append(toolMsgs, m.Content)
		}
	}
	if len(toolMsgs) != 2 || toolMsgs[0] != "one" || toolMsgs[1] != "two" {
		t.Fatalf("tool messages = %v, want [one two]", toolMsgs)
	}
}

func TestAgentRunsMutatingToolCallsSequentially(t *testing.T) {
	started := make(chan string, 2)
	release := make(chan struct{})

	tool1 := &mutatingSlowTool{slowTool{name: "mut_one", started: started, release: release}}
	tool2 := &mutatingSlowTool{slowTool{name: "mut_two", started: started, release: release}}

	provider := &parallelProvider{
		responses: [][]llm.ToolCallDelta{
			{
				{Index: 0, ID: "call_1", Type: "function", Name: "mut_one", Arguments: `{"text":"one"}`},
				{Index: 1, ID: "call_2", Type: "function", Name: "mut_two", Arguments: `{"text":"two"}`},
			},
		},
		finalContent: "done",
	}

	a := New(Options{
		Provider:     provider,
		Model:        "test",
		SystemPrompt: "system",
		ToolRegistry: tools.NewRegistry(tool1, tool2),
	})

	var out strings.Builder
	done := make(chan error, 1)
	go func() {
		done <- a.RunTurn(context.Background(), "run both", &out, nil)
	}()

	// Only the first tool should start.
	select {
	case name := <-started:
		if name != "mut_one" {
			t.Fatalf("expected mut_one to start first, got %q", name)
		}
	case <-make(chan struct{}):
		t.Fatal("timed out waiting for first tool to start")
	}

	// Second tool should NOT have started yet.
	select {
	case name := <-started:
		t.Fatalf("second tool started too early: %q", name)
	default:
		// Expected: second tool hasn't started.
	}

	close(release)
	if err := <-done; err != nil {
		t.Fatal(err)
	}
}

// parallelProvider emits a batch of tool calls on the first request, then returns finalContent.
type parallelProvider struct {
	responses    [][]llm.ToolCallDelta
	finalContent string
	calls        int
}

func (p *parallelProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
	p.calls++
	if p.calls <= len(p.responses) {
		return emit(llm.StreamEvent{ToolCalls: p.responses[p.calls-1]})
	}
	return emit(llm.StreamEvent{Content: p.finalContent})
}