AutoZero/tools/builtin_test.go

package tools

import (
	"context"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"
	"unicode/utf8"
)

func lookupTool(t *testing.T, name string) *Tool {
	t.Helper()
	reg := NewRegistry()
	for _, tl := range Builtin() {
		if err := reg.Register(tl); err != nil {
			t.Fatalf("register %s: %v", tl.Name, err)
		}
	}
	tl, ok := reg.Get(name)
	if !ok {
		t.Fatalf("tool %q not registered", name)
	}
	return tl
}

// TestBashTimeoutIsEnforced：bash 的超时必须真的终止命令。
//
// 命令里 `sleep 60 &` 会留下一个继承 stdout 管道写端的子进程。修复前
// exec.CommandContext 只杀直接子进程，cmd.Run() 要等管道写端全部关闭才返回
// ——于是"120 秒超时"实测挂了 23 分钟（session 20261008-214717），整个回合静止。
func TestBashTimeoutIsEnforced(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Skip("依赖 POSIX 进程组语义")
	}
	old := bashTimeout
	bashTimeout = 400 * time.Millisecond
	defer func() { bashTimeout = old }()

	tl := lookupTool(t, "bash")
	start := time.Now()
	res, err := tl.Run(context.Background(), `{"command":"sleep 60 & sleep 60"}`)
	elapsed := time.Since(start)
	if err != nil {
		t.Fatalf("工具不该返回 error（失败走 Result.Err）：%v", err)
	}
	if res.Err == nil {
		t.Errorf("超时必须成为显式失败信号，text=%q", res.Text)
	}
	// WaitGrace 给 2 秒余量；再宽就是抓"完全没生效"的量级了。
	if elapsed > 10*time.Second {
		t.Errorf("超时没有生效：预算 %v，实际 %v", old, elapsed)
	}
}

// TestBashCancelStopsTool：调用方取消（用户按停止）后不该再留下长时间阻塞。
func TestBashCancelStopsTool(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Skip("依赖 POSIX 进程组语义")
	}
	tl := lookupTool(t, "bash")
	ctx, cancel := context.WithCancel(context.Background())
	go func() {
		time.Sleep(300 * time.Millisecond)
		cancel()
	}()
	start := time.Now()
	res, _ := tl.Run(ctx, `{"command":"sleep 60 & sleep 60"}`)
	if elapsed := time.Since(start); elapsed > 10*time.Second {
		t.Errorf("取消后仍阻塞 %v（res.text=%q）", elapsed, res.Text)
	}
}

// TestBashOutputTruncationKeepsUTF8：截断不能把一个汉字劈开。
// 中文一字 3 字节，16000 字节的硬切几乎必然落在字符中间，尾部非法 UTF-8
// 进 JSON 会变替换字符，也可能被模型网关直接拒掉。
func TestBashOutputTruncationKeepsUTF8(t *testing.T) {
	old := bashMaxOutput
	bashMaxOutput = 100
	defer func() { bashMaxOutput = old }()

	tl := lookupTool(t, "bash")
	res, err := tl.Run(context.Background(),
		`{"command":"awk 'BEGIN{for(i=0;i<2000;i++)printf \"中文测试\"}'"}`)
	if err != nil {
		t.Fatalf("run: %v", err)
	}
	if res.Err != nil {
		t.Fatalf("命令失败：%s %v", res.Text, res.Err)
	}
	const marker = "\n...[truncated]"
	body := strings.TrimSuffix(res.Text, marker)
	if body == res.Text {
		t.Fatalf("没有截断标记，len=%d text=%q", len(res.Text), res.Text)
	}
	if !utf8.ValidString(body) {
		tail := len(body) - 9
		if tail < 0 {
			tail = 0
		}
		t.Errorf("截断处切坏了 UTF-8，尾部=%q", body[tail:])
	}
	if len(body) > bashMaxOutput {
		t.Errorf("截断没生效：%d > %d", len(body), bashMaxOutput)
	}
}

