AutoZero/agent/subagent_approval_test.go

package agent

import (
	"bytes"
	"context"
	"fmt"
	"strings"
	"testing"

	"git.gopher.ink/mirrors2/AutoZero/llm"
	"git.gopher.ink/mirrors2/AutoZero/session"
	"git.gopher.ink/mirrors2/AutoZero/tools"
)

// 子 agent 的审批通道必须通向**同一个宿主**。修之前只继承了 Stdin：WebUI 那条路
// stdin 传的就是 nil，于是子 agent 撞上危险门时既没回调也没 stdin，询问直接按
// "无人可问"拒绝，弹框根本不会出现——模型只会说"没有审批人，你手动来吧"。

// subApproverAgent 造一个只有假 bash 的父 agent 并挂上 agent 工具（runs 记录 bash
// 真的跑了几次，绝不让真命令落地）。
func subApproverAgent(t *testing.T, runs *int) (*Agent, *tools.Tool) {
	t.Helper()
	t.Setenv("HOME", t.TempDir())
	reg := tools.NewRegistry()
	if err := reg.Register(&tools.Tool{
		Name: "bash", SideEffect: true, Kind: tools.KindExec,
		Run: func(context.Context, string) (tools.Result, error) {
			*runs++
			return tools.TextResult("ran"), nil
		},
	}); err != nil {
		t.Fatal(err)
	}
	sess, err := session.New()
	if err != nil {
		t.Fatal(err)
	}
	var out bytes.Buffer
	a := &Agent{
		Tools:    reg,
		Session:  sess,
		System:   "test",
		MaxDepth: 1,
		Opt:      Options{Stdout: &out, Stderr: &out},
		trusted:  map[string]bool{},
	}
	if err := a.Tools.Register(a.SubAgentTool()); err != nil {
		t.Fatal(err)
	}
	tl, ok := a.Tools.Get("agent")
	if !ok {
		t.Fatal("agent tool not registered")
	}
	return a, tl
}

// childToolProvider 驱动子 agent：第一步调 bash 跑给定命令，第二步收尾。
// 用例直接调用 agent 工具，所以父级不会消耗 provider 的 call 1。
func childToolProvider(cmd string) *scriptedProvider {
	return &scriptedProvider{fn: func(call int, _ string, _ []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		if call == 1 {
			return llm.Message{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{
				{ID: "b1", Name: "bash", Input: fmt.Sprintf(`{"command":%q}`, cmd)},
			}}, nil
		}
		return llm.Message{Role: llm.RoleAssistant, Content: "SUB-DONE"}, nil
	}}
}

const (
	sshKeyRead = `cat ~/.ssh/id_rsa`
	rmRFCmd    = `rm -rf /tmp/x` // 注意：danger_gate_test.go 里的 rmRF 是整段 JSON 入参，不是命令
)

// TestSubAgentDangerReachesParentApprover：子 agent 的不可逆动作要走到父级的审批
// 回调，人点了放行就真的执行。
func TestSubAgentDangerReachesParentApprover(t *testing.T) {
	runs := 0
	a, tl := subApproverAgent(t, &runs)
	a.Provider = childToolProvider(rmRFCmd)
	a.Mode = ModeAsk

	var asked []ApprovalRequest
	a.Opt.Approval = func(req ApprovalRequest) ApprovalVerdict {
		asked = append(asked, req)
		return ApprovalOnce
	}
	res, err := tl.Run(context.Background(), `{"prompt":"child"}`)
	if err != nil {
		t.Fatal(err)
	}
	if len(asked) != 1 {
		t.Fatalf("子 agent 的危险动作应当问一次，实际 %d 次", len(asked))
	}
	if asked[0].Tool != "bash" || !asked[0].Danger {
		t.Errorf("询问内容不对：tool=%q danger=%v", asked[0].Tool, asked[0].Danger)
	}
	if !strings.Contains(asked[0].Preview, "rm-recursive") && !strings.Contains(asked[0].Preview, "递归删除") {
		t.Errorf("审批预览里要点名危险原因：%q", asked[0].Preview)
	}
	if runs != 1 {
		t.Errorf("人已放行，bash 却没跑（runs=%d）", runs)
	}
	if !strings.Contains(res.Text, "SUB-DONE") {
		t.Errorf("子 agent 收尾回复没回来：%q", res.Text)
	}
}

// TestSubAgentDangerDenyStaysFailClosed：接上通道不等于放松门——父级拒绝时子
// agent 依然不执行，只是终于知道"是人不批"而不是"没人可问"。
func TestSubAgentDangerDenyStaysFailClosed(t *testing.T) {
	runs := 0
	a, tl := subApproverAgent(t, &runs)
	a.Provider = childToolProvider(rmRFCmd)
	a.Mode = ModeAsk

	asked := 0
	a.Opt.Approval = func(ApprovalRequest) ApprovalVerdict {
		asked++
		return ApprovalDeny
	}
	if _, err := tl.Run(context.Background(), `{"prompt":"child"}`); err != nil {
		t.Fatal(err)
	}
	if asked != 1 {
		t.Fatalf("应当问一次，实际 %d 次", asked)
	}
	if runs != 0 {
		t.Errorf("人被拒了还执行（runs=%d）", runs)
	}
}

// TestSubAgentSecretReadPiercesBypass：bypass 收回的是"别问我"，不是"凭据可以念"。
// 子 agent 里读私钥仍要过审批——这条以前根本问不到人，所以从来没被验证过。
func TestSubAgentSecretReadPiercesBypass(t *testing.T) {
	runs := 0
	a, tl := subApproverAgent(t, &runs)
	a.Provider = childToolProvider(sshKeyRead)
	a.Mode = ModeBypass
	a.Opt.ApproveAll = true // 连 --yes 一起叠上：这两样都不该让凭据读取静默通过

	asked := 0
	var deny bool
	a.Opt.Approval = func(ApprovalRequest) ApprovalVerdict {
		asked++
		deny = true
		return ApprovalDeny
	}
	if _, err := tl.Run(context.Background(), `{"prompt":"child"}`); err != nil {
		t.Fatal(err)
	}
	if !deny {
		t.Fatal("子 agent 没有把凭据读取交给父级审批回调")
	}
	if asked != 1 {
		t.Fatalf("应当只问一次，实际 %d 次", asked)
	}
	if runs != 0 {
		t.Errorf("拒绝后仍执行了（runs=%d）", runs)
	}
}

// TestSubAgentNoApproverStillFailsClosed：反向对照——父级确实没人可问（-p / cron /
// 渠道：回调与 stdin 都是 nil）时，子 agent 保持原有语义，直接拒，不静默放行。
func TestSubAgentNoApproverStillFailsClosed(t *testing.T) {
	runs := 0
	a, tl := subApproverAgent(t, &runs)
	a.Provider = childToolProvider(rmRFCmd)
	a.Mode = ModeAsk
	a.Opt.Approval = nil
	a.Opt.Stdin = nil
	if _, err := tl.Run(context.Background(), `{"prompt":"child"}`); err != nil {
		t.Fatal(err)
	}
	if runs != 0 {
		t.Errorf("无人可问时必须拒绝，却执行了（runs=%d）", runs)
	}
}