diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 4292c8dda..fad5406f4 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -39,6 +39,7 @@ type AgentLoop struct { blackboards sync.Map // sessionKey -> *multiagent.Blackboard fallback *providers.FallbackChain channelManager *channels.Manager + runRegistry *multiagent.RunRegistry // tracks active handoff/spawn runs for cascade stop } // processOptions configures how a message is processed @@ -55,9 +56,10 @@ type processOptions struct { func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers.LLMProvider) *AgentLoop { registry := NewAgentRegistry(cfg, provider) + runRegistry := multiagent.NewRunRegistry() // Register shared tools to all agents - registerSharedTools(cfg, msgBus, registry, provider) + registerSharedTools(cfg, msgBus, registry, provider, runRegistry) // Set up shared fallback chain cooldown := providers.NewCooldownTracker() @@ -77,6 +79,7 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers state: stateManager, summarizing: sync.Map{}, fallback: fallbackChain, + runRegistry: runRegistry, } } @@ -122,7 +125,7 @@ func (r *registryResolver) ListAgents() []multiagent.AgentInfo { } // registerSharedTools registers tools that are shared across all agents (web, message, spawn). -func registerSharedTools(cfg *config.Config, msgBus *bus.MessageBus, registry *AgentRegistry, provider providers.LLMProvider) { +func registerSharedTools(cfg *config.Config, msgBus *bus.MessageBus, registry *AgentRegistry, provider providers.LLMProvider, runReg *multiagent.RunRegistry) { for _, agentID := range registry.ListAgentIDs() { agent, ok := registry.GetAgent(agentID) if !ok { @@ -197,6 +200,7 @@ func registerSharedTools(cfg *config.Config, msgBus *bus.MessageBus, registry *A } return registry.CanSpawnSubagent(currentAgentIDForHandoff, to) })) + handoffTool.SetRunRegistry(runReg, "") agent.Tools.Register(handoffTool) // List agents tool: discover available agents diff --git a/pkg/multiagent/handoff.go b/pkg/multiagent/handoff.go index 47859f6c9..f90794543 100644 --- a/pkg/multiagent/handoff.go +++ b/pkg/multiagent/handoff.go @@ -58,13 +58,14 @@ const DefaultMaxHandoffDepth = 3 // HandoffRequest describes a delegation from one agent to another. type HandoffRequest struct { - FromAgentID string - ToAgentID string - Task string - Context map[string]string // k-v to write to blackboard before handoff - Depth int // current depth level (0 = top-level) - Visited []string // agent IDs already in the call chain - MaxDepth int // max allowed depth (0 = use DefaultMaxHandoffDepth) + FromAgentID string + ToAgentID string + Task string + Context map[string]string // k-v to write to blackboard before handoff + Depth int // current depth level (0 = top-level) + Visited []string // agent IDs already in the call chain + MaxDepth int // max allowed depth (0 = use DefaultMaxHandoffDepth) + ParentRunKey string // parent run session key for cascade stop tracking } // HandoffResult contains the outcome of a handoff execution. @@ -139,6 +140,7 @@ func ExecuteHandoff(ctx context.Context, resolver AgentResolver, board *Blackboa ht.depth = req.Depth + 1 ht.visited = newVisited ht.maxDepth = maxDepth + ht.parentSessionKey = req.ParentRunKey } } } diff --git a/pkg/multiagent/handoff_tool.go b/pkg/multiagent/handoff_tool.go index 098991a27..1f4ed23e0 100644 --- a/pkg/multiagent/handoff_tool.go +++ b/pkg/multiagent/handoff_tool.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "strings" + "time" "github.com/sipeed/picoclaw/pkg/tools" ) @@ -19,6 +20,8 @@ type HandoffTool struct { visited []string // agent IDs already in the call chain maxDepth int // max allowed depth (0 = use DefaultMaxHandoffDepth) allowlistChecker AllowlistChecker // optional; nil = allow all + registry *RunRegistry // optional; nil = no run tracking + parentSessionKey string // session key of the parent run } // NewHandoffTool creates a handoff tool bound to a specific source agent. @@ -101,6 +104,12 @@ func (t *HandoffTool) SetAllowlistChecker(checker AllowlistChecker) { t.allowlistChecker = checker } +// SetRunRegistry sets the registry for tracking active runs (cascade cancellation). +func (t *HandoffTool) SetRunRegistry(registry *RunRegistry, parentSessionKey string) { + t.registry = registry + t.parentSessionKey = parentSessionKey +} + // SetContext updates the origin channel and chat ID for handoff routing. func (t *HandoffTool) SetContext(channel, chatID string) { t.originChannel = channel @@ -143,14 +152,33 @@ func (t *HandoffTool) Execute(ctx context.Context, args map[string]any) *tools.T } } - result := ExecuteHandoff(ctx, t.resolver, t.board, HandoffRequest{ - FromAgentID: t.fromAgentID, - ToAgentID: agentID, - Task: task, - Context: contextMap, - Depth: t.depth, - Visited: t.visited, - MaxDepth: t.maxDepth, + // Create cancellable context for cascade stop support. + // If the parent context is cancelled, this handoff is also cancelled. + childCtx, cancel := context.WithCancel(ctx) + defer cancel() + + // Register this run in the registry for cascade cancellation. + childSessionKey := fmt.Sprintf("handoff:%s:%s:%d:%d", t.fromAgentID, agentID, t.depth, time.Now().UnixNano()) + if t.registry != nil { + t.registry.Register(&ActiveRun{ + SessionKey: childSessionKey, + AgentID: agentID, + ParentKey: t.parentSessionKey, + Cancel: cancel, + StartedAt: time.Now(), + }) + defer t.registry.Deregister(childSessionKey) + } + + result := ExecuteHandoff(childCtx, t.resolver, t.board, HandoffRequest{ + FromAgentID: t.fromAgentID, + ToAgentID: agentID, + Task: task, + Context: contextMap, + Depth: t.depth, + Visited: t.visited, + MaxDepth: t.maxDepth, + ParentRunKey: childSessionKey, }, t.originChannel, t.originChatID) if !result.Success {