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#"{