diff --git a/pkg/agent/loop_reporter_test.go b/pkg/agent/loop_reporter_test.go new file mode 100644 index 000000000..ef7ea0905 --- /dev/null +++ b/pkg/agent/loop_reporter_test.go @@ -0,0 +1,132 @@ +package agent + +import ( + "context" + "os" + "testing" + "time" + + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/orch" +) + +// makeOrchTestLoop creates a minimal AgentLoop with a temp workspace and +// a real Broadcaster wired as the reporter. +// Returns the loop, the broadcaster, and a cleanup function. +func makeOrchTestLoop(t *testing.T) (*AgentLoop, *orch.Broadcaster) { + t.Helper() + tmpDir, err := os.MkdirTemp("", "agent-orch-test-*") + if err != nil { + t.Fatalf("MkdirTemp: %v", err) + } + t.Cleanup(func() { os.RemoveAll(tmpDir) }) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 512, + MaxToolIterations: 5, + }, + }, + } + al := NewAgentLoop(cfg, bus.NewMessageBus(), &mockProvider{}) + b := orch.NewBroadcaster() + al.SetOrchReporter(b) + return al, b +} + +// collectOrchEvents drains the subscriber channel until an agent_gc event +// arrives or the deadline is exceeded. +func collectOrchEvents(t *testing.T, ch <-chan orch.Event, timeout time.Duration) []orch.Event { + t.Helper() + var events []orch.Event + deadline := time.After(timeout) + for { + select { + case ev := <-ch: + events = append(events, ev) + if ev.Type == "agent_gc" { + return events + } + case <-deadline: + t.Fatalf("timed out waiting for agent_gc; events so far: %+v", events) + } + } +} + +// TestAgentLoop_ProcessDirect_EmitsSpawnWaitingGC verifies that a main +// session processed via ProcessDirect emits the full lifecycle: +// +// agent_spawn(sessionKey) → agent_state(waiting) → agent_gc(completed) +// +// and that the Broadcaster snapshot is empty after the call returns. +func TestAgentLoop_ProcessDirect_EmitsSpawnWaitingGC(t *testing.T) { + al, b := makeOrchTestLoop(t) + sub := b.Subscribe() + defer b.Unsubscribe(sub) + + const sessionKey = "orch-test-session" + _, err := al.ProcessDirect(context.Background(), "hello", sessionKey) + if err != nil { + t.Fatalf("ProcessDirect: %v", err) + } + + events := collectOrchEvents(t, sub.Ch, 5*time.Second) + + // First event: agent_spawn with correct ID. + if events[0].Type != "agent_spawn" || events[0].ID != sessionKey { + t.Errorf("first event must be agent_spawn(%s), got: %+v", sessionKey, events[0]) + } + + // At least one agent_state(waiting) for this session. + var hasWaiting bool + for _, ev := range events { + if ev.Type == "agent_state" && ev.ID == sessionKey && ev.State == "waiting" { + hasWaiting = true + break + } + } + if !hasWaiting { + t.Errorf("missing agent_state(waiting) for %s; events: %+v", sessionKey, events) + } + + // Last event: agent_gc(completed) for this session. + last := events[len(events)-1] + if last.Type != "agent_gc" || last.ID != sessionKey || last.Reason != "completed" { + t.Errorf("last event must be agent_gc(completed,%s), got: %+v", sessionKey, last) + } + + // Snapshot must be empty — session removed on GC. + if snap := b.Snapshot(); len(snap) != 0 { + t.Errorf("snapshot must be empty after GC, got: %v", snap) + } +} + +// TestAgentLoop_ProcessHeartbeat_EmitsSpawnAndGC verifies that heartbeat +// sessions appear on canvas with sessionKey = "heartbeat". +func TestAgentLoop_ProcessHeartbeat_EmitsSpawnAndGC(t *testing.T) { + al, b := makeOrchTestLoop(t) + sub := b.Subscribe() + defer b.Unsubscribe(sub) + + _, err := al.ProcessHeartbeat(context.Background(), "check system", "heartbeat-chan", "none") + if err != nil { + t.Fatalf("ProcessHeartbeat: %v", err) + } + + events := collectOrchEvents(t, sub.Ch, 5*time.Second) + + // ProcessHeartbeat always uses sessionKey = "heartbeat". + const want = "heartbeat" + if events[0].Type != "agent_spawn" || events[0].ID != want { + t.Errorf("first event must be agent_spawn(%s), got: %+v", want, events[0]) + } + + last := events[len(events)-1] + if last.Type != "agent_gc" || last.ID != want || last.Reason != "completed" { + t.Errorf("last event must be agent_gc(completed,%s), got: %+v", want, last) + } +} diff --git a/pkg/tools/subagent_reporter_test.go b/pkg/tools/subagent_reporter_test.go index 5ff68f4e4..aaf3bd7a4 100644 --- a/pkg/tools/subagent_reporter_test.go +++ b/pkg/tools/subagent_reporter_test.go @@ -7,8 +7,28 @@ import ( "time" "github.com/sipeed/picoclaw/pkg/orch" + "github.com/sipeed/picoclaw/pkg/providers" ) +// blockingProvider blocks inside Chat until the context is cancelled. +// The ready channel is closed the moment Chat is entered, so callers can +// synchronise before cancelling the context. +type blockingProvider struct { + ready chan struct{} +} + +func newBlockingProvider() *blockingProvider { + return &blockingProvider{ready: make(chan struct{})} +} + +func (p *blockingProvider) Chat(ctx context.Context, _ []providers.Message, _ []providers.ToolDefinition, _ string, _ map[string]any) (*providers.LLMResponse, error) { + close(p.ready) // signal: we are now blocking + <-ctx.Done() + return nil, ctx.Err() +} + +func (p *blockingProvider) GetDefaultModel() string { return "test" } + // TestSubagentManager_Spawn_EmitsLifecycleEvents verifies that Spawn() fires // the correct sequence of orchestration events through a real Broadcaster: // @@ -151,3 +171,73 @@ func TestSubagentManager_Spawn_SnapshotLiveDuringExecution(t *testing.T) { } } } + +// TestSubagentManager_Spawn_CancelledDuringExecution verifies that when the +// context is cancelled while a subagent's LLM call is in progress, the +// Broadcaster receives agent_gc with reason="cancelled" and the agent is +// removed from the snapshot. +// +// Synchronisation: +// 1. blockingProvider.ready is closed when Chat() is entered (goroutine is +// now blocked inside the LLM call). +// 2. Only then is the context cancelled, so there is no race between spawn +// and cancellation. +func TestSubagentManager_Spawn_CancelledDuringExecution(t *testing.T) { + b := orch.NewBroadcaster() + sub := b.Subscribe() + defer b.Unsubscribe(sub) + + bp := newBlockingProvider() + mgr := NewSubagentManager(bp, "test-model", "/tmp/test", nil, b) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + _, err := mgr.Spawn(ctx, "long task", "cancel-me", "", "cli", "direct", nil) + if err != nil { + t.Fatalf("Spawn() error: %v", err) + } + + // Wait until the subagent goroutine is inside Chat (blocking on ctx). + select { + case <-bp.ready: + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for blockingProvider to enter Chat") + } + + // Now cancel — the LLM call unblocks with ctx.Err(). + cancel() + + // Collect events until agent_gc. + var events []orch.Event + deadline := time.After(3 * time.Second) +loop: + for { + select { + case ev := <-sub.Ch: + events = append(events, ev) + if ev.Type == "agent_gc" { + break loop + } + case <-deadline: + t.Fatalf("timed out waiting for agent_gc; events so far: %+v", events) + } + } + + // Locate agent_gc and verify reason = "cancelled". + var gcEv orch.Event + for _, ev := range events { + if ev.Type == "agent_gc" { + gcEv = ev + break + } + } + if gcEv.Reason != "cancelled" { + t.Errorf("agent_gc reason = %q, want %q; events: %+v", gcEv.Reason, "cancelled", events) + } + + // Snapshot must be empty after the GC event. + if snap := b.Snapshot(); len(snap) != 0 { + t.Errorf("snapshot must be empty after agent_gc(cancelled), got: %v", snap) + } +}