agentu/internal/agent/agent_test.go

package agent

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"sync"
	"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 streamToolProvider struct {
	calls int
	t     *testing.T
}

func (p *streamToolProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
	p.calls++
	switch p.calls {
	case 1:
		return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{
			{Index: 0, ID: "call_1", Type: "function", Name: "test_stream", Arguments: `{}`},
		}})
	case 2:
		last := req.Messages[len(req.Messages)-1]
		if last.Role != llm.RoleTool || last.Content != "full result" {
			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 streamEchoTool struct{}

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

func (streamEchoTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
	return "full result", nil
}

func (streamEchoTool) ExecuteStream(ctx context.Context, raw json.RawMessage, emit func(string) error) (string, error) {
	if err := emit("one\n"); err != nil {
		return "", err
	}
	if err := emit("two\n"); err != nil {
		return "", err
	}
	return "full result", nil
}

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

	var out strings.Builder
	var logs strings.Builder
	if err := a.RunTurn(context.Background(), "stream", &out, &logs); err != nil {
		t.Fatal(err)
	}
	if out.String() != "done" {
		t.Fatalf("out = %q", out.String())
	}
	logText := logs.String()
	if strings.Count(logText, "[output]") != 2 {
		t.Fatalf("logs = %q", logText)
	}
	if !strings.Contains(logText, "one\n") || !strings.Contains(logText, "two\n") {
		t.Fatalf("logs missing streamed chunks: %q", logText)
	}
	if strings.Contains(logText, "[output]\nfull result") {
		t.Fatalf("logs contain duplicate full output: %q", logText)
	}
}

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 btwProvider struct {
	req llm.ChatRequest
	err error
}

func (p *btwProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
	p.req = req
	if p.err != nil {
		return p.err
	}
	return emit(llm.StreamEvent{Content: "side answer", Usage: &llm.Usage{TotalTokens: 999}})
}

func TestBTWUsesContextWithoutMutatingConversation(t *testing.T) {
	provider := &btwProvider{}
	a := New(Options{Provider: provider, Model: "test", Thinking: "high", ThinkingEnabled: true})
	a.SetMessages([]llm.Message{
		{Role: llm.RoleUser, Content: "main question"},
		{Role: llm.RoleAssistant, Content: "main answer"},
	})
	a.SetLastUsage(&llm.Usage{TotalTokens: 7})
	before := a.Messages()

	var out strings.Builder
	if err := a.BTW(context.Background(), "what does that mean?", &out); err != nil {
		t.Fatal(err)
	}
	if out.String() != "side answer" {
		t.Fatalf("output = %q", out.String())
	}
	if !reflect.DeepEqual(a.Messages(), before) {
		t.Fatalf("conversation changed:\nbefore: %+v\nafter:  %+v", before, a.Messages())
	}
	if usage := a.LastUsage(); usage == nil || usage.TotalTokens != 7 {
		t.Fatalf("usage changed: %+v", usage)
	}
	if len(provider.req.Tools) != 0 || provider.req.ToolChoice != "" {
		t.Fatalf("btw exposed tools: %+v / %q", provider.req.Tools, provider.req.ToolChoice)
	}
	last := provider.req.Messages[len(provider.req.Messages)-1]
	if last.Role != llm.RoleUser || last.Content != "what does that mean?" {
		t.Fatalf("last request message = %+v", last)
	}
	if provider.req.Extra["thinking"] != "high" {
		t.Fatalf("thinking extra = %+v", provider.req.Extra)
	}
}

func TestBTWRejectsEmptyQuestionAndReturnsProviderError(t *testing.T) {
	provider := &btwProvider{err: errors.New("boom")}
	a := New(Options{Provider: provider, Model: "test"})
	if err := a.BTW(context.Background(), "  ", nil); err == nil {
		t.Fatal("empty question did not fail")
	}
	if err := a.BTW(context.Background(), "question", nil); err == nil || !strings.Contains(err.Error(), "boom") {
		t.Fatalf("provider error = %v", err)
	}
}

func TestDotContinuesCurrentTaskWithoutRewritingHistory(t *testing.T) {
	provider := &btwProvider{}
	a := New(Options{Provider: provider, Model: "test"})
	if err := a.RunTurn(context.Background(), "  .  ", nil, nil); err != nil {
		t.Fatal(err)
	}

	requestMessage := provider.req.Messages[len(provider.req.Messages)-1]
	if requestMessage.Content != continuationPrompt {
		t.Fatalf("request content = %q, want %q", requestMessage.Content, continuationPrompt)
	}
	history := a.Messages()
	if history[0].Content != "." {
		t.Fatalf("stored input = %q, want literal dot", history[0].Content)
	}
}

