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.
This commit is contained in:
parent
276523b10a
commit
cdc8c1457e
1 changed files with 135 additions and 0 deletions
135
pkg/multiagent/cascade.go
Normal file
135
pkg/multiagent/cascade.go
Normal file
|
|
@ -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
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue