feat(permissions): add environment variable permission handling and related tests

Signed-off-by: Jiaxiao Zhou <duibao55328@gmail.com>

Author: Mossaka

crates/mcp-server/src/tools.rs

@@ -50,6 +50,9 @@ pub async fn handle_tools_call(
         "grant-network-permission" => {
             handle_grant_network_permission(&req, lifecycle_manager).await
         }
+        "grant-environment-variable-permission" => {
+            handle_grant_environment_variable_permission(&req, lifecycle_manager).await
+        }
         _ => handle_component_call(&req, lifecycle_manager).await,
     };
 
@@ -214,6 +217,37 @@ fn get_builtin_tools() -> Vec<Tool> {
             ),
             annotations: None,
         },
+        Tool {
+            name: Cow::Borrowed("grant-environment-variable-permission"),
+            description: Some(Cow::Borrowed(
+                "Grants environment variable access permission to a component, allowing it to access specific environment variables."
+            )),
+            input_schema: Arc::new(
+                serde_json::from_value(json!({
+                    "type": "object",
+                    "properties": {
+                      "component_id": {
+                        "type": "string",
+                        "description": "ID of the component to grant environment variable permission to"
+                      },
+                      "details": {
+                        "type": "object",
+                        "properties": {
+                          "key": { 
+                            "type": "string",
+                            "description": "Environment variable key to grant access to"
+                          }
+                        },
+                        "required": ["key"],
+                        "additionalProperties": false
+                      }
+                    },
+                    "required": ["component_id", "details"]
+                  }))
+                .unwrap_or_default(),
+            ),
+            annotations: None,
+        },
     ]
 }
 
@@ -358,20 +392,75 @@ async fn handle_grant_network_permission(
     }
 }
 
+#[instrument(skip(lifecycle_manager))]
+async fn handle_grant_environment_variable_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 environment variable permission to component {}",
+        component_id
+    );
+
+    let result = lifecycle_manager
+        .grant_permission(component_id, "environment", details)
+        .await;
+
+    match result {
+        Ok(()) => {
+            let status_text = serde_json::to_string(&json!({
+                "status": "permission granted",
+                "component_id": component_id,
+                "permission_type": "environment",
+                "details": details
+            }))?;
+
+            let contents = vec![Content::text(status_text)];
+
+            Ok(CallToolResult {
+                content: contents,
+                is_error: None,
+            })
+        }
+        Err(e) => {
+            error!("Failed to grant environment variable permission: {}", e);
+            Err(anyhow::anyhow!(
+                "Failed to grant environment variable permission to component {}: {}",
+                component_id,
+                e
+            ))
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
     fn test_get_builtin_tools() {
         let tools = get_builtin_tools();
-        assert_eq!(tools.len(), 6);
+        assert_eq!(tools.len(), 7);
         assert!(tools.iter().any(|t| t.name == "load-component"));
         assert!(tools.iter().any(|t| t.name == "unload-component"));
         assert!(tools.iter().any(|t| t.name == "list-components"));
         assert!(tools.iter().any(|t| t.name == "get-policy"));
         assert!(tools.iter().any(|t| t.name == "grant-storage-permission"));
         assert!(tools.iter().any(|t| t.name == "grant-network-permission"));
+        assert!(tools
+            .iter()
+            .any(|t| t.name == "grant-environment-variable-permission"));
     }
 
     #[tokio::test]

crates/wassette/src/lib.rs

@@ -41,6 +41,9 @@ pub enum PermissionRule {
         uri: String,
         access: Vec<AccessType>,
     },
+    Environment {
+        key: String,
+    },
 }
 
 /// Access types for storage permissions
@@ -1027,6 +1030,15 @@ impl LifecycleManager {
                     access: access_types?,
                 })
             }
+            "environment" => {
+                let key = details
+                    .get("key")
+                    .and_then(|v| v.as_str())
+                    .ok_or_else(|| anyhow!("Missing 'key' field for environment permission"))?;
+                Ok(PermissionRule::Environment {
+                    key: key.to_string(),
+                })
+            }
             _ => Err(anyhow!("Unknown permission type: {}", permission_type)),
         }
     }
@@ -1151,6 +1163,29 @@ impl LifecycleManager {
                     }
                 }
             }
+            PermissionRule::Environment { key } => {
+                // For environment permissions, we need to create a struct with key field
+                let env_perms = policy
+                    .permissions
+                    .environment
+                    .get_or_insert_with(Default::default);
+                let allow_list = env_perms.allow.get_or_insert_with(Vec::new);
+
+                // Create a simple struct with the key field
+                let env_allow = serde_json::json!({ "key": key });
+                if let Ok(env_allow_struct) = serde_json::from_value(env_allow) {
+                    // Avoid duplicates by checking if key already exists
+                    if !allow_list.iter().any(|existing| {
+                        if let Ok(existing_json) = serde_json::to_value(existing) {
+                            existing_json.get("key").and_then(|k| k.as_str()) == Some(&key)
+                        } else {
+                            false
+                        }
+                    }) {
+                        allow_list.push(env_allow_struct);
+                    }
+                }
+            }
         }
         Ok(())
     }
