From cdc8c1457e10efbd4db94d148e5f51f73ed44212 Mon Sep 17 00:00:00 2001
From: Leandro Barbosa <leandrobar93@gmail.com>
Date: Wed, 18 Feb 2026 17:44:42 -0300
Subject: [PATCH] feat: add RunRegistry with cascade stop for agent run
 cancellation

RunRegistry tracks active handoff/spawn runs with parent-child
relationships. CascadeStop recursively cancels a run and all its
descendants with cycle-safe seen-set protection.

Supports: Register, Deregister, CascadeStop, StopAll, GetChildren.
---
 pkg/multiagent/cascade.go | 135 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 135 insertions(+)
 create mode 100644 pkg/multiagent/cascade.go

diff --git a/pkg/multiagent/cascade.go b/pkg/multiagent/cascade.go
new file mode 100644
index 000000000..338b0c8fb
--- /dev/null
+++ b/pkg/multiagent/cascade.go
@@ -0,0 +1,135 @@
+package multiagent
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	"github.com/sipeed/picoclaw/pkg/logger"
+)
+
+// ActiveRun represents a running handoff or spawn that can be cancelled.
+type ActiveRun struct {
+	SessionKey string
+	AgentID    string
+	ParentKey  string             // parent session key ("" for top-level)
+	Cancel     context.CancelFunc // cancels this run's context
+	StartedAt  time.Time
+}
+
+// RunRegistry tracks active agent runs for cascade cancellation.
+// Thread-safe via sync.Map.
+type RunRegistry struct {
+	runs sync.Map // sessionKey -> *ActiveRun
+}
+
+// NewRunRegistry creates an empty run registry.
+func NewRunRegistry() *RunRegistry {
+	return &RunRegistry{}
+}
+
+// Register adds an active run to the registry.
+func (r *RunRegistry) Register(run *ActiveRun) {
+	r.runs.Store(run.SessionKey, run)
+	logger.DebugCF("cascade", "Run registered",
+		map[string]interface{}{
+			"session_key": run.SessionKey,
+			"agent_id":    run.AgentID,
+			"parent_key":  run.ParentKey,
+		})
+}
+
+// Deregister removes a run from the registry (normal completion).
+func (r *RunRegistry) Deregister(sessionKey string) {
+	r.runs.Delete(sessionKey)
+	logger.DebugCF("cascade", "Run deregistered",
+		map[string]interface{}{
+			"session_key": sessionKey,
+		})
+}
+
+// CascadeStop cancels a run and all its descendants.
+// Returns the number of runs cancelled. Uses a seen-set to prevent infinite loops.
+func (r *RunRegistry) CascadeStop(sessionKey string) int {
+	seen := make(map[string]bool)
+	killed := r.cascadeStop(sessionKey, seen)
+	if killed > 0 {
+		logger.InfoCF("cascade", "Cascade stop completed",
+			map[string]interface{}{
+				"root_key": sessionKey,
+				"killed":   killed,
+			})
+	}
+	return killed
+}
+
+func (r *RunRegistry) cascadeStop(sessionKey string, seen map[string]bool) int {
+	if seen[sessionKey] {
+		return 0
+	}
+	seen[sessionKey] = true
+	killed := 0
+
+	// Cancel and remove this run
+	if v, ok := r.runs.LoadAndDelete(sessionKey); ok {
+		run := v.(*ActiveRun)
+		run.Cancel()
+		killed++
+		logger.DebugCF("cascade", "Run cancelled",
+			map[string]interface{}{
+				"session_key": sessionKey,
+				"agent_id":    run.AgentID,
+			})
+	}
+
+	// Find and cascade-stop all children (runs whose ParentKey == sessionKey)
+	r.runs.Range(func(key, value interface{}) bool {
+		childRun := value.(*ActiveRun)
+		if childRun.ParentKey == sessionKey {
+			killed += r.cascadeStop(key.(string), seen)
+		}
+		return true
+	})
+
+	return killed
+}
+
+// StopAll cancels every active run. Returns the number cancelled.
+func (r *RunRegistry) StopAll() int {
+	killed := 0
+	r.runs.Range(func(key, value interface{}) bool {
+		run := value.(*ActiveRun)
+		run.Cancel()
+		r.runs.Delete(key)
+		killed++
+		return true
+	})
+	if killed > 0 {
+		logger.InfoCF("cascade", "Stop all completed",
+			map[string]interface{}{"killed": killed})
+	}
+	return killed
+}
+
+// ActiveCount returns the number of currently active runs.
+func (r *RunRegistry) ActiveCount() int {
+	count := 0
+	r.runs.Range(func(_, _ interface{}) bool {
+		count++
+		return true
+	})
+	return count
+}
+
+// GetChildren returns session keys of all direct children of the given parent.
+func (r *RunRegistry) GetChildren(parentKey string) []string {
+	var children []string
+	r.runs.Range(func(key, value interface{}) bool {
+		run := value.(*ActiveRun)
+		if run.ParentKey == parentKey {
+			children = append(children, key.(string))
+		}
+		return true
+	})
+	return children
+}