// TestBashRunsWithinWorkspace：workDir 语义不能被换掉时丢掉。
func TestBashRunsWithinWorkspace(t *testing.T) {
	dir := t.TempDir()
	reg := NewRegistry()
	for _, tl := range BuiltinIn(dir, nil, nil) {
		if err := reg.Register(tl); err != nil {
			t.Fatal(err)
		}
	}
	tl, ok := reg.Get("bash")
	if !ok {
		t.Fatal("bash not registered")
	}
	res, err := tl.Run(context.Background(), `{"command":"pwd"}`)
	if err != nil || res.Err != nil {
		t.Fatalf("run: err=%v res=%+v", err, res)
	}
	// macOS 的临时目录带 /private 前缀，两侧都解引用再比。
	got, want := eval(res.Text), dir
	if want2 := eval(dir); want2 != "" {
		want = want2
	}
	if got != want {
		t.Errorf("没有在工作区里执行：got=%q want=%q", got, want)
	}
}

func eval(s string) string {
	p, err := filepath.EvalSymlinks(strings.TrimSpace(s))
	if err != nil {
		return ""
	}
	return p
}

// recordingSnap 收集留档钩子收到的改动。
type recordingSnap struct{ got []Change }

func (r *recordingSnap) Record(c Change) { r.got = append(r.got, c) }

// TestWriteEditFeedTheSnapshotHook：write/edit 必须在**落盘之后**把
// 路径/改前内容/是否新建/改后内容交给留档层。少一份或报错都算回归。
func TestWriteEditFeedTheSnapshotHook(t *testing.T) {
	dir := t.TempDir()
	snap := &recordingSnap{}
	reg := NewRegistry()
	for _, tl := range BuiltinIn(dir, NewFence(), snap) {
		if err := reg.Register(tl); err != nil {
			t.Fatal(err)
		}
	}
	w, _ := reg.Get("write")
	e, _ := reg.Get("edit")
	p := filepath.Join(dir, "s.txt")

	// 新建：Existed=false，Before 为空
	if _, err := w.Run(context.Background(),
		`{"path":"`+p+`","content":"第一版"}`); err != nil {
		t.Fatal(err)
	}
	if len(snap.got) != 1 || snap.got[0].Existed || string(snap.got[0].After) != "第一版" {
		t.Fatalf("新建的留档不对：%+v", snap.got)
	}
	// 改动：Existed=true，Before 是旧内容
	if _, err := w.Run(context.Background(),
		`{"path":"`+p+`","content":"第二版"}`); err != nil {
		t.Fatal(err)
	}
	g := snap.got[1]
	if !g.Existed || string(g.Before) != "第一版" || string(g.After) != "第二版" {
		t.Errorf("覆盖写的留档不对：%+v", g)
	}
	// edit：同样要交一份，且 Before 是编辑前全文
	if _, err := e.Run(context.Background(),
		`{"path":"`+p+`","old_string":"第二版","new_string":"第三版"}`); err != nil {
		t.Fatal(err)
	}
	g = snap.got[2]
	if !g.Existed || string(g.Before) != "第二版" || string(g.After) != "第三版" {
		t.Errorf("edit 的留档不对：%+v", g)
	}
	// 失败的写不该留档（否则索引里会出现现实中没发生的改动）
	if _, err := e.Run(context.Background(),
		`{"path":"`+p+`","old_string":"不存在","new_string":"x"}`); err == nil {
		t.Fatal("edit 应报错")
	}
	if len(snap.got) != 3 {
		t.Errorf("失败的改动也留档了：%d 条", len(snap.got))
	}
}

// TestSnapshotHookDisabledIsNoop：不传留档层时行为与从前完全一致。
func TestSnapshotHookDisabledIsNoop(t *testing.T) {
	reg := NewRegistry()
	for _, tl := range BuiltinIn(t.TempDir(), NewFence(), nil) {
		if err := reg.Register(tl); err != nil {
			t.Fatal(err)
		}
	}
	w, _ := reg.Get("write")
	res, err := w.Run(context.Background(), `{"path":"x.txt","content":"hi"}`)
	if err != nil || res.Err != nil {
		t.Fatalf("write: %v %+v", err, res)
	}
}