Cyrene_v2/internal/runtime/orchestration/spawnflow_test.go

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)
		}
	}
}