From 8c875ef3ccbd061d4655d0d8e0b25f60b6752cb1 Mon Sep 17 00:00:00 2001 From: Lifei Zhou Date: Wed, 30 Jul 2025 11:34:09 +1000 Subject: [PATCH 1/2] disabled route strategy for recipe --- crates/goose-cli/src/session/builder.rs | 1 + crates/goose/src/agents/agent.rs | 38 +++++++++++++++++++------ crates/goose/src/agents/reply_parts.rs | 24 ++++++++++------ 3 files changed, 46 insertions(+), 17 deletions(-) diff --git a/crates/goose-cli/src/session/builder.rs b/crates/goose-cli/src/session/builder.rs index d4b1b812a1ed..6b44c2f17a3d 100644 --- a/crates/goose-cli/src/session/builder.rs +++ b/crates/goose-cli/src/session/builder.rs @@ -342,6 +342,7 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> Session { // Extensions need to be added after the session is created because we change directory when resuming a session // If we get extensions_override, only run those extensions and none other let extensions_to_run: Vec<_> = if let Some(extensions) = session_config.extensions_override { + agent.disable_router_for_recipe().await; extensions.into_iter().collect() } else { ExtensionConfigManager::get_all() diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index d412517ecfdd..99ce7e4d1d72 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -69,7 +69,7 @@ pub struct ReplyContext { pub system_prompt: String, pub goose_mode: String, pub initial_messages: Vec, - pub config: &'static Config, + pub max_turns: u32, } /// Result of processing tool requests @@ -97,6 +97,7 @@ pub struct Agent { pub(super) tool_result_rx: ToolResultReceiver, pub(super) tool_monitor: Arc>>, pub(super) router_tool_selector: Mutex>>>, + pub(super) router_disabled_override: Mutex, pub(super) scheduler_service: Mutex>>, pub(super) retry_manager: RetryManager, } @@ -172,6 +173,7 @@ impl Agent { tool_result_rx: Arc::new(Mutex::new(tool_rx)), tool_monitor, router_tool_selector: Mutex::new(None), + router_disabled_override: Mutex::new(false), scheduler_service: Mutex::new(None), retry_manager, } @@ -234,6 +236,14 @@ impl Agent { let (tools, toolshim_tools, system_prompt) = self.prepare_tools_and_prompt().await?; let goose_mode = Self::determine_goose_mode(session.as_ref(), config); + let max_turns = session + .as_ref() + .and_then(|s| s.max_turns) + .unwrap_or_else(|| { + config + .get_param("GOOSE_MAX_TURNS") + .unwrap_or(DEFAULT_MAX_TURNS) + }); Ok(ReplyContext { messages, @@ -243,6 +253,7 @@ impl Agent { goose_mode, initial_messages, config, + max_turns, }) } @@ -336,6 +347,13 @@ impl Agent { *scheduler_service = Some(scheduler); } + /// Disable router tool selector for recipe execution + /// This prevents the router from being reinitialized even if config changes + pub async fn disable_router_for_recipe(&self) { + *self.router_disabled_override.lock().await = true; + *self.router_tool_selector.lock().await = None; + } + /// Get a reference count clone to the provider pub async fn provider(&self) -> Result, anyhow::Error> { match &*self.provider.lock().await { @@ -750,6 +768,11 @@ impl Agent { &self, strategy: Option, ) -> Vec { + // If router is disabled for recipe execution, return empty list + if *self.router_disabled_override.lock().await { + return vec![]; + } + let mut prefixed_tools = vec![]; match strategy { Some(RouterToolSelectionStrategy::Vector) => { @@ -842,7 +865,7 @@ impl Agent { mut system_prompt, goose_mode, initial_messages, - config, + max_turns, } = context; let reply_span = tracing::Span::current(); @@ -859,12 +882,6 @@ impl Agent { Ok(Box::pin(async_stream::try_stream! { let _ = reply_span.enter(); let mut turns_taken = 0u32; - let max_turns = session - .as_ref() - .and_then(|s| s.max_turns) - .unwrap_or_else(|| { - config.get_param("GOOSE_MAX_TURNS").unwrap_or(DEFAULT_MAX_TURNS) - }); loop { if is_token_cancelled(&cancel_token) { @@ -1151,6 +1168,11 @@ impl Agent { provider: Option>, reindex_all: Option, ) -> Result<()> { + // Check if router is disabled for recipe execution + if *self.router_disabled_override.lock().await { + return Ok(()); + } + let config = Config::global(); let _extension_manager = self.extension_manager.read().await; let provider = match provider { diff --git a/crates/goose/src/agents/reply_parts.rs b/crates/goose/src/agents/reply_parts.rs index 61452c01eac2..7b52d8b60dc7 100644 --- a/crates/goose/src/agents/reply_parts.rs +++ b/crates/goose/src/agents/reply_parts.rs @@ -48,16 +48,22 @@ impl Agent { }; // Get tools from extension manager - let mut tools = match tool_selection_strategy { - Some(RouterToolSelectionStrategy::Vector) => { - self.list_tools_for_router(Some(RouterToolSelectionStrategy::Vector)) - .await - } - Some(RouterToolSelectionStrategy::Llm) => { - self.list_tools_for_router(Some(RouterToolSelectionStrategy::Llm)) - .await + // Check if router is disabled for recipe execution first + let mut tools = if *self.router_disabled_override.lock().await { + // If router is disabled, use regular tools + self.list_tools(None).await + } else { + match tool_selection_strategy { + Some(RouterToolSelectionStrategy::Vector) => { + self.list_tools_for_router(Some(RouterToolSelectionStrategy::Vector)) + .await + } + Some(RouterToolSelectionStrategy::Llm) => { + self.list_tools_for_router(Some(RouterToolSelectionStrategy::Llm)) + .await + } + _ => self.list_tools(None).await, } - _ => self.list_tools(None).await, }; // Add frontend tools let frontend_tools = self.frontend_tools.lock().await; From b5db1e494b34148560b610323616d414b6ff0c71 Mon Sep 17 00:00:00 2001 From: Lifei Zhou Date: Wed, 30 Jul 2025 13:13:39 +1000 Subject: [PATCH 2/2] clean up --- crates/goose/src/agents/agent.rs | 37 +++++++------------------- crates/goose/src/agents/reply_parts.rs | 1 - 2 files changed, 9 insertions(+), 29 deletions(-) diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 99ce7e4d1d72..4d14f926b44b 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -69,7 +69,7 @@ pub struct ReplyContext { pub system_prompt: String, pub goose_mode: String, pub initial_messages: Vec, - pub max_turns: u32, + pub config: &'static Config, } /// Result of processing tool requests @@ -236,14 +236,6 @@ impl Agent { let (tools, toolshim_tools, system_prompt) = self.prepare_tools_and_prompt().await?; let goose_mode = Self::determine_goose_mode(session.as_ref(), config); - let max_turns = session - .as_ref() - .and_then(|s| s.max_turns) - .unwrap_or_else(|| { - config - .get_param("GOOSE_MAX_TURNS") - .unwrap_or(DEFAULT_MAX_TURNS) - }); Ok(ReplyContext { messages, @@ -253,7 +245,6 @@ impl Agent { goose_mode, initial_messages, config, - max_turns, }) } @@ -347,8 +338,6 @@ impl Agent { *scheduler_service = Some(scheduler); } - /// Disable router tool selector for recipe execution - /// This prevents the router from being reinitialized even if config changes pub async fn disable_router_for_recipe(&self) { *self.router_disabled_override.lock().await = true; *self.router_tool_selector.lock().await = None; @@ -495,7 +484,7 @@ impl Agent { || tool_call.name == ROUTER_LLM_SEARCH_TOOL_NAME { let selector = self.router_tool_selector.lock().await.clone(); - let mut selected_tools = match selector.as_ref() { + let selected_tools = match selector.as_ref() { Some(selector) => match selector.select_tools(tool_call.arguments.clone()).await { Ok(tools) => tools, Err(e) => { @@ -518,18 +507,6 @@ impl Agent { } }; - // Append final_output tool if present (for structured output recipes, [Issue #3700](https://github.com/block/goose/issues/3700) - if let Some(final_output_tool) = self.final_output_tool.lock().await.as_ref() { - let tool = final_output_tool.tool(); - let tool_content = Content::text(format!( - "Tool: {}\nDescription: {}\nSchema: {}", - tool.name, - tool.description.unwrap_or_default(), - serde_json::to_string_pretty(&tool.input_schema).unwrap_or_default() - )); - selected_tools.push(tool_content); - } - ToolCallResult::from(Ok(selected_tools)) } else { // Clone the result to ensure no references to extension_manager are returned @@ -768,7 +745,6 @@ impl Agent { &self, strategy: Option, ) -> Vec { - // If router is disabled for recipe execution, return empty list if *self.router_disabled_override.lock().await { return vec![]; } @@ -865,7 +841,7 @@ impl Agent { mut system_prompt, goose_mode, initial_messages, - max_turns, + config, } = context; let reply_span = tracing::Span::current(); @@ -882,6 +858,12 @@ impl Agent { Ok(Box::pin(async_stream::try_stream! { let _ = reply_span.enter(); let mut turns_taken = 0u32; + let max_turns = session + .as_ref() + .and_then(|s| s.max_turns) + .unwrap_or_else(|| { + config.get_param("GOOSE_MAX_TURNS").unwrap_or(DEFAULT_MAX_TURNS) + }); loop { if is_token_cancelled(&cancel_token) { @@ -1168,7 +1150,6 @@ impl Agent { provider: Option>, reindex_all: Option, ) -> Result<()> { - // Check if router is disabled for recipe execution if *self.router_disabled_override.lock().await { return Ok(()); } diff --git a/crates/goose/src/agents/reply_parts.rs b/crates/goose/src/agents/reply_parts.rs index 7b52d8b60dc7..9ef7a18f93d0 100644 --- a/crates/goose/src/agents/reply_parts.rs +++ b/crates/goose/src/agents/reply_parts.rs @@ -48,7 +48,6 @@ impl Agent { }; // Get tools from extension manager - // Check if router is disabled for recipe execution first let mut tools = if *self.router_disabled_override.lock().await { // If router is disabled, use regular tools self.list_tools(None).await