AskaEth
b99268d019
- 11.4.2：三类节点首次真进图 —— memory.recall/memory.write（loop/memory.go 召回与写入点）、 guardrail.input/output（流式主路径两侧 + agent.RunStep 静音路径）、orchestration.spawn （scheduler 每次派发/收束）；新增 trace.ChildSpan 让不持有 Tracer 的叶子包也能开子 span - 11.4.3（设计先行，规格先写进 design/19 §1.1）：tool 节点参数摘要 —— 放行表 argValueKeys + trace.IsSecretKey 统一判据（密钥类键不给值）+ 单值 80 rune / ≤24 条 + 出图前再清洗； 前端 NodeDetail 新增「参数摘要」区，未放行键显示「值未展示」 - 13.4.1：流式主路径输出侧护栏改后台 goroutine（脱离 context、30s 兜底、不阻塞回合）， 输入侧仍在模型调用前同步拦下；拦截闭环到 guardrail 节点 + 审计事件 - 拆文件：scheduler.go 加 span 后 546 行 → 编排 span 助手拆出 scheduler_spans.go（回到 492） - 契约：无帧名/事件形状变化；graph_update 节点 props 新增**可选**键（tool args_keys/args_summary、 guardrail action/allow/elapsed_ms、memory action/scope/hits/…、orchestration subagent/persona）， 文档同步由主会话统一落
275 lines
10 KiB
Go
275 lines
10 KiB
Go
package orchestration
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"io"
|
||
"log/slog"
|
||
"strings"
|
||
"sync"
|
||
"testing"
|
||
|
||
"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
|
||
"git.yeij.top/AskaEth/Cyrene_v2/internal/db"
|
||
"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/loop"
|
||
"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/tool"
|
||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/trace"
|
||
)
|
||
|
||
// roundExtractor 是确定性的记忆提取器：一轮结束后必然产生一次记忆写入，
|
||
// 使「真实跑一轮」里 memory 的写侧也有节点（不实现 asyncExtractor，因此
|
||
// 提取留在同步路径上，测试不用等后台）。
|
||
type roundExtractor struct{}
|
||
|
||
func (roundExtractor) Extract(context.Context, []model.Message) ([]memory.MemoryItem, error) {
|
||
return []memory.MemoryItem{{Content: "子代理记住了这次委托"}}, nil
|
||
}
|
||
|
||
// delegateTool 模拟「派活」工具：它把前台会话当作 ParentSession 派一个后台
|
||
// 子代理（与生产里的 subagent_spawn 同一个归属口径：会话来自 ctx，绝不从
|
||
// 模型输出推断），并记下被派发的子代理 id 供测试等待。
|
||
type delegateTool struct {
|
||
sched *Scheduler
|
||
mu sync.Mutex
|
||
spawn []string
|
||
err error
|
||
}
|
||
|
||
func (d *delegateTool) Name() string { return "delegate" }
|
||
func (d *delegateTool) Description() string { return "dispatch a background sub-agent" }
|
||
func (d *delegateTool) Parameters() (map[string]any, error) { return nil, nil }
|
||
|
||
func (d *delegateTool) Execute(ctx context.Context, _ any) (any, error) {
|
||
session := loop.SessionFromContext(ctx)
|
||
if session == "" {
|
||
d.setErr(context.Canceled)
|
||
return "no session in context", nil
|
||
}
|
||
// 与生产的 subagent_spawn 同一口径：后台子代理必须活得比前台回合久，
|
||
// 因此派发用一个脱离取消但保留会话值的 context。
|
||
spawnCtx := context.WithoutCancel(ctx)
|
||
sub, err := d.sched.Spawn(spawnCtx, AgentSpec{Goal: "调研一下", ParentSession: session})
|
||
if err != nil {
|
||
d.setErr(err)
|
||
return "dispatch failed", nil
|
||
}
|
||
d.mu.Lock()
|
||
d.spawn = append(d.spawn, sub.ID)
|
||
d.mu.Unlock()
|
||
return "已派发 " + sub.ID, nil
|
||
}
|
||
|
||
func (d *delegateTool) setErr(err error) {
|
||
d.mu.Lock()
|
||
defer d.mu.Unlock()
|
||
if d.err == nil {
|
||
d.err = err
|
||
}
|
||
}
|
||
|
||
func (d *delegateTool) spawned() []string {
|
||
d.mu.Lock()
|
||
defer d.mu.Unlock()
|
||
return append([]string(nil), d.spawn...)
|
||
}
|
||
|
||
// roundDiscardLogger 返回一个不写任何地方的 logger，保持测试输出干净。
|
||
func roundDiscardLogger() *slog.Logger {
|
||
return slog.New(slog.NewTextHandler(io.Discard, nil))
|
||
}
|
||
|
||
// flowRound 是一轮「前台 + 后台子代理」的真实运行装置：它把两条真实循环
|
||
// （前台流式路径、子代理模板）、调度器与派活工具串起来。
|
||
type flowRound struct {
|
||
front *loop.Loop
|
||
flow *loop.FlowPublisher
|
||
capture *graphCapture
|
||
sched *Scheduler
|
||
delegate *delegateTool
|
||
ctx context.Context
|
||
}
|
||
|
||
// newFlowRound 组装装置：真实 SQLite 内存库上的记忆栈、真实护栏、实时
|
||
// tracer/flow，前台走流式路径（生产 chat 路径），子代理模板同样带记忆。
|
||
func newFlowRound(t *testing.T) *flowRound {
|
||
t.Helper()
|
||
capture := &graphCapture{}
|
||
flow := loop.NewFlowPublisher(capture.publish, nil)
|
||
tracer := trace.New(nil)
|
||
|
||
ctx := context.Background()
|
||
sqlDB, oerr := db.Open(ctx, ":memory:", roundDiscardLogger())
|
||
if oerr != nil {
|
||
t.Fatalf("db.Open: %v", oerr)
|
||
}
|
||
t.Cleanup(func() { _ = sqlDB.Close() })
|
||
memStore, merr := memory.NewStore(ctx, sqlDB, roundDiscardLogger())
|
||
if merr != nil {
|
||
t.Fatalf("memory.NewStore: %v", merr)
|
||
}
|
||
retriever := memory.NewRetriever(memStore)
|
||
|
||
front, frontDispatch := newFrontRoundLoop(t, flow, tracer, memStore, retriever)
|
||
subLoop := newSubRoundLoop(t, flow, tracer, memStore, retriever)
|
||
sched := newScheduler(t, Options{MaxParallel: 1, Loop: subLoop, Tracer: tracer, Flow: flow, Memory: memStore})
|
||
|
||
delegate := &delegateTool{sched: sched}
|
||
if rerr := frontDispatch.Register(delegate); rerr != nil {
|
||
t.Fatalf("Register delegate: %v", rerr)
|
||
}
|
||
return &flowRound{front: front, flow: flow, capture: capture, sched: sched, delegate: delegate, ctx: ctx}
|
||
}
|
||
|
||
// newFrontRoundLoop 构建前台循环：派活工具 + 流式路径（生产 chat 路径）+ 真实
|
||
// 记忆/护栏。它把 dispatcher 一并返回，调用方才能挂上需要调度器的派活工具。
|
||
func newFrontRoundLoop(t *testing.T, flow *loop.FlowPublisher, tracer *trace.Tracer, memStore *memory.Store, retriever *memory.Retriever) (*loop.Loop, *tool.Dispatcher) {
|
||
t.Helper()
|
||
dispatch, derr := tool.New(registry.New())
|
||
if derr != nil {
|
||
t.Fatalf("tool.New: %v", derr)
|
||
}
|
||
adapter := &toolThenFinalAdapter{toolName: "delegate", args: map[string]any{"goal": "调研"}, final: "已派发，稍后回填"}
|
||
ag, aerr := agent.New("front", dispatch, nil,
|
||
agent.WithModel(adapter, nil), agent.WithOptions(agent.Options{Model: "mock-1"}))
|
||
if aerr != nil {
|
||
t.Fatalf("agent.New: %v", aerr)
|
||
}
|
||
front := loop.New(loop.Options{
|
||
Agent: ag, Tools: dispatch, Tracer: tracer, Flow: flow, MaxRounds: 8,
|
||
Guards: guardrail.New(), Memory: memStore, Retriever: retriever,
|
||
Extractor: roundExtractor{}, PersonaID: "cyrene",
|
||
OnStream: func(model.StreamEvent) {},
|
||
})
|
||
return front, dispatch
|
||
}
|
||
|
||
// newSubRoundLoop 构建子代理模板：真实循环 + 真实记忆检索，因此泳道里也有
|
||
// memory 节点，且它的 llm/tool 步骤会折进派生它的会话的图。
|
||
func newSubRoundLoop(t *testing.T, flow *loop.FlowPublisher, tracer *trace.Tracer, memStore *memory.Store, retriever *memory.Retriever) *loop.Loop {
|
||
t.Helper()
|
||
dispatch, serr := tool.New(registry.New())
|
||
if serr != nil {
|
||
t.Fatalf("tool.New: %v", serr)
|
||
}
|
||
if rerr := dispatch.Register(&echoTool{name: "echo", result: map[string]any{"ok": true}}); rerr != nil {
|
||
t.Fatalf("Register: %v", rerr)
|
||
}
|
||
adapter := &toolThenFinalAdapter{toolName: "echo", args: map[string]any{
|
||
"path": "/tmp/sub-report.txt",
|
||
"api_key": secretArg,
|
||
}, final: "子代理完成"}
|
||
ag, aerr := agent.New("be brief", dispatch, nil,
|
||
agent.WithModel(adapter, nil), agent.WithOptions(agent.Options{Model: "mock-1"}))
|
||
if aerr != nil {
|
||
t.Fatalf("agent.New: %v", aerr)
|
||
}
|
||
return loop.New(loop.Options{
|
||
Agent: ag, Tools: dispatch, Tracer: tracer, Flow: flow, MaxRounds: 8,
|
||
Memory: memStore, Retriever: retriever, Extractor: roundExtractor{}, PersonaID: "cyrene",
|
||
})
|
||
}
|
||
|
||
// run 跑完前台一轮、等子代理收尾，并返回前台会话的最后一份图。
|
||
func (r *flowRound) run(t *testing.T) (*trace.FlowGraph, string) {
|
||
t.Helper()
|
||
r.flow.StartRound(parentSession, "帮我派个活")
|
||
res, err := r.front.Run(r.ctx, parentSession, []model.Message{model.TextMessage(model.RoleUser, "帮我派个活")})
|
||
if err != nil {
|
||
t.Fatalf("front Run: %v", err)
|
||
}
|
||
r.flow.EndRound(parentSession, res.Text)
|
||
if len(r.delegate.spawned()) != 1 {
|
||
t.Fatalf("派发次数 = %v，want 1（delegate 错误：%v）", r.delegate.spawned(), r.delegate.err)
|
||
}
|
||
subID := r.delegate.spawned()[0]
|
||
merged, werr := r.sched.Wait(r.ctx, []string{subID})
|
||
if werr != nil {
|
||
t.Fatalf("Wait: %v", werr)
|
||
}
|
||
if merged.Succeeded != 1 {
|
||
t.Fatalf("子代理未成功：%+v（summary=%s）", merged, merged.Summary)
|
||
}
|
||
return r.capture.graph(t, parentSession), subID
|
||
}
|
||
|
||
// TestRealRound_FlowGraphHasAllThreeNodeKinds 是一条端到端的「真实跑一轮」用例
|
||
// （裁决 11.4.2）：前台会话带着真实记忆栈、真实护栏与实时 tracer/flow 跑完一轮，
|
||
// 期间经工具派发一个后台子代理 —— 同一张图里必须同时出现 memory（记忆召回/写入）、
|
||
// guardrail（输入/输出护栏）与 orchestration（派发标记）三类节点，且工具参数里的
|
||
// 密钥值绝不出图（裁决 11.4.3 红线）。
|
||
func TestRealRound_FlowGraphHasAllThreeNodeKinds(t *testing.T) {
|
||
t.Parallel()
|
||
r := newFlowRound(t)
|
||
g, subID := r.run(t)
|
||
|
||
kinds := make(map[string]int)
|
||
for _, n := range g.Nodes {
|
||
kinds[n.Kind]++
|
||
}
|
||
for _, want := range []string{"memory", "guardrail", "orchestration"} {
|
||
if kinds[want] == 0 {
|
||
t.Fatalf("图中缺 %q 节点（实得 %v）", want, kinds)
|
||
}
|
||
}
|
||
assertOrchestrationNode(t, g, subID)
|
||
assertSubagentLane(t, g, subID)
|
||
assertRoundGraphRedaction(t, g)
|
||
}
|
||
|
||
// assertOrchestrationNode 断言派发标记：带被派发的子代理 id、收尾为 done。
|
||
func assertOrchestrationNode(t *testing.T, g *trace.FlowGraph, subID string) {
|
||
t.Helper()
|
||
orch := flowKindNodes(g, "orchestration")
|
||
if len(orch) != 1 {
|
||
t.Fatalf("orchestration 节点 = %d 个，want 1", len(orch))
|
||
}
|
||
if got := orch[0].Props["subagent"]; got != subID {
|
||
t.Fatalf("orchestration 节点 subagent = %v，want %s（props：%+v）", got, subID, orch[0].Props)
|
||
}
|
||
if orch[0].Status != trace.FlowStatusDone {
|
||
t.Fatalf("orchestration 节点状态 = %q，want done（props：%+v）", orch[0].Status, orch[0].Props)
|
||
}
|
||
}
|
||
|
||
// assertSubagentLane 断言泳道仍在：子代理根 + 它自己的步骤链（含记忆召回与工具调用）。
|
||
func assertSubagentLane(t *testing.T, g *trace.FlowGraph, subID string) {
|
||
t.Helper()
|
||
subNode := flowSubagentNode(t, g, subID)
|
||
lane := flowReachable(g, subNode.ID)
|
||
hasMemoryInLane, hasToolInLane := false, false
|
||
for _, n := range g.Nodes {
|
||
if !lane[n.ID] {
|
||
continue
|
||
}
|
||
if n.Kind == "memory" {
|
||
hasMemoryInLane = true
|
||
}
|
||
if n.Kind == "tool" {
|
||
hasToolInLane = true
|
||
}
|
||
}
|
||
if !hasMemoryInLane || !hasToolInLane {
|
||
t.Fatalf("子代理泳道缺 memory/tool 步骤（memory=%v tool=%v）", hasMemoryInLane, hasToolInLane)
|
||
}
|
||
}
|
||
|
||
// assertRoundGraphRedaction 断言红线：工具参数里的密钥不得出现在发布的图里，
|
||
// 而放行的结构类键仍要可读（可视化要能看参数）。
|
||
func assertRoundGraphRedaction(t *testing.T, g *trace.FlowGraph) {
|
||
t.Helper()
|
||
blob, jerr := json.Marshal(g)
|
||
if jerr != nil {
|
||
t.Fatalf("marshal graph: %v", jerr)
|
||
}
|
||
if strings.Contains(string(blob), secretArg) {
|
||
t.Fatalf("发布的图泄漏了工具参数：%s", blob)
|
||
}
|
||
if !strings.Contains(string(blob), "/tmp/sub-report.txt") {
|
||
t.Fatalf("放行键 path 的值没出图：%s", blob)
|
||
}
|
||
}
|