func TestExpandContinuationMessagesOnlyExpandsSingleDot(t *testing.T) {
	messages := []llm.Message{
		{Role: llm.RoleUser, Content: "."},
		{Role: llm.RoleUser, Content: " . "},
		{Role: llm.RoleUser, Content: ".."},
		{Role: llm.RoleUser, Content: ". continue"},
		{Role: llm.RoleAssistant, Content: "."},
	}
	expanded := expandContinuationMessages(messages)
	if expanded[0].Content != continuationPrompt || expanded[1].Content != continuationPrompt {
		t.Fatalf("single dots were not expanded: %+v", expanded)
	}
	for _, index := range []int{2, 3, 4} {
		if expanded[index].Content != messages[index].Content {
			t.Fatalf("message %d changed from %q to %q", index, messages[index].Content, expanded[index].Content)
		}
	}
	if messages[0].Content != "." || messages[1].Content != " . " {
		t.Fatalf("source messages were mutated: %+v", messages)
	}
}

func TestCompleteConversationPrefixDropsPendingToolRound(t *testing.T) {
	messages := []llm.Message{
		{Role: llm.RoleUser, Content: "first"},
		{Role: llm.RoleAssistant, Content: "answer"},
		{Role: llm.RoleUser, Content: "main"},
		{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{{ID: "call_1"}, {ID: "call_2"}}},
		{Role: llm.RoleTool, ToolCallID: "call_1", Content: "one"},
	}
	got := completeConversationPrefix(messages)
	if len(got) != 3 || got[len(got)-1].Content != "main" {
		t.Fatalf("complete prefix = %+v", got)
	}
}

type concurrentBTWProvider struct {
	mainStarted chan struct{}
	releaseMain chan struct{}
	once        sync.Once
	mu          sync.Mutex
	btwRequest  llm.ChatRequest
}

func (p *concurrentBTWProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
	last := req.Messages[len(req.Messages)-1].Content
	if last == "main task" {
		p.once.Do(func() { close(p.mainStarted) })
		select {
		case <-p.releaseMain:
			return emit(llm.StreamEvent{Content: "main answer"})
		case <-ctx.Done():
			return ctx.Err()
		}
	}
	p.mu.Lock()
	p.btwRequest = req
	p.mu.Unlock()
	return emit(llm.StreamEvent{Content: "side answer"})
}

func TestBTWCanRunDuringMainTurn(t *testing.T) {
	provider := &concurrentBTWProvider{mainStarted: make(chan struct{}), releaseMain: make(chan struct{})}
	a := New(Options{Provider: provider, Model: "test"})
	mainDone := make(chan error, 1)
	go func() { mainDone <- a.RunTurn(context.Background(), "main task", nil, nil) }()
	<-provider.mainStarted

	var side strings.Builder
	if err := a.BTW(context.Background(), "side question", &side); err != nil {
		t.Fatal(err)
	}
	if side.String() != "side answer" {
		t.Fatalf("side output = %q", side.String())
	}
	provider.mu.Lock()
	req := provider.btwRequest
	provider.mu.Unlock()
	if len(req.Messages) < 2 || req.Messages[len(req.Messages)-2].Content != "main task" {
		t.Fatalf("btw request did not include active main turn: %+v", req.Messages)
	}

	close(provider.releaseMain)
	if err := <-mainDone; err != nil {
		t.Fatal(err)
	}
	if len(a.Messages()) != 2 {
		t.Fatalf("main conversation contains btw messages: %+v", a.Messages())
	}
}

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
}

type authorizationProvider struct {
	call       llm.ToolCallDelta
	calls      int
	toolOutput string
}

func (p *authorizationProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
	p.calls++
	if p.calls == 1 {
		return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{p.call}})
	}
	p.toolOutput = req.Messages[len(req.Messages)-1].Content
	return emit(llm.StreamEvent{Content: "done"})
}

type recordingAuthorizer struct {
	decision ToolAuthorizationDecision
	requests []ToolAuthorizationRequest
}

