Cyrene_v2/internal/runtime/loop/flow_spans_test.go

package loop

import (
	"context"
	"encoding/json"
	"strings"
	"sync"
	"testing"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/guardrail"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/memory"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/trace"
)

// 本文件钉住三条裁决在循环侧的行为：
//   - 11.4.2：真实跑一轮后，图里真的出现 memory / guardrail 节点（各有生产者）；
//   - 11.4.3：tool 节点的参数摘要出图、密钥值绝不出图；
//   - 13.4.1：输出侧检查异步（不阻塞回合）、输入侧检查仍在模型调用前同步拦下。

// flowRecorder 收集 FlowPublisher 发布的每一帧图，按会话保留最后一份快照。
// 后台的输出侧检查会从 goroutine 里发布，因此读取必须加锁。
type flowRecorder struct {
	mu     sync.Mutex
	graphs map[string]*trace.FlowGraph
	frames int
}

func newFlowRecorder() *flowRecorder {
	return &flowRecorder{graphs: make(map[string]*trace.FlowGraph)}
}

// publisher 返回接在该记录器上的实时图发布者。
func (r *flowRecorder) publisher() *FlowPublisher {
	return NewFlowPublisher(func(_ string, payload any) {
		g, ok := payload.(*trace.FlowGraph)
		if !ok || g == nil {
			return
		}
		r.mu.Lock()
		defer r.mu.Unlock()
		r.graphs[g.SessionID] = g
		r.frames++
	}, func(string) bool { return true })
}

// last 返回某会话的最后一份图快照；一帧都没发过时判失败。
func (r *flowRecorder) last(t *testing.T, session string) *trace.FlowGraph {
	t.Helper()
	r.mu.Lock()
	defer r.mu.Unlock()
	g := r.graphs[session]
	if g == nil {
		t.Fatalf("会话 %q 一帧图都没发布（frames=%d）", session, r.frames)
	}
	return g
}

// fixedExtractor 是确定性的记忆提取器：永远返回脚本化的条目，因此一轮结束后
// 必然发生一次记忆写入（用来证明写侧有生产者）。它不实现 asyncExtractor，所以
// 提取留在同步路径上，测试无需等后台。
type fixedExtractor struct{ items []memory.MemoryItem }

func (e fixedExtractor) Extract(context.Context, []model.Message) ([]memory.MemoryItem, error) {
	return e.items, nil
}

// nodesOfKind 返回图里某一类 kind 的节点，顺序即出现顺序。
func nodesOfKind(g *trace.FlowGraph, kind string) []trace.FlowNode {
	out := make([]trace.FlowNode, 0, len(g.Nodes))
	for _, n := range g.Nodes {
		if n.Kind == kind {
			out = append(out, n)
		}
	}
	return out
}

// nodeByKind 返回图里第一个某类 kind 的节点。
func nodeByKind(t *testing.T, g *trace.FlowGraph, kind string) trace.FlowNode {
	t.Helper()
	nodes := nodesOfKind(g, kind)
	if len(nodes) == 0 {
		t.Fatalf("图中没有 %q 节点（实得 kind：%v）", kind, graphKinds(g))
	}
	return nodes[0]
}

// graphKinds 返回图里出现的 kind 列表，仅用于失败信息。
func graphKinds(g *trace.FlowGraph) []string {
	out := make([]string, 0, len(g.Nodes))
	for _, n := range g.Nodes {
		out = append(out, n.Kind)
	}
	return out
}

// waitUntil 轮询 cond 直到成立或超时（后台检查落定需要一点时间）。
func waitUntil(t *testing.T, what string, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Fatalf("等待 %s 超时", what)
}

// TestRun_RealRound_FlowGraphHasMemoryAndGuardrailNodes 真实跑一轮（流式路径，
// 带真实记忆栈、真实护栏管线与实时 tracer/flow），断言图里出现 memory 与
// guardrail 两类真节点（裁决 11.4.2），并顺带钉住 tool 节点的参数摘要
// （裁决 11.4.3）：放行的 path 出图、密钥值绝不出图。
func TestRun_RealRound_FlowGraphHasMemoryAndGuardrailNodes(t *testing.T) {
	t.Parallel()
	const secret = "sk-live-0123456789abcdef"
	rec := newFlowRecorder()

	call := `{"tool_calls":[{"name":"echo","args":{"path":"/tmp/report.txt","api_key":"` + secret + `"}}]}`
	ad := &mockAdapter{stream: []mgResp{
		{resp: model.Response{Model: "mock-1", Content: call}},
		{resp: model.Response{Model: "mock-1", Content: "报告读完了"}},
	}}
	td := dispatcherWith(t, &mockTool{name: "echo", result: "报告内容"})
	a := newTestAgent(t, ad, td)
	store := newTestMemoryStore(t)
	l := New(Options{
		Agent: a, Tools: td, MaxRounds: 8,
		Tracer:    trace.New(nil),
		Flow:      rec.publisher(),
		Guards:    guardrail.New(),
		Memory:    store,
		Retriever: memory.NewRetriever(store),
		Extractor: fixedExtractor{items: []memory.MemoryItem{{Content: "主人喜欢甜食"}}},
		PersonaID: "cyrene",
		// 流式路径是生产路径（gateway 总会装 OnStream），护栏 span
		// 也在这一侧产生，因此这条「真实跑一轮」用例走它。
		OnStream: func(model.StreamEvent) {},
	})

	res, err := l.Run(context.Background(), "web", []model.Message{model.TextMessage(model.RoleUser, "看看报告")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.Text != "报告读完了" {
		t.Fatalf("Text = %q, want 报告读完了", res.Text)
	}

	g := rec.last(t, "web")
	assertMemoryNodes(t, g)
	assertGuardrailNodes(t, rec)
	assertToolArgsSummary(t, rec.last(t, "web"), secret)
}

// assertMemoryNodes 断言读侧（召回）与写侧（提取落库）各留下一个真节点。
func assertMemoryNodes(t *testing.T, g *trace.FlowGraph) {
	t.Helper()
	memoryNodes := nodesOfKind(g, "memory")
	if len(memoryNodes) < 2 {
		t.Fatalf("memory 节点 = %d 个，want ≥2（召回 + 写入）；图中 kind：%v", len(memoryNodes), graphKinds(g))
	}
	recall := memoryNodes[0]
	if got := recall.Props["action"]; got != "recall" {
		t.Fatalf("memory 节点 action = %v，want recall（props：%+v）", got, recall.Props)
	}
	if got := recall.Props["scope"]; got != "owner" {
		t.Fatalf("memory 召回 scope = %v，want owner（props：%+v）", got, recall.Props)
	}
}

// assertGuardrailNodes 断言输入/输出两侧都留下节点且都落下裁决。输出侧检查是
// 异步的（裁决 13.4.1），因此先等它的裁决落定。
func assertGuardrailNodes(t *testing.T, rec *flowRecorder) {
	t.Helper()
	waitUntil(t, "两类 guardrail 节点都落下裁决", func() bool {
		guardNodes := nodesOfKind(rec.last(t, "web"), "guardrail")
		if len(guardNodes) < 2 {
			return false
		}
		for _, n := range guardNodes {
			if _, ok := n.Props["allow"].(bool); !ok {
				return false
			}
		}
		return true
	})
	g := rec.last(t, "web")
	actions := make(map[string]bool)
	allowed := 0
	for _, n := range nodesOfKind(g, "guardrail") {
		action, _ := n.Props["action"].(string)
		actions[action] = true
		if allow, _ := n.Props["allow"].(bool); allow {
			allowed++
		}
	}
	if !actions[spanNameGuardrailInput] || !actions[spanNameGuardrailOutput] {
		t.Fatalf("guardrail 节点缺输入/输出任一侧（实得 %v；图中 kind：%v）", actions, graphKinds(g))
	}
	if allowed < 2 {
		t.Fatalf("allow=true 的 guardrail 节点 = %d 个，want ≥2（props：%+v）", allowed, nodesOfKind(g, "guardrail"))
	}
}

// assertToolArgsSummary 钉住参数摘要的红线：放行的结构类键有值、密钥类键连条目
// 都没有，且发布的图里不含密钥值。
func assertToolArgsSummary(t *testing.T, g *trace.FlowGraph, secret string) {
	t.Helper()
	toolNode := nodeByKind(t, g, "tool")
	summary, ok := toolNode.Props["args_summary"].(map[string]any)
	if !ok {
		t.Fatalf("tool 节点缺 args_summary（props：%+v）", toolNode.Props)
	}
	if summary["path"] != "/tmp/report.txt" {
		t.Fatalf("args_summary[path] = %v，want /tmp/report.txt", summary["path"])
	}
	if _, leaked := summary["api_key"]; leaked {
		t.Fatalf("密钥类键进了摘要：%+v", summary)
	}
	keys, ok := toolNode.Props["args_keys"].([]any)
	if !ok || len(keys) != 2 {
		t.Fatalf("args_keys = %v，want 两个参数名", toolNode.Props["args_keys"])
	}
	blob, merr := json.Marshal(g)
	if merr != nil {
		t.Fatalf("marshal graph: %v", merr)
	}
	if strings.Contains(string(blob), secret) {
		t.Fatalf("发布的图泄漏了工具参数里的密钥：%s", blob)
	}
}

// gatedOutputGuard 的输出侧检查会一直等到 release 关闭（或 context 结束）才返回
// 拦截裁决，用来证明该检查真的落在关键路径之外。
type gatedOutputGuard struct {
	entered chan struct{}
	release chan struct{}
	once    sync.Once
}

func (g *gatedOutputGuard) CheckInput(context.Context, *model.Request) (*guardrail.Verdict, error) {
	return &guardrail.Verdict{Allow: true}, nil
}

func (g *gatedOutputGuard) CheckOutput(ctx context.Context, _ *model.Response) (*guardrail.Verdict, error) {
	g.once.Do(func() { close(g.entered) })
	select {
	case <-g.release:
		return &guardrail.Verdict{Allow: false, Reason: "输出含敏感内容"}, nil
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

// TestStreamStep_OutputGuardrailDoesNotBlockTurn 是裁决 13.4.1 的输出侧用例：
// 输出侧检查挂住不放时，回合仍必须返回（真并发、不阻塞）；检查放行后它的拦截
// 结果仍要闭环 —— 落进 guardrail 节点并发出审计事件（异步 ≠ 可绕过）。
func TestStreamStep_OutputGuardrailDoesNotBlockTurn(t *testing.T) {
	t.Parallel()
	rec := newFlowRecorder()
	guard := &gatedOutputGuard{entered: make(chan struct{}), release: make(chan struct{})}
	ad := &mockAdapter{stream: []mgResp{{resp: model.Response{Model: "mock-1", Content: "回答正文"}}}}
	td := dispatcherWith(t)
	a := newTestAgent(t, ad, td)

	var mu sync.Mutex
	var audits []AuditEvent
	l := New(Options{
		Agent: a, Tools: td, MaxRounds: 8,
		Tracer: trace.New(nil), Flow: rec.publisher(),
		Guards:   guardrail.New(guardrail.WithGuard(guard)),
		OnStream: func(model.StreamEvent) {},
		OnAudit: func(ev AuditEvent) {
			mu.Lock()
			audits = append(audits, ev)
			mu.Unlock()
		},
	})

	type runOut struct {
		res *Result
		err error
	}
	done := make(chan runOut, 1)
	go func() {
		res, err := l.Run(context.Background(), "web", []model.Message{model.TextMessage(model.RoleUser, "你好")})
		done <- runOut{res: res, err: err}
	}()

	select {
	case <-guard.entered:
	case <-time.After(2 * time.Second):
		t.Fatal("输出侧检查始终没有启动")
	}
	var out runOut
	select {
	case out = <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("回合被输出侧护栏阻塞了：检查未放行时 Run 必须已经返回（裁决 13.4.1 真并发）")
	}
	if out.err != nil {
		t.Fatalf("Run: %v", out.err)
	}
	if out.res.Text != "回答正文" {
		t.Fatalf("Text = %q，want 模型产出（输出侧检查不得改写已流转的回复）", out.res.Text)
	}

	close(guard.release)
	waitUntil(t, "输出侧拦截闭环", func() bool {
		mu.Lock()
		audited := len(audits) == 1 && audits[0].What == auditGuardrailOutputBlock
		mu.Unlock()
		if !audited {
			return false
		}
		g := rec.last(t, "web")
		for _, n := range nodesOfKind(g, "guardrail") {
			if n.Props["action"] == spanNameGuardrailOutput && n.Props["allow"] == false {
				return n.Status == trace.FlowStatusError
			}
		}
		return false
	})
}

// denyingInputGuard 直接拒掉输入侧检查。
type denyingInputGuard struct{}

func (denyingInputGuard) CheckInput(context.Context, *model.Request) (*guardrail.Verdict, error) {
	return &guardrail.Verdict{Allow: false, Reason: "输入含敏感内容"}, nil
}

func (denyingInputGuard) CheckOutput(context.Context, *model.Response) (*guardrail.Verdict, error) {
	return &guardrail.Verdict{Allow: true}, nil
}

// TestStreamStep_InputGuardrailBlocksBeforeModel 是裁决 13.4.1 的输入侧用例：
// 输入侧检查仍是模型调用之前的同步前置门 —— 拒绝时模型一次都不许被调用，回合
// 直接带着拦截理由返回，且图上留下一个显红的 guardrail 节点。
func TestStreamStep_InputGuardrailBlocksBeforeModel(t *testing.T) {
	t.Parallel()
	rec := newFlowRecorder()
	ad := &mockAdapter{stream: []mgResp{{resp: model.Response{Model: "mock-1", Content: "不该出现"}}}}
	td := dispatcherWith(t)
	a := newTestAgent(t, ad, td)
	l := New(Options{
		Agent: a, Tools: td, MaxRounds: 8,
		Tracer: trace.New(nil), Flow: rec.publisher(),
		Guards:   guardrail.New(guardrail.WithGuard(denyingInputGuard{})),
		OnStream: func(model.StreamEvent) {},
	})

	res, err := l.Run(context.Background(), "web", []model.Message{model.TextMessage(model.RoleUser, "危险输入")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.Text != "输入含敏感内容" || res.StopReason != agent.StopReasonGuardrailInput {
		t.Fatalf("res = %+v，want 输入侧拦截理由与 guardrail_input", res)
	}
	if ad.streamCalls != 0 {
		t.Fatalf("模型被调用了 %d 次，输入侧前置门必须在模型之前拦下", ad.streamCalls)
	}

	g := rec.last(t, "web")
	node := nodeByKind(t, g, "guardrail")
	if node.Props["action"] != spanNameGuardrailInput {
		t.Fatalf("guardrail 节点 action = %v，want %s", node.Props["action"], spanNameGuardrailInput)
	}
	if allow, _ := node.Props["allow"].(bool); allow {
		t.Fatalf("被拦下的 guardrail 节点 allow = %v，want false（props：%+v）", node.Props["allow"], node.Props)
	}
	if node.Status != trace.FlowStatusError {
		t.Fatalf("guardrail 节点状态 = %q，want error", node.Status)
	}
}