feat(agent): wire SubTurn into AgentLoop and Spawn Tool

- Add subTurnResults sync.Map to AgentLoop for per-session channel tracking
- Add register/unregister/dequeue methods in steering.go
- Poll SubTurn results in runLLMIteration at loop start and after each tool,
  injecting results as [SubTurn Result] messages into parent conversation
- Initialize root turnState in runAgentLoop, propagate via context
  (withTurnState/turnStateFromContext), call rootTS.Finish() on completion
- Wire Spawn Tool to spawnSubTurn via SetSpawner in registerSharedTools,
  recovering parentTS from context for proper turn hierarchy
- Refactor subagent.go to use SetSpawner pattern
- Add TestSubTurnResultChannelRegistration and TestDequeuePendingSubTurnResults
This commit is contained in:
Administrator
2026-03-16 17:27:04 +08:00
parent ae23193295
commit ceeae15d8a
5 changed files with 348 additions and 73 deletions
+104 -4
View File
@@ -49,6 +49,7 @@ type AgentLoop struct {
cmdRegistry *commands.Registry
mcp mcpRuntime
steering *steeringQueue
subTurnResults sync.Map
mu sync.RWMutex
// Track active requests for safe provider cleanup
activeRequests sync.WaitGroup
@@ -85,9 +86,6 @@ func NewAgentLoop(
) *AgentLoop {
registry := NewAgentRegistry(cfg, provider)
// Register shared tools to all agents
registerSharedTools(cfg, msgBus, registry, provider)
// Set up shared fallback chain
cooldown := providers.NewCooldownTracker()
fallbackChain := providers.NewFallbackChain(cooldown)
@@ -110,11 +108,15 @@ func NewAgentLoop(
steering: newSteeringQueue(parseSteeringMode(cfg.Agents.Defaults.SteeringMode)),
}
// Register shared tools to all agents (now that al is created)
registerSharedTools(al, cfg, msgBus, registry, provider)
return al
}
// registerSharedTools registers tools that are shared across all agents (web, message, spawn).
func registerSharedTools(
al *AgentLoop,
cfg *config.Config,
msgBus *bus.MessageBus,
registry *AgentRegistry,
@@ -230,12 +232,76 @@ func registerSharedTools(
if cfg.Tools.IsToolEnabled("subagent") {
subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Workspace)
subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature)
// Set the spawner that links into AgentLoop's turnState
subagentManager.SetSpawner(func(
ctx context.Context,
task, label, targetAgentID string,
tls *tools.ToolRegistry,
maxTokens int,
temperature float64,
hasMaxTokens, hasTemperature bool,
) (*tools.ToolResult, error) {
// 1. Recover parent Turn State from Context
parentTS := turnStateFromContext(ctx)
if parentTS == nil {
// Fallback: If no turnState exists in context, create an isolated ad-hoc root turn state
// so that the tool can still function outside of an agent loop (e.g. tests, raw invocations).
parentTS = &turnState{
ctx: ctx,
turnID: "adhoc-root",
depth: 0,
session: newEphemeralSession(nil),
pendingResults: make(chan *tools.ToolResult, 16),
}
}
// 2. Build Tools slice from registry
var tlSlice []tools.Tool
for _, name := range tls.List() {
if t, ok := tls.Get(name); ok {
tlSlice = append(tlSlice, t)
}
}
// 3. System Prompt
systemPrompt := "You are a subagent. Complete the given task independently and report the result.\n" +
"You have access to tools - use them as needed to complete your task.\n" +
"After completing the task, provide a clear summary of what was done.\n\n" +
"Task: " + task
// 4. Resolve Model
modelToUse := agent.Model
if targetAgentID != "" {
if targetAgent, ok := al.GetRegistry().GetAgent(targetAgentID); ok {
modelToUse = targetAgent.Model
}
}
// 5. Build SubTurnConfig
cfg := SubTurnConfig{
Model: modelToUse,
Tools: tlSlice,
SystemPrompt: systemPrompt,
}
if hasMaxTokens {
cfg.MaxTokens = maxTokens
}
// 6. Spawn SubTurn
return spawnSubTurn(ctx, al, parentTS, cfg)
})
spawnTool := tools.NewSpawnTool(subagentManager)
currentAgentID := agentID
spawnTool.SetAllowlistChecker(func(targetAgentID string) bool {
return registry.CanSpawnSubagent(currentAgentID, targetAgentID)
})
agent.Tools.Register(spawnTool)
// Also register the synchronous subagent tool
subagentTool := tools.NewSubagentTool(subagentManager)
agent.Tools.Register(subagentTool)
} else {
logger.WarnCF("agent", "spawn tool requires subagent to be enabled", nil)
}
@@ -450,7 +516,7 @@ func (al *AgentLoop) ReloadProviderAndConfig(
}
// Ensure shared tools are re-registered on the new registry
registerSharedTools(cfg, al.bus, registry, provider)
registerSharedTools(al, cfg, al.bus, registry, provider)
// Atomically swap the config and registry under write lock
// This ensures readers see a consistent pair
@@ -896,6 +962,20 @@ func (al *AgentLoop) runAgentLoop(
agent *AgentInstance,
opts processOptions,
) (string, error) {
// Initialize a root TurnState for this iteration, allowing sub-turns to be spawned.
rootTS := &turnState{
ctx: ctx,
turnID: opts.SessionKey, // Associate this turn graph with the current session key
depth: 0,
session: agent.Sessions,
pendingResults: make(chan *tools.ToolResult, 16),
}
ctx = withTurnState(ctx, rootTS)
// Ensure the parent's pending results channel is cleaned up when this root turn finishes
defer al.unregisterSubTurnResultChannel(rootTS.turnID)
al.registerSubTurnResultChannel(rootTS.turnID, rootTS.pendingResults)
// 0. Record last channel for heartbeat notifications (skip internal channels and cli)
if opts.Channel != "" && opts.ChatID != "" {
if !constants.IsInternalChannel(opts.Channel) {
@@ -940,6 +1020,9 @@ func (al *AgentLoop) runAgentLoop(
return "", err
}
// Signal completion to rootTS so it knows it is finished, terminating any active sub-turns
rootTS.Finish()
// If last tool had ForUser content and we already sent it, we might not need to send final response
// This is controlled by the tool's Silent flag and ForUser content
@@ -1055,6 +1138,14 @@ func (al *AgentLoop) runLLMIteration(
}
}
// Poll for any pending SubTurn results and inject them as assistant context.
if subResults := al.dequeuePendingSubTurnResults(opts.SessionKey); len(subResults) > 0 {
for _, r := range subResults {
msg := providers.Message{Role: "user", Content: fmt.Sprintf("[SubTurn Result] %s", r.ForLLM)}
pendingMessages = append(pendingMessages, msg)
}
}
// Determine effective model tier for this conversation turn.
// selectCandidates evaluates routing once and the decision is sticky for
// all tool-follow-up iterations within the same turn so that a multi-step
@@ -1459,6 +1550,15 @@ func (al *AgentLoop) runLLMIteration(
steeringAfterTools = steerMsgs
break
}
// Also poll for any SubTurn results that arrived during tool execution.
if subResults := al.dequeuePendingSubTurnResults(opts.SessionKey); len(subResults) > 0 {
for _, r := range subResults {
msg := providers.Message{Role: "user", Content: fmt.Sprintf("[SubTurn Result] %s", r.ForLLM)}
messages = append(messages, msg)
agent.Sessions.AddFullMessage(opts.SessionKey, msg)
}
}
}
// If steering messages were captured during tool execution, they
+41
View File
@@ -8,6 +8,7 @@ import (
"github.com/sipeed/picoclaw/pkg/logger"
"github.com/sipeed/picoclaw/pkg/providers"
"github.com/sipeed/picoclaw/pkg/tools"
)
// SteeringMode controls how queued steering messages are dequeued.
@@ -186,3 +187,43 @@ func (al *AgentLoop) Continue(ctx context.Context, sessionKey, channel, chatID s
SkipInitialSteeringPoll: true,
})
}
// ====================== SubTurn Result Polling ======================
// dequeuePendingSubTurnResults polls the SubTurn result channel for the given
// session and returns all available results without blocking.
// Returns nil if no channel is registered for this session.
func (al *AgentLoop) dequeuePendingSubTurnResults(sessionKey string) []*tools.ToolResult {
chInterface, ok := al.subTurnResults.Load(sessionKey)
if !ok {
return nil
}
ch, ok := chInterface.(chan *tools.ToolResult)
if !ok {
return nil
}
var results []*tools.ToolResult
for {
select {
case result := <-ch:
if result != nil {
results = append(results, result)
}
default:
return results
}
}
}
// registerSubTurnResultChannel registers a SubTurn result channel for the given session.
// This allows the parent loop to poll for results from child SubTurns.
func (al *AgentLoop) registerSubTurnResultChannel(sessionKey string, ch chan *tools.ToolResult) {
al.subTurnResults.Store(sessionKey, ch)
}
// unregisterSubTurnResultChannel removes the SubTurn result channel for the given session.
func (al *AgentLoop) unregisterSubTurnResultChannel(sessionKey string) {
al.subTurnResults.Delete(sessionKey)
}
+22 -5
View File
@@ -54,7 +54,20 @@ type SubTurnOrphanResultEvent struct {
Result *tools.ToolResult
}
// ====================== turnState (Simplified, reusable with existing structs) ======================
// ====================== turnState ======================
type turnStateKeyType struct{}
var turnStateKey = turnStateKeyType{}
func withTurnState(ctx context.Context, ts *turnState) context.Context {
return context.WithValue(ctx, turnStateKey, ts)
}
func turnStateFromContext(ctx context.Context) *turnState {
ts, _ := ctx.Value(turnStateKey).(*turnState)
return ts
}
type turnState struct {
ctx context.Context
cancelFunc context.CancelFunc // Used to cancel all children when this turn finishes
@@ -189,14 +202,18 @@ func spawnSubTurn(ctx context.Context, al *AgentLoop, parentTS *turnState, cfg S
parentTS.childTurnIDs = append(parentTS.childTurnIDs, childID)
parentTS.mu.Unlock()
// 5. Emit Spawn event (currently using Mock, will be replaced by real EventBus)
// 5. Register the parent's pendingResults channel so the parent loop can poll it
al.registerSubTurnResultChannel(parentTS.turnID, parentTS.pendingResults)
defer al.unregisterSubTurnResultChannel(parentTS.turnID)
// 6. Emit Spawn event (currently using Mock, will be replaced by real EventBus)
MockEventBus.Emit(SubTurnSpawnEvent{
ParentID: parentTS.turnID,
ChildID: childID,
Config: cfg,
})
// 6. Defer emitting End event, and recover from panics to ensure it's always fired
// 7. Defer emitting End event, and recover from panics to ensure it's always fired
defer func() {
if r := recover(); r != nil {
err = fmt.Errorf("subturn panicked: %v", r)
@@ -209,11 +226,11 @@ func spawnSubTurn(ctx context.Context, al *AgentLoop, parentTS *turnState, cfg S
})
}()
// 7. Execute sub-turn via the real agent loop.
// 8. Execute sub-turn via the real agent loop.
// Build a child AgentInstance from SubTurnConfig, inheriting defaults from the parent agent.
result, err = runTurn(childCtx, al, childTS, cfg)
// 8. Deliver result back to parent Turn
// 9. Deliver result back to parent Turn
deliverSubTurnResult(parentTS, childID, result)
return result, err
+70
View File
@@ -253,3 +253,73 @@ func TestSpawnSubTurn_OrphanResultRouting(t *testing.T) {
t.Error("Parent history was polluted by orphan result")
}
}
// ====================== Extra Independent Test: Result Channel Registration ======================
func TestSubTurnResultChannelRegistration(t *testing.T) {
al, _, _, _, cleanup := newTestAgentLoop(t)
defer cleanup()
parent := &turnState{
ctx: context.Background(),
turnID: "parent-reg-1",
depth: 0,
pendingResults: make(chan *tools.ToolResult, 4),
session: &ephemeralSessionStore{},
}
cfg := SubTurnConfig{Model: "gpt-4o-mini", Tools: []tools.Tool{}}
// Before spawn: channel should not be registered
if results := al.dequeuePendingSubTurnResults(parent.turnID); results != nil {
t.Error("expected no channel before spawnSubTurn")
}
_, _ = spawnSubTurn(context.Background(), al, parent, cfg)
// After spawn completes: channel should be unregistered (defer cleanup in spawnSubTurn)
if _, ok := al.subTurnResults.Load(parent.turnID); ok {
t.Error("channel should be unregistered after spawnSubTurn completes")
}
}
// ====================== Extra Independent Test: Dequeue Pending SubTurn Results ======================
func TestDequeuePendingSubTurnResults(t *testing.T) {
al, _, _, _, cleanup := newTestAgentLoop(t)
defer cleanup()
sessionKey := "test-session-dequeue"
ch := make(chan *tools.ToolResult, 4)
// Register channel manually
al.registerSubTurnResultChannel(sessionKey, ch)
defer al.unregisterSubTurnResultChannel(sessionKey)
// Empty channel returns nil
if results := al.dequeuePendingSubTurnResults(sessionKey); len(results) != 0 {
t.Errorf("expected empty results, got %d", len(results))
}
// Put 3 results in
ch <- &tools.ToolResult{ForLLM: "result-1"}
ch <- &tools.ToolResult{ForLLM: "result-2"}
ch <- &tools.ToolResult{ForLLM: "result-3"}
results := al.dequeuePendingSubTurnResults(sessionKey)
if len(results) != 3 {
t.Errorf("expected 3 results, got %d", len(results))
}
if results[0].ForLLM != "result-1" || results[2].ForLLM != "result-3" {
t.Error("results order or content mismatch")
}
// Channel should be drained now
if results := al.dequeuePendingSubTurnResults(sessionKey); len(results) != 0 {
t.Errorf("expected empty after drain, got %d", len(results))
}
// Unregistered session returns nil
al.unregisterSubTurnResultChannel(sessionKey)
if results := al.dequeuePendingSubTurnResults(sessionKey); results != nil {
t.Error("expected nil for unregistered session")
}
}