Add CPU resource limits support using Wasmtime fuel API

CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 
 ### Added
 
+- CPU resource limits support for WebAssembly components using Wasmtime's fuel API
 - Comprehensive Docker documentation and Dockerfile for running Wassette in containers with enhanced security isolation, including examples for mounting components, secrets, configuration files, and production deployment patterns with Docker Compose
 - `rust-toolchain.toml` file specifying Rust 1.90 as the stable toolchain version, ensuring consistent Rust version across development environments and CI/CD pipelines
 - AI agent development guides (`AGENTS.md` and `Claude.md`) that consolidate development guidelines from `.github/instructions/` into accessible documentation for AI agents working on the project

crates/mcp-server/src/tools.rs

@@ -723,6 +723,60 @@ pub async fn handle_grant_memory_permission(
     }
 }
 
+pub async fn handle_grant_cpu_permission(
+    req: &CallToolRequestParam,
+    lifecycle_manager: &LifecycleManager,
+) -> Result<CallToolResult> {
+    let args = extract_args_from_request(req)?;
+
+    let component_id = args
+        .get("component_id")
+        .and_then(|v| v.as_str())
+        .ok_or_else(|| anyhow::anyhow!("Missing required argument: 'component_id'"))?;
+
+    let details = args
+        .get("details")
+        .ok_or_else(|| anyhow::anyhow!("Missing required argument: 'details'"))?;
+
+    info!("Granting CPU permission to component {}", component_id);
+
+    lifecycle_manager
+        .ensure_component_loaded(component_id)
+        .await
+        .map_err(|e| anyhow::anyhow!("Component not found: {} ({})", component_id, e))?;
+
+    let result = lifecycle_manager
+        .grant_permission(component_id, "resource", details)
+        .await;
+
+    match result {
+        Ok(()) => {
+            let status_text = serde_json::to_string(&json!({
+                "status": "permission granted successfully",
+                "component_id": component_id,
+                "permission_type": "cpu",
+                "details": details
+            }))?;
+
+            let contents = vec![Content::text(status_text)];
+
+            Ok(CallToolResult {
+                content: Some(contents),
+                structured_content: None,
+                is_error: None,
+            })
+        }
+        Err(e) => {
+            error!("Failed to grant CPU permission: {}", e);
+            Err(anyhow::anyhow!(
+                "Failed to grant CPU permission to component {}: {}",
+                component_id,
+                e
+            ))
+        }
+    }
+}
+
 #[instrument(skip(lifecycle_manager))]
 pub async fn handle_revoke_storage_permission(
     req: &CallToolRequestParam,

crates/wassette/src/lib.rs

@@ -935,7 +935,11 @@ impl LifecycleManager {
     async fn get_wasi_state_for_component(
         &self,
         component_id: &str,
-    ) -> Result<(WassetteWasiState<WasiState>, Option<CustomResourceLimiter>)> {
+    ) -> Result<(
+        WassetteWasiState<WasiState>,
+        Option<CustomResourceLimiter>,
+        Option<f64>,
+    )> {
         let policy_template = self
             .policy_manager
             .template_for_component(component_id)
@@ -944,9 +948,10 @@ impl LifecycleManager {
         let wasi_state = policy_template.build()?;
         let allowed_hosts = policy_template.allowed_hosts.clone();
         let resource_limiter = wasi_state.resource_limiter.clone();
+        let cpu_limit = policy_template.cpu_limit;
 
         let wassette_wasi_state = WassetteWasiState::new(wasi_state, allowed_hosts)?;
-        Ok((wassette_wasi_state, resource_limiter))
+        Ok((wassette_wasi_state, resource_limiter, cpu_limit))
     }
 
     /// Executes a function call on a WebAssembly component
@@ -962,7 +967,8 @@ impl LifecycleManager {
             .await
             .ok_or_else(|| anyhow!("Component not found: {}", component_id))?;
 
-        let (state, resource_limiter) = self.get_wasi_state_for_component(component_id).await?;
+        let (state, resource_limiter, cpu_limit) =
+            self.get_wasi_state_for_component(component_id).await?;
 
         let mut store = Store::new(self.runtime.as_ref(), state);
 
@@ -979,6 +985,16 @@ impl LifecycleManager {
             });
         }
 
+        // Apply CPU limits if configured in the policy using fuel
+        // Convert CPU cores to fuel units (1 core ≈ 10 billion instructions)
+        if let Some(cpu_cores) = cpu_limit {
+            let fuel = (cpu_cores * 10_000_000_000.0) as u64;
+            store.set_fuel(fuel)?;
+        } else {
+            // Set unlimited fuel when no CPU limit is configured
+            store.set_fuel(u64::MAX)?;
+        }
+
         let instance = component.instance_pre.instantiate_async(&mut store).await?;
 
         // Use the new function identifier lookup instead of dot-splitting

crates/wassette/src/policy_internal.rs

@@ -426,28 +426,50 @@
             }
             "resource" => {
                 // Handle both direct memory field and nested resources.limits.memory structure
-                let memory = if let Some(memory_str) =
-                    details.get("memory").and_then(|v| v.as_str())
-                {
-                    // Direct memory field (for backward compatibility or direct API calls)
-                    memory_str
-                } else if let Some(memory_str) = details
+                let memory =
+                    if let Some(memory_str) = details.get("memory").and_then(|v| v.as_str()) {
+                        // Direct memory field (for backward compatibility or direct API calls)
+                        Some(memory_str)
+                    } else if let Some(memory_str) = details
+                        .get("resources")
+                        .and_then(|r| r.get("limits"))
+                        .and_then(|l| l.get("memory"))
+                        .and_then(|m| m.as_str())
+                    {
+                        // Nested structure from CLI (resources.limits.memory)
+                        Some(memory_str)
+                    } else {
+                        None
+                    };
+
+                // Handle CPU limits
+                let cpu = if let Some(cpu_str) = details.get("cpu").and_then(|v| v.as_str()) {
+                    // Direct cpu field
+                    Some(cpu_str)
+                } else if let Some(cpu_str) = details
                     .get("resources")
                     .and_then(|r| r.get("limits"))
-                    .and_then(|l| l.get("memory"))
-                    .and_then(|m| m.as_str())
+                    .and_then(|l| l.get("cpu"))
+                    .and_then(|c| c.as_str())
                 {
-                    // Nested structure from CLI (resources.limits.memory)
-                    memory_str
+                    // Nested structure from CLI (resources.limits.cpu)
+                    Some(cpu_str)
                 } else {
-                    return Err(anyhow!("Missing 'memory' field for resource permission. Expected either 'memory' or 'resources.limits.memory'"));
+                    None
                 };
 
+                // At least one of memory or cpu must be provided
+                if memory.is_none() && cpu.is_none() {
+                    return Err(anyhow!(
+                        "Missing resource field. Expected 'memory', 'cpu', or both"
+                    ));
+                }
+
                 // Create structured resource limits instead of hardcoded JSON
                 let resource_limits = policy::ResourceLimits {
                     limits: Some(policy::ResourceLimitValues::new(
-                        None,
-                        Some(policy::MemoryLimit::String(memory.to_string())),
+                        cpu.map(|c| policy::CpuLimit::String(c.to_string())),
+                        memory.map(|m| policy::MemoryLimit::String(m.to_string())),
                     )),
                     ..Default::default()
                 };
@@ -620,19 +642,44 @@
             .and_then(|m| m.as_str())
         {
             // Original CLI format: {"resources": {"limits": {"memory": "512Mi"}}}
-            memory_str
+            Some(memory_str)
         } else if let Some(memory_str) = details
             .get("limits")
             .and_then(|l| l.get("memory"))
             .and_then(|m| m.as_str())
         {
             // Converted ResourceLimits format: {"limits": {"memory": "512Mi"}}
-            memory_str
+            Some(memory_str)
+        } else {
+            None
+        };
+
+        // Extract the CPU limit from the details - handle both formats
+        let cpu_str = if let Some(cpu_str) = details
+            .get("resources")
+            .and_then(|r| r.get("limits"))
+            .and_then(|l| l.get("cpu"))
+            .and_then(|c| c.as_str())
+        {
+            // Original CLI format: {"resources": {"limits": {"cpu": "500m"}}}
+            Some(cpu_str)
+        } else if let Some(cpu_str) = details
+            .get("limits")
+            .and_then(|l| l.get("cpu"))
+            .and_then(|c| c.as_str())
+        {
+            // Converted ResourceLimits format: {"limits": {"cpu": "500m"}}
+            Some(cpu_str)
         } else {
+            None
+        };
+
+        // At least one of memory or cpu must be provided
+        if memory_str.is_none() && cpu_str.is_none() {
             return Err(anyhow!(
-                "Invalid resource permission format: missing memory field"
+                "Invalid resource permission format: missing memory or cpu field"
             ));
-        };
+        }
 
         // Initialize resources if not present
         let resources = policy
@@ -645,8 +692,15 @@
             .limits
             .get_or_insert_with(|| policy::ResourceLimitValues::new(None, None));
 
-        // Set the memory limit
-        limits.memory = Some(policy::MemoryLimit::String(memory_str.to_string()));
+        // Set the memory limit if provided
+        if let Some(memory) = memory_str {
+            limits.memory = Some(policy::MemoryLimit::String(memory.to_string()));
+        }
+
+        // Set the CPU limit if provided
+        if let Some(cpu) = cpu_str {
+            limits.cpu = Some(policy::CpuLimit::String(cpu.to_string()));
+        }
 
         Ok(())
     }
