package agent import ( "bytes" "context" "net/http" "net/http/httptest" "os" "path/filepath" "strings" "testing" "git.gopher.ink/mirrors2/AutoZero/llm" "git.gopher.ink/mirrors2/AutoZero/session" "git.gopher.ink/mirrors2/AutoZero/tools" ) // mockOpenAI 返回两段 SSE:第一轮发起 read 工具调用,第二轮给最终回复。 func mockOpenAI(t *testing.T, workspace string) *httptest.Server { t.Helper() turn := 0 return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") turn++ if turn == 1 { w.Write([]byte("data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"c1\",\"function\":{\"name\":\"read\",\"arguments\":\"{\\\"path\\\":\\\"" + workspace + "/note.txt\\\"}\"}}]}}]}\n\n")) w.Write([]byte("data: {\"choices\":[{\"finish_reason\":\"tool_calls\"}]}\n\n")) w.Write([]byte("data: [DONE]\n\n")) return } // 校验工具结果确实被回填 body := readAll(t, r) if !strings.Contains(body, "tool_call_id") || !strings.Contains(body, "hello from file") { t.Errorf("tool result not fed back: %s", body) } w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"done\"}}]}\n\n")) w.Write([]byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":2}}\n\n")) w.Write([]byte("data: [DONE]\n\n")) })) } func readAll(t *testing.T, r *http.Request) string { t.Helper() var buf bytes.Buffer buf.ReadFrom(r.Body) return buf.String() } func TestAgentLoopWithToolCall(t *testing.T) { ws := t.TempDir() os.WriteFile(filepath.Join(ws, "note.txt"), []byte("hello from file"), 0o644) srv := mockOpenAI(t, ws) defer srv.Close() dir := t.TempDir() t.Setenv("HOME", dir) // 会话文件写进临时 HOME reg := tools.NewRegistry() for _, tl := range tools.Builtin() { if err := reg.Register(tl); err != nil { t.Fatal(err) } } sess, err := session.New() if err != nil { t.Fatal(err) } var out bytes.Buffer a := &Agent{ Provider: &llm.OpenAIProvider{BaseURL: srv.URL, APIKey: "test", Model: "mock"}, Tools: reg, Session: sess, System: "test system", Opt: Options{Stdout: &out, Stderr: &out, ApproveAll: true}, } if err := a.Run(context.Background(), "read note.txt and tell me"); err != nil { t.Fatal(err) } if got := out.String(); !strings.Contains(got, "done") { t.Fatalf("unexpected output: %q", got) } // 会话应包含 user / assistant(toolcall) / tool / assistant 四条 if len(sess.Messages) != 4 { t.Fatalf("want 4 messages, got %d", len(sess.Messages)) } // 重新加载验证 JSONL 持久化 reloaded, err := session.Load(sess.Path) if err != nil { t.Fatal(err) } if len(reloaded.Messages) != 4 { t.Fatalf("reload: want 4 messages, got %d", len(reloaded.Messages)) } } // 审批门控:deny/once/always 三种裁决与信任缓存、policy=always 的全工具询问。 func TestExecOneApprovalVerdicts(t *testing.T) { newAgent := func(policy string, verdicts []ApprovalVerdict, calls *int) *Agent { reg := tools.NewRegistry() reg.Register(&tools.Tool{ Name: "writer", Description: "side effect", SideEffect: true, Run: func(ctx context.Context, input string) (tools.Result, error) { *calls++ return tools.TextResult("ran"), nil }, }) reg.Register(&tools.Tool{ Name: "reader", Description: "read only", ReadOnly: true, ConcurrencySafe: true, Run: func(ctx context.Context, input string) (tools.Result, error) { *calls++ return tools.TextResult("read"), nil }, }) idx := 0 a := &Agent{Tools: reg, ApprovalPolicy: policy} a.Opt.Approval = func(ApprovalRequest) ApprovalVerdict { v := verdicts[min(idx, len(verdicts)-1)] idx++ return v } return a } t.Run("deny", func(t *testing.T) { calls := 0 a := newAgent("dangerous-only", []ApprovalVerdict{ApprovalDeny}, &calls) if _, err := a.execOne(context.Background(), llm.ToolCall{Name: "writer", Input: "{}"}); err == nil { t.Fatal("denied tool should error") } if calls != 0 { t.Error("tool should not run") } }) t.Run("once asks again", func(t *testing.T) { calls := 0 a := newAgent("dangerous-only", []ApprovalVerdict{ApprovalOnce, ApprovalOnce}, &calls) for i := 0; i < 2; i++ { if _, err := a.execOne(context.Background(), llm.ToolCall{Name: "writer", Input: "{}"}); err != nil { t.Fatal(err) } } if calls != 2 { t.Errorf("calls = %d, want 2", calls) } }) t.Run("always caches trust", func(t *testing.T) { calls := 0 a := newAgent("dangerous-only", []ApprovalVerdict{ApprovalAlways}, &calls) for i := 0; i < 2; i++ { if _, err := a.execOne(context.Background(), llm.ToolCall{Name: "writer", Input: "{}"}); err != nil { t.Fatal(err) } } if calls != 2 { t.Errorf("calls = %d, want 2", calls) } if !a.trusted["writer"] { t.Error("writer should be trusted after ApprovalAlways") } }) t.Run("policy always asks read-only tools", func(t *testing.T) { calls := 0 a := newAgent("always", []ApprovalVerdict{ApprovalOnce}, &calls) if _, err := a.execOne(context.Background(), llm.ToolCall{Name: "reader", Input: "{}"}); err != nil { t.Fatal(err) } if calls != 1 { t.Error("reader should run after approval") } }) t.Run("ApproveAll skips approval", func(t *testing.T) { calls := 0 a := newAgent("dangerous-only", nil, &calls) a.Opt.ApproveAll = true if _, err := a.execOne(context.Background(), llm.ToolCall{Name: "writer", Input: "{}"}); err != nil { t.Fatal(err) } if calls != 1 { t.Error("writer should run under ApproveAll") } }) }