Add inline python extension (#3107)

Co-authored-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: Michael Neale <michael.neale@gmail.com>

crates/goose/Cargo.toml

@@ -81,6 +81,7 @@ jsonwebtoken = "9.3.1"
 blake3 = "1.5"
 fs2 = "0.4.3"
 tokio-stream = "0.1.17"
+tempfile = "3.15.0"
 dashmap = "6.1"
 ahash = "0.8"
 tokio-util = "0.7.15"
@@ -94,7 +95,6 @@ winapi = { version = "0.3", features = ["wincred"] }
 
 [dev-dependencies]
 criterion = "0.5"
-tempfile = "3.15.0"
 serial_test = "3.2.0"
 mockall = "0.13.1"
 wiremock = "0.6.0"

crates/goose/src/agents/extension.rs

@@ -28,6 +28,8 @@ pub enum ExtensionError {
     SetupError(String),
     #[error("Join error occurred during task execution: {0}")]
     TaskJoinError(#[from] tokio::task::JoinError),
+    #[error("IO error: {0}")]
+    IoError(#[from] std::io::Error),
 }
 
 pub type ExtensionResult<T> = Result<T, ExtensionError>;
@@ -202,6 +204,21 @@ pub enum ExtensionConfig {
         #[serde(default)]
         bundled: Option<bool>,
     },
+    /// Inline Python code that will be executed using uvx
+    #[serde(rename = "inline_python")]
+    InlinePython {
+        /// The name used to identify this extension
+        name: String,
+        /// The Python code to execute
+        code: String,
+        /// Description of what the extension does
+        description: Option<String>,
+        /// Timeout in seconds
+        timeout: Option<u64>,
+        /// Python package dependencies required by this extension
+        #[serde(default)]
+        dependencies: Option<Vec<String>>,
+    },
 }
 
 impl Default for ExtensionConfig {
@@ -265,6 +282,21 @@ impl ExtensionConfig {
         }
     }
 
+    pub fn inline_python<S: Into<String>, T: Into<u64>>(
+        name: S,
+        code: S,
+        description: S,
+        timeout: T,
+    ) -> Self {
+        Self::InlinePython {
+            name: name.into(),
+            code: code.into(),
+            description: Some(description.into()),
+            timeout: Some(timeout.into()),
+            dependencies: None,
+        }
+    }
+
     pub fn with_args<I, S>(self, args: I) -> Self
     where
         I: IntoIterator<Item = S>,
@@ -307,6 +339,7 @@ impl ExtensionConfig {
             Self::Stdio { name, .. } => name,
             Self::Builtin { name, .. } => name,
             Self::Frontend { name, .. } => name,
+            Self::InlinePython { name, .. } => name,
         }
         .to_string()
     }
@@ -328,6 +361,9 @@ impl std::fmt::Display for ExtensionConfig {
             ExtensionConfig::Frontend { name, tools, .. } => {
                 write!(f, "Frontend({}: {} tools)", name, tools.len())
             }
+            ExtensionConfig::InlinePython { name, code, .. } => {
+                write!(f, "InlinePython({}: {} chars)", name, code.len())
+            }
         }
     }
 }

crates/goose/src/agents/extension_manager.rs

@@ -7,6 +7,7 @@ use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 use std::sync::LazyLock;
 use std::time::Duration;
+use tempfile::tempdir;
 use tokio::sync::Mutex;
 use tokio::task;
 use tokio_stream::wrappers::ReceiverStream;
@@ -35,6 +36,7 @@ pub struct ExtensionManager {
     clients: HashMap<String, McpClientBox>,
     instructions: HashMap<String, String>,
     resource_capable_extensions: HashSet<String>,
+    temp_dirs: HashMap<String, tempfile::TempDir>,
 }
 
 /// A flattened representation of a resource used by the agent to prepare inference
@@ -105,6 +107,7 @@ impl ExtensionManager {
             clients: HashMap::new(),
             instructions: HashMap::new(),
             resource_capable_extensions: HashSet::new(),
+            temp_dirs: HashMap::new(),
         }
     }
 
@@ -267,6 +270,49 @@ impl ExtensionManager {
                     .await?,
                 )
             }
