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 := ¶llelProvider{ 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 := ¶llelProvider{ 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}) }