package agent import ( "context" "encoding/json" "strings" "testing" "agentu/internal/llm" "agentu/internal/tools" ) 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}) } 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()) } }