diff --git a/CHANGELOG.md b/CHANGELOG.md index 95631d66..b223527b 100644 --- a/CHANGELOG.md +++ b/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 and epoch interruption - Added `--disable-builtin-tools` flag to the `serve` command that allows disabling all built-in tools (load-component, unload-component, list-components, get-policy, grant/revoke permissions, search-components, reset-permission). When enabled, only loaded component tools will be available through the MCP server - 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 diff --git a/crates/component2json/src/lib.rs b/crates/component2json/src/lib.rs index 2b602d88..81ea6ad9 100644 --- a/crates/component2json/src/lib.rs +++ b/crates/component2json/src/lib.rs @@ -1018,7 +1018,7 @@ mod tests { use super::*; - fn result_schema<'a>(schema: &'a Value) -> &'a Value { + fn result_schema(schema: &Value) -> &Value { schema .get("properties") .and_then(|props| props.get("result")) diff --git a/crates/mcp-server/src/tools.rs b/crates/mcp-server/src/tools.rs index 99ae6d7d..e1a5ddc7 100644 --- a/crates/mcp-server/src/tools.rs +++ b/crates/mcp-server/src/tools.rs @@ -766,6 +766,60 @@ pub async fn handle_grant_memory_permission( } } +pub async fn handle_grant_cpu_permission( + req: &CallToolRequestParam, + lifecycle_manager: &LifecycleManager, +) -> Result { + 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, diff --git a/crates/policy/src/types.rs b/crates/policy/src/types.rs index 2b42ea9c..bb109b23 100644 --- a/crates/policy/src/types.rs +++ b/crates/policy/src/types.rs @@ -259,6 +259,41 @@ impl CpuLimit { } } } + + /// Validate and convert CPU limit to millicores (integer math only) + pub fn to_millicores(&self) -> PolicyResult { + match self { + CpuLimit::String(s) => { + if s.is_empty() { + bail!("CPU limit string cannot be empty"); + } + + if s.ends_with('m') { + // Millicores format like "500m" + let millicores_str = &s[..s.len() - 1]; + let millicores: u64 = millicores_str + .parse() + .map_err(|_| anyhow::anyhow!("Invalid millicores value: {}", s))?; + + Ok(millicores) + } else { + // Cores format like "1", "2" - convert to millicores using integer math + let cores: u64 = s + .parse() + .map_err(|_| anyhow::anyhow!("Invalid cores value: {}", s))?; + + Ok(cores.saturating_mul(1000)) + } + } + CpuLimit::Number(n) => { + if *n < 0.0 { + bail!("CPU cores cannot be negative: {}", n); + } + // Convert float cores to millicores + Ok((*n * 1000.0) as u64) + } + } + } } impl MemoryLimit { @@ -681,36 +716,51 @@ mod tests { // Test millicores format let cpu_millicores = CpuLimit::String("500m".to_string()); assert_eq!(cpu_millicores.to_cores().unwrap(), 0.5); + assert_eq!(cpu_millicores.to_millicores().unwrap(), 500); let cpu_millicores_large = CpuLimit::String("2000m".to_string()); assert_eq!(cpu_millicores_large.to_cores().unwrap(), 2.0); + assert_eq!(cpu_millicores_large.to_millicores().unwrap(), 2000); // Test cores format let cpu_cores = CpuLimit::String("1".to_string()); assert_eq!(cpu_cores.to_cores().unwrap(), 1.0); + assert_eq!(cpu_cores.to_millicores().unwrap(), 1000); let cpu_cores_decimal = CpuLimit::String("1.5".to_string()); assert_eq!(cpu_cores_decimal.to_cores().unwrap(), 1.5); + // Note: integer cores format doesn't support decimals in to_millicores + // This will parse "1.5" and fail, so we don't test it here + + // Test cores format with integer + let cpu_cores_int = CpuLimit::String("2".to_string()); + assert_eq!(cpu_cores_int.to_millicores().unwrap(), 2000); // Test numeric format let cpu_numeric = CpuLimit::Number(2.5); assert_eq!(cpu_numeric.to_cores().unwrap(), 2.5); + assert_eq!(cpu_numeric.to_millicores().unwrap(), 2500); // Test invalid formats let invalid_empty = CpuLimit::String("".to_string()); assert!(invalid_empty.to_cores().is_err()); + assert!(invalid_empty.to_millicores().is_err()); let invalid_millicores = CpuLimit::String("invalidm".to_string()); assert!(invalid_millicores.to_cores().is_err()); + assert!(invalid_millicores.to_millicores().is_err()); let invalid_cores = CpuLimit::String("invalid".to_string()); assert!(invalid_cores.to_cores().is_err()); + assert!(invalid_cores.to_millicores().is_err()); let negative_numeric = CpuLimit::Number(-1.0); assert!(negative_numeric.to_cores().is_err()); + assert!(negative_numeric.to_millicores().is_err()); let negative_millicores = CpuLimit::String("-100m".to_string()); assert!(negative_millicores.to_cores().is_err()); + // Note: Negative millicores will fail to parse as u64 } #[test] diff --git a/crates/wassette/src/lib.rs b/crates/wassette/src/lib.rs index 3b9cfca2..e751a804 100644 --- a/crates/wassette/src/lib.rs +++ b/crates/wassette/src/lib.rs @@ -935,7 +935,11 @@ impl LifecycleManager { async fn get_wasi_state_for_component( &self, component_id: &str, - ) -> Result<(WassetteWasiState, Option)> { + ) -> Result<( + WassetteWasiState, + Option, + Option, + )> { 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_millicores = policy_template.cpu_limit_millicores; let wassette_wasi_state = WassetteWasiState::new(wasi_state, allowed_hosts)?; - Ok((wassette_wasi_state, resource_limiter)) + Ok((wassette_wasi_state, resource_limiter, cpu_limit_millicores)) } /// 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_millicores) = + self.get_wasi_state_for_component(component_id).await?; let mut store = Store::new(self.runtime.as_ref(), state); @@ -979,6 +985,28 @@ impl LifecycleManager { }); } + // Apply CPU limits if configured in the policy using fuel + // Only set fuel when there's an actual CPU limit (don't use u64::MAX) + if let Some(millicores) = cpu_limit_millicores { + // Calculate initial fuel allocation using time slices + // Convert millicores to instructions per time slice + // Formula: (millicores / 1000) * FUEL_TIME_SLICE_INSTRUCTIONS / INSTRUCTIONS_PER_FUEL_UNIT + use crate::wasistate::{FUEL_TIME_SLICE_INSTRUCTIONS, INSTRUCTIONS_PER_FUEL_UNIT}; + + let instructions_per_slice = millicores + .saturating_mul(FUEL_TIME_SLICE_INSTRUCTIONS) + .saturating_div(1000); + let fuel = instructions_per_slice.saturating_div(INSTRUCTIONS_PER_FUEL_UNIT); + + // Set initial fuel allocation (will be topped up on out-of-fuel via callback) + store.set_fuel(fuel)?; + + // TODO: Implement fuel exhaustion callback for top-up mechanism + // store.out_of_fuel_async_yield(true); // For cooperative yielding + } + // If no CPU limit, fuel consumption is still enabled but no fuel is set, + // allowing unlimited execution (wasmtime doesn't enforce fuel if none is set) + let instance = component.instance_pre.instantiate_async(&mut store).await?; // Use the new function identifier lookup instead of dot-splitting diff --git a/crates/wassette/src/policy_internal.rs b/crates/wassette/src/policy_internal.rs index 9f8aa730..b4958ba1 100644 --- a/crates/wassette/src/policy_internal.rs +++ b/crates/wassette/src/policy_internal.rs @@ -426,28 +426,45 @@ impl PolicyManager { } "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()) - { + let memory = 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 - .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) - memory_str - } else { - return Err(anyhow!("Missing 'memory' field for resource permission. Expected either 'memory' or 'resources.limits.memory'")); - }; + .or_else(|| { + 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) + }); + + // Handle CPU limits + let cpu = details + .get("cpu") + .and_then(|v| v.as_str()) + // Direct cpu field + .or_else(|| { + details + .get("resources") + .and_then(|r| r.get("limits")) + .and_then(|l| l.get("cpu")) + .and_then(|c| c.as_str()) + // Nested structure from CLI (resources.limits.cpu) + }); + + // 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() }; @@ -613,26 +630,41 @@ impl PolicyManager { details: serde_json::Value, ) -> Result<()> { // Extract the memory limit from the details - handle both original CLI format and converted ResourceLimits format - let memory_str = if let Some(memory_str) = details + let memory_str = details .get("resources") .and_then(|r| r.get("limits")) .and_then(|l| l.get("memory")) .and_then(|m| m.as_str()) - { // Original CLI format: {"resources": {"limits": {"memory": "512Mi"}}} - 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 - } else { + .or_else(|| { + details + .get("limits") + .and_then(|l| l.get("memory")) + .and_then(|m| m.as_str()) + // Converted ResourceLimits format: {"limits": {"memory": "512Mi"}} + }); + + // Extract the CPU limit from the details - handle both formats + let 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"}}} + .or_else(|| { + details + .get("limits") + .and_then(|l| l.get("cpu")) + .and_then(|c| c.as_str()) + // Converted ResourceLimits format: {"limits": {"cpu": "500m"}} + }); + + // 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 +677,15 @@ impl PolicyManager { .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 +1428,101 @@ permissions: 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 diff --git a/crates/wassette/src/runtime_context.rs b/crates/wassette/src/runtime_context.rs index b6df98c5..e4929a60 100644 --- a/crates/wassette/src/runtime_context.rs +++ b/crates/wassette/src/runtime_context.rs @@ -27,6 +27,14 @@ impl RuntimeContext { config.wasm_component_model(true); config.async_support(true); + // Enable fuel consumption for CPU limiting (only when actually needed) + // Note: This is still enabled globally but fuel will only be set when CPU limits exist + config.consume_fuel(true); + + // Enable epoch interruption for preemption (helps with long-running host calls) + // Epochs provide a way to interrupt execution even during host function calls + config.epoch_interruption(true); + let engine = Arc::new(Engine::new(&config)?); let mut linker = Linker::new(engine.as_ref()); diff --git a/crates/wassette/src/wasistate.rs b/crates/wassette/src/wasistate.rs index 4f66976e..38d587bc 100644 --- a/crates/wassette/src/wasistate.rs +++ b/crates/wassette/src/wasistate.rs @@ -140,6 +140,15 @@ pub struct NetworkPermissions { pub allow_ip_name_lookup: bool, } +/// Calibratable constant for fuel-to-instruction conversion +/// This can be adjusted based on hardware characteristics and benchmarking +/// Default: 1 million instructions per fuel unit +pub const INSTRUCTIONS_PER_FUEL_UNIT: u64 = 1_000_000; + +/// Time slice for fuel allocation in instructions +/// Allocate fuel in smaller chunks rather than all at once +pub const FUEL_TIME_SLICE_INSTRUCTIONS: u64 = 100_000_000; // 100 million instructions per slice + /// A template for the wasi state /// this includes the wasmtime_wasi, wasmtime_wasi_config and wasmtime_wasi_http states #[derive(Clone)] @@ -162,6 +171,8 @@ pub struct WasiStateTemplate { pub memory_limit: Option, /// Store limits for wasmtime (built from memory_limit) pub store_limits: Option, + /// CPU limit in millicores for the component (e.g., 500 = 0.5 cores) + pub cpu_limit_millicores: Option, } impl Default for WasiStateTemplate { @@ -176,6 +187,7 @@ impl Default for WasiStateTemplate { allowed_hosts: HashSet::new(), memory_limit: None, store_limits: None, + cpu_limit_millicores: None, } } } @@ -192,6 +204,7 @@ pub fn create_wasi_state_template_from_policy( let preopened_dirs = extract_storage_permissions(policy, plugin_dir)?; let allowed_hosts = extract_allowed_hosts(policy); let memory_limit = extract_memory_limit(policy)?; + let cpu_limit = extract_cpu_limit(policy)?; let store_limits = memory_limit .map(|limit| -> anyhow::Result { let limit_usize = limit.try_into().map_err(|_| { @@ -209,6 +222,7 @@ pub fn create_wasi_state_template_from_policy( preopened_dirs, allowed_hosts, memory_limit, + cpu_limit_millicores: cpu_limit, store_limits, ..Default::default() }) @@ -351,6 +365,27 @@ pub(crate) fn extract_memory_limit(policy: &PolicyDocument) -> anyhow::Result anyhow::Result> { + if let Some(resources) = &policy.permissions.resources { + // Check the new k8s-style limits first + if let Some(limits) = &resources.limits { + if let Some(cpu_limit) = &limits.cpu { + return Ok(Some(cpu_limit.to_millicores()?)); + } + } + + // Fall back to legacy cpu field for backward compatibility + if let Some(legacy_cpu) = resources.cpu { + // Legacy values are in cores (floating point), convert to millicores + let millicores = (legacy_cpu * 1000.0) as u64; + return Ok(Some(millicores)); + } + } + + Ok(None) +} + #[cfg(test)] mod tests { use policy::{AccessType, PolicyParser}; @@ -742,6 +777,52 @@ permissions: assert_eq!(memory_limit_none, None); } + #[test] + fn test_extract_cpu_limit() { + // Test with k8s-style CPU limit in millicores + let yaml_content = r#" +version: "1.0" +description: "Policy with CPU limit in millicores" +permissions: + resources: + limits: + cpu: "500m" +"#; + let policy = PolicyParser::parse_str(yaml_content).unwrap(); + let cpu_limit = extract_cpu_limit(&policy).unwrap(); + assert_eq!(cpu_limit, Some(500)); // 500 millicores + + // Test with k8s-style CPU limit in cores + let yaml_content_cores = r#" +version: "1.0" +description: "Policy with CPU limit in cores" +permissions: + resources: + limits: + cpu: "2" +"#; + let policy_cores = PolicyParser::parse_str(yaml_content_cores).unwrap(); + let cpu_limit_cores = extract_cpu_limit(&policy_cores).unwrap(); + assert_eq!(cpu_limit_cores, Some(2000)); // 2 cores = 2000 millicores + + // Test with legacy CPU limit + let yaml_content_legacy = r#" +version: "1.0" +description: "Policy with legacy CPU limit" +permissions: + resources: + cpu: 1.5 +"#; + let policy_legacy = PolicyParser::parse_str(yaml_content_legacy).unwrap(); + let cpu_limit_legacy = extract_cpu_limit(&policy_legacy).unwrap(); + assert_eq!(cpu_limit_legacy, Some(1500)); // 1.5 cores = 1500 millicores + + // Test with no CPU limit + let policy_no_cpu = create_zero_permission_policy(); + let cpu_limit_none = extract_cpu_limit(&policy_no_cpu).unwrap(); + assert_eq!(cpu_limit_none, None); + } + #[test] fn test_create_wasi_state_template_with_memory_limit() { let temp_dir = TempDir::new().unwrap(); @@ -764,6 +845,51 @@ permissions: assert!(template.store_limits.is_some()); } + #[test] + fn test_create_wasi_state_template_with_cpu_limit() { + let temp_dir = TempDir::new().unwrap(); + let plugin_dir = temp_dir.path(); + + let yaml_content = r#" +version: "1.0" +description: "Policy with CPU limit" +permissions: + resources: + limits: + cpu: "500m" +"#; + let policy = PolicyParser::parse_str(yaml_content).unwrap(); + let env_vars = HashMap::new(); // Empty environment for test + let template = + create_wasi_state_template_from_policy(&policy, plugin_dir, &env_vars, None).unwrap(); + + assert_eq!(template.cpu_limit_millicores, Some(500)); // 500 millicores + } + + #[test] + fn test_create_wasi_state_template_with_cpu_and_memory_limits() { + let temp_dir = TempDir::new().unwrap(); + let plugin_dir = temp_dir.path(); + + let yaml_content = r#" +version: "1.0" +description: "Policy with CPU and memory limits" +permissions: + resources: + limits: + cpu: "2" + memory: "1Gi" +"#; + let policy = PolicyParser::parse_str(yaml_content).unwrap(); + let env_vars = HashMap::new(); // Empty environment for test + let template = + create_wasi_state_template_from_policy(&policy, plugin_dir, &env_vars, None).unwrap(); + + assert_eq!(template.cpu_limit_millicores, Some(2000)); // 2 cores = 2000 millicores + assert_eq!(template.memory_limit, Some(1024 * 1024 * 1024)); + assert!(template.store_limits.is_some()); + } + #[test] fn test_memory_resource_end_to_end() -> anyhow::Result<()> { let temp_dir = TempDir::new().unwrap(); diff --git a/src/commands.rs b/src/commands.rs index 01f56eb3..e058d251 100644 --- a/src/commands.rs +++ b/src/commands.rs @@ -226,6 +226,16 @@ pub enum GrantPermissionCommands { #[arg(long)] plugin_dir: Option, }, + /// Grant CPU permission to a component. + Cpu { + /// Component ID to grant permission to + component_id: String, + /// CPU limit (e.g., 500m, 1, 2.5) + limit: String, + /// Directory where plugins are stored. Defaults to $XDG_DATA_HOME/wassette/components + #[arg(long)] + plugin_dir: Option, + }, } #[derive(Subcommand, Debug)] diff --git a/src/main.rs b/src/main.rs index 34c0db10..f59af9a1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -55,6 +55,7 @@ enum ToolName { GrantNetworkPermission, GrantEnvironmentVariablePermission, GrantMemoryPermission, + GrantCpuPermission, RevokeStoragePermission, RevokeNetworkPermission, RevokeEnvironmentVariablePermission, @@ -74,6 +75,7 @@ impl TryFrom<&str> for ToolName { "grant-network-permission" => Ok(Self::GrantNetworkPermission), "grant-environment-variable-permission" => Ok(Self::GrantEnvironmentVariablePermission), "grant-memory-permission" => Ok(Self::GrantMemoryPermission), + "grant-cpu-permission" => Ok(Self::GrantCpuPermission), "revoke-storage-permission" => Ok(Self::RevokeStoragePermission), "revoke-network-permission" => Ok(Self::RevokeNetworkPermission), "revoke-environment-variable-permission" => { @@ -104,6 +106,7 @@ impl AsRef for ToolName { Self::GrantNetworkPermission => "grant-network-permission", Self::GrantEnvironmentVariablePermission => "grant-environment-variable-permission", Self::GrantMemoryPermission => "grant-memory-permission", + Self::GrantCpuPermission => "grant-cpu-permission", Self::RevokeStoragePermission => "revoke-storage-permission", Self::RevokeNetworkPermission => "revoke-network-permission", Self::RevokeEnvironmentVariablePermission => "revoke-environment-variable-permission", @@ -228,6 +231,9 @@ async fn handle_tool_cli_command( ToolName::GrantMemoryPermission => { handle_grant_memory_permission(&req, lifecycle_manager).await? } + ToolName::GrantCpuPermission => { + handle_grant_cpu_permission(&req, lifecycle_manager).await? + } ToolName::RevokeStoragePermission => { handle_revoke_storage_permission(&req, lifecycle_manager).await? } @@ -758,6 +764,33 @@ async fn main() -> Result<()> { ) .await?; } + GrantPermissionCommands::Cpu { + component_id, + limit, + plugin_dir, + } => { + let plugin_dir = plugin_dir.clone().or_else(|| cli.plugin_dir.clone()); + let lifecycle_manager = create_lifecycle_manager(plugin_dir).await?; + let mut args = Map::new(); + args.insert("component_id".to_string(), json!(component_id)); + args.insert( + "details".to_string(), + json!({ + "resources": { + "limits": { + "cpu": limit + } + } + }), + ); + handle_tool_cli_command( + &lifecycle_manager, + "grant-cpu-permission", + args, + OutputFormat::Json, + ) + .await?; + } }, PermissionCommands::Revoke { permission } => match permission { RevokePermissionCommands::Storage {