feat(orchestration,loop,flow): 接线子代理编排 — 派发入口 + 分支归父会话 + 节点日志(含脱敏)
修的是「文档说 ✅、生产不可达」:此前 Scheduler.Spawn 只有测试调用方、无 bus 入口、无派发工具;Span.Record 无生产调用方 → 泳道画不出、节点日志恒空。
- 入口:subagent_spawn 工具(internal/runtime/orchestration/spawntool.go,同包便于用真 Scheduler/loop/Tracer/Flow 做集成测试;注册在 cmd/cyrene/subagent_tool.go 照 terminal_tool 范式)。默认 owner-only,不设 AllowInGroup/AllowInStranger;安全仍由 restrictSchemas + applyToolSecurityGate 双层保证,未绕代审
- 父会话归属:AgentSpec.ParentSession(请求语义)→ Scheduler 转运行期 ctx 标记 loop.WithFlowBranch → 分支 span 以父 session 发布;**allow 谓词一行未改**;单图语义靠「分支 Run 不调 StartRound/EndRound」+ FlowPublisher 不新建图;未归属 spawn 一张图都不发(有测试锁)
- 前台最多等 2s(对齐 task.autoBackgroundThreshold)→ 合并结果作工具结果;转后台则广播 bus subagent.result(design/05 §9 回填 hook);新增 Scheduler.Wait(ids) 按批等待(Join 会串并发批次,故不重用)
- 子代理跑在 WithoutCancel(ctx):保留 span 归属、不被工具超时/回合结束顺手杀掉;生命周期由 PerAgentTimeout + CancelAll 约束;内部 WithMutedBroadcast 静音(不进前台 stream)
- 节点日志:loop/spanlog.go 新增 llm.done / tool.done / subagent.done;工具参数**只记键名不记值**;仅步骤真正结束才追加(跳过/被拒路径不发,不伪造)
- 脱敏(红线):新增 trace/redact.go(形状规则 sk-/Bearer/JWT/key:value + 键名规则 api_key/token/authorization/password…,递归嵌套、不改调用方事件),graph.go 发布前必经;**顺带修掉一个既有泄漏**:result_summary prop 原先未脱敏(工具结果里的密钥会进图),由图级变异测试抓出
- 测试:新增 22 个(工具级 / 图级 / 日志级),含 4 次变异验证(去掉 WithFlowBranch、去掉 record*、还原未脱敏、去掉静音 各自 FAIL);-race 过 orchestration+trace+loop
- 已知限制:真实模型端到端未验(无凭据/不起服务);工具白名单是 design gap 未做;subagent.result 暂无消费者(design/15 未实现);后台子代理跨轮次边界迟到 span 会丢泳道根;PerAgentTimeout=10min 待确认
This commit is contained in:
1 parent
3f79159fbd
commit
de1fa83b79
25 files changed
+2173
-97
No files matched your search
@@ -0,0 +1,43 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/orchestration"
|
||||
)
|
||||
|
||||
// defaultSubagentTimeout bounds one sub-agent's whole run. DevDocs/design/05-子代理.md
|
||||
// §7 requires a timeout and an anti-runaway bound; the scheduler already
|
||||
// implements both (PerAgentTimeout, Abort/CancelAll), so production must supply a
|
||||
// value or the per-agent timeout stays inert (the scheduler's zero means "no
|
||||
// timeout"). The number is deliberately generous — a sub-agent runs up to 8 tool
|
||||
// rounds like any loop — and is a single knob to revisit when design/05 fixes a
|
||||
// concrete figure.
|
||||
const defaultSubagentTimeout = 10 * time.Minute
|
||||
|
||||
// wireSubagentTool registers the sub-agent dispatch tool with the dispatcher, so
|
||||
// the agent can finally reach the orchestration layer at runtime (before this,
|
||||
// (*orchestration.Scheduler).Spawn had no production caller and design/05 was
|
||||
// library-only).
|
||||
//
|
||||
// The capability is registered with its default reachability markers (both
|
||||
// false), which keeps it owner-only: the design/11 §3 policy engine hides the
|
||||
// schema from a group/stranger session and the loop's security gate rejects the
|
||||
// call there. Design/05 §6/§7 give no group/stranger allowance for spawning
|
||||
// background agents, and being able to spawn one is a heavy operation, so the
|
||||
// safe default is chosen deliberately rather than enabled by accident.
|
||||
func (a *app) wireSubagentTool() {
|
||||
log := logger.WithSource(a.log, "cmd/cyrene")
|
||||
if a.scheduler == nil {
|
||||
log.Warn("sub-agent tool skipped: no scheduler")
|
||||
return
|
||||
}
|
||||
spawner := orchestration.NewSpawnTool(a.scheduler, orchestration.WithSpawnLogger(a.log))
|
||||
if err := a.dispatcher.Register(spawner); err != nil {
|
||||
log.Warn("sub-agent tool register skipped", "error", err)
|
||||
return
|
||||
}
|
||||
log.Info("sub-agent tool registered", "tool", orchestration.SpawnToolName,
|
||||
"owner_only", true)
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/restrict"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/orchestration"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/tool"
|
||||
)
|
||||
|
||||
// TestWireSubagentToolRegistersOwnerOnly locks the composition-root wiring: the
|
||||
// dispatch tool must actually land in the capability registry (otherwise
|
||||
// orchestration stays library-only), and it must carry the default reachability
|
||||
// markers so a group/stranger session can neither see nor call it.
|
||||
func TestWireSubagentToolRegistersOwnerOnly(t *testing.T) {
|
||||
reg := registry.New()
|
||||
d, err := tool.New(reg)
|
||||
if err != nil {
|
||||
t.Fatalf("tool.New: %v", err)
|
||||
}
|
||||
a := &app{log: testLog(), dispatcher: d, scheduler: orchestration.New(orchestration.Options{})}
|
||||
|
||||
a.wireSubagentTool()
|
||||
|
||||
if _, gerr := reg.Get(orchestration.SpawnToolName); gerr != nil {
|
||||
t.Fatalf("capability %q is not registered: %v", orchestration.SpawnToolName, gerr)
|
||||
}
|
||||
names := make([]string, 0)
|
||||
for _, s := range d.ToolSchemas() {
|
||||
names = append(names, s.Name)
|
||||
}
|
||||
found := false
|
||||
for _, n := range names {
|
||||
if n == orchestration.SpawnToolName {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatalf("tool schemas = %v, want the spawn tool advertised", names)
|
||||
}
|
||||
if !restrict.SetContext(reg, restrict.ActorOwner).Allow(orchestration.SpawnToolName) {
|
||||
t.Fatal("owner was denied the spawn tool")
|
||||
}
|
||||
if restrict.SetContext(reg, restrict.ActorGroupMember).Allow(orchestration.SpawnToolName) {
|
||||
t.Fatal("group member may invoke the spawn tool; want owner-only")
|
||||
}
|
||||
if restrict.SetContext(reg, restrict.ActorStranger).Allow(orchestration.SpawnToolName) {
|
||||
t.Fatal("stranger may invoke the spawn tool; want owner-only")
|
||||
}
|
||||
// The registered capability must be the dispatch tool itself, not a placeholder.
|
||||
cap, _ := reg.Get(orchestration.SpawnToolName)
|
||||
sp, ok := cap.Value.(*orchestration.SpawnTool)
|
||||
if !ok {
|
||||
t.Fatalf("capability value = %T, want *orchestration.SpawnTool", cap.Value)
|
||||
}
|
||||
if _, err := sp.Execute(context.Background(), map[string]any{}); err == nil {
|
||||
t.Fatal("the registered tool accepted a goal-less call")
|
||||
}
|
||||
}
|
||||
|
||||
// TestWireSubagentToolNilScheduler asserts a degraded boot (no scheduler) logs and
|
||||
// skips the registration instead of panicking.
|
||||
func TestWireSubagentToolNilScheduler(t *testing.T) {
|
||||
reg := registry.New()
|
||||
d, err := tool.New(reg)
|
||||
if err != nil {
|
||||
t.Fatalf("tool.New: %v", err)
|
||||
}
|
||||
a := &app{log: testLog(), dispatcher: d}
|
||||
|
||||
a.wireSubagentTool()
|
||||
|
||||
if _, gerr := reg.Get(orchestration.SpawnToolName); gerr == nil {
|
||||
t.Fatal("the tool was registered without a scheduler")
|
||||
}
|
||||
}
|
||||
+14
-8
@@ -435,14 +435,19 @@ func (a *app) wireRuntime() error {
|
||||
})
|
||||
|
||||
a.scheduler = orchestration.New(orchestration.Options{
|
||||
MaxParallel: defaultMaxParallel,
|
||||
Loop: a.loop,
|
||||
Bus: a.bus,
|
||||
Memory: store,
|
||||
Tracer: tracer,
|
||||
Approver: audit.NewApprover(a.log),
|
||||
Logger: logger.WithSource(a.log, "orchestration"),
|
||||
AgentID: defaultAgentID,
|
||||
MaxParallel: defaultMaxParallel,
|
||||
PerAgentTimeout: defaultSubagentTimeout,
|
||||
Loop: a.loop,
|
||||
Bus: a.bus,
|
||||
Memory: store,
|
||||
Tracer: tracer,
|
||||
// The sub-agent's span and its inner llm/tool steps fold into the spawning
|
||||
// session's live graph, so a branch shows up as a swimlane of that session's
|
||||
// flow instead of opening a second graph (design/19 §2).
|
||||
Flow: flowPub,
|
||||
Approver: audit.NewApprover(a.log),
|
||||
Logger: logger.WithSource(a.log, "orchestration"),
|
||||
AgentID: defaultAgentID,
|
||||
})
|
||||
return nil
|
||||
}
|
||||
@@ -575,6 +580,7 @@ func (a *app) wireAuxiliary() error {
|
||||
a.wireMCP()
|
||||
a.wireSkills()
|
||||
a.wireTerminalTool()
|
||||
a.wireSubagentTool()
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
@@ -23,6 +23,11 @@ type allowedSessions func(sessionID string) bool
|
||||
// a realtime input→think→tool→output graph (design/19 §1). It is nil-tolerant
|
||||
// (a nil publisher is a no-op) and safe for concurrent use; session graphs are
|
||||
// serialized under one mutex because parallel tool spans finish concurrently.
|
||||
//
|
||||
// Ownership: one session owns one live round. A Run that declares itself a branch
|
||||
// of another session (loop.WithFlowBranch) never calls StartRound/EndRound and
|
||||
// publishes its spans under the owning session instead, so a sub-agent renders as
|
||||
// a swimlane of the parent's graph rather than as a second graph (design/19 §2).
|
||||
type FlowPublisher struct {
|
||||
mu sync.Mutex
|
||||
publish func(subject string, payload any)
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
package loop
|
||||
|
||||
import "context"
|
||||
|
||||
// flowBranchKey is the context key marking a Run as a branch of another
|
||||
// session's flow graph. It is set by the orchestration layer (through
|
||||
// WithFlowBranch) before it drives a sub-agent's Run.
|
||||
type flowBranchKey struct{}
|
||||
|
||||
// WithFlowBranch returns ctx marking the current Run as a branch: its spans fold
|
||||
// into parentSession's live flow graph and the Run neither starts nor ends a
|
||||
// round of its own — the owning session's round does. It is how the orchestration
|
||||
// layer attributes a sub-agent's work to the session that spawned it, so parallel
|
||||
// sub-agents render as swimlanes of one graph instead of replacing the frontend's
|
||||
// single flowGraph with a graph of their own (design/19 §2).
|
||||
//
|
||||
// An empty parentSession still marks the Run as a branch: it owns no round and
|
||||
// publishes nothing. That is the state of a background sub-agent nobody attributed
|
||||
// to a session, and it must not silently become a graph of its own. A nil ctx is
|
||||
// returned unchanged.
|
||||
func WithFlowBranch(ctx context.Context, parentSession string) context.Context {
|
||||
if ctx == nil {
|
||||
return ctx
|
||||
}
|
||||
return context.WithValue(ctx, flowBranchKey{}, parentSession)
|
||||
}
|
||||
|
||||
// flowBranch reports the parent session recorded by WithFlowBranch and whether the
|
||||
// Run is a branch at all. A non-branch Run owns its own round and publishes under
|
||||
// its own session id.
|
||||
func flowBranch(ctx context.Context) (string, bool) {
|
||||
if ctx == nil {
|
||||
return "", false
|
||||
}
|
||||
v := ctx.Value(flowBranchKey{})
|
||||
if v == nil {
|
||||
return "", false
|
||||
}
|
||||
s, _ := v.(string)
|
||||
return s, true
|
||||
}
|
||||
|
||||
// flowOwner returns the session whose flow graph should receive this Run's spans:
|
||||
// the branch parent when the Run is a sub-agent branch (empty when the branch is
|
||||
// unattributed, which suppresses publication), otherwise the Run's own sessionID.
|
||||
func flowOwner(ctx context.Context, sessionID string) string {
|
||||
if parent, isBranch := flowBranch(ctx); isBranch {
|
||||
return parent
|
||||
}
|
||||
return sessionID
|
||||
}
|
||||
|
||||
// FlowOwnerFromContext returns the session whose flow graph owns work started
|
||||
// from ctx: the branch parent when ctx runs inside a sub-agent branch, otherwise
|
||||
// the session id the enclosing Run was driven with (empty when neither is set).
|
||||
//
|
||||
// A tool that dispatches background work — the sub-agent spawn tool — passes this
|
||||
// as the parent session, so a sub-agent spawned from inside another sub-agent
|
||||
// stays attributed to the same top-level graph instead of to an id that owns no
|
||||
// graph.
|
||||
func FlowOwnerFromContext(ctx context.Context) string {
|
||||
if parent, isBranch := flowBranch(ctx); isBranch {
|
||||
return parent
|
||||
}
|
||||
return SessionFromContext(ctx)
|
||||
}
|
||||
@@ -55,7 +55,10 @@ func (l *Loop) Run(ctx context.Context, sessionID string, messages []model.Messa
|
||||
msgs = l.resolveImages(ctx, msgs)
|
||||
toolState := &toolLoopState{prevFP: make(map[string]bool)}
|
||||
res := &Result{}
|
||||
if l.flow != nil {
|
||||
// A branch Run (a sub-agent) folds its spans into the parent session's live
|
||||
// graph; only the owning Run starts and ends a round. Without this a sub-agent
|
||||
// would reset the frontend's single flowGraph and gap the parent's own nodes.
|
||||
if _, isBranch := flowBranch(ctx); l.flow != nil && !isBranch {
|
||||
l.flow.StartRound(sessionID, lastUserText(msgs))
|
||||
defer func() { l.flow.EndRound(sessionID, res.Text) }()
|
||||
}
|
||||
|
||||
@@ -20,6 +20,12 @@ func loopSession(ctx context.Context) string {
|
||||
return s
|
||||
}
|
||||
|
||||
// SessionFromContext returns the session id the enclosing Run is driving, or ""
|
||||
// when none was attached. It is exported so a tool executing inside the loop can
|
||||
// attribute the work it starts (a spawned sub-agent, a background task) to the
|
||||
// conversation that launched it — never inferred from model output.
|
||||
func SessionFromContext(ctx context.Context) string { return loopSession(ctx) }
|
||||
|
||||
// loopTurnIDKey is the context key for the assistant turn the current Run is
|
||||
// producing. The gateway attaches it (WithTurnID) before starting a turn so the
|
||||
// streaming events the loop emits can be correlated to that turn by the client.
|
||||
|
||||
@@ -0,0 +1,118 @@
|
||||
package loop
|
||||
|
||||
import (
|
||||
"sort"
|
||||
"time"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/trace"
|
||||
)
|
||||
|
||||
// Log event types recorded on a span so the flow graph's node log[] shows what
|
||||
// happened at each step (design/19 §1 节点实时日志). An event is appended only when
|
||||
// the step really finished, so the log never announces work that did not run.
|
||||
const (
|
||||
// EventLLMDone is recorded on an llm_call span once the model call returned.
|
||||
EventLLMDone = "llm.done"
|
||||
// EventToolDone is recorded on a tool_call span once the tool call settled.
|
||||
EventToolDone = "tool.done"
|
||||
)
|
||||
|
||||
// maxLogSummaryRunes bounds a log value before publication. The graph redacts and
|
||||
// truncates again on the way out; this only keeps an individual span small.
|
||||
const maxLogSummaryRunes = 240
|
||||
|
||||
// recordLLMDone appends a desensitized summary of a finished model call to its
|
||||
// span: the configured model, both stop reasons, how many tool calls were
|
||||
// requested and how long the call took. No prompt or completion text is recorded,
|
||||
// so an llm node's log can never carry conversation content.
|
||||
func (l *Loop) recordLLMDone(span *trace.Span, out *stepOutcome) {
|
||||
if span == nil || out == nil {
|
||||
return
|
||||
}
|
||||
props := map[string]any{
|
||||
"stop_reason": out.stopReason,
|
||||
"tool_calls": len(out.toolCalls),
|
||||
"elapsed_ms": elapsedMS(span),
|
||||
}
|
||||
if m := l.modelName(out); m != "" {
|
||||
props["model"] = m
|
||||
}
|
||||
if out.modelStopReason != "" {
|
||||
props["model_stop_reason"] = out.modelStopReason
|
||||
}
|
||||
_ = span.Record(trace.Event{Type: EventLLMDone, Props: props})
|
||||
}
|
||||
|
||||
// recordToolDone appends a desensitized summary of one tool call to its span: the
|
||||
// tool name, whether it succeeded, the argument NAMES (never their values), the
|
||||
// error kind on failure, and a truncated result summary. Argument values are
|
||||
// deliberately omitted — a tool argument is exactly where a credential, a message
|
||||
// body or a file path with a secret would appear — and everything recorded here is
|
||||
// scrubbed again by the graph before publication.
|
||||
func (l *Loop) recordToolDone(span *trace.Span, tc agent.ToolCall, out ToolOutcome) {
|
||||
if span == nil {
|
||||
return
|
||||
}
|
||||
props := map[string]any{
|
||||
"tool": tc.Name,
|
||||
"ok": out.Err == nil,
|
||||
"args_keys": argKeys(tc.Args),
|
||||
"elapsed_ms": elapsedMS(span),
|
||||
}
|
||||
if out.Err != nil {
|
||||
if out.ErrKind != "" {
|
||||
props["error_kind"] = out.ErrKind
|
||||
}
|
||||
props["error"] = trimSummary(out.Err.Error())
|
||||
} else {
|
||||
props["result_summary"] = trimSummary(formatToolResult(out.Result))
|
||||
}
|
||||
_ = span.Record(trace.Event{Type: EventToolDone, Props: props})
|
||||
}
|
||||
|
||||
// modelName reports the model that served a step: the routed request model when
|
||||
// the streaming path knows it, else the loop's configured model. It is empty when
|
||||
// neither is available.
|
||||
func (l *Loop) modelName(out *stepOutcome) string {
|
||||
if out.model != "" {
|
||||
return out.model
|
||||
}
|
||||
if l.agent == nil {
|
||||
return ""
|
||||
}
|
||||
return l.agent.Model()
|
||||
}
|
||||
|
||||
// elapsedMS reports how long a span has been open so far. A span's Duration is
|
||||
// only stamped by Finish, so a log event recorded before Finish measures it here.
|
||||
func elapsedMS(span *trace.Span) int64 {
|
||||
if span == nil || span.StartedAt.IsZero() {
|
||||
return 0
|
||||
}
|
||||
return time.Since(span.StartedAt).Milliseconds()
|
||||
}
|
||||
|
||||
// argKeys returns a tool call's argument names in sorted order. Only the names are
|
||||
// logged; the values stay out of the trace (see recordToolDone).
|
||||
func argKeys(args map[string]any) []string {
|
||||
if len(args) == 0 {
|
||||
return nil
|
||||
}
|
||||
keys := make([]string, 0, len(args))
|
||||
for k := range args {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
return keys
|
||||
}
|
||||
|
||||
// trimSummary shortens a log value to maxLogSummaryRunes runes without splitting a
|
||||
// multi-byte rune.
|
||||
func trimSummary(s string) string {
|
||||
r := []rune(s)
|
||||
if len(r) <= maxLogSummaryRunes {
|
||||
return s
|
||||
}
|
||||
return string(r[:maxLogSummaryRunes]) + "…"
|
||||
}
|
||||
@@ -21,6 +21,10 @@ type stepOutcome struct {
|
||||
// call (Anthropic max_tokens / end_turn / tool_use), recorded on the llm_call
|
||||
// span so a truncation is visible even when the step itself ended normally.
|
||||
modelStopReason string
|
||||
// model is the routed model that served the call, known on the streaming path
|
||||
// (which builds the request itself) and empty on the agent path. It is
|
||||
// recorded on the llm_call span's log event.
|
||||
model string
|
||||
}
|
||||
|
||||
// step runs a single model call. When an OnStream callback is configured it
|
||||
@@ -53,6 +57,7 @@ func (l *Loop) step(ctx context.Context, sessionID string, messages []model.Mess
|
||||
l.spanProp(llmSpan, "model_stop_reason", out.modelStopReason)
|
||||
}
|
||||
l.spanProp(llmSpan, "tool_calls", len(out.toolCalls))
|
||||
l.recordLLMDone(llmSpan, out)
|
||||
l.finishSpan(llmSpan, trace.SpanStatusOK)
|
||||
return out, nil
|
||||
}
|
||||
|
||||
@@ -59,7 +59,9 @@ func (l *Loop) streamStep(ctx context.Context, messages []model.Message) (*stepO
|
||||
}
|
||||
|
||||
if calls, ok := decodeToolCalls(resp); ok {
|
||||
return &stepOutcome{final: false, toolCalls: calls, stopReason: agent.StopReasonToolCalls, modelStopReason: resp.StopReason}, nil
|
||||
return &stepOutcome{final: false, toolCalls: calls, stopReason: agent.StopReasonToolCalls,
|
||||
modelStopReason: resp.StopReason, model: resp.Model}, nil
|
||||
}
|
||||
return &stepOutcome{final: true, text: resp.Content, stopReason: agent.StopReasonStop, modelStopReason: resp.StopReason}, nil
|
||||
return &stepOutcome{final: true, text: resp.Content, stopReason: agent.StopReasonStop,
|
||||
modelStopReason: resp.StopReason, model: resp.Model}, nil
|
||||
}
|
||||
@@ -293,6 +293,7 @@ func (l *Loop) invokeOne(ctx context.Context, d *tool.Dispatcher, tc agent.ToolC
|
||||
}
|
||||
l.emitToolExecution(ctx, tc.Name, phase)
|
||||
out := ToolOutcome{Result: result, Err: err, ErrKind: errorKind(err)}
|
||||
l.recordToolDone(toolSpan, tc, out)
|
||||
if err != nil {
|
||||
l.spanProp(toolSpan, "error", err.Error())
|
||||
l.spanProp(toolSpan, "error_kind", out.ErrKind)
|
||||
|
||||
@@ -9,14 +9,15 @@ import (
|
||||
// span starts a named span of the given kind when a Tracer is configured, else
|
||||
// returns a nil no-op span. The returned context carries the span as the current
|
||||
// one so nested spans chain as children. When a FlowPublisher is configured the
|
||||
// span's session (threaded via ctx) is stamped as a prop and a running placeholder
|
||||
// span's flow owner (the branch parent for a sub-agent Run, else the Run's own
|
||||
// session, both threaded via ctx) is stamped as a prop and a running placeholder
|
||||
// node is published immediately, so the flow graph shows work in flight.
|
||||
func (l *Loop) span(ctx context.Context, name string, kind trace.SpanKind) (context.Context, *trace.Span) {
|
||||
if l.tracer == nil {
|
||||
return ctx, nil
|
||||
}
|
||||
sctx, sp := l.tracer.StartSpan(ctx, name, kind)
|
||||
if sid := loopSession(ctx); sid != "" {
|
||||
if sid := flowOwner(ctx, loopSession(ctx)); sid != "" {
|
||||
_ = sp.SetProp("session", sid)
|
||||
if l.flow != nil {
|
||||
l.flow.SpanStarted(sid, sp)
|
||||
|
||||
@@ -56,6 +56,14 @@ type Options struct {
|
||||
// Tracer emits a per-sub-agent span (kind subagent) around each run. Optional.
|
||||
Tracer *trace.Tracer
|
||||
|
||||
// Flow, when non-nil, receives the sub-agent span transitions so a spawned
|
||||
// branch appears in the flow graph of the session named by
|
||||
// AgentSpec.ParentSession (design/19 §2). *loop.FlowPublisher is the
|
||||
// production implementation. Optional: a nil sink leaves sub-agent branches
|
||||
// unvisualized (the sub-agent's own inner spans still fold into the parent
|
||||
// graph through the loop's flow branch).
|
||||
Flow FlowSpanSink
|
||||
|
||||
// Approver records sensitive-operation review-and-approve (代审) triggered by
|
||||
// Approve. Policy judgement belongs to the upper layer. Optional.
|
||||
Approver *audit.Approver
|
||||
@@ -68,3 +76,15 @@ type Options struct {
|
||||
// from Bus, "orchestrator" is used.
|
||||
AgentID string
|
||||
}
|
||||
|
||||
// FlowSpanSink folds a span transition into a session's live flow graph so the
|
||||
// realtime visualization shows the work in flight (design/19 §1). *loop.FlowPublisher
|
||||
// is the production implementation; tests substitute a capturing fake. The
|
||||
// interface keeps orchestration from depending on the concrete publisher beyond
|
||||
// the loop template it already takes.
|
||||
type FlowSpanSink interface {
|
||||
// SpanStarted registers a running placeholder node for a span that just began.
|
||||
SpanStarted(sessionID string, sp *trace.Span)
|
||||
// SpanFinished folds a just-finished span into the session graph.
|
||||
SpanFinished(sessionID string, sp *trace.Span)
|
||||
}
|
||||
@@ -2,11 +2,8 @@ package orchestration
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"sync"
|
||||
@@ -14,7 +11,6 @@ import (
|
||||
"time"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/bus"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
|
||||
"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"
|
||||
@@ -47,6 +43,7 @@ type Scheduler struct {
|
||||
comm *comm.Probe
|
||||
memory *memory.Store
|
||||
tracer *trace.Tracer
|
||||
flow FlowSpanSink
|
||||
approver *audit.Approver
|
||||
log *slog.Logger
|
||||
agentID string
|
||||
@@ -92,6 +89,7 @@ func New(opts Options) *Scheduler {
|
||||
bus: opts.Bus,
|
||||
memory: opts.Memory,
|
||||
tracer: opts.Tracer,
|
||||
flow: opts.Flow,
|
||||
approver: opts.Approver,
|
||||
log: sourcedLog(opts.Logger),
|
||||
agentID: agentID,
|
||||
@@ -136,11 +134,13 @@ func (s *Scheduler) Spawn(ctx context.Context, spec AgentSpec) (*SubAgent, error
|
||||
RoleID: spec.RoleID,
|
||||
Mode: spec.Mode,
|
||||
AttachInfo: spec.AttachInfo,
|
||||
ParentSession: spec.ParentSession,
|
||||
Data: spec.Data,
|
||||
Context: cloneMessages(spec.Context),
|
||||
WorkingMemory: memory.NewWorkingMemory(0),
|
||||
status: StatusPending,
|
||||
done: make(chan struct{}),
|
||||
finishedCh: make(chan struct{}),
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
@@ -318,17 +318,24 @@ func (s *Scheduler) run(ctx context.Context, sub *SubAgent) {
|
||||
s.publish(EventStarted, sub.progress())
|
||||
|
||||
runCtx, cancel := s.subCtx(ctx)
|
||||
// A sub-agent must not leak its inner monologue into the caller's chat stream
|
||||
// (design/05 §8: 少说话、多做事; design/04 §2: 不污染前台): muting suppresses the
|
||||
// UI stream/capacity broadcast while progress still reaches the frontend over
|
||||
// the bus and the flow lane. The flow branch attributes every span this Run
|
||||
// emits to the spawning session's graph instead of opening a second graph.
|
||||
runCtx = loop.WithMutedBroadcast(runCtx)
|
||||
runCtx = loop.WithFlowBranch(runCtx, sub.ParentSession)
|
||||
runCtx, span := s.subSpan(runCtx, sub)
|
||||
sub.setCancel(cancel)
|
||||
defer cancel()
|
||||
|
||||
res, err := s.runner.Run(runCtx, sub.ID, sub.Context)
|
||||
if err != nil {
|
||||
s.spanErr(span, err)
|
||||
s.spanErr(sub, span, err)
|
||||
s.abortOrFail(sub, err)
|
||||
return
|
||||
}
|
||||
s.spanOK(span, res)
|
||||
s.spanOK(sub, span, res)
|
||||
sub.setResult(res)
|
||||
s.finish(sub, StatusDone, doneSummary(res), nil)
|
||||
}
|
||||
@@ -354,6 +361,7 @@ func (s *Scheduler) finish(sub *SubAgent, status Status, summary string, err err
|
||||
sub.finishOnce.Do(func() {
|
||||
sub.setOutcome(status, summary, err)
|
||||
s.publish(eventFor(status), sub.progress())
|
||||
sub.markFinished()
|
||||
|
||||
s.mu.Lock()
|
||||
s.active--
|
||||
@@ -373,33 +381,71 @@ func (s *Scheduler) subCtx(parent context.Context) (context.Context, context.Can
|
||||
}
|
||||
|
||||
// subSpan starts a sub-agent span when a tracer is configured, chaining it as a
|
||||
// child of any span already in ctx. It returns the ctx carrying the span and the
|
||||
// span itself (nil when the tracer is nil).
|
||||
// child of any span already in ctx. When the sub-agent is attributed to a flow
|
||||
// session (AgentSpec.ParentSession) the span is also folded into that session's
|
||||
// live graph as a subagent lane root, so the branch renders inside the graph of
|
||||
// the session that spawned it rather than in a graph of its own (design/19 §2). It
|
||||
// returns the ctx carrying the span and the span itself (nil when the tracer is
|
||||
// nil).
|
||||
func (s *Scheduler) subSpan(ctx context.Context, sub *SubAgent) (context.Context, *trace.Span) {
|
||||
if s.tracer == nil {
|
||||
return ctx, nil
|
||||
}
|
||||
return s.tracer.StartSpan(ctx, "subagent."+sub.ID, trace.SpanKindSubagent)
|
||||
sctx, span := s.tracer.StartSpan(ctx, "subagent."+sub.ID, trace.SpanKindSubagent)
|
||||
if sess := sub.ParentSession; sess != "" && s.flow != nil {
|
||||
_ = span.SetProp("session", sess)
|
||||
s.flow.SpanStarted(sess, span)
|
||||
}
|
||||
return sctx, span
|
||||
}
|
||||
|
||||
// spanOK finishes a sub-agent span as ok, recording the result text.
|
||||
func (s *Scheduler) spanOK(span *trace.Span, res *loop.Result) {
|
||||
// subagentEventDone is the log event type recorded on a sub-agent span when its
|
||||
// run settles, so the dashboard node carries the outcome summary (design/19 §1
|
||||
// 节点实时日志).
|
||||
const subagentEventDone = "subagent.done"
|
||||
|
||||
// spanOK finishes a sub-agent span as ok, recording the result text and a
|
||||
// desensitized outcome event, then folds the closed span into the parent flow.
|
||||
func (s *Scheduler) spanOK(sub *SubAgent, span *trace.Span, res *loop.Result) {
|
||||
if span == nil {
|
||||
return
|
||||
}
|
||||
if res != nil {
|
||||
_ = span.SetProp("text", res.Text)
|
||||
_ = span.Record(trace.Event{Type: subagentEventDone, Props: map[string]any{
|
||||
"status": string(StatusDone),
|
||||
"steps": res.Steps,
|
||||
"summary": compactText(res.Text, summaryTextMax),
|
||||
}})
|
||||
}
|
||||
span.Finish(trace.SpanStatusOK)
|
||||
s.foldSpanFinish(sub, span)
|
||||
}
|
||||
|
||||
// spanErr finishes a sub-agent span as an error, recording the cause.
|
||||
func (s *Scheduler) spanErr(span *trace.Span, err error) {
|
||||
// spanErr finishes a sub-agent span as an error, recording the cause, then folds
|
||||
// the closed span into the parent flow.
|
||||
func (s *Scheduler) spanErr(sub *SubAgent, span *trace.Span, err error) {
|
||||
if span == nil {
|
||||
return
|
||||
}
|
||||
_ = span.SetProp("error", err.Error())
|
||||
_ = span.Record(trace.Event{Type: subagentEventDone, Props: map[string]any{
|
||||
"status": string(StatusFailed),
|
||||
"error": compactText(err.Error(), summaryTextMax),
|
||||
}})
|
||||
span.Finish(trace.SpanStatusError)
|
||||
s.foldSpanFinish(sub, span)
|
||||
}
|
||||
|
||||
// foldSpanFinish folds a just-closed sub-agent span into its parent session's
|
||||
// flow graph, flipping the lane root to its final status. The span already
|
||||
// carries the parent session as its "session" prop, so the publisher knows which
|
||||
// graph owns it.
|
||||
func (s *Scheduler) foldSpanFinish(sub *SubAgent, span *trace.Span) {
|
||||
if span == nil || s.flow == nil || sub.ParentSession == "" {
|
||||
return
|
||||
}
|
||||
s.flow.SpanFinished(sub.ParentSession, span)
|
||||
}
|
||||
|
||||
// publish broadcasts a progress event over the bus. It is a no-op when no comm
|
||||
@@ -420,33 +466,6 @@ func (s *Scheduler) nextID() string {
|
||||
return fmt.Sprintf("%s-%d", s.prefix, s.seq.Add(1))
|
||||
}
|
||||
|
||||
// merge aggregates every registered sub-agent into a MergedResult. It is called by
|
||||
// Join once all sub-agents are terminal.
|
||||
func (s *Scheduler) merge() *MergedResult {
|
||||
s.mu.Lock()
|
||||
subs := make([]*SubAgent, 0, len(s.agents))
|
||||
for _, sub := range s.agents {
|
||||
subs = append(subs, sub)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
res := &MergedResult{Total: len(subs), Agents: make([]SubAgentResult, 0, len(subs))}
|
||||
for _, sub := range subs {
|
||||
snap := sub.Snapshot()
|
||||
res.Agents = append(res.Agents, snap)
|
||||
switch snap.Status {
|
||||
case StatusDone:
|
||||
res.Succeeded++
|
||||
case StatusFailed:
|
||||
res.Failed++
|
||||
case StatusAborted:
|
||||
res.Aborted++
|
||||
}
|
||||
}
|
||||
res.Summary = summarizeMerged(res)
|
||||
return res
|
||||
}
|
||||
|
||||
// eventFor maps a terminal status to its progress event name.
|
||||
func eventFor(status Status) string {
|
||||
switch status {
|
||||
@@ -469,26 +488,3 @@ func wrapLoop(l *loop.Loop) runner {
|
||||
}
|
||||
return l
|
||||
}
|
||||
|
||||
// sourcedLog returns the base logger tagged with source "orchestration",
|
||||
// replacing a nil logger with a discard one so the scheduler can always log.
|
||||
func sourcedLog(l *slog.Logger) *slog.Logger {
|
||||
if l == nil {
|
||||
l = slog.New(slog.NewTextHandler(io.Discard, nil))
|
||||
}
|
||||
return logger.WithSource(l, "orchestration")
|
||||
}
|
||||
|
||||
// prefixBytes is the number of random bytes used to build the sub-agent id
|
||||
// prefix.
|
||||
const prefixBytes = 4
|
||||
|
||||
// randPrefix returns a short random prefix for sub-agent ids, falling back to
|
||||
// "sub" if the system entropy source is unavailable.
|
||||
func randPrefix() string {
|
||||
var b [prefixBytes]byte
|
||||
if _, err := rand.Read(b[:]); err != nil {
|
||||
return "sub"
|
||||
}
|
||||
return "sub-" + hex.EncodeToString(b[:])
|
||||
}
|
||||
@@ -0,0 +1,287 @@
|
||||
package orchestration
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop"
|
||||
"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"
|
||||
)
|
||||
|
||||
// parentSession is the flow session the branch tests attribute their sub-agents
|
||||
// to; it mirrors the web chat session the gateway publishes under.
|
||||
const parentSession = "web"
|
||||
|
||||
// secretArg and secretValue are credential-shaped strings planted in a tool call's
|
||||
// arguments and in its result. Neither may appear in a published flow graph.
|
||||
const (
|
||||
secretArg = "sk-arg-0123456789abcdef"
|
||||
secretValue = "sk-live-0123456789abcdef"
|
||||
)
|
||||
|
||||
// graphCapture records every flow graph the publisher emits, per session.
|
||||
type graphCapture struct {
|
||||
mu sync.Mutex
|
||||
bySession map[string]*trace.FlowGraph
|
||||
roundIDs map[string][]string
|
||||
}
|
||||
|
||||
func (c *graphCapture) publish(_ string, payload any) {
|
||||
g, ok := payload.(*trace.FlowGraph)
|
||||
if !ok || g == nil {
|
||||
return
|
||||
}
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if c.bySession == nil {
|
||||
c.bySession = make(map[string]*trace.FlowGraph)
|
||||
c.roundIDs = make(map[string][]string)
|
||||
}
|
||||
c.bySession[g.SessionID] = g
|
||||
if ids := c.roundIDs[g.SessionID]; len(ids) == 0 || ids[len(ids)-1] != g.ID {
|
||||
c.roundIDs[g.SessionID] = append(ids, g.ID)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *graphCapture) graph(t *testing.T, session string) *trace.FlowGraph {
|
||||
t.Helper()
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
g := c.bySession[session]
|
||||
if g == nil {
|
||||
t.Fatalf("no flow graph published for session %q (captured sessions: %v)", session, c.sessionsLocked())
|
||||
}
|
||||
return g
|
||||
}
|
||||
|
||||
func (c *graphCapture) ids(session string) []string {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return append([]string(nil), c.roundIDs[session]...)
|
||||
}
|
||||
|
||||
func (c *graphCapture) has(session string) bool {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
_, ok := c.bySession[session]
|
||||
return ok
|
||||
}
|
||||
|
||||
func (c *graphCapture) sessionsLocked() []string {
|
||||
out := make([]string, 0, len(c.bySession))
|
||||
for s := range c.bySession {
|
||||
out = append(out, s)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// toolThenFinalAdapter behaves like a model that asks for one tool call and then
|
||||
// answers: it returns the scripted call until the request already carries a tool
|
||||
// result. Deciding from the request (rather than a call counter) keeps it correct
|
||||
// when two sub-agents drive the same instance concurrently.
|
||||
type toolThenFinalAdapter struct {
|
||||
toolName string
|
||||
args map[string]any
|
||||
final string
|
||||
}
|
||||
|
||||
func (a *toolThenFinalAdapter) Name() string { return "mock" }
|
||||
|
||||
func (a *toolThenFinalAdapter) Chat(_ context.Context, req model.Request, _ model.Options) (model.Response, error) {
|
||||
for _, m := range req.Messages {
|
||||
if m.Role == model.RoleTool {
|
||||
return model.Response{Model: "mock-1", Content: a.final, StopReason: "end_turn"}, nil
|
||||
}
|
||||
}
|
||||
return model.Response{Model: "mock-1", ToolCalls: []model.ToolCall{
|
||||
{ID: "call-1", Name: a.toolName, Args: a.args},
|
||||
}}, nil
|
||||
}
|
||||
|
||||
func (a *toolThenFinalAdapter) ChatStream(ctx context.Context, req model.Request, on func(model.StreamEvent)) (model.Response, error) {
|
||||
// Emit a delta so a test can detect whether the loop forwarded the sub-agent's
|
||||
// stream to the caller's UI channel.
|
||||
if on != nil {
|
||||
on(model.StreamEvent{Kind: model.StreamKindTextDelta, Delta: "内部草稿"})
|
||||
}
|
||||
return a.Chat(ctx, req, model.Options{})
|
||||
}
|
||||
|
||||
// echoTool is a minimal tool that reports back a scripted result after an optional
|
||||
// delay, so two branches' steps overlap in time.
|
||||
type echoTool struct {
|
||||
name string
|
||||
result any
|
||||
delay time.Duration
|
||||
}
|
||||
|
||||
func (e *echoTool) Name() string { return e.name }
|
||||
func (e *echoTool) Description() string { return "echo the scripted result" }
|
||||
func (e *echoTool) Parameters() (map[string]any, error) { return nil, nil }
|
||||
|
||||
func (e *echoTool) Execute(ctx context.Context, _ any) (any, error) {
|
||||
if e.delay > 0 {
|
||||
select {
|
||||
case <-time.After(e.delay):
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
return e.result, nil
|
||||
}
|
||||
|
||||
// branchLoop builds a real loop over the given dispatcher with tracing and flow
|
||||
// publication wired, plus a scheduler driving it — the production shape, minus the
|
||||
// model provider.
|
||||
func branchLoop(t *testing.T, flow *loop.FlowPublisher, tracer *trace.Tracer, d *tool.Dispatcher, opts Options) *Scheduler {
|
||||
t.Helper()
|
||||
ad := &toolThenFinalAdapter{toolName: "echo", args: map[string]any{
|
||||
"command": "date",
|
||||
"api_key": secretArg,
|
||||
}, final: "子代理完成"}
|
||||
ag, err := agent.New("be brief", d, nil,
|
||||
agent.WithModel(ad, nil),
|
||||
agent.WithOptions(agent.Options{Model: "mock-1"}))
|
||||
if err != nil {
|
||||
t.Fatalf("agent.New: %v", err)
|
||||
}
|
||||
lp := loop.New(loop.Options{Agent: ag, Tools: d, Tracer: tracer, Flow: flow, MaxRounds: 8})
|
||||
opts.Loop = lp
|
||||
opts.Tracer = tracer
|
||||
opts.Flow = flow
|
||||
return newScheduler(t, opts)
|
||||
}
|
||||
|
||||
// flowNodeByID returns the node with the given id, failing the test when absent.
|
||||
func flowNodeByID(t *testing.T, g *trace.FlowGraph, id string) trace.FlowNode {
|
||||
t.Helper()
|
||||
for _, n := range g.Nodes {
|
||||
if n.ID == id {
|
||||
return n
|
||||
}
|
||||
}
|
||||
t.Fatalf("no node %q in graph %s (nodes: %v)", id, g.ID, nodeIDs(g))
|
||||
return trace.FlowNode{}
|
||||
}
|
||||
|
||||
func nodeIDs(g *trace.FlowGraph) []string {
|
||||
out := make([]string, 0, len(g.Nodes))
|
||||
for _, n := range g.Nodes {
|
||||
out = append(out, n.ID)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// flowKindNodes returns every node of the given kind, in graph insertion order
|
||||
// (which is the order the spans started, i.e. the lane's step order).
|
||||
func flowKindNodes(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
|
||||
}
|
||||
|
||||
// flowLaneNodes returns the nodes inside a lane (a reachable set) in insertion
|
||||
// order.
|
||||
func flowLaneNodes(g *trace.FlowGraph, lane map[string]bool) []trace.FlowNode {
|
||||
out := make([]trace.FlowNode, 0, len(lane))
|
||||
for _, n := range g.Nodes {
|
||||
if lane[n.ID] {
|
||||
out = append(out, n)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// flowLaneChain returns the branch's steps ordered by insertion time and asserts
|
||||
// they form one chain: the subagent root feeds the first step and each step feeds
|
||||
// the next. A cross-linked lane (one step reached from two parents) or a stray node
|
||||
// breaks the chain assertion.
|
||||
func flowLaneChain(t *testing.T, g *trace.FlowGraph, sub trace.FlowNode) []trace.FlowNode {
|
||||
t.Helper()
|
||||
lane := flowReachable(g, sub.ID)
|
||||
if len(lane) == 0 {
|
||||
t.Fatalf("subagent node %s (label %s) produced no steps", sub.ID, sub.Label)
|
||||
}
|
||||
chain := flowLaneNodes(g, lane)
|
||||
if !hasFlowEdge(g, sub.ID, chain[0].ID) {
|
||||
t.Fatalf("edges = %+v, want the lane to start at %s (root %s)", g.Edges, sub.ID, chain[0].ID)
|
||||
}
|
||||
for i := 1; i < len(chain); i++ {
|
||||
if !hasFlowEdge(g, chain[i-1].ID, chain[i].ID) {
|
||||
t.Fatalf("lane chain broken between %s and %s (edges: %+v)", chain[i-1].ID, chain[i].ID, g.Edges)
|
||||
}
|
||||
}
|
||||
return chain
|
||||
}
|
||||
|
||||
// chainKinds renders a lane chain as "llm>tool>llm" for failure messages.
|
||||
func chainKinds(chain []trace.FlowNode) string {
|
||||
parts := make([]string, 0, len(chain))
|
||||
for _, n := range chain {
|
||||
parts = append(parts, n.Kind)
|
||||
}
|
||||
return strings.Join(parts, ">")
|
||||
}
|
||||
|
||||
// flowSubagentNode returns the branch root whose label names the given sub-agent
|
||||
// id. A node is keyed by its span id, so the label is what identifies the owner.
|
||||
func flowSubagentNode(t *testing.T, g *trace.FlowGraph, subID string) trace.FlowNode {
|
||||
t.Helper()
|
||||
want := "subagent." + subID
|
||||
for _, n := range flowKindNodes(g, "subagent") {
|
||||
if n.Label == want || n.ID == want {
|
||||
return n
|
||||
}
|
||||
}
|
||||
t.Fatalf("no subagent node labelled %q in graph %s (nodes: %v)", want, g.ID, nodeIDs(g))
|
||||
return trace.FlowNode{}
|
||||
}
|
||||
|
||||
// flowReachable returns the set of node ids reachable from id, excluding id.
|
||||
func flowReachable(g *trace.FlowGraph, id string) map[string]bool {
|
||||
out := make(map[string]bool)
|
||||
queue := []string{id}
|
||||
for len(queue) > 0 {
|
||||
from := queue[0]
|
||||
queue = queue[1:]
|
||||
for _, e := range g.Edges {
|
||||
if e.From != from || out[e.To] || e.To == id {
|
||||
continue
|
||||
}
|
||||
out[e.To] = true
|
||||
queue = append(queue, e.To)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func hasFlowEdge(g *trace.FlowGraph, from, to string) bool {
|
||||
for _, e := range g.Edges {
|
||||
if e.From == from && e.To == to {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// flowLogEvent returns the node's log entry of the given type.
|
||||
func flowLogEvent(t *testing.T, n trace.FlowNode, typ string) trace.Event {
|
||||
t.Helper()
|
||||
for _, ev := range n.Log {
|
||||
if ev.Type == typ {
|
||||
return ev
|
||||
}
|
||||
}
|
||||
t.Fatalf("node %s has no %q log entry (log: %+v)", n.ID, typ, n.Log)
|
||||
return trace.Event{}
|
||||
}
|
||||
@@ -0,0 +1,275 @@
|
||||
package orchestration
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop"
|
||||
"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"
|
||||
)
|
||||
|
||||
// TestSubagentBranchFoldsIntoParentGraph is the end-to-end wiring test for the
|
||||
// branch visualization: a real loop driven by the scheduler under a parent
|
||||
// session must appear in that session's graph as a subagent lane root whose inner
|
||||
// llm/tool steps hang off it, must not start a round (so the frontend's single
|
||||
// flowGraph is never replaced), and must not open a graph under its own id.
|
||||
func TestSubagentBranchFoldsIntoParentGraph(t *testing.T) {
|
||||
t.Parallel()
|
||||
capture := &graphCapture{}
|
||||
flow := loop.NewFlowPublisher(capture.publish, nil)
|
||||
tracer := trace.New(nil)
|
||||
|
||||
d, err := tool.New(registry.New())
|
||||
if err != nil {
|
||||
t.Fatalf("tool.New: %v", err)
|
||||
}
|
||||
if rerr := d.Register(&echoTool{name: "echo", result: map[string]any{"key": secretValue}}); rerr != nil {
|
||||
t.Fatalf("Register: %v", rerr)
|
||||
}
|
||||
s := branchLoop(t, flow, tracer, d, Options{MaxParallel: 2})
|
||||
|
||||
flow.StartRound(parentSession, "帮我调研")
|
||||
sub, serr := s.Spawn(context.Background(), AgentSpec{Goal: "调研", ParentSession: parentSession})
|
||||
if serr != nil {
|
||||
t.Fatalf("Spawn: %v", serr)
|
||||
}
|
||||
res, werr := s.Wait(context.Background(), []string{sub.ID})
|
||||
if werr != nil {
|
||||
t.Fatalf("Wait: %v", werr)
|
||||
}
|
||||
if res.Succeeded != 1 || res.Total != 1 {
|
||||
t.Fatalf("merged = %d total / %d ok, want 1/1 (%+v)", res.Total, res.Succeeded, res.Agents)
|
||||
}
|
||||
flow.EndRound(parentSession, "结果如下")
|
||||
|
||||
if ids := capture.ids(parentSession); len(ids) != 1 || ids[0] != parentSession+"-1" {
|
||||
t.Fatalf("published rounds for %q = %v, want exactly one round: the branch must not start its own",
|
||||
parentSession, ids)
|
||||
}
|
||||
if capture.has(sub.ID) {
|
||||
t.Fatalf("sub-agent %s opened a flow graph of its own; the branch must fold into %s", sub.ID, parentSession)
|
||||
}
|
||||
|
||||
g := capture.graph(t, parentSession)
|
||||
subNode := flowSubagentNode(t, g, sub.ID)
|
||||
if subNode.Kind != "subagent" || subNode.Status != trace.FlowStatusDone {
|
||||
t.Fatalf("subagent node = %+v, want kind subagent / status done", subNode)
|
||||
}
|
||||
// The branch's steps chain inside its own lane: llm (asks for the tool), tool,
|
||||
// llm (answers).
|
||||
chain := flowLaneChain(t, g, subNode)
|
||||
if len(chain) != 3 || chain[0].Kind != "llm" || chain[1].Kind != "tool" || chain[2].Kind != "llm" {
|
||||
t.Fatalf("lane chain = %s, want llm>tool>llm", chainKinds(chain))
|
||||
}
|
||||
llmNode, toolNode, finalNode := chain[0], chain[1], chain[2]
|
||||
if in := flowNodeByID(t, g, "input"); in.Label != "帮我调研" {
|
||||
t.Fatalf("input node label = %q, want the parent round's seed text", in.Label)
|
||||
}
|
||||
|
||||
// Node logs: every executed step must carry a real event, and nothing else.
|
||||
llmEvent := flowLogEvent(t, llmNode, loop.EventLLMDone)
|
||||
if got := fmt.Sprint(llmEvent.Props["model"]); got != "mock-1" {
|
||||
t.Fatalf("llm log model = %q, want mock-1 (props: %+v)", got, llmEvent.Props)
|
||||
}
|
||||
if got := fmt.Sprint(llmEvent.Props["stop_reason"]); got != "tool_calls" {
|
||||
t.Fatalf("llm log stop_reason = %q, want tool_calls", got)
|
||||
}
|
||||
if got, ok := llmEvent.Props["elapsed_ms"].(int64); !ok || got < 0 {
|
||||
t.Fatalf("llm log elapsed_ms = %v, want a duration in ms", llmEvent.Props["elapsed_ms"])
|
||||
}
|
||||
if got := fmt.Sprint(flowLogEvent(t, finalNode, loop.EventLLMDone).Props["stop_reason"]); got != "stop" {
|
||||
t.Fatalf("final llm log stop_reason = %q, want stop", got)
|
||||
}
|
||||
toolEvent := flowLogEvent(t, toolNode, loop.EventToolDone)
|
||||
if got := fmt.Sprint(toolEvent.Props["tool"]); got != "echo" {
|
||||
t.Fatalf("tool log tool = %q, want echo", got)
|
||||
}
|
||||
if got := fmt.Sprint(toolEvent.Props["args_keys"]); !strings.Contains(got, "api_key") || !strings.Contains(got, "command") {
|
||||
t.Fatalf("tool log args_keys = %q, want the argument names", got)
|
||||
}
|
||||
|
||||
// Red line: no credential from the tool call or its result may be published.
|
||||
blob, merr := json.Marshal(g)
|
||||
if merr != nil {
|
||||
t.Fatalf("marshal graph: %v", merr)
|
||||
}
|
||||
if strings.Contains(string(blob), secretArg) {
|
||||
t.Fatalf("published graph leaked a tool argument: %s", blob)
|
||||
}
|
||||
if strings.Contains(string(blob), secretValue) {
|
||||
t.Fatalf("published graph leaked a tool result value: %s", blob)
|
||||
}
|
||||
if !strings.Contains(string(blob), "[redacted]") {
|
||||
t.Fatalf("published graph carries no redaction marker; result summary was not scrubbed: %s", blob)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSpawnParallelBranchesStaySeparate asserts two sub-agents running at the same
|
||||
// time land in two disjoint swimlanes of the parent graph: each lane has its own
|
||||
// llm→tool chain and nothing in one lane is reachable from the other (design/19 §2
|
||||
// 泳道, the cross-linking failure the lane cursor exists to prevent).
|
||||
func TestSpawnParallelBranchesStaySeparate(t *testing.T) {
|
||||
t.Parallel()
|
||||
capture := &graphCapture{}
|
||||
flow := loop.NewFlowPublisher(capture.publish, nil)
|
||||
tracer := trace.New(nil)
|
||||
|
||||
d, err := tool.New(registry.New())
|
||||
if err != nil {
|
||||
t.Fatalf("tool.New: %v", err)
|
||||
}
|
||||
if rerr := d.Register(&echoTool{name: "echo", result: "ok", delay: 15 * time.Millisecond}); rerr != nil {
|
||||
t.Fatalf("Register: %v", rerr)
|
||||
}
|
||||
s := branchLoop(t, flow, tracer, d, Options{MaxParallel: 2})
|
||||
|
||||
ctx := context.Background()
|
||||
flow.StartRound(parentSession, "并行调研")
|
||||
subA, aerr := s.Spawn(ctx, AgentSpec{Goal: "a", ParentSession: parentSession})
|
||||
if aerr != nil {
|
||||
t.Fatalf("Spawn a: %v", aerr)
|
||||
}
|
||||
subB, berr := s.Spawn(ctx, AgentSpec{Goal: "b", ParentSession: parentSession})
|
||||
if berr != nil {
|
||||
t.Fatalf("Spawn b: %v", berr)
|
||||
}
|
||||
res, werr := s.Wait(ctx, []string{subA.ID, subB.ID})
|
||||
if werr != nil {
|
||||
t.Fatalf("Wait: %v", werr)
|
||||
}
|
||||
if res.Succeeded != 2 {
|
||||
t.Fatalf("merged = %d ok, want 2 (%+v)", res.Succeeded, res.Agents)
|
||||
}
|
||||
|
||||
g := capture.graph(t, parentSession)
|
||||
nodeA := flowSubagentNode(t, g, subA.ID)
|
||||
nodeB := flowSubagentNode(t, g, subB.ID)
|
||||
if nodeA.ID == nodeB.ID {
|
||||
t.Fatal("both sub-agents resolved to one node")
|
||||
}
|
||||
setA := flowReachable(g, nodeA.ID)
|
||||
setB := flowReachable(g, nodeB.ID)
|
||||
for _, sub := range []trace.FlowNode{nodeA, nodeB} {
|
||||
chain := flowLaneChain(t, g, sub)
|
||||
if len(chain) != 3 || chain[0].Kind != "llm" || chain[1].Kind != "tool" || chain[2].Kind != "llm" {
|
||||
t.Fatalf("lane %s chain = %s, want llm>tool>llm", sub.Label, chainKinds(chain))
|
||||
}
|
||||
}
|
||||
for id := range setA {
|
||||
if setB[id] {
|
||||
t.Fatalf("node %s is reachable from both lanes (cross-linked edges: %+v)", id, g.Edges)
|
||||
}
|
||||
}
|
||||
// The trunk cursor must still be the parent's seed, not the last branch node.
|
||||
if !hasFlowEdge(g, "input", nodeA.ID) || !hasFlowEdge(g, "input", nodeB.ID) {
|
||||
t.Fatalf("edges = %+v, want both branch roots hanging off the parent's input node", g.Edges)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSubagentRunDoesNotStreamToFrontSession asserts a sub-agent's inner stream is
|
||||
// muted: the loop still drives the model (its llm step appears in the parent's flow
|
||||
// lane), but no delta reaches the caller's OnStream, so a background sub-agent's
|
||||
// monologue never leaks into the front conversation (design/05 §8, design/04 §2
|
||||
// 不污染前台).
|
||||
func TestSubagentRunDoesNotStreamToFrontSession(t *testing.T) {
|
||||
t.Parallel()
|
||||
capture := &graphCapture{}
|
||||
flow := loop.NewFlowPublisher(capture.publish, nil)
|
||||
tracer := trace.New(nil)
|
||||
|
||||
d, err := tool.New(registry.New())
|
||||
if err != nil {
|
||||
t.Fatalf("tool.New: %v", err)
|
||||
}
|
||||
if rerr := d.Register(&echoTool{name: "echo", result: "ok"}); rerr != nil {
|
||||
t.Fatalf("Register: %v", rerr)
|
||||
}
|
||||
ad := &toolThenFinalAdapter{toolName: "echo", args: map[string]any{"command": "date"}, final: "完成"}
|
||||
ag, aerr := agent.New("be brief", d, nil, agent.WithModel(ad, nil),
|
||||
agent.WithOptions(agent.Options{Model: "mock-1"}))
|
||||
if aerr != nil {
|
||||
t.Fatalf("agent.New: %v", aerr)
|
||||
}
|
||||
|
||||
var mu sync.Mutex
|
||||
var streamed []model.StreamEvent
|
||||
lp := loop.New(loop.Options{
|
||||
Agent: ag, Tools: d, Tracer: tracer, Flow: flow, MaxRounds: 8,
|
||||
OnStream: func(ev model.StreamEvent) {
|
||||
mu.Lock()
|
||||
streamed = append(streamed, ev)
|
||||
mu.Unlock()
|
||||
},
|
||||
})
|
||||
s := newScheduler(t, Options{MaxParallel: 1, Loop: lp, Tracer: tracer, Flow: flow})
|
||||
|
||||
ctx := context.Background()
|
||||
flow.StartRound(parentSession, "调研")
|
||||
sub, serr := s.Spawn(ctx, AgentSpec{Goal: "x", ParentSession: parentSession})
|
||||
if serr != nil {
|
||||
t.Fatalf("Spawn: %v", serr)
|
||||
}
|
||||
if _, werr := s.Wait(ctx, []string{sub.ID}); werr != nil {
|
||||
t.Fatalf("Wait: %v", werr)
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
got := append([]model.StreamEvent(nil), streamed...)
|
||||
mu.Unlock()
|
||||
if len(got) != 0 {
|
||||
t.Fatalf("sub-agent pushed %d stream events into the front session: %+v", len(got), got)
|
||||
}
|
||||
// The work is still visible where it belongs: as a lane of the parent's graph.
|
||||
g := capture.graph(t, parentSession)
|
||||
subNode := flowSubagentNode(t, g, sub.ID)
|
||||
if chain := flowLaneChain(t, g, subNode); len(chain) == 0 {
|
||||
t.Fatal("the muted sub-agent left no steps in the parent's flow lane")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSpawnWithoutParentSessionStaysOutOfEveryGraph asserts the flow attribution is
|
||||
// opt-in: a spawn with no ParentSession publishes no graph at all, so an
|
||||
// unattributed background agent can never surface as a stray lane.
|
||||
func TestSpawnWithoutParentSessionStaysOutOfEveryGraph(t *testing.T) {
|
||||
t.Parallel()
|
||||
capture := &graphCapture{}
|
||||
flow := loop.NewFlowPublisher(capture.publish, nil)
|
||||
tracer := trace.New(nil)
|
||||
|
||||
d, err := tool.New(registry.New())
|
||||
if err != nil {
|
||||
t.Fatalf("tool.New: %v", err)
|
||||
}
|
||||
if rerr := d.Register(&echoTool{name: "echo", result: "ok"}); rerr != nil {
|
||||
t.Fatalf("Register: %v", rerr)
|
||||
}
|
||||
s := branchLoop(t, flow, tracer, d, Options{MaxParallel: 1})
|
||||
|
||||
flow.StartRound(parentSession, "调研")
|
||||
sub, serr := s.Spawn(context.Background(), AgentSpec{Goal: "unattributed"})
|
||||
if serr != nil {
|
||||
t.Fatalf("Spawn: %v", serr)
|
||||
}
|
||||
if _, werr := s.Wait(context.Background(), []string{sub.ID}); werr != nil {
|
||||
t.Fatalf("Wait: %v", werr)
|
||||
}
|
||||
|
||||
if capture.has(sub.ID) {
|
||||
t.Fatal("an unattributed sub-agent opened a flow graph of its own")
|
||||
}
|
||||
g := capture.graph(t, parentSession)
|
||||
for _, n := range g.Nodes {
|
||||
if n.Kind == "subagent" {
|
||||
t.Fatalf("node %s appeared in %q although the spawn declared no parent session", n.ID, parentSession)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,387 @@
|
||||
package orchestration
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop"
|
||||
)
|
||||
|
||||
// Constants for the sub-agent dispatch tool (design/05 §1 创建, §7 并行与安全).
|
||||
const (
|
||||
// SpawnToolName is the registered name of the tool the agent calls to
|
||||
// dispatch background sub-agents.
|
||||
SpawnToolName = "subagent_spawn"
|
||||
|
||||
// DefaultSpawnWait is how long a foreground dispatch waits for its batch
|
||||
// before handing the turn back to the model. It mirrors the runtime's
|
||||
// established "前台运行,超过 2 秒自动移入后台" threshold
|
||||
// (DevDocs/design/04-运行时.md §8.5, internal/runtime/task). The sub-agents keep
|
||||
// running either way; only whether the tool result already carries their
|
||||
// output changes.
|
||||
DefaultSpawnWait = 2 * time.Second
|
||||
|
||||
// maxSpawnBatch caps how many sub-agents one dispatch may create. It bounds the
|
||||
// queue a single model call can build (design/05 §7 防失控); the parallel
|
||||
// execution limit itself stays with the scheduler's MaxParallel.
|
||||
maxSpawnBatch = 8
|
||||
|
||||
// maxSpawnTextRunes caps how many runes of one sub-agent's result text are
|
||||
// handed back to the spawning model, so a runaway sub-agent cannot flood the
|
||||
// front session's context. The scheduler keeps the full text on its snapshot.
|
||||
maxSpawnTextRunes = 4000
|
||||
|
||||
// SubagentResultSubject is the bus subject a finished background batch is
|
||||
// announced on, carrying the *MergedResult. It is the design/05 §9 结果回填 hook
|
||||
// for whichever component owns the session message queue (design/15); the
|
||||
// WS-visible per-agent status already travels on the scheduler's own progress
|
||||
// events, so this subject has no consumer until that queue lands.
|
||||
SubagentResultSubject = "subagent.result"
|
||||
)
|
||||
|
||||
// SpawnTool is the runtime entry point that lets the running agent dispatch
|
||||
// sub-agents (design/05 §1). It is the only production caller of
|
||||
// (*Scheduler).Spawn: the model invokes it like any other tool, the scheduler
|
||||
// applies the parallel limit, per-agent timeout and abort machinery, and the tool
|
||||
// hands back either the merged result or — when the batch is still running after
|
||||
// the foreground wait — a dispatch acknowledgement.
|
||||
//
|
||||
// It is safe for concurrent use and holds no per-call state.
|
||||
type SpawnTool struct {
|
||||
sched *Scheduler
|
||||
wait time.Duration
|
||||
log *slog.Logger
|
||||
}
|
||||
|
||||
// SpawnOption customizes a SpawnTool built by NewSpawnTool.
|
||||
type SpawnOption func(*SpawnTool)
|
||||
|
||||
// WithSpawnWait overrides how long a foreground dispatch waits for its batch. A
|
||||
// non-positive value keeps DefaultSpawnWait.
|
||||
func WithSpawnWait(d time.Duration) SpawnOption {
|
||||
return func(t *SpawnTool) {
|
||||
if d > 0 {
|
||||
t.wait = d
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// WithSpawnLogger replaces the tool's logger (tagged source "orchestration/spawn").
|
||||
func WithSpawnLogger(l *slog.Logger) SpawnOption {
|
||||
return func(t *SpawnTool) { t.log = spawnLog(l) }
|
||||
}
|
||||
|
||||
// NewSpawnTool builds the dispatch tool over sched. A nil scheduler is tolerated:
|
||||
// Execute then fails fast with ErrNoScheduler instead of panicking.
|
||||
func NewSpawnTool(sched *Scheduler, opts ...SpawnOption) *SpawnTool {
|
||||
t := &SpawnTool{sched: sched, wait: DefaultSpawnWait, log: spawnLog(nil)}
|
||||
for _, opt := range opts {
|
||||
if opt != nil {
|
||||
opt(t)
|
||||
}
|
||||
}
|
||||
return t
|
||||
}
|
||||
|
||||
// Name returns the tool identifier.
|
||||
func (t *SpawnTool) Name() string { return SpawnToolName }
|
||||
|
||||
// Description explains the tool to the model: what a sub-agent is, that it runs in
|
||||
// an isolated context in parallel, and that a slow batch is handed to the
|
||||
// background rather than blocking the conversation (design/05 §1/§2).
|
||||
func (t *SpawnTool) Description() string {
|
||||
return "Dispatch one or more sub-agents to work in parallel on independent parts of a task. " +
|
||||
"Each sub-agent runs in its own isolated context and returns a summary plus its result. " +
|
||||
"A batch that does not finish within 2 seconds keeps running in the background and reports its " +
|
||||
"results later; prefer several sub-agents over one serial chain for independent work."
|
||||
}
|
||||
|
||||
// Parameters returns the input JSON Schema: either a single goal shorthand or a
|
||||
// tasks array for a parallel batch, plus an optional background flag.
|
||||
func (t *SpawnTool) Parameters() (map[string]any, error) {
|
||||
return map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"goal": spawnStringProp("Objective the sub-agent works toward alone. Required unless tasks is given."),
|
||||
"role": spawnStringProp("Role/persona id to bind the sub-agent to (optional)."),
|
||||
"mode": spawnStringProp("Behavior mode for the role, e.g. engineer (optional)."),
|
||||
"attach_info": spawnStringProp("Extra context or data to hand the sub-agent (optional)."),
|
||||
"tasks": spawnTasksProp(),
|
||||
"background": map[string]any{
|
||||
"type": "boolean",
|
||||
"description": "Dispatch and return immediately without waiting for the foreground window (optional, default false).",
|
||||
},
|
||||
},
|
||||
"anyOf": []any{
|
||||
map[string]any{"required": []any{"goal"}},
|
||||
map[string]any{"required": []any{"tasks"}},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Execute dispatches the requested sub-agents, then either returns their merged
|
||||
// results (the batch settled inside the foreground window) or a dispatch
|
||||
// acknowledgement (it moved to the background, reported later on the bus).
|
||||
//
|
||||
// The sub-agents are spawned on a detached context: their lifetime is their own
|
||||
// (bounded by the scheduler's per-agent timeout and Abort/CancelAll), never the
|
||||
// dispatching turn's, so a slow batch is not killed the moment the tool returns
|
||||
// (design/04 §2 后台不阻塞前台). Context values are kept so the sub-agent span —
|
||||
// and the flow lane it roots — still chains under the step that dispatched it.
|
||||
func (t *SpawnTool) Execute(ctx context.Context, args any) (any, error) {
|
||||
m, ok := args.(map[string]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("%s: args must be an object", SpawnToolName)
|
||||
}
|
||||
if t.sched == nil {
|
||||
return nil, fmt.Errorf("%s: %w", SpawnToolName, ErrNoScheduler)
|
||||
}
|
||||
specs, err := spawnSpecs(m, loop.FlowOwnerFromContext(ctx))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%s: %w", SpawnToolName, err)
|
||||
}
|
||||
|
||||
spawnCtx := context.WithoutCancel(ctx)
|
||||
ids := make([]string, 0, len(specs))
|
||||
for _, spec := range specs {
|
||||
sub, serr := t.sched.Spawn(spawnCtx, spec)
|
||||
if serr != nil {
|
||||
return nil, fmt.Errorf("%s: %w", SpawnToolName, serr)
|
||||
}
|
||||
ids = append(ids, sub.ID)
|
||||
}
|
||||
t.log.Info("sub-agents dispatched", "count", len(ids), "role", specs[0].RoleID,
|
||||
"mode", specs[0].Mode, "session", specs[0].ParentSession)
|
||||
|
||||
if boolArg(m, "background") {
|
||||
go t.await(ids)
|
||||
return runningSpawnResult(ids, specs), nil
|
||||
}
|
||||
return t.foreground(ids, specs), nil
|
||||
}
|
||||
|
||||
// foreground waits up to t.wait for the dispatched batch to settle. A settled
|
||||
// batch is merged into the tool result; a batch still running when the window
|
||||
// closes is announced in the background and acknowledged immediately.
|
||||
func (t *SpawnTool) foreground(ids []string, specs []AgentSpec) any {
|
||||
waitCtx, cancel := context.WithTimeout(context.Background(), t.wait)
|
||||
defer cancel()
|
||||
res, err := t.sched.Wait(waitCtx, ids)
|
||||
if err != nil {
|
||||
t.log.Info("sub-agent batch moved to background", "count", len(ids), "waited", t.wait.String())
|
||||
go t.await(ids)
|
||||
return runningSpawnResult(ids, specs)
|
||||
}
|
||||
return doneSpawnResult(res)
|
||||
}
|
||||
|
||||
// await blocks until the batch settles and broadcasts the merged outcome as the
|
||||
// design/05 §2 "结束后写执行总结 → 返回主会话" hand-off. It runs on a detached
|
||||
// goroutine; a nil bus or an unknown id ends it quietly.
|
||||
func (t *SpawnTool) await(ids []string) {
|
||||
res, err := t.sched.Wait(context.Background(), ids)
|
||||
if err != nil {
|
||||
t.log.Warn("sub-agent batch wait failed", "count", len(ids), "error", err)
|
||||
if res == nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
b := t.sched.Bus()
|
||||
if b == nil {
|
||||
return
|
||||
}
|
||||
if perr := b.Publish(context.Background(), SubagentResultSubject, res); perr != nil {
|
||||
t.log.Warn("publish sub-agent results", "error", perr)
|
||||
}
|
||||
}
|
||||
|
||||
// SpawnResult is the payload the dispatch tool returns to the spawning model: the
|
||||
// batch's merged outcome when it settled inside the foreground window, or a
|
||||
// dispatch acknowledgement when it moved to the background.
|
||||
type SpawnResult struct {
|
||||
// Status is "done" (settled) or "running" (dispatched to the background).
|
||||
Status string `json:"status"`
|
||||
// Summary is a one-line aggregate over the batch.
|
||||
Summary string `json:"summary"`
|
||||
// Agents holds one entry per dispatched sub-agent, in dispatch order.
|
||||
Agents []SpawnAgentResult `json:"agents"`
|
||||
}
|
||||
|
||||
// SpawnAgentResult is one sub-agent's entry in a SpawnResult.
|
||||
type SpawnAgentResult struct {
|
||||
// ID is the sub-agent id (also the id the WS subagent_event frames carry).
|
||||
ID string `json:"id"`
|
||||
// Goal is the objective the sub-agent was given.
|
||||
Goal string `json:"goal"`
|
||||
// Role is the bound role/persona, if any.
|
||||
Role string `json:"role,omitempty"`
|
||||
// Mode is the bound behavior mode, if any.
|
||||
Mode string `json:"mode,omitempty"`
|
||||
// Status is the lifecycle status at return time.
|
||||
Status string `json:"status"`
|
||||
// Summary is the execution summary written when the sub-agent settled.
|
||||
Summary string `json:"summary,omitempty"`
|
||||
// Text is the sub-agent's result text (capped at maxSpawnTextRunes).
|
||||
Text string `json:"text,omitempty"`
|
||||
// Steps is how many loop steps the sub-agent ran.
|
||||
Steps int `json:"steps,omitempty"`
|
||||
}
|
||||
|
||||
// doneSpawnResult maps a settled batch onto the tool result.
|
||||
func doneSpawnResult(res *MergedResult) *SpawnResult {
|
||||
out := &SpawnResult{Status: "done", Agents: make([]SpawnAgentResult, 0)}
|
||||
if res == nil {
|
||||
return out
|
||||
}
|
||||
out.Summary = res.Summary
|
||||
for _, a := range res.Agents {
|
||||
out.Agents = append(out.Agents, SpawnAgentResult{
|
||||
ID: a.ID,
|
||||
Goal: a.Goal,
|
||||
Role: a.RoleID,
|
||||
Mode: a.Mode,
|
||||
Status: string(a.Status),
|
||||
Summary: a.Summary,
|
||||
Text: compactText(a.Text, maxSpawnTextRunes),
|
||||
Steps: a.Steps,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// runningSpawnResult acknowledges a batch that is still running in the background.
|
||||
func runningSpawnResult(ids []string, specs []AgentSpec) *SpawnResult {
|
||||
out := &SpawnResult{
|
||||
Status: "running",
|
||||
Summary: fmt.Sprintf("已派发 %d 个子代理,正在后台运行;完成后结果回填", len(ids)),
|
||||
Agents: make([]SpawnAgentResult, 0, len(ids)),
|
||||
}
|
||||
for i, id := range ids {
|
||||
spec := specs[i]
|
||||
out.Agents = append(out.Agents, SpawnAgentResult{
|
||||
ID: id,
|
||||
Goal: spec.Goal,
|
||||
Role: spec.RoleID,
|
||||
Mode: spec.Mode,
|
||||
Status: string(StatusRunning),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// spawnSpecs decodes the tool arguments into one AgentSpec per sub-agent. A
|
||||
// non-empty tasks array is a parallel batch and takes precedence over the single
|
||||
// goal shorthand.
|
||||
func spawnSpecs(m map[string]any, parentSession string) ([]AgentSpec, error) {
|
||||
tasks, err := taskList(m["tasks"])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(tasks) > maxSpawnBatch {
|
||||
return nil, fmt.Errorf("too many tasks: %d (max %d)", len(tasks), maxSpawnBatch)
|
||||
}
|
||||
if len(tasks) == 0 {
|
||||
spec, serr := specFromMap(m, parentSession)
|
||||
if serr != nil {
|
||||
return nil, serr
|
||||
}
|
||||
return []AgentSpec{spec}, nil
|
||||
}
|
||||
specs := make([]AgentSpec, 0, len(tasks))
|
||||
for i, task := range tasks {
|
||||
spec, serr := specFromMap(task, parentSession)
|
||||
if serr != nil {
|
||||
return nil, fmt.Errorf("tasks[%d]: %w", i, serr)
|
||||
}
|
||||
specs = append(specs, spec)
|
||||
}
|
||||
return specs, nil
|
||||
}
|
||||
|
||||
// specFromMap reads one AgentSpec out of a decoded arguments object (the top-level
|
||||
// object or one tasks entry). A blank goal wraps ErrMissingGoal.
|
||||
func specFromMap(m map[string]any, parentSession string) (AgentSpec, error) {
|
||||
goal := strings.TrimSpace(stringArg(m, "goal"))
|
||||
if goal == "" {
|
||||
return AgentSpec{}, ErrMissingGoal
|
||||
}
|
||||
return AgentSpec{
|
||||
Goal: goal,
|
||||
RoleID: strings.TrimSpace(stringArg(m, "role")),
|
||||
Mode: strings.TrimSpace(stringArg(m, "mode")),
|
||||
AttachInfo: strings.TrimSpace(stringArg(m, "attach_info")),
|
||||
ParentSession: parentSession,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// taskList reads a tasks array of objects. A missing or empty array yields nil; a
|
||||
// non-array or an entry that is not an object is rejected so a malformed batch
|
||||
// never silently degrades into a single spawn.
|
||||
func taskList(v any) ([]map[string]any, error) {
|
||||
if v == nil {
|
||||
return nil, nil
|
||||
}
|
||||
raw, ok := v.([]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("tasks must be an array")
|
||||
}
|
||||
out := make([]map[string]any, 0, len(raw))
|
||||
for i, item := range raw {
|
||||
m, ok := item.(map[string]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("tasks[%d] must be an object", i)
|
||||
}
|
||||
out = append(out, m)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// spawnStringProp builds a string property schema.
|
||||
func spawnStringProp(desc string) map[string]any {
|
||||
return map[string]any{"type": "string", "description": desc}
|
||||
}
|
||||
|
||||
// spawnTasksProp builds the tasks array property schema for a parallel batch.
|
||||
func spawnTasksProp() map[string]any {
|
||||
return map[string]any{
|
||||
"type": "array",
|
||||
"description": "Sub-agents to dispatch in parallel. When given, it overrides goal/role/mode/attach_info.",
|
||||
"maxItems": maxSpawnBatch,
|
||||
"items": map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"goal": spawnStringProp("Objective this sub-agent works toward alone."),
|
||||
"role": spawnStringProp("Role/persona id to bind the sub-agent to (optional)."),
|
||||
"mode": spawnStringProp("Behavior mode for the role (optional)."),
|
||||
"attach_info": spawnStringProp("Extra context or data for this sub-agent (optional)."),
|
||||
},
|
||||
"required": []any{"goal"},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// stringArg reads a string argument, tolerating a missing or non-string value.
|
||||
func stringArg(m map[string]any, key string) string {
|
||||
s, _ := m[key].(string)
|
||||
return s
|
||||
}
|
||||
|
||||
// boolArg reads a boolean argument, tolerating a missing or non-boolean value.
|
||||
func boolArg(m map[string]any, key string) bool {
|
||||
b, _ := m[key].(bool)
|
||||
return b
|
||||
}
|
||||
|
||||
// spawnLog returns l tagged with source "orchestration/spawn", replacing a nil
|
||||
// logger with a discard one so the tool can always log.
|
||||
func spawnLog(l *slog.Logger) *slog.Logger {
|
||||
if l == nil {
|
||||
l = slog.New(slog.NewTextHandler(io.Discard, nil))
|
||||
}
|
||||
return logger.WithSource(l, "orchestration/spawn")
|
||||
}
|
||||
@@ -0,0 +1,367 @@
|
||||
package orchestration
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/bus"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/restrict"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/tool"
|
||||
)
|
||||
|
||||
// countingRunner is a scheduler runner that records how many invocations overlap,
|
||||
// so a test can prove the batch still goes through the scheduler's parallel gate.
|
||||
type countingRunner struct {
|
||||
delay time.Duration
|
||||
result *loop.Result
|
||||
|
||||
mu sync.Mutex
|
||||
active int
|
||||
maxActive int
|
||||
total int
|
||||
}
|
||||
|
||||
func (r *countingRunner) Run(ctx context.Context, _ string, _ []model.Message) (*loop.Result, error) {
|
||||
r.mu.Lock()
|
||||
r.active++
|
||||
r.total++
|
||||
if r.active > r.maxActive {
|
||||
r.maxActive = r.active
|
||||
}
|
||||
r.mu.Unlock()
|
||||
|
||||
if r.delay > 0 {
|
||||
select {
|
||||
case <-time.After(r.delay):
|
||||
case <-ctx.Done():
|
||||
}
|
||||
}
|
||||
|
||||
r.mu.Lock()
|
||||
r.active--
|
||||
r.mu.Unlock()
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if r.result != nil {
|
||||
return r.result, nil
|
||||
}
|
||||
return &loop.Result{Text: "ok", Steps: 1}, nil
|
||||
}
|
||||
|
||||
func (r *countingRunner) peak() int {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.maxActive
|
||||
}
|
||||
|
||||
// spawnToolFor builds a dispatch tool over a scheduler with the given options and
|
||||
// foreground window.
|
||||
func spawnToolFor(t *testing.T, opts Options, wait time.Duration) (*Scheduler, *SpawnTool) {
|
||||
t.Helper()
|
||||
s := newScheduler(t, opts)
|
||||
return s, NewSpawnTool(s, WithSpawnWait(wait))
|
||||
}
|
||||
|
||||
// decodeSpawnResult re-reads the tool result through JSON, so a test asserts the
|
||||
// shape the model actually receives rather than the in-memory struct.
|
||||
func decodeSpawnResult(t *testing.T, raw any) SpawnResult {
|
||||
t.Helper()
|
||||
b, err := json.Marshal(raw)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal tool result: %v", err)
|
||||
}
|
||||
var out SpawnResult
|
||||
if err := json.Unmarshal(b, &out); err != nil {
|
||||
t.Fatalf("unmarshal tool result: %v", err)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// statusOf returns the status the scheduler reports for id, or "" when unknown.
|
||||
func statusOf(s *Scheduler, id string) Status {
|
||||
for _, a := range s.Agents() {
|
||||
if a.ID == id {
|
||||
return a.Status
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func TestSpawnToolReturnsMergedResult(t *testing.T) {
|
||||
t.Parallel()
|
||||
s, sp := spawnToolFor(t, Options{MaxParallel: 2}, time.Second)
|
||||
s.runner = &mockRunner{result: &loop.Result{Text: "answer", Steps: 3, StopReason: "stop"}}
|
||||
|
||||
raw, err := sp.Execute(context.Background(), map[string]any{
|
||||
"goal": "research X", "role": "cyrene", "mode": "engineer", "attach_info": "ctx",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Execute: %v", err)
|
||||
}
|
||||
got := decodeSpawnResult(t, raw)
|
||||
if got.Status != "done" {
|
||||
t.Fatalf("Status = %q, want done", got.Status)
|
||||
}
|
||||
if len(got.Agents) != 1 {
|
||||
t.Fatalf("Agents = %d, want 1", len(got.Agents))
|
||||
}
|
||||
a := got.Agents[0]
|
||||
if a.Status != string(StatusDone) || a.Text != "answer" || a.Steps != 3 {
|
||||
t.Fatalf("agent = %+v, want a done agent with text answer/3 steps", a)
|
||||
}
|
||||
if a.Goal != "research X" || a.Role != "cyrene" || a.Mode != "engineer" {
|
||||
t.Fatalf("agent = %+v, want the goal/role/mode to round-trip", a)
|
||||
}
|
||||
if a.Summary == "" || !strings.Contains(got.Summary, "1 成功") {
|
||||
t.Fatalf("summary = %q / %q, want a settled aggregate", got.Summary, a.Summary)
|
||||
}
|
||||
if statusOf(s, a.ID) != StatusDone {
|
||||
t.Fatalf("scheduler status = %q, want done", statusOf(s, a.ID))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSpawnToolParallelBatchHonorsParallelLimit(t *testing.T) {
|
||||
t.Parallel()
|
||||
s := newScheduler(t, Options{MaxParallel: 1})
|
||||
r := &countingRunner{delay: 20 * time.Millisecond, result: &loop.Result{Text: "ok", Steps: 1}}
|
||||
s.runner = r
|
||||
sp := NewSpawnTool(s, WithSpawnWait(5*time.Second))
|
||||
|
||||
raw, err := sp.Execute(context.Background(), map[string]any{
|
||||
"tasks": []any{
|
||||
map[string]any{"goal": "a"},
|
||||
map[string]any{"goal": "b"},
|
||||
map[string]any{"goal": "c"},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Execute: %v", err)
|
||||
}
|
||||
got := decodeSpawnResult(t, raw)
|
||||
if got.Status != "done" || len(got.Agents) != 3 {
|
||||
t.Fatalf("status/agents = %q/%d, want done/3", got.Status, len(got.Agents))
|
||||
}
|
||||
if got.Agents[0].Goal != "a" || got.Agents[1].Goal != "b" || got.Agents[2].Goal != "c" {
|
||||
t.Fatalf("agents = %+v, want dispatch order a/b/c", got.Agents)
|
||||
}
|
||||
for _, a := range got.Agents {
|
||||
if a.Status != string(StatusDone) {
|
||||
t.Fatalf("agent %s status = %q, want done", a.ID, a.Status)
|
||||
}
|
||||
}
|
||||
// The batch went through the scheduler's semaphore: with MaxParallel=1 the
|
||||
// three sub-agents can never run at the same time.
|
||||
if peak := r.peak(); peak != 1 {
|
||||
t.Fatalf("max concurrent runs = %d, want 1 (the parallel limit)", peak)
|
||||
}
|
||||
if !strings.Contains(got.Summary, "3 成功") {
|
||||
t.Fatalf("summary = %q, want 3 成功", got.Summary)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSpawnToolMovesSlowBatchToBackground(t *testing.T) {
|
||||
t.Parallel()
|
||||
s := newScheduler(t, Options{MaxParallel: 2})
|
||||
gate := make(chan struct{})
|
||||
s.runner = &mockRunner{gate: gate, result: &loop.Result{Text: "late", Steps: 2}}
|
||||
sp := NewSpawnTool(s, WithSpawnWait(40*time.Millisecond))
|
||||
|
||||
results := make(chan *MergedResult, 1)
|
||||
unsub, err := s.Bus().Subscribe(context.Background(), SubagentResultSubject,
|
||||
func(_ context.Context, e bus.Event) error {
|
||||
if res, ok := e.Payload.(*MergedResult); ok {
|
||||
select {
|
||||
case results <- res:
|
||||
default:
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Subscribe: %v", err)
|
||||
}
|
||||
defer unsub()
|
||||
|
||||
start := time.Now()
|
||||
raw, err := sp.Execute(context.Background(), map[string]any{"goal": "slow"})
|
||||
elapsed := time.Since(start)
|
||||
if err != nil {
|
||||
t.Fatalf("Execute: %v", err)
|
||||
}
|
||||
got := decodeSpawnResult(t, raw)
|
||||
if got.Status != "running" {
|
||||
t.Fatalf("Status = %q, want running (the batch outlived the foreground window)", got.Status)
|
||||
}
|
||||
if elapsed > time.Second {
|
||||
t.Fatalf("Execute blocked for %v, want it to hand the turn back at ~40ms", elapsed)
|
||||
}
|
||||
if len(got.Agents) != 1 || got.Agents[0].Status != string(StatusRunning) || got.Agents[0].Goal != "slow" {
|
||||
t.Fatalf("agents = %+v, want a running acknowledgement", got.Agents)
|
||||
}
|
||||
|
||||
close(gate)
|
||||
select {
|
||||
case res := <-results:
|
||||
if res.Total != 1 || res.Succeeded != 1 {
|
||||
t.Fatalf("background merge = %d total / %d ok, want 1/1", res.Total, res.Succeeded)
|
||||
}
|
||||
if res.Agents[0].Text != "late" {
|
||||
t.Fatalf("background result text = %q, want late", res.Agents[0].Text)
|
||||
}
|
||||
case <-time.After(waitTimeout):
|
||||
t.Fatal("a background batch never published its merged result")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSpawnToolKeepsPerAgentTimeout(t *testing.T) {
|
||||
t.Parallel()
|
||||
s := newScheduler(t, Options{MaxParallel: 1, PerAgentTimeout: 40 * time.Millisecond})
|
||||
s.runner = &mockRunner{gate: make(chan struct{})} // never closed: block until the timeout
|
||||
sp := NewSpawnTool(s, WithSpawnWait(2*time.Second))
|
||||
|
||||
raw, err := sp.Execute(context.Background(), map[string]any{"goal": "hang"})
|
||||
if err != nil {
|
||||
t.Fatalf("Execute: %v", err)
|
||||
}
|
||||
got := decodeSpawnResult(t, raw)
|
||||
if got.Status != "done" {
|
||||
t.Fatalf("Status = %q, want done (the batch settled by timing out)", got.Status)
|
||||
}
|
||||
a := got.Agents[0]
|
||||
if a.Status != string(StatusAborted) {
|
||||
t.Fatalf("agent status = %q, want aborted", a.Status)
|
||||
}
|
||||
if !strings.Contains(a.Summary, "超时") {
|
||||
t.Fatalf("summary = %q, want a timeout notice", a.Summary)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSpawnToolAbortStillCancels(t *testing.T) {
|
||||
t.Parallel()
|
||||
s := newScheduler(t, Options{MaxParallel: 1})
|
||||
s.runner = &mockRunner{gate: make(chan struct{})}
|
||||
sp := NewSpawnTool(s, WithSpawnWait(20*time.Millisecond))
|
||||
|
||||
raw, err := sp.Execute(context.Background(), map[string]any{"goal": "hang", "background": true})
|
||||
if err != nil {
|
||||
t.Fatalf("Execute: %v", err)
|
||||
}
|
||||
got := decodeSpawnResult(t, raw)
|
||||
if got.Status != "running" || len(got.Agents) != 1 {
|
||||
t.Fatalf("status/agents = %q/%d, want running/1", got.Status, len(got.Agents))
|
||||
}
|
||||
id := got.Agents[0].ID
|
||||
|
||||
waitFor(t, func() bool { return statusOf(s, id) == StatusRunning }, waitTimeout)
|
||||
if err := s.Abort(context.Background(), id); err != nil {
|
||||
t.Fatalf("Abort: %v", err)
|
||||
}
|
||||
waitFor(t, func() bool { return statusOf(s, id) == StatusAborted }, waitTimeout)
|
||||
}
|
||||
|
||||
func TestSpawnToolRejectsBadArgs(t *testing.T) {
|
||||
t.Parallel()
|
||||
s, sp := spawnToolFor(t, Options{}, time.Millisecond)
|
||||
s.runner = &mockRunner{result: &loop.Result{Text: "x"}}
|
||||
|
||||
tooMany := make([]any, 0, maxSpawnBatch+1)
|
||||
for i := 0; i <= maxSpawnBatch; i++ {
|
||||
tooMany = append(tooMany, map[string]any{"goal": fmt.Sprintf("g%d", i)})
|
||||
}
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
args any
|
||||
want error
|
||||
}{
|
||||
{"args are not an object", "nope", nil},
|
||||
{"no goal", map[string]any{}, ErrMissingGoal},
|
||||
{"blank goal", map[string]any{"goal": " "}, ErrMissingGoal},
|
||||
{"empty tasks", map[string]any{"tasks": []any{}}, ErrMissingGoal},
|
||||
{"task without goal", map[string]any{"tasks": []any{map[string]any{"role": "x"}}}, ErrMissingGoal},
|
||||
{"tasks is not an array", map[string]any{"tasks": "x"}, nil},
|
||||
{"tasks entry is not an object", map[string]any{"tasks": []any{"x"}}, nil},
|
||||
{"batch over the cap", map[string]any{"tasks": tooMany}, nil},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
_, err := sp.Execute(context.Background(), tc.args)
|
||||
if err == nil {
|
||||
t.Fatalf("%s: Execute succeeded, want an error", tc.name)
|
||||
}
|
||||
if tc.want != nil && !errors.Is(err, tc.want) {
|
||||
t.Fatalf("%s: err = %v, want %v", tc.name, err, tc.want)
|
||||
}
|
||||
}
|
||||
if got := len(s.Agents()); got != 0 {
|
||||
t.Fatalf("rejected calls spawned %d sub-agents, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSpawnToolWithoutScheduler(t *testing.T) {
|
||||
t.Parallel()
|
||||
sp := NewSpawnTool(nil)
|
||||
_, err := sp.Execute(context.Background(), map[string]any{"goal": "x"})
|
||||
if !errors.Is(err, ErrNoScheduler) {
|
||||
t.Fatalf("err = %v, want ErrNoScheduler", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSpawnToolThroughDispatcher drives the tool through the real dispatcher so the
|
||||
// advertised JSON Schema, the argument validation and the registry's reachability
|
||||
// markers are all exercised — not just Execute.
|
||||
func TestSpawnToolThroughDispatcher(t *testing.T) {
|
||||
t.Parallel()
|
||||
s, sp := spawnToolFor(t, Options{MaxParallel: 2}, time.Second)
|
||||
s.runner = &mockRunner{result: &loop.Result{Text: "sub answer", Steps: 1}}
|
||||
|
||||
reg := registry.New()
|
||||
d, err := tool.New(reg)
|
||||
if err != nil {
|
||||
t.Fatalf("tool.New: %v", err)
|
||||
}
|
||||
if err := d.Register(sp); err != nil {
|
||||
t.Fatalf("Register: %v", err)
|
||||
}
|
||||
|
||||
schemas := d.ToolSchemas()
|
||||
if len(schemas) != 1 || schemas[0].Name != SpawnToolName {
|
||||
t.Fatalf("schemas = %+v, want the spawn tool advertised", schemas)
|
||||
}
|
||||
|
||||
wrongArgs := map[string]any{"role": "cyrene"}
|
||||
if _, err := d.Invoke(context.Background(), SpawnToolName, wrongArgs); err == nil {
|
||||
t.Fatal("Invoke with neither goal nor tasks succeeded, want a schema/validation error")
|
||||
} else if !errors.Is(err, tool.ErrInvalidArgs) {
|
||||
t.Fatalf("err = %v, want ErrInvalidArgs", err)
|
||||
}
|
||||
|
||||
raw, err := d.Invoke(context.Background(), SpawnToolName, map[string]any{"goal": "do it"})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if got := decodeSpawnResult(t, raw); got.Status != "done" || len(got.Agents) != 1 {
|
||||
t.Fatalf("result = %+v, want one settled agent", got)
|
||||
}
|
||||
|
||||
// Registered with the default (both false) markers: owner-only, so a group or
|
||||
// stranger session can neither see nor invoke it (design/11 §3).
|
||||
if !restrict.SetContext(reg, restrict.ActorOwner).Allow(SpawnToolName) {
|
||||
t.Fatal("owner was denied the spawn tool")
|
||||
}
|
||||
if restrict.SetContext(reg, restrict.ActorGroupMember).Allow(SpawnToolName) {
|
||||
t.Fatal("group member may invoke the spawn tool; want owner-only")
|
||||
}
|
||||
if restrict.SetContext(reg, restrict.ActorStranger).Allow(SpawnToolName) {
|
||||
t.Fatal("stranger may invoke the spawn tool; want owner-only")
|
||||
}
|
||||
}
|
||||
@@ -47,6 +47,13 @@ type AgentSpec struct {
|
||||
Context []model.Message
|
||||
// AttachInfo is free-form attachment/injection text supplied at spawn.
|
||||
AttachInfo string
|
||||
// ParentSession is the session whose live flow graph owns the branch, so the
|
||||
// sub-agent's span and its inner steps render as a swimlane of the spawning
|
||||
// session's graph instead of as a graph of their own (design/19 §2). The
|
||||
// spawning tool takes it from loop.FlowOwnerFromContext. Blank disables the
|
||||
// flow attribution; it never affects which session the sub-agent runs in, so
|
||||
// its checkpoint, memory and stream isolation are unchanged.
|
||||
ParentSession string
|
||||
// Data carries an optional structured payload for the caller's use.
|
||||
Data any
|
||||
}
|
||||
@@ -99,6 +106,10 @@ var (
|
||||
|
||||
// ErrAgentNotFound is returned by Abort when the sub-agent id is unknown.
|
||||
ErrAgentNotFound = errors.New("orchestration: sub-agent not found")
|
||||
|
||||
// ErrNoScheduler is returned by the spawn tool when it was built without a
|
||||
// scheduler (a wiring mistake, reported instead of panicking).
|
||||
ErrNoScheduler = errors.New("orchestration: no scheduler configured")
|
||||
)
|
||||
|
||||
// cloneMessages returns a deep copy of msgs: the message slice and each message's
|
||||
|
||||
@@ -11,12 +11,12 @@ import (
|
||||
)
|
||||
|
||||
// SubAgent is a single spawned sub-agent instance. The identity fields (ID, Goal,
|
||||
// RoleID, Mode, AttachInfo, Data, Context, WorkingMemory) are set once at Spawn
|
||||
// and never mutated. The lifecycle fields (Status, Summary, Result, Err) are
|
||||
// written by the sub-agent's background goroutine and read through the
|
||||
// RoleID, Mode, AttachInfo, ParentSession, Data, Context, WorkingMemory) are set
|
||||
// once at Spawn and never mutated. The lifecycle fields (Status, Summary, Result,
|
||||
// Err) are written by the sub-agent's background goroutine and read through the
|
||||
// synchronized accessors below, so reading them before the sub-agent finishes is
|
||||
// safe. A sub-agent is terminal exactly once its Done reports true; Join returns
|
||||
// once every sub-agent is terminal.
|
||||
// once every sub-agent is terminal, Wait once the named ids are.
|
||||
//
|
||||
// SubAgent is created by (*Scheduler).Spawn; it should not be built directly.
|
||||
type SubAgent struct {
|
||||
@@ -30,6 +30,9 @@ type SubAgent struct {
|
||||
Mode string
|
||||
// AttachInfo is the free-form attachment text supplied at spawn.
|
||||
AttachInfo string
|
||||
// ParentSession is the session whose flow graph owns this sub-agent's branch
|
||||
// (blank when the branch is not visualized).
|
||||
ParentSession string
|
||||
// Data is the optional structured payload supplied at spawn.
|
||||
Data any
|
||||
// Context is the isolated conversation context. It is a deep copy of the
|
||||
@@ -52,6 +55,10 @@ type SubAgent struct {
|
||||
abortReq bool
|
||||
doneClosed bool
|
||||
|
||||
// finishedCh is closed exactly once when the sub-agent reaches a terminal
|
||||
// state, so Wait can block per-id without polling.
|
||||
finishedCh chan struct{}
|
||||
|
||||
launchOnce sync.Once
|
||||
finishOnce sync.Once
|
||||
}
|
||||
@@ -106,6 +113,14 @@ func (s *SubAgent) FinishedAt() time.Time {
|
||||
// Done reports whether the sub-agent reached a terminal state.
|
||||
func (s *SubAgent) Done() bool { return s.Status().Terminal() }
|
||||
|
||||
// waitDone returns the channel closed when the sub-agent reaches a terminal
|
||||
// state. It is what (*Scheduler).Wait blocks on, per-id, without polling.
|
||||
func (s *SubAgent) waitDone() <-chan struct{} { return s.finishedCh }
|
||||
|
||||
// markFinished closes the completion channel exactly once. It is called from
|
||||
// finish under finishOnce, so a sub-agent wakes every Wait caller exactly once.
|
||||
func (s *SubAgent) markFinished() { close(s.finishedCh) }
|
||||
|
||||
// Snapshot returns a point-in-time copy of the sub-agent's state, safe to read
|
||||
// while the sub-agent is still running.
|
||||
func (s *SubAgent) Snapshot() SubAgentResult {
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
package orchestration
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"io"
|
||||
"log/slog"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
|
||||
)
|
||||
|
||||
// sourcedLog returns the base logger tagged with source "orchestration",
|
||||
// replacing a nil logger with a discard one so the scheduler can always log.
|
||||
func sourcedLog(l *slog.Logger) *slog.Logger {
|
||||
if l == nil {
|
||||
l = slog.New(slog.NewTextHandler(io.Discard, nil))
|
||||
}
|
||||
return logger.WithSource(l, "orchestration")
|
||||
}
|
||||
|
||||
// prefixBytes is the number of random bytes used to build the sub-agent id
|
||||
// prefix.
|
||||
const prefixBytes = 4
|
||||
|
||||
// randPrefix returns a short random prefix for sub-agent ids, falling back to
|
||||
// "sub" if the system entropy source is unavailable.
|
||||
func randPrefix() string {
|
||||
var b [prefixBytes]byte
|
||||
if _, err := rand.Read(b[:]); err != nil {
|
||||
return "sub"
|
||||
}
|
||||
return "sub-" + hex.EncodeToString(b[:])
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
package orchestration
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// Wait blocks until every sub-agent named by ids reaches a terminal state (or ctx
|
||||
// is canceled) and returns the merged outcome of exactly those sub-agents, in the
|
||||
// order the ids were given.
|
||||
//
|
||||
// Unlike Join — which covers the whole scheduler — Wait scopes both the wait and
|
||||
// the merge to one caller's batch, so concurrent spawns from different sessions
|
||||
// never wait on, or merge, each other's sub-agents. It is what the spawn tool uses
|
||||
// to collect the results it is about to hand back.
|
||||
//
|
||||
// It returns ErrAgentNotFound (wrapped with %w) for an unknown id, and the wrapped
|
||||
// ctx error together with the partial merged outcome when ctx is canceled first.
|
||||
func (s *Scheduler) Wait(ctx context.Context, ids []string) (*MergedResult, error) {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
subs, err := s.pick(ids)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, sub := range subs {
|
||||
select {
|
||||
case <-sub.waitDone():
|
||||
case <-ctx.Done():
|
||||
return mergeSubs(subs), fmt.Errorf("orchestration: wait: %w", ctx.Err())
|
||||
}
|
||||
}
|
||||
return mergeSubs(subs), nil
|
||||
}
|
||||
|
||||
// pick resolves ids to their registered sub-agents in the given order. It returns
|
||||
// ErrAgentNotFound wrapped for the first unknown id.
|
||||
func (s *Scheduler) pick(ids []string) ([]*SubAgent, error) {
|
||||
subs := make([]*SubAgent, 0, len(ids))
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
for _, id := range ids {
|
||||
sub, ok := s.agents[id]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("orchestration: wait %q: %w", id, ErrAgentNotFound)
|
||||
}
|
||||
subs = append(subs, sub)
|
||||
}
|
||||
return subs, nil
|
||||
}
|
||||
|
||||
// merge aggregates every registered sub-agent into a MergedResult. It is called by
|
||||
// Join once all sub-agents are terminal.
|
||||
func (s *Scheduler) merge() *MergedResult {
|
||||
s.mu.Lock()
|
||||
subs := make([]*SubAgent, 0, len(s.agents))
|
||||
for _, sub := range s.agents {
|
||||
subs = append(subs, sub)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
return mergeSubs(subs)
|
||||
}
|
||||
|
||||
// mergeSubs aggregates the given sub-agents into a MergedResult, counting the
|
||||
// terminal statuses and appending each snapshot in slice order. A sub-agent that
|
||||
// is still running contributes its current snapshot and is counted in neither
|
||||
// Succeeded/Failed/Aborted.
|
||||
func mergeSubs(subs []*SubAgent) *MergedResult {
|
||||
res := &MergedResult{Total: len(subs), Agents: make([]SubAgentResult, 0, len(subs))}
|
||||
for _, sub := range subs {
|
||||
snap := sub.Snapshot()
|
||||
res.Agents = append(res.Agents, snap)
|
||||
switch snap.Status {
|
||||
case StatusDone:
|
||||
res.Succeeded++
|
||||
case StatusFailed:
|
||||
res.Failed++
|
||||
case StatusAborted:
|
||||
res.Aborted++
|
||||
}
|
||||
}
|
||||
res.Summary = summarizeMerged(res)
|
||||
return res
|
||||
}
|
||||
@@ -256,11 +256,12 @@ func (g *RunGraph) addEdge(from, to string) {
|
||||
}
|
||||
|
||||
// SeedInput creates the "input" node for the round. It returns the node id so a
|
||||
// caller can reference it.
|
||||
// caller can reference it. The label is redacted (design/19 §3), so a message that
|
||||
// happens to carry a credential cannot publish one.
|
||||
func (g *RunGraph) SeedInput(text string) string {
|
||||
g.Upsert(NodeSpec{
|
||||
ID: "input",
|
||||
Label: truncate(text, maxSummaryRunes),
|
||||
Label: RedactText(text),
|
||||
Kind: "input",
|
||||
Status: FlowStatusDone,
|
||||
StartedAt: g.started,
|
||||
@@ -269,12 +270,12 @@ func (g *RunGraph) SeedInput(text string) string {
|
||||
}
|
||||
|
||||
// SeedOutput creates the "output" node for the round, wired to the last node so
|
||||
// the trunk ends at the answer.
|
||||
// the trunk ends at the answer. The label is redacted like SeedInput's.
|
||||
func (g *RunGraph) SeedOutput(text string) string {
|
||||
g.ended = time.Now()
|
||||
g.Upsert(NodeSpec{
|
||||
ID: "output",
|
||||
Label: truncate(text, maxSummaryRunes),
|
||||
Label: RedactText(text),
|
||||
Kind: "output",
|
||||
Status: FlowStatusDone,
|
||||
StartedAt: g.ended,
|
||||
@@ -300,7 +301,7 @@ func (g *RunGraph) AddSpan(rec SpanRecord) {
|
||||
StartedAt: rec.StartedAt,
|
||||
EndedAt: rec.EndedAt,
|
||||
Props: safeSpanProps(rec),
|
||||
Log: rec.Events,
|
||||
Log: redactEvents(rec.Events),
|
||||
ParentID: rec.ParentID,
|
||||
})
|
||||
}
|
||||
@@ -392,8 +393,9 @@ func labelForStart(name string, kind SpanKind) string {
|
||||
}
|
||||
|
||||
// safeSpanProps carries only a whitelist of span props to the frontend, plus a
|
||||
// desensitized tool result summary. It never forwards raw arguments or secrets,
|
||||
// and truncates string text so long content cannot leak a full sensitive payload.
|
||||
// desensitized tool result summary. It never forwards raw arguments or secrets, and
|
||||
// every string is scrubbed of credential shapes and truncated, so long or sensitive
|
||||
// content cannot leak a full payload (design/19 §3).
|
||||
func safeSpanProps(rec SpanRecord) map[string]any {
|
||||
var props map[string]any
|
||||
for _, k := range safePropKeys {
|
||||
@@ -411,21 +413,12 @@ func safeSpanProps(rec SpanRecord) map[string]any {
|
||||
if props == nil {
|
||||
props = make(map[string]any)
|
||||
}
|
||||
props["result_summary"] = truncate(summarize(v), maxSummaryRunes)
|
||||
props["result_summary"] = RedactText(summarize(v))
|
||||
}
|
||||
}
|
||||
return props
|
||||
}
|
||||
|
||||
// redactValue shortens string props so a long error/step value cannot leak
|
||||
// sensitive content through the dashboard; non-string values pass through.
|
||||
func redactValue(v any) any {
|
||||
if s, ok := v.(string); ok {
|
||||
return truncate(s, maxSummaryRunes)
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
// safePropKeys are the span prop keys forwarded to a node's props.
|
||||
var safePropKeys = []string{"stop_reason", "step", "error", "error_kind", "tool"}
|
||||
|
||||
|
||||
@@ -0,0 +1,141 @@
|
||||
package trace
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Redaction for everything the flow graph publishes (design/19 §3 红线: 不泄露敏感 ——
|
||||
// 工具参数/内容脱敏;design/11 密钥不落日志). The graph is relayed to the admin panel
|
||||
// and every connected client, so every string that reaches it — span props and
|
||||
// span log events alike — is scrubbed and truncated here. Deciding to redact at
|
||||
// publish time (rather than only at record time) means a runtime that records a
|
||||
// raw tool result still cannot leak it.
|
||||
|
||||
// redactedMarker replaces anything that looks like a credential.
|
||||
const redactedMarker = "[redacted]"
|
||||
|
||||
// secretPatterns match credential shapes that must never reach the dashboard: an
|
||||
// API-key style assignment (plain or JSON), a standalone provider key, a bearer
|
||||
// token, and a JWT. They are deliberately shallow (shape, not semantics): a false
|
||||
// positive only over-redacts a dashboard string, while a false negative publishes a
|
||||
// live credential.
|
||||
var secretPatterns = []*regexp.Regexp{
|
||||
regexp.MustCompile(`(?i)\b[a-z0-9_]*(?:api[_-]?key|access[_-]?token|token|secret|password|passwd|pwd)[a-z0-9_]*"?\s*[:=]\s*"?[^\s"',;}\]]{6,}`),
|
||||
regexp.MustCompile(`\bsk-[A-Za-z0-9_\-]{8,}`),
|
||||
regexp.MustCompile(`(?i)\bbearer\s+[A-Za-z0-9_\-\.=]{10,}`),
|
||||
regexp.MustCompile(`\beyJ[A-Za-z0-9_\-]+\.[A-Za-z0-9_\-]+\.[A-Za-z0-9_\-]+`),
|
||||
}
|
||||
|
||||
// secretKeyNames are prop names whose value is a credential whatever its shape, so
|
||||
// a structured payload that names its secret ({api_key: "..."}) is redacted even
|
||||
// when the value itself looks like ordinary text. Keys are compared after
|
||||
// lower-casing and normalizing -/. to _.
|
||||
var secretKeyNames = map[string]bool{
|
||||
"api_key": true, "apikey": true, "api_token": true, "access_token": true,
|
||||
"auth_token": true, "authorization": true, "token": true, "secret": true,
|
||||
"client_secret": true, "password": true, "passwd": true, "pwd": true,
|
||||
"credential": true, "credentials": true, "private_key": true, "session_key": true,
|
||||
"raw_api_key": true, "x_api_key": true, "bearer": true,
|
||||
}
|
||||
|
||||
// secretKeySuffixes catch prefixed variants such as openai_api_key.
|
||||
var secretKeySuffixes = []string{"_api_key", "_access_token", "_auth_token", "_secret", "_token"}
|
||||
|
||||
// isSecretKey reports whether a map key names a credential.
|
||||
func isSecretKey(key string) bool {
|
||||
k := strings.NewReplacer("-", "_", ".", "_").Replace(strings.ToLower(strings.TrimSpace(key)))
|
||||
if secretKeyNames[k] {
|
||||
return true
|
||||
}
|
||||
for _, suffix := range secretKeySuffixes {
|
||||
if strings.HasSuffix(k, suffix) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// RedactText scrubs credential-looking substrings from s and truncates the result
|
||||
// to maxSummaryRunes, so a published string can never carry a secret or an
|
||||
// unbounded blob.
|
||||
func RedactText(s string) string {
|
||||
return truncate(redactSecrets(s), maxSummaryRunes)
|
||||
}
|
||||
|
||||
// RedactEvent returns ev with every string value — including those nested in maps
|
||||
// and slices, and including any value under a credential-named key — scrubbed and
|
||||
// truncated. Timestamps, types and non-string scalars pass through unchanged.
|
||||
func RedactEvent(ev Event) Event {
|
||||
out := Event{Time: ev.Time, Type: ev.Type}
|
||||
if ev.Props == nil {
|
||||
return out
|
||||
}
|
||||
out.Props = redactProps(ev.Props)
|
||||
return out
|
||||
}
|
||||
|
||||
// redactProps scrubs a property map, applying the key-name rule at the top level
|
||||
// as well as the shape rule to each value.
|
||||
func redactProps(props map[string]any) map[string]any {
|
||||
out := make(map[string]any, len(props))
|
||||
for k, v := range props {
|
||||
if isSecretKey(k) {
|
||||
out[k] = redactedMarker
|
||||
continue
|
||||
}
|
||||
out[k] = redactValue(v)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// redactEvents returns a redacted copy of events for publication. A nil or empty
|
||||
// slice yields nil so a node carries no empty log array.
|
||||
func redactEvents(events []Event) []Event {
|
||||
if len(events) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]Event, len(events))
|
||||
for i, ev := range events {
|
||||
out[i] = RedactEvent(ev)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// redactSecrets replaces every credential-shaped substring with the marker.
|
||||
func redactSecrets(s string) string {
|
||||
if s == "" {
|
||||
return s
|
||||
}
|
||||
for _, re := range secretPatterns {
|
||||
s = re.ReplaceAllString(s, redactedMarker)
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// redactValue scrubs one prop/event value: a string is redacted and truncated, a
|
||||
// map or slice is copied and redacted recursively (with credential-named keys
|
||||
// replaced outright), and any other value (number, bool, time) passes through
|
||||
// unchanged.
|
||||
func redactValue(v any) any {
|
||||
switch x := v.(type) {
|
||||
case string:
|
||||
return RedactText(x)
|
||||
case map[string]any:
|
||||
return redactProps(x)
|
||||
case []any:
|
||||
out := make([]any, len(x))
|
||||
for i, item := range x {
|
||||
out[i] = redactValue(item)
|
||||
}
|
||||
return out
|
||||
case []string:
|
||||
out := make([]any, len(x))
|
||||
for i, item := range x {
|
||||
out[i] = RedactText(item)
|
||||
}
|
||||
return out
|
||||
default:
|
||||
return v
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,132 @@
|
||||
package trace
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TestRedactTextScrubsCredentials asserts every credential shape the dashboard
|
||||
// could be handed is replaced, while ordinary text passes through unchanged.
|
||||
func TestRedactTextScrubsCredentials(t *testing.T) {
|
||||
t.Parallel()
|
||||
secrets := []string{
|
||||
"api_key=sk-abcdef0123456789",
|
||||
`{"token":"abcdef1234567890"}`,
|
||||
"password=hunter2hunter2",
|
||||
"Authorization: Bearer eyJhbGciOiJIUzI1NiJ9.abc.def",
|
||||
"eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.SflKxwRJSMeKKF2QT4",
|
||||
"access_token: ghp_0123456789abcdefghij",
|
||||
}
|
||||
for _, in := range secrets {
|
||||
got := RedactText(in)
|
||||
if !strings.Contains(got, redactedMarker) {
|
||||
t.Fatalf("RedactText(%q) = %q, want the redaction marker", in, got)
|
||||
}
|
||||
for _, leak := range []string{"sk-abcdef0123456789", "hunter2hunter2", "abcdef1234567890",
|
||||
"eyJhbGciOiJIUzI1NiJ9", "ghp_0123456789abcdefghij"} {
|
||||
if strings.Contains(got, leak) {
|
||||
t.Fatalf("RedactText(%q) = %q, still carries %q", in, got, leak)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
benign := "the tool returned 3 rows; stop_reason=tool_calls"
|
||||
if got := RedactText(benign); got != benign {
|
||||
t.Fatalf("RedactText(%q) = %q, want it unchanged", benign, got)
|
||||
}
|
||||
if got := truncate("", maxSummaryRunes); got != "" {
|
||||
t.Fatalf("empty text = %q, want empty", got)
|
||||
}
|
||||
long := strings.Repeat("x", maxSummaryRunes+10)
|
||||
if got := RedactText(long); !strings.HasSuffix(got, "…") || len([]rune(got)) != maxSummaryRunes+1 {
|
||||
t.Fatalf("long text was not truncated to the summary bound: %d runes", len([]rune(got)))
|
||||
}
|
||||
}
|
||||
|
||||
// TestRedactEventScrubsNestedValues asserts a secret nested in a map or a slice is
|
||||
// scrubbed too — the shapes a tool result or error payload actually has.
|
||||
func TestRedactEventScrubsNestedValues(t *testing.T) {
|
||||
t.Parallel()
|
||||
ev := Event{Type: "tool.done", Props: map[string]any{
|
||||
"tool": "http",
|
||||
"headers": map[string]any{
|
||||
"Authorization": "Bearer sk-live-0123456789abcdef",
|
||||
},
|
||||
"args_keys": []string{"api_key", "url"},
|
||||
"nested": []any{map[string]any{"password": "hunter2hunter2"}},
|
||||
"count": 2,
|
||||
}}
|
||||
|
||||
got := RedactEvent(ev)
|
||||
if got.Type != ev.Type || got.Time != ev.Time {
|
||||
t.Fatalf("event envelope changed: %+v", got)
|
||||
}
|
||||
blob, err := json.Marshal(got)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal: %v", err)
|
||||
}
|
||||
for _, leak := range []string{"sk-live-0123456789abcdef", "hunter2hunter2"} {
|
||||
if strings.Contains(string(blob), leak) {
|
||||
t.Fatalf("redacted event still carries %q: %s", leak, blob)
|
||||
}
|
||||
}
|
||||
if count, ok := got.Props["count"].(int); !ok || count != 2 {
|
||||
t.Fatalf("non-string prop = %v, want it passed through", got.Props["count"])
|
||||
}
|
||||
// The original event must not be mutated (a caller may still need it).
|
||||
if orig, _ := ev.Props["headers"].(map[string]any)["Authorization"].(string); !strings.Contains(orig, "sk-live-") {
|
||||
t.Fatalf("RedactEvent mutated the caller's event: %v", orig)
|
||||
}
|
||||
}
|
||||
|
||||
// TestNodePropsAndLogAreRedacted asserts the published graph — both a tool node's
|
||||
// result_summary prop and its log events — never carries a credential. This is the
|
||||
// design/19 §3 red line at the last point before publication.
|
||||
func TestNodePropsAndLogAreRedacted(t *testing.T) {
|
||||
t.Parallel()
|
||||
rec := spanRec("t1", "step1", "tool.http", SpanKindTool, SpanStatusOK,
|
||||
map[string]any{"result": map[string]any{"body": `{"api_key":"sk-abcdef0123456789"}`}})
|
||||
rec.Events = []Event{{
|
||||
Type: "tool.done",
|
||||
Props: map[string]any{
|
||||
"tool": "http",
|
||||
"args_keys": []string{"api_key"},
|
||||
"result_summary": `{"api_key":"sk-abcdef0123456789"}`,
|
||||
},
|
||||
}}
|
||||
|
||||
g := graphFor([]SpanRecord{rec}, "token=abcdef1234567890", "reply")
|
||||
node := nodeByID(g, "t1")
|
||||
if len(node.Log) != 1 || node.Log[0].Type != "tool.done" {
|
||||
t.Fatalf("node log = %+v, want the recorded event", node.Log)
|
||||
}
|
||||
blob, err := json.Marshal(g)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal graph: %v", err)
|
||||
}
|
||||
if strings.Contains(string(blob), "sk-abcdef0123456789") {
|
||||
t.Fatalf("published graph leaked the tool result: %s", blob)
|
||||
}
|
||||
if !strings.Contains(string(blob), redactedMarker) {
|
||||
t.Fatalf("published graph carries no redaction marker: %s", blob)
|
||||
}
|
||||
if got := node.Props["result_summary"]; !strings.Contains(got.(string), redactedMarker) {
|
||||
t.Fatalf("result_summary prop = %v, want it scrubbed", got)
|
||||
}
|
||||
// The round's own input/output labels are scrubbed as well.
|
||||
if in := nodeByID(g, "input").Label; strings.Contains(in, "abcdef1234567890") {
|
||||
t.Fatalf("input label leaked a credential: %q", in)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRedactEventKeepsEmptyPropsNil guards the wire shape: an event without props
|
||||
// must not gain an empty map.
|
||||
func TestRedactEventKeepsEmptyPropsNil(t *testing.T) {
|
||||
t.Parallel()
|
||||
ev := Event{Type: "llm.done", Time: time.Now()}
|
||||
if got := RedactEvent(ev); got.Props != nil {
|
||||
t.Fatalf("props = %v, want nil", got.Props)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user