+            ExtensionConfig::InlinePython {
+                name,
+                code,
+                timeout,
+                dependencies,
+                ..
+            } => {
+                let temp_dir = tempdir()?;
+                let file_path = temp_dir.path().join(format!("{}.py", name));
+                std::fs::write(&file_path, code)?;
+
+                let mut args = vec![];
+
+                let mut all_deps = vec!["mcp".to_string()];
+
+                if let Some(deps) = dependencies.as_ref() {
+                    all_deps.extend(deps.iter().cloned());
+                }
+
+                for dep in all_deps {
+                    args.push("--with".to_string());
+                    args.push(dep);
+                }
+
+                args.push("python".to_string());
+                args.push(file_path.to_str().unwrap().to_string());
+
+                let transport = StdioTransport::new("uvx", args, HashMap::new());
+                let handle = transport.start().await?;
+                let client = Box::new(
+                    McpClient::connect(
+                        handle,
+                        Duration::from_secs(
+                            timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT),
+                        ),
+                    )
+                    .await?,
+                );
+
+                self.temp_dirs.insert(sanitized_name.clone(), temp_dir);
+
+                client
+            }
             _ => unreachable!(),
         };
 
@@ -317,6 +363,7 @@ impl ExtensionManager {
         self.clients.remove(&sanitized_name);
         self.instructions.remove(&sanitized_name);
         self.resource_capable_extensions.remove(&sanitized_name);
+        self.temp_dirs.remove(&sanitized_name);
         Ok(())
     }
 
@@ -777,8 +824,11 @@ impl ExtensionManager {
                     }
                     | ExtensionConfig::Stdio {
                         description, name, ..
+                    }
+                    | ExtensionConfig::InlinePython {
+                        description, name, ..
                     } => {
-                        // For SSE/StreamableHttp/Stdio, use description if available
+                        // For SSE/StreamableHttp/Stdio/InlinePython, use description if available
                         description
                             .as_ref()
                             .map(|s| s.to_string())

crates/goose/src/recipe/mod.rs

@@ -432,6 +432,7 @@ impl RecipeBuilder {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use std::fs;
 
     #[test]
     fn test_from_content_with_json() {
@@ -653,6 +654,52 @@ sub_recipes:
         assert_eq!(author.contact, Some("test@example.com".to_string()));
     }
 
+    #[test]
+    fn test_inline_python_extension() {
+        let content = r#"{
+            "version": "1.0.0",
+            "title": "Test Recipe",
+            "description": "A test recipe",
+            "instructions": "Test instructions",
+            "extensions": [
+                {
+                    "type": "inline_python",
+                    "name": "test_python",
+                    "code": "print('hello world')",
+                    "timeout": 300,
+                    "description": "Test python extension",
+                    "dependencies": ["numpy", "matplotlib"]
+                }
+            ]
+        }"#;
+
+        let recipe = Recipe::from_content(content).unwrap();
+
+        assert!(recipe.extensions.is_some());
+        let extensions = recipe.extensions.unwrap();
+        assert_eq!(extensions.len(), 1);
+
+        match &extensions[0] {
+            ExtensionConfig::InlinePython {
+                name,
+                code,
+                description,
+                timeout,
+                dependencies,
+            } => {
+                assert_eq!(name, "test_python");
+                assert_eq!(code, "print('hello world')");
+                assert_eq!(description.as_deref(), Some("Test python extension"));
+                assert_eq!(timeout, &Some(300));
+                assert!(dependencies.is_some());
+                let deps = dependencies.as_ref().unwrap();
+                assert!(deps.contains(&"numpy".to_string()));
+                assert!(deps.contains(&"matplotlib".to_string()));
+            }
+            _ => panic!("Expected InlinePython extension"),
+        }
+    }
+
     #[test]
     fn test_from_content_with_activities() {
         let content = r#"{

recipe-inline-python.yaml

@@ -0,0 +1,87 @@
+version: "1.0.0"
+title: "Word Counter"
+description: "A recipe that counts words in text using an inline Python extension"
+
+instructions: |
+  This recipe provides a simple word counting tool. You can:
+  1. Count words in a text string
+  2. Get statistics about the text (unique words, average word length)
+  3. Generate word clouds from the text
+
+parameters:
+  - key: text
+    input_type: string
+    requirement: required
+    description: "The text to analyze"
+
+extensions:
+  - type: inline_python
+    name: word_counter
+    code: |
+      from mcp.server.fastmcp import FastMCP
+      import json
+      from collections import Counter
+      import numpy as np
+      from wordcloud import WordCloud
+      import matplotlib.pyplot as plt
+      import base64
+      from io import BytesIO
+      
+      mcp = FastMCP("word_counter")
+      
+      @mcp.tool()
+      def count_words(text: str) -> str:
+          """Count the number of words in a text string and generate statistics"""
+          words = text.split()
+          word_count = len(words)
+          unique_words = len(set(words))
+          avg_length = sum(len(word) for word in words) / word_count if word_count > 0 else 0
+          
+          # Generate word cloud
+          wordcloud = WordCloud(width=800, height=400, background_color='white').generate(text)
+          
+          # Save word cloud to base64
+          plt.figure(figsize=(10, 5))
+          plt.imshow(wordcloud, interpolation='bilinear')
+          plt.axis('off')
+          
+          # Save to bytes
+          buf = BytesIO()
+          plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
+          buf.seek(0)
+          img_base64 = base64.b64encode(buf.getvalue()).decode()
+          plt.close()
+          
+          result = {
+              "total_words": word_count,
+              "unique_words": unique_words,
+              "average_word_length": round(avg_length, 2),
+              "most_common": dict(Counter(words).most_common(5)),
+              "wordcloud_base64": img_base64
+          }
+          
+          return json.dumps(result, indent=2)
+          
+      if __name__ == "__main__":
+          mcp.run()
+    timeout: 300
+    description: "Count words and provide text statistics with visualization"
+    dependencies:
+      - "mcp"
+      - "numpy"
+      - "matplotlib"
+      - "wordcloud"
+      - "beautifulsoup4"
+      - "html2text"
+      - "requests"
+
+prompt: |
+  You are a helpful assistant that can analyze text using word counting tools.
+  The user has provided the following text to analyze: {{ text }}
+  
+  Use the word_counter__count_words tool to analyze this text and provide insights about:
+  - Total word count
+  - Number of unique words
+  - Average word length
+  - Most commonly used words
+  - A word cloud visualization of the text

ui/desktop/openapi.json

@@ -1702,6 +1702,51 @@
                 ]
               }
             }
