Agent-Native/runtime/loop_guardrails_test.go

package runtime

import (
	"context"
	"strings"
	"testing"

	"git-ssh.u1680832.nyat.app/weichen/Agent-Native.git/runtime/guardrails"
	"git-ssh.u1680832.nyat.app/weichen/Agent-Native.git/runtime/hooks"
	"git-ssh.u1680832.nyat.app/weichen/Agent-Native.git/runtime/llm"
	"git-ssh.u1680832.nyat.app/weichen/Agent-Native.git/runtime/tools"
)

// TestInputGuardrailDenySkipsLoop：输入护栏 deny 时不进循环、不调模型，
// 以理由直接收尾；原文已进转录（留痕）。
func TestInputGuardrailDenySkipsLoop(t *testing.T) {
	calls := 0
	bus := hooks.New()
	guardrails.Register(bus, guardrails.Set{
		Input: func(ctx context.Context, text string) guardrails.Verdict {
			return guardrails.Verdict{Decision: hooks.Deny, Reason: "含敏感指令"}
		},
	}, guardrails.Scope{Entry: true})
	res, err := Run(context.Background(), Config{
		Model: "fake",
		Hooks: bus,
		Stream: func(ctx context.Context, model string, messages []llm.Message, decls []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
			calls++
			ch := make(chan llm.Event, 1)
			ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}}
			close(ch)
			return ch
		},
	}, llm.Message{Role: llm.RoleUser, Content: "bad input"})
	if err != nil {
		t.Fatal(err)
	}
	if calls != 0 {
		t.Fatalf("model called %d times after input deny", calls)
	}
	if res.Text != "含敏感指令" {
		t.Fatalf("result text = %q", res.Text)
	}
	if !strings.Contains(transcriptText(res.Messages), "bad input") {
		t.Fatal("original input missing from transcript")
	}
}

// TestInputGuardrailModifyAppendsRewrite：modify 时改写内容追加为新的用户消息，
// 模型最终看到的是改写后的输入；原文保留在转录里。
func TestInputGuardrailModifyAppendsRewrite(t *testing.T) {
	var saw string
	bus := hooks.New()
	guardrails.Register(bus, guardrails.Set{
		Input: func(ctx context.Context, text string) guardrails.Verdict {
			return guardrails.Verdict{Decision: hooks.Modify, Replacement: "改写后的输入"}
		},
	}, guardrails.Scope{Entry: true})
	_, err := Run(context.Background(), Config{
		Model: "fake",
		Hooks: bus,
		Stream: func(ctx context.Context, model string, messages []llm.Message, decls []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
			saw = transcriptText(messages)
			ch := make(chan llm.Event, 1)
			ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}}
			close(ch)
			return ch
		},
	}, llm.Message{Role: llm.RoleUser, Content: "原始输入"})
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(saw, "原始输入") || !strings.Contains(saw, "改写后的输入") {
		t.Fatalf("request messages = %q", saw)
	}
}

// TestOutputGuardrailDenyContinuesThenCaps：输出护栏 deny 视同续跑，
// 理由进转录；达内核续跑上限后收尾。
func TestOutputGuardrailDenyContinuesThenCaps(t *testing.T) {
	calls := 0
	bus := hooks.New()
	guardrails.Register(bus, guardrails.Set{
		Output: func(ctx context.Context, text string) guardrails.Verdict {
			return guardrails.Verdict{Decision: hooks.Deny, Reason: "缺少来源标注"}
		},
	}, guardrails.Scope{Final: true})
	res, err := Run(context.Background(), Config{
		Model:    "fake",
		MaxTurns: 8,
		Hooks:    bus,
		Stream: script(
			llm.Message{Role: llm.RoleAssistant, Content: "v1", StopReason: llm.StopEnd},
			llm.Message{Role: llm.RoleAssistant, Content: "v2", StopReason: llm.StopEnd},
			llm.Message{Role: llm.RoleAssistant, Content: "v3", StopReason: llm.StopEnd},
			llm.Message{Role: llm.RoleAssistant, Content: "v4", StopReason: llm.StopEnd},
		),
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	_ = calls
	if strings.Count(transcriptText(res.Messages), "缺少来源标注") != 3 {
		t.Fatalf("reason count = %s", transcriptText(res.Messages))
	}
	if res.Text != "v4" {
		t.Fatalf("final text = %q", res.Text)
	}
}

// TestToolGuardrailDenyPairsResult：工具护栏 deny 写配对错误结果且不执行工具。
func TestToolGuardrailDenyPairsResult(t *testing.T) {
	hits := 0
	reg := tools.NewRegistry(t.TempDir())
	reg.Register(hitTool{hits: &hits})
	bus := hooks.New()
	guardrails.Register(bus, guardrails.Set{
		Tool: func(ctx context.Context, name, args string) guardrails.Verdict {
			return guardrails.Verdict{Decision: hooks.Deny, Reason: "护栏禁用"}
		},
	}, guardrails.Scope{})
	res, err := Run(context.Background(), Config{
		Model: "fake",
		Tools: reg,
		Hooks: bus,
		Stream: script(
			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
				ID: "c1", Name: "echo", Arguments: `{}`,
			}}},
			llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd},
		),
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	if hits != 0 {
		t.Fatalf("tool executed %d times", hits)
	}
	paired := false
	for _, m := range res.Messages {
		if m.Role == llm.RoleTool && m.ToolCallID == "c1" && m.IsError && m.Content == "护栏禁用" {
			paired = true
		}
	}
	if !paired {
		t.Fatalf("missing paired result: %s", transcriptText(res.Messages))
	}
}

// TestSubRunScopeSkipsInputOutput：子代理 run（Entry/Final=false）注册同名护栏也不触发。
func TestSubRunScopeSkipsInputOutput(t *testing.T) {
	bus := hooks.New()
	guardrails.Register(bus, guardrails.Set{
		Input:  func(ctx context.Context, text string) guardrails.Verdict { return guardrails.Verdict{Decision: hooks.Deny, Reason: "x"} },
		Output: func(ctx context.Context, text string) guardrails.Verdict { return guardrails.Verdict{Decision: hooks.Deny, Reason: "y"} },
	}, guardrails.Scope{})
	res, err := Run(context.Background(), Config{
		Model: "fake",
		Hooks: bus,
		Stream: script(llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}),
	}, llm.Message{Role: llm.RoleUser, Content: "go"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Text != "ok" {
		t.Fatalf("sub-scope run text = %q", res.Text)
	}
}