func (a *recordingAuthorizer) AuthorizeToolCall(ctx context.Context, req ToolAuthorizationRequest) (ToolAuthorizationDecision, error) {
	a.requests = append(a.requests, req)
	return a.decision, ctx.Err()
}

func TestAgentAuthorizesMutationAndCreatesRollbackSnapshot(t *testing.T) {
	dir := t.TempDir()
	provider := &authorizationProvider{call: llm.ToolCallDelta{
		Index: 0, ID: "write_1", Type: "function", Name: "file_write",
		Arguments: `{"path":"new.txt","content":"hello"}`,
	}}
	authorizer := &recordingAuthorizer{decision: ToolAuthorizationAllow}
	rollback := tools.NewRollbackStore()
	a := New(Options{
		Provider:      provider,
		Model:         "test",
		ToolRegistry:  tools.Builtins(dir),
		Authorizer:    authorizer,
		RollbackStore: rollback,
	})

	if err := a.RunTurn(context.Background(), "write it", nil, nil); err != nil {
		t.Fatal(err)
	}
	if len(authorizer.requests) != 1 || !authorizer.requests[0].RollbackAvailable {
		t.Fatalf("authorization requests = %+v", authorizer.requests)
	}
	if !strings.Contains(authorizer.requests[0].Preview, "write 5 bytes to new.txt") {
		t.Fatalf("preview = %q", authorizer.requests[0].Preview)
	}
	rollbackID := rollback.LastID()
	if rollbackID == "" || !strings.Contains(provider.toolOutput, rollbackID) {
		t.Fatalf("tool output %q missing rollback ID %q", provider.toolOutput, rollbackID)
	}
	if _, err := rollback.Restore(rollbackID); err != nil {
		t.Fatal(err)
	}
	if _, err := os.Stat(filepath.Join(dir, "new.txt")); !os.IsNotExist(err) {
		t.Fatalf("new file still exists after rollback: %v", err)
	}
}

func TestAgentDeniedMutationDoesNotExecute(t *testing.T) {
	dir := t.TempDir()
	provider := &authorizationProvider{call: llm.ToolCallDelta{
		Index: 0, ID: "write_1", Type: "function", Name: "file_write",
		Arguments: `{"path":"denied.txt","content":"no"}`,
	}}
	authorizer := &recordingAuthorizer{decision: ToolAuthorizationDeny}
	a := New(Options{Provider: provider, Model: "test", ToolRegistry: tools.Builtins(dir), Authorizer: authorizer})

	if err := a.RunTurn(context.Background(), "write it", nil, nil); err != nil {
		t.Fatal(err)
	}
	if _, err := os.Stat(filepath.Join(dir, "denied.txt")); !os.IsNotExist(err) {
		t.Fatalf("denied file exists: %v", err)
	}
	if !strings.Contains(provider.toolOutput, "denied by user") {
		t.Fatalf("tool output = %q", provider.toolOutput)
	}
}

func TestAgentDoesNotAuthorizeApplyPatchDryRun(t *testing.T) {
	dir := t.TempDir()
	provider := &authorizationProvider{call: llm.ToolCallDelta{
		Index: 0, ID: "patch_1", Type: "function", Name: "apply_patch",
		Arguments: `{"patch":"*** Begin Patch\n*** Add File: dry.txt\n+dry\n*** End Patch","dry_run":true}`,
	}}
	authorizer := &recordingAuthorizer{decision: ToolAuthorizationDeny}
	a := New(Options{Provider: provider, Model: "test", ToolRegistry: tools.Builtins(dir), Authorizer: authorizer})

	if err := a.RunTurn(context.Background(), "preview", nil, nil); err != nil {
		t.Fatal(err)
	}
	if len(authorizer.requests) != 0 {
		t.Fatalf("dry run requested authorization: %+v", authorizer.requests)
	}
	if _, err := os.Stat(filepath.Join(dir, "dry.txt")); !os.IsNotExist(err) {
		t.Fatalf("dry run created a file: %v", err)
	}
}

func TestSetMessagesRefreshesCurrentSystemInstructions(t *testing.T) {
	a := New(Options{SystemPrompt: "current instructions"})
	a.SetMessages([]llm.Message{
		{Role: llm.RoleSystem, Content: "old instructions"},
		{Role: llm.RoleUser, Content: "hello"},
	})
	messages := a.Messages()
	if messages[0].Content != "current instructions" || messages[1].Content != "hello" {
		t.Fatalf("messages = %+v", messages)
	}
}

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