diff --git a/cmd/cyrene/subagent_tool.go b/cmd/cyrene/subagent_tool.go new file mode 100644 index 0000000..f34a01b --- /dev/null +++ b/cmd/cyrene/subagent_tool.go @@ -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) +} diff --git a/cmd/cyrene/subagent_tool_test.go b/cmd/cyrene/subagent_tool_test.go new file mode 100644 index 0000000..7502218 --- /dev/null +++ b/cmd/cyrene/subagent_tool_test.go @@ -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") + } +} diff --git a/cmd/cyrene/wire.go b/cmd/cyrene/wire.go index ec179d0..413d83c 100644 --- a/cmd/cyrene/wire.go +++ b/cmd/cyrene/wire.go @@ -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 } diff --git a/internal/runtime/loop/flow.go b/internal/runtime/loop/flow.go index 420e521..ef7ba63 100644 --- a/internal/runtime/loop/flow.go +++ b/internal/runtime/loop/flow.go @@ -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) diff --git a/internal/runtime/loop/flow_branch.go b/internal/runtime/loop/flow_branch.go new file mode 100644 index 0000000..be84c84 --- /dev/null +++ b/internal/runtime/loop/flow_branch.go @@ -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) +} diff --git a/internal/runtime/loop/run.go b/internal/runtime/loop/run.go index 7424e69..2ba6ea0 100644 --- a/internal/runtime/loop/run.go +++ b/internal/runtime/loop/run.go @@ -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) }() } diff --git a/internal/runtime/loop/session_ctx.go b/internal/runtime/loop/session_ctx.go index 5d11191..a2d8238 100644 --- a/internal/runtime/loop/session_ctx.go +++ b/internal/runtime/loop/session_ctx.go @@ -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. diff --git a/internal/runtime/loop/spanlog.go b/internal/runtime/loop/spanlog.go new file mode 100644 index 0000000..314795d --- /dev/null +++ b/internal/runtime/loop/spanlog.go @@ -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]) + "…" +} diff --git a/internal/runtime/loop/step.go b/internal/runtime/loop/step.go index 95dce2b..48100e1 100644 --- a/internal/runtime/loop/step.go +++ b/internal/runtime/loop/step.go @@ -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 } diff --git a/internal/runtime/loop/stream.go b/internal/runtime/loop/stream.go index 5210952..726f52f 100644 --- a/internal/runtime/loop/stream.go +++ b/internal/runtime/loop/stream.go @@ -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 } diff --git a/internal/runtime/loop/tools.go b/internal/runtime/loop/tools.go index ac3b623..165d6cd 100644 --- a/internal/runtime/loop/tools.go +++ b/internal/runtime/loop/tools.go @@ -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) diff --git a/internal/runtime/loop/trace.go b/internal/runtime/loop/trace.go index 38e58b1..f705634 100644 --- a/internal/runtime/loop/trace.go +++ b/internal/runtime/loop/trace.go @@ -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) diff --git a/internal/runtime/orchestration/options.go b/internal/runtime/orchestration/options.go index 4407fc7..4086016 100644 --- a/internal/runtime/orchestration/options.go +++ b/internal/runtime/orchestration/options.go @@ -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) +} diff --git a/internal/runtime/orchestration/scheduler.go b/internal/runtime/orchestration/scheduler.go index 78f166f..ce65d29 100644 --- a/internal/runtime/orchestration/scheduler.go +++ b/internal/runtime/orchestration/scheduler.go @@ -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[:]) -} diff --git a/internal/runtime/orchestration/spawnflow_helpers_test.go b/internal/runtime/orchestration/spawnflow_helpers_test.go new file mode 100644 index 0000000..5d0d88f --- /dev/null +++ b/internal/runtime/orchestration/spawnflow_helpers_test.go @@ -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{} +} diff --git a/internal/runtime/orchestration/spawnflow_test.go b/internal/runtime/orchestration/spawnflow_test.go new file mode 100644 index 0000000..14b6844 --- /dev/null +++ b/internal/runtime/orchestration/spawnflow_test.go @@ -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) + } + } +} diff --git a/internal/runtime/orchestration/spawntool.go b/internal/runtime/orchestration/spawntool.go new file mode 100644 index 0000000..718edf5 --- /dev/null +++ b/internal/runtime/orchestration/spawntool.go @@ -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") +} diff --git a/internal/runtime/orchestration/spawntool_test.go b/internal/runtime/orchestration/spawntool_test.go new file mode 100644 index 0000000..fb56a85 --- /dev/null +++ b/internal/runtime/orchestration/spawntool_test.go @@ -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") + } +} diff --git a/internal/runtime/orchestration/spec.go b/internal/runtime/orchestration/spec.go index 65f1759..22edf9e 100644 --- a/internal/runtime/orchestration/spec.go +++ b/internal/runtime/orchestration/spec.go @@ -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 diff --git a/internal/runtime/orchestration/subagent.go b/internal/runtime/orchestration/subagent.go index d732835..3853e62 100644 --- a/internal/runtime/orchestration/subagent.go +++ b/internal/runtime/orchestration/subagent.go @@ -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 { diff --git a/internal/runtime/orchestration/support.go b/internal/runtime/orchestration/support.go new file mode 100644 index 0000000..ad4c64c --- /dev/null +++ b/internal/runtime/orchestration/support.go @@ -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[:]) +} diff --git a/internal/runtime/orchestration/wait.go b/internal/runtime/orchestration/wait.go new file mode 100644 index 0000000..973fd14 --- /dev/null +++ b/internal/runtime/orchestration/wait.go @@ -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 +} diff --git a/internal/runtime/trace/graph.go b/internal/runtime/trace/graph.go index 73661af..ede6531 100644 --- a/internal/runtime/trace/graph.go +++ b/internal/runtime/trace/graph.go @@ -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"} diff --git a/internal/runtime/trace/redact.go b/internal/runtime/trace/redact.go new file mode 100644 index 0000000..46621b3 --- /dev/null +++ b/internal/runtime/trace/redact.go @@ -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 + } +} diff --git a/internal/runtime/trace/redact_test.go b/internal/runtime/trace/redact_test.go new file mode 100644 index 0000000..329bfbb --- /dev/null +++ b/internal/runtime/trace/redact_test.go @@ -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) + } +}