+          },
+          {
+            "type": "object",
+            "description": "Inline Python code that will be executed using uvx",
+            "required": [
+              "name",
+              "code",
+              "type"
+            ],
+            "properties": {
+              "code": {
+                "type": "string",
+                "description": "The Python code to execute"
+              },
+              "dependencies": {
+                "type": "array",
+                "items": {
+                  "type": "string"
+                },
+                "description": "Python package dependencies required by this extension",
+                "nullable": true
+              },
+              "description": {
+                "type": "string",
+                "description": "Description of what the extension does",
+                "nullable": true
+              },
+              "name": {
+                "type": "string",
+                "description": "The name used to identify this extension"
+              },
+              "timeout": {
+                "type": "integer",
+                "format": "int64",
+                "description": "Timeout in seconds",
+                "nullable": true,
+                "minimum": 0
+              },
+              "type": {
+                "type": "string",
+                "enum": [
+                  "inline_python"
+                ]
+              }
+            }
           }
         ],
         "description": "Represents the different types of MCP extensions that can be added to the manager",

ui/desktop/src/api/types.gen.ts

@@ -217,6 +217,28 @@ export type ExtensionConfig = {
      */
     tools: Array<Tool>;
     type: 'frontend';
+} | {
+    /**
+     * The Python code to execute
+     */
+    code: string;
+    /**
+     * Python package dependencies required by this extension
+     */
+    dependencies?: Array<string> | null;
+    /**
+     * Description of what the extension does
+     */
+    description?: string | null;
+    /**
+     * The name used to identify this extension
+     */
+    name: string;
+    /**
+     * Timeout in seconds
+     */
+    timeout?: number | null;
+    type: 'inline_python';
 };
 
 export type ExtensionEntry = ExtensionConfig & {

ui/desktop/src/components/settings/extensions/ExtensionsSection.tsx

@@ -55,8 +55,8 @@ export default function ExtensionsSection({
         if (a.type !== 'builtin' && b.type === 'builtin') return 1;
 
         // Then sort by bundled (handle null/undefined cases)
-        const aBundled = a.bundled === true;
-        const bBundled = b.bundled === true;
+        const aBundled = 'bundled' in a && a.bundled === true;
+        const bBundled = 'bundled' in b && b.bundled === true;
         if (aBundled && !bBundled) return -1;
         if (!aBundled && bBundled) return 1;
 

ui/desktop/src/components/settings/extensions/bundled-extensions.ts

@@ -43,7 +43,7 @@ export async function syncBundledExtensions(
       const existingExt = existingExtensions.find((ext) => nameToKey(ext.name) === bundledExt.id);
 
       // Skip if extension exists and is already marked as bundled
-      if (existingExt?.bundled) continue;
+      if (existingExt && 'bundled' in existingExt && existingExt.bundled) continue;
 
       // Create the config for this extension
       let extConfig: ExtensionConfig;

ui/desktop/src/components/settings/extensions/subcomponents/ExtensionItem.tsx

@@ -68,7 +68,8 @@ export default function ExtensionItem({
   // Over time we can take the first part of the conditional away as people have bundled: true in their config.yaml entries
 
   // allow configuration editing if extension is not a builtin/bundled extension AND isStatic = false
-  const editable = !(extension.type === 'builtin' || extension.bundled) && !isStatic;
+  const editable =
+    !(extension.type === 'builtin' || ('bundled' in extension && extension.bundled)) && !isStatic;
 
   return (
     <Card

ui/desktop/src/components/settings/extensions/subcomponents/ExtensionList.tsx

@@ -45,7 +45,12 @@ export function getFriendlyTitle(extension: FixedExtensionEntry): string {
   let name = '';
 
   // if it's a builtin, check if there's a display_name (old configs didn't have this field)
-  if (extension.bundled === true && 'display_name' in extension && extension.display_name) {
+  if (
+    'bundled' in extension &&
+    extension.bundled === true &&
+    'display_name' in extension &&
+    extension.display_name
+  ) {
     // If we have a display_name for a builtin, use it directly
     return extension.display_name;
   } else {

ui/desktop/src/components/settings/extensions/utils.ts

@@ -100,7 +100,10 @@ export function extensionToFormData(extension: FixedExtensionEntry): ExtensionFo
       extension.type === 'stdio' || extension.type === 'sse' || extension.type === 'streamable_http'
         ? extension.description || ''
         : '',
-    type: extension.type === 'frontend' ? 'stdio' : extension.type,
+    type:
+      extension.type === 'frontend' || extension.type === 'inline_python'
+        ? 'stdio'
+        : extension.type,
     cmd: extension.type === 'stdio' ? combineCmdAndArgs(extension.cmd, extension.args) : undefined,
     endpoint:
       extension.type === 'sse' || extension.type === 'streamable_http' ? extension.uri : undefined,

ui/desktop/src/components/settings/permission/PermissionRulesModal.tsx

@@ -66,8 +66,8 @@ export default function PermissionRulesModal({ isOpen, onClose }: PermissionRule
       if (a.type !== 'builtin' && b.type === 'builtin') return 1;
 
       // Then sort by bundled (handle null/undefined cases)
-      const aBundled = a.bundled === true;
-      const bBundled = b.bundled === true;
+      const aBundled = 'bundled' in a && a.bundled === true;
+      const bBundled = 'bundled' in b && b.bundled === true;
       if (aBundled && !bBundled) return -1;
       if (!aBundled && bBundled) return 1;
 

ui/desktop/src/components/settings/permission/PermissionSetting.tsx

@@ -62,8 +62,8 @@ export default function PermissionSettingsView({ onClose }: { onClose: () => voi
       if (a.type !== 'builtin' && b.type === 'builtin') return 1;
 
       // Then sort by bundled (handle null/undefined cases)
-      const aBundled = a.bundled === true;
-      const bBundled = b.bundled === true;
+      const aBundled = 'bundled' in a && a.bundled === true;
+      const bBundled = 'bundled' in b && b.bundled === true;
       if (aBundled && !bBundled) return -1;
       if (!aBundled && bBundled) return 1;