From b7e63ca0bd57604a88e111e43b5b64e9e463bcdc Mon Sep 17 00:00:00 2001
From: rabi <ramishra@redhat.com>
Date: Tue, 6 Jan 2026 13:52:00 +0530
Subject: [PATCH] fix(code_execution): handle model quirks with tool calls

- Auto-prefix unprefixed tool names (read_module -> code_execution__read_module)
- Parse stringified JSON arrays in search_modules terms parameter
- Add note that all tool calls are synchronous (no async/await)
- Improve search_modules documentation to clarify array parameter format

Fixes issues where some models (e.g. Nemotron) strip tool prefixes or
stringify array parameters.

Signed-off-by: rabi <ramishra@redhat.com>
---
 .../src/agents/code_execution_extension.rs    | 20 ++++++++++++------
 crates/goose/src/agents/extension_manager.rs  | 21 +++++++++++++++----
 2 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/crates/goose/src/agents/code_execution_extension.rs b/crates/goose/src/agents/code_execution_extension.rs
index 2d66599201ff..664169e7bdce 100644
--- a/crates/goose/src/agents/code_execution_extension.rs
+++ b/crates/goose/src/agents/code_execution_extension.rs
@@ -466,6 +466,8 @@ impl CodeExecutionClient {
                 - WRONG: Multiple execute_code calls, each with one tool
                 - RIGHT: One execute_code call with a script that calls all needed tools
 
+                IMPORTANT: All tool calls are SYNCHRONOUS. Do NOT use async/await.
+
                 Workflow:
                     1. Use the read_module tool to discover tools and signatures
                     2. Write ONE script that imports and calls ALL tools needed for the task
@@ -572,12 +574,16 @@ impl CodeExecutionClient {
             .and_then(|a| a.get("terms"))
             .ok_or("Missing required parameter: terms")?;
 
-        let terms_vec = if let Some(s) = terms.as_str() {
-            vec![s.to_string()]
-        } else if let Some(arr) = terms.as_array() {
+        let terms_vec = if let Some(arr) = terms.as_array() {
             arr.iter()
                 .filter_map(|v| v.as_str().map(String::from))
                 .collect()
+        } else if let Some(s) = terms.as_str() {
+            if s.starts_with('[') && s.ends_with(']') {
+                serde_json::from_str::<Vec<String>>(s).unwrap_or_else(|_| vec![s.to_string()])
+            } else {
+                vec![s.to_string()]
+            }
         } else {
             return Err("Parameter 'terms' must be a string or array of strings".to_string());
         };
@@ -830,9 +836,11 @@ impl McpClientTrait for CodeExecutionClient {
                         Search for tools by name or description across all available modules.
 
                         USAGE:
-                        - Single term: search_modules with terms="file"
-                        - Multiple terms: search_modules with terms=["git", "shell"]
-                        - Regex patterns: search_modules with terms="sh.*", regex=true
+                        - Single term: terms="github" (just a plain string)
+                        - Multiple terms: terms=["git", "shell"] (a JSON array, NOT a string)
+                        - Regex patterns: terms="sh.*", regex=true
+
+                        IMPORTANT: Do NOT stringify arrays. Use terms=["a","b"] not terms="[\"a\",\"b\"]"
 
                         Returns matching servers and tools with descriptions.
                         Use this when you don't know which module contains the tool you need.
diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs
index 206b873c89cb..2682b9da6f0e 100644
--- a/crates/goose/src/agents/extension_manager.rs
+++ b/crates/goose/src/agents/extension_manager.rs
@@ -1022,17 +1022,30 @@ impl ExtensionManager {
         tool_call: CallToolRequestParam,
         cancellation_token: CancellationToken,
     ) -> Result<ToolCallResult> {
+        // Some models strip the tool prefix, so auto-add it for known code_execution tools
+        let tool_name_str = tool_call.name.to_string();
+        let prefixed_name = if !tool_name_str.contains("__") {
+            let code_exec_tools = ["execute_code", "read_module", "search_modules"];
+            if code_exec_tools.contains(&tool_name_str.as_str())
+                && self.extensions.lock().await.contains_key("code_execution")
+            {
+                format!("code_execution__{}", tool_name_str)
+            } else {
+                tool_name_str
+            }
+        } else {
+            tool_name_str
+        };
+
         // Dispatch tool call based on the prefix naming convention
         let (client_name, client) =
-            self.get_client_for_tool(&tool_call.name)
+            self.get_client_for_tool(&prefixed_name)
                 .await
                 .ok_or_else(|| {
                     ErrorData::new(ErrorCode::RESOURCE_NOT_FOUND, tool_call.name.clone(), None)
                 })?;
 
-        // rsplit returns the iterator in reverse, tool_name is then at 0
-        let tool_name = tool_call
-            .name
+        let tool_name = prefixed_name
             .strip_prefix(client_name.as_str())
             .and_then(|s| s.strip_prefix("__"))
             .ok_or_else(|| {