From 710b703f35aee5be8acbc14eb71b6c3f9bdf6191 Mon Sep 17 00:00:00 2001 From: dj-oyu <68707227+dj-oyu@users.noreply.github.com> Date: Sun, 15 Mar 2026 18:38:32 +0900 Subject: [PATCH] refactor: extract orchestration tool registration to loop_ext.go Move the 67-line spawn/subagent/answer/review_plan tool registration block from registerSharedTools into registerOrchestrationTools() in loop_ext.go. The upstream-shared registerSharedTools now has a single 1-line call to the extracted function. Co-Authored-By: Claude Opus 4.6 (1M context) --- pkg/agent/loop.go | 79 ++----------------------------------------- pkg/agent/loop_ext.go | 72 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+), 77 deletions(-) diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index c9377e1ac..413b060a8 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -298,86 +298,11 @@ func registerSharedTools( } } - // Spawn tool — only registered when orchestration is explicitly enabled. - - if agent.Subagents != nil && agent.Subagents.Enabled { - webSearchOpts := tools.WebSearchToolOptions{ - BraveAPIKeys: config.MergeAPIKeys(cfg.Tools.Web.Brave.APIKey, cfg.Tools.Web.Brave.APIKeys), - BraveMaxResults: cfg.Tools.Web.Brave.MaxResults, - BraveEnabled: cfg.Tools.Web.Brave.Enabled, - TavilyAPIKeys: config.MergeAPIKeys(cfg.Tools.Web.Tavily.APIKey, cfg.Tools.Web.Tavily.APIKeys), - TavilyBaseURL: cfg.Tools.Web.Tavily.BaseURL, - TavilyMaxResults: cfg.Tools.Web.Tavily.MaxResults, - TavilyEnabled: cfg.Tools.Web.Tavily.Enabled, - DuckDuckGoMaxResults: cfg.Tools.Web.DuckDuckGo.MaxResults, - DuckDuckGoEnabled: cfg.Tools.Web.DuckDuckGo.Enabled, - PerplexityAPIKeys: config.MergeAPIKeys( - cfg.Tools.Web.Perplexity.APIKey, - cfg.Tools.Web.Perplexity.APIKeys, - ), - PerplexityMaxResults: cfg.Tools.Web.Perplexity.MaxResults, - PerplexityEnabled: cfg.Tools.Web.Perplexity.Enabled, - } - - subagentManager := tools.NewSubagentManager( - - provider, - - agent.Model, - - agent.Workspace, - - msgBus, - - al.reporter(), - - webSearchOpts, - ) - - subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature) - - // Wire session recorder for DAG persistence. - - recorder := newSessionRecorder(agent.Sessions) - - conductorKey := routing.BuildAgentMainSessionKey(agent.ID) - - subagentManager.SetSessionRecorder(recorder, conductorKey) - - agent.SubagentMgr = subagentManager - - spawnTool := tools.NewSpawnTool(subagentManager) - - currentAgentID := agentID - - spawnTool.SetAllowlistChecker(func(targetAgentID string) bool { - return registry.CanSpawnSubagent(currentAgentID, targetAgentID) - }) - - agent.Tools.Register(spawnTool) - - // Register blocking subagent tool alongside spawn - - subagentTool := tools.NewSubagentTool(subagentManager) - - agent.Tools.Register(subagentTool) - - // Register conductor-side escalation tools (answer questions, review plans) - - agent.Tools.Register(tools.NewAnswerSubagentTool(subagentManager)) - - agent.Tools.Register(tools.NewReviewSubagentPlanTool(subagentManager)) - } + // Orchestration tools (spawn, subagent, answer, review_plan) + registerOrchestrationTools(cfg, agent, agentID, registry, provider, msgBus, al) // Update context builder with the complete tools registry - agent.ContextBuilder.SetToolsRegistry(agent.Tools) - - // Set orchestration mode if enabled - - if agent.Subagents != nil && agent.Subagents.Enabled { - agent.ContextBuilder.SetOrchestrationEnabled(true) - } } } diff --git a/pkg/agent/loop_ext.go b/pkg/agent/loop_ext.go index 1dbf3122c..89dbd9a24 100644 --- a/pkg/agent/loop_ext.go +++ b/pkg/agent/loop_ext.go @@ -8,7 +8,10 @@ import ( "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/orch" + "github.com/sipeed/picoclaw/pkg/providers" + "github.com/sipeed/picoclaw/pkg/routing" "github.com/sipeed/picoclaw/pkg/stats" + "github.com/sipeed/picoclaw/pkg/tools" "github.com/sipeed/picoclaw/pkg/utils" ) @@ -82,6 +85,75 @@ func (al *AgentLoop) SetHeartbeatThreadUpdater(fn func(int)) { al.onHeartbeatThreadUpdate = fn } +// registerOrchestrationTools registers spawn, subagent, answer, and review_plan +// tools for agents with orchestration enabled. +func registerOrchestrationTools( + cfg *config.Config, + agent *AgentInstance, + agentID string, + registry *AgentRegistry, + provider providers.LLMProvider, + msgBus *bus.MessageBus, + al *AgentLoop, +) { + if agent.Subagents == nil || !agent.Subagents.Enabled { + return + } + + webSearchOpts := tools.WebSearchToolOptions{ + BraveAPIKeys: config.MergeAPIKeys(cfg.Tools.Web.Brave.APIKey, cfg.Tools.Web.Brave.APIKeys), + BraveMaxResults: cfg.Tools.Web.Brave.MaxResults, + BraveEnabled: cfg.Tools.Web.Brave.Enabled, + TavilyAPIKeys: config.MergeAPIKeys(cfg.Tools.Web.Tavily.APIKey, cfg.Tools.Web.Tavily.APIKeys), + TavilyBaseURL: cfg.Tools.Web.Tavily.BaseURL, + TavilyMaxResults: cfg.Tools.Web.Tavily.MaxResults, + TavilyEnabled: cfg.Tools.Web.Tavily.Enabled, + DuckDuckGoMaxResults: cfg.Tools.Web.DuckDuckGo.MaxResults, + DuckDuckGoEnabled: cfg.Tools.Web.DuckDuckGo.Enabled, + PerplexityAPIKeys: config.MergeAPIKeys( + cfg.Tools.Web.Perplexity.APIKey, + cfg.Tools.Web.Perplexity.APIKeys, + ), + PerplexityMaxResults: cfg.Tools.Web.Perplexity.MaxResults, + PerplexityEnabled: cfg.Tools.Web.Perplexity.Enabled, + } + + subagentManager := tools.NewSubagentManager( + provider, + agent.Model, + agent.Workspace, + msgBus, + al.reporter(), + webSearchOpts, + ) + + subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature) + + // Wire session recorder for DAG persistence. + recorder := newSessionRecorder(agent.Sessions) + conductorKey := routing.BuildAgentMainSessionKey(agent.ID) + subagentManager.SetSessionRecorder(recorder, conductorKey) + + agent.SubagentMgr = subagentManager + + spawnTool := tools.NewSpawnTool(subagentManager) + currentAgentID := agentID + spawnTool.SetAllowlistChecker(func(targetAgentID string) bool { + return registry.CanSpawnSubagent(currentAgentID, targetAgentID) + }) + agent.Tools.Register(spawnTool) + + // Register blocking subagent tool alongside spawn + agent.Tools.Register(tools.NewSubagentTool(subagentManager)) + + // Register conductor-side escalation tools (answer questions, review plans) + agent.Tools.Register(tools.NewAnswerSubagentTool(subagentManager)) + agent.Tools.Register(tools.NewReviewSubagentPlanTool(subagentManager)) + + // Set orchestration mode on context builder + agent.ContextBuilder.SetOrchestrationEnabled(true) +} + // handleTaskIntervention checks if a message is a reply to an active task and // either cancels the task or injects a user intervention. Returns (response, handled). func (al *AgentLoop) handleTaskIntervention(msg bus.InboundMessage) (string, bool) {