agentu/internal/agent/agent_test.go

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