feat: wire cascade stop into handoff pipeline

- HandoffTool wraps context with cancel, registers run in RunRegistry
- Deregisters on completion (normal or error) via defer
- Propagates ParentRunKey to nested handoffs for correct tree structure
- AgentLoop creates shared RunRegistry, passes to all HandoffTools
This commit is contained in:
Leandro Barbosa 2026-02-18 17:44:49 -03:00
parent cdc8c1457e
commit 03b95932ee
3 changed files with 51 additions and 17 deletions

View file

@ -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

View file

@ -65,6 +65,7 @@ type HandoffRequest struct {
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
}
}
}

View file

@ -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,7 +152,25 @@ func (t *HandoffTool) Execute(ctx context.Context, args map[string]any) *tools.T
}
}
result := ExecuteHandoff(ctx, t.resolver, t.board, HandoffRequest{
// 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,
@ -151,6 +178,7 @@ func (t *HandoffTool) Execute(ctx context.Context, args map[string]any) *tools.T
Depth: t.depth,
Visited: t.visited,
MaxDepth: t.maxDepth,
ParentRunKey: childSessionKey,
}, t.originChannel, t.originChatID)
if !result.Success {