@@ -1200,6 +1235,11 @@ impl LifecycleManager {
                     return Err(anyhow!("Storage access cannot be empty"));
                 }
             }
+            PermissionRule::Environment { key } => {
+                if key.is_empty() {
+                    return Err(anyhow!("Environment variable key cannot be empty"));
+                }
+            }
         }
         Ok(())
     }

tests/grant_permission_integration_test.rs

@@ -635,3 +635,115 @@ async fn test_grant_permission_sequential_grants() -> Result<()> {
 
     Ok(())
 }
+
+#[cfg(any(target_os = "linux", target_os = "macos"))]
+#[test(tokio::test)]
+async fn test_grant_permission_environment_variable_basic() -> Result<()> {
+    let (manager, _tempdir) = setup_lifecycle_manager().await?;
+    let component_path = build_fetch_component().await?;
+
+    let (component_id, _) = manager
+        .load_component(&format!("file://{}", component_path.to_str().unwrap()))
+        .await?;
+
+    // Test granting environment variable permission
+    let result = manager
+        .grant_permission(
+            &component_id,
+            "environment",
+            &serde_json::json!({"key": "API_KEY"}),
+        )
+        .await;
+
+    assert!(result.is_ok());
+
+    // Verify policy file was created and contains the permission
+    let policy_info = manager.get_policy_info(&component_id).await;
+    assert!(policy_info.is_some());
+    let policy_info = policy_info.unwrap();
+
+    // Verify policy contains the permission
+    let policy_content = tokio::fs::read_to_string(&policy_info.local_path).await?;
+    assert!(policy_content.contains("API_KEY"));
+    assert!(policy_content.contains("environment"));
+
+    Ok(())
+}
+
+#[cfg(any(target_os = "linux", target_os = "macos"))]
+#[test(tokio::test)]
+async fn test_grant_permission_environment_variable_multiple() -> Result<()> {
+    let (manager, _tempdir) = setup_lifecycle_manager().await?;
+    let component_path = build_fetch_component().await?;
+
+    let (component_id, _) = manager
+        .load_component(&format!("file://{}", component_path.to_str().unwrap()))
+        .await?;
+
+    // Grant multiple environment variable permissions
+    let api_key_result = manager
+        .grant_permission(
+            &component_id,
+            "environment",
+            &serde_json::json!({"key": "API_KEY"}),
+        )
+        .await;
+
+    let config_url_result = manager
+        .grant_permission(
+            &component_id,
+            "environment",
+            &serde_json::json!({"key": "CONFIG_URL"}),
+        )
+        .await;
+
+    assert!(api_key_result.is_ok());
+    assert!(config_url_result.is_ok());
+
+    // Verify policy file contains all permissions
+    let policy_info = manager.get_policy_info(&component_id).await;
+    assert!(policy_info.is_some());
+    let policy_info = policy_info.unwrap();
+    let policy_content = tokio::fs::read_to_string(&policy_info.local_path).await?;
+
+    assert!(policy_content.contains("API_KEY"));
+    assert!(policy_content.contains("CONFIG_URL"));
+    assert!(policy_content.contains("environment"));
+
+    Ok(())
+}
+
+#[cfg(any(target_os = "linux", target_os = "macos"))]
+#[test(tokio::test)]
+async fn test_grant_permission_environment_variable_duplicate_prevention() -> Result<()> {
+    let (manager, _tempdir) = setup_lifecycle_manager().await?;
+    let component_path = build_fetch_component().await?;
+
+    let (component_id, _) = manager
+        .load_component(&format!("file://{}", component_path.to_str().unwrap()))
+        .await?;
+
+    // Grant the same environment variable permission twice
+    let details = serde_json::json!({"key": "API_KEY"});
+    let first_result = manager
+        .grant_permission(&component_id, "environment", &details)
+        .await;
+    let second_result = manager
+        .grant_permission(&component_id, "environment", &details)
+        .await;
+
+    assert!(first_result.is_ok());
+    assert!(second_result.is_ok());
+
+    // Verify policy file contains only one instance
+    let policy_info = manager.get_policy_info(&component_id).await;
+    assert!(policy_info.is_some());
+    let policy_info = policy_info.unwrap();
+    let policy_content = tokio::fs::read_to_string(&policy_info.local_path).await?;
+
+    // Count occurrences of the environment key - should be exactly 1
+    let occurrences = policy_content.matches("API_KEY").count();
+    assert_eq!(occurrences, 1);
+
+    Ok(())
+}