AutoZero/agent/agent_test.go

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