Agent-System/internal/agent/pinative_test.go

package agent

import (
	"context"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"chat-agent-system/internal/agentx/llm"
	"chat-agent-system/internal/config"
)

func TestPiNativeWorkspaceWritesTarget(t *testing.T) {
	dir := t.TempDir()
	b := &PiNativeBackend{cfg: &config.Config{}, stream: scriptedNative(
		llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
			ID: "1", Name: "write_file", Arguments: `{"path":"output.md","content":"# hello"}`,
		}}},
		llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd},
	)}
	res, err := b.Execute(context.Background(), Request{
		WorkDir:      dir,
		Prompt:       "write the doc",
		Instructions: "be precise",
		Targets:      []string{"output.md"},
		ToolProfile:  ToolProfileWorkspace,
	})
	if err != nil {
		t.Fatal(err)
	}
	if !res.Success || string(res.Patches["output.md"]) != "# hello" {
		t.Fatalf("%#v", res)
	}
	if _, err := os.Stat(filepath.Join(dir, "execution.log")); err != nil {
		t.Fatal(err)
	}
}

func TestPiNativeReasoningMaterializesText(t *testing.T) {
	dir := t.TempDir()
	b := &PiNativeBackend{cfg: &config.Config{}, stream: scriptedNative(
		llm.Message{Role: llm.RoleAssistant, Content: "reply body", StopReason: llm.StopEnd},
	)}
	res, err := b.Execute(context.Background(), Request{
		WorkDir:     dir,
		Prompt:      "say hello",
		Targets:     []string{"reply.md"},
		ToolProfile: ToolProfileReasoning,
	})
	if err != nil {
		t.Fatal(err)
	}
	if !res.Success || !strings.Contains(string(res.Patches["reply.md"]), "reply body") {
		t.Fatalf("%#v", res)
	}
}

func TestPiNativeRejectsUnknownProfileAndMissingKey(t *testing.T) {
	b := &PiNativeBackend{cfg: &config.Config{}}
	res, err := b.Execute(context.Background(), Request{WorkDir: t.TempDir(), ToolProfile: "shell"})
	if err != nil || res.Success || !strings.Contains(res.ErrorReason, "unknown tool profile") {
		t.Fatalf("%v %#v", err, res)
	}
	res, err = b.Execute(context.Background(), Request{WorkDir: t.TempDir(), ToolProfile: ToolProfileWorkspace})
	if err != nil || res.Success || !strings.Contains(res.ErrorReason, "api key") {
		t.Fatalf("%v %#v", err, res)
	}
}

func scriptedNative(turns ...llm.Message) llm.StreamFunc {
	i := 0
	return func(ctx context.Context, model string, messages []llm.Message, tools []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
		ch := make(chan llm.Event, 2)
		go func() {
			defer close(ch)
			msg := turns[i]
			i++
			ch <- llm.Event{Kind: llm.KindMessageEnd, Message: msg}
		}()
		return ch
	}
}