@@ -1389,6 +1443,101 @@
         Ok(())
     }
 
+    #[test]
+    fn test_add_cpu_resource_permission_to_policy() -> Result<()> {
+        let manager_result = tokio::runtime::Runtime::new()
+            .unwrap()
+            .block_on(async { create_test_manager().await });
+        let manager = manager_result.unwrap();
+
+        // Create a minimal policy
+        let mut policy = policy::PolicyDocument {
+            version: "1.0".to_string(),
+            description: Some("Test policy".to_string()),
+            permissions: policy::Permissions::default(),
+        };
+
+        // Test adding CPU resource permission
+        let resource_details = serde_json::json!({
+            "resources": {
+                "limits": {
+                    "cpu": "500m"
+                }
+            }
+        });
+
+        manager
+            .manager
+            .policy_manager
+            .add_resource_permission_to_policy(&mut policy, resource_details)?;
+
+        // Verify the policy has the CPU resource permission
+        let resources = policy.permissions.resources.expect("Should have resources");
+        let limits = resources.limits.expect("Should have limits");
+        let cpu = limits.cpu.expect("Should have CPU limit");
+
+        match cpu {
+            policy::CpuLimit::String(cpu_str) => {
+                assert_eq!(cpu_str, "500m");
+            }
+            _ => panic!("Expected string CPU limit"),
+        }
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_add_cpu_and_memory_resource_permissions_to_policy() -> Result<()> {
+        let manager_result = tokio::runtime::Runtime::new()
+            .unwrap()
+            .block_on(async { create_test_manager().await });
+        let manager = manager_result.unwrap();
+
+        // Create a minimal policy
+        let mut policy = policy::PolicyDocument {
+            version: "1.0".to_string(),
+            description: Some("Test policy".to_string()),
+            permissions: policy::Permissions::default(),
+        };
+
+        // Test adding both CPU and memory resource permissions at once
+        let resource_details = serde_json::json!({
+            "resources": {
+                "limits": {
+                    "cpu": "2",
+                    "memory": "1Gi"
+                }
+            }
+        });
+
+        manager
+            .manager
+            .policy_manager
+            .add_resource_permission_to_policy(&mut policy, resource_details)?;
+
+        // Verify the policy has both resource permissions
+        let resources = policy.permissions.resources.expect("Should have resources");
+        let limits = resources.limits.expect("Should have limits");
+
+        let cpu = limits.cpu.as_ref().expect("Should have CPU limit");
+        match cpu {
+            policy::CpuLimit::String(cpu_str) => {
+                assert_eq!(cpu_str, "2");
+            }
+            _ => panic!("Expected string CPU limit"),
+        }
+
+        let memory = limits.memory.as_ref().expect("Should have memory limit");
+        match memory {
+            policy::MemoryLimit::String(mem_str) => {
+                assert_eq!(mem_str, "1Gi");
+            }
+            _ => panic!("Expected string memory limit"),
+        }
+
+        Ok(())
+    }
+
     #[test]
     fn test_access_type_serialization() -> Result<()> {
         // Test serialization of AccessType

crates/wassette/src/runtime_context.rs

@@ -26,6 +26,8 @@ impl RuntimeContext {
         let mut config = wasmtime::Config::new();
         config.wasm_component_model(true);
         config.async_support(true);
+        // Enable fuel consumption for CPU limiting
+        config.consume_fuel(true);
 
         let engine = Arc::new(Engine::new(&config)?);