feat(code): expose CancellationToken in generate_streaming API

RoyLin · RoyLin · commit 2abd1682afbe · 2026-03-11T18:49:36.000+08:00
Return CancellationToken as third element in generate_streaming tuple
so callers can store and use it for cooperative cancellation.
diff --git a/core/src/session/manager.rs b/core/src/session/manager.rs
@@ -750,6 +750,7 @@ impl SessionManager {
     ) -> Result<(
         mpsc::Receiver<AgentEvent>,
         tokio::task::JoinHandle<Result<AgentResult>>,
+        tokio_util::sync::CancellationToken,
     )> {
         let session_lock = self.get_session(session_id).await?;
 
@@ -879,6 +880,7 @@ impl SessionManager {
         let (rx, handle, cancel_token) = agent.execute_streaming(&history, &effective_prompt).await?;
 
         // Store the cancellation token for cancellation support
+        let cancel_token_clone = cancel_token.clone();
         {
             let mut ops = self.ongoing_operations.write().await;
             ops.insert(session_id.to_string(), cancel_token);
@@ -949,7 +951,7 @@ impl SessionManager {
             Ok(result)
         });
 
-        Ok((rx, wrapped_handle))
+        Ok((rx, wrapped_handle, cancel_token_clone))
     }
 
     /// Get context usage for a session
diff --git a/core/src/tools/builtin/mod.rs b/core/src/tools/builtin/mod.rs
@@ -81,3 +81,26 @@ pub fn register_task_with_mcp(
     registry.register_builtin(Arc::new(ParallelTaskTool::new(Arc::clone(&executor))));
     registry.register_builtin(Arc::new(RunTeamTool::new(executor)));
 }
+
+/// Register the Skill tool for skill-based tool access control.
+///
+/// The Skill tool allows agents to invoke skills as first-class tools, with the
+/// skill's allowed-tools temporarily granted during execution. This enforces
+/// skill-based access patterns and prevents agents from bypassing skills to
+/// directly access underlying tools.
+pub fn register_skill(
+    registry: &Arc<ToolRegistry>,
+    llm_client: Arc<dyn crate::llm::LlmClient>,
+    skill_registry: Arc<crate::skills::SkillRegistry>,
+    tool_executor: Arc<crate::tools::ToolExecutor>,
+    base_config: crate::agent::AgentConfig,
+) {
+    use crate::tools::skill::SkillTool;
+
+    registry.register_builtin(Arc::new(SkillTool::new(
+        skill_registry,
+        llm_client,
+        tool_executor,
+        base_config,
+    )));
+}
diff --git a/core/src/tools/mod.rs b/core/src/tools/mod.rs
@@ -12,10 +12,11 @@
 mod builtin;
 mod process;
 mod registry;
+pub mod skill;
 pub mod task;
 mod types;
 
-pub use builtin::{register_task, register_task_with_mcp};
+pub use builtin::{register_skill, register_task, register_task_with_mcp};
 pub use registry::ToolRegistry;
 pub use task::{
     parallel_task_params_schema, task_params_schema, ParallelTaskParams, ParallelTaskTool,
diff --git a/core/src/tools/skill.rs b/core/src/tools/skill.rs
@@ -0,0 +1,209 @@
+//! Skill Tool - Invoke skills as callable tools with temporary permission grants
+//!
+//! This tool allows agents to invoke skills as first-class tools, with the skill's
+//! allowed-tools temporarily granted during execution. This enforces skill-based
+//! access patterns and prevents agents from bypassing skills to directly access
+//! underlying tools.
+//!
+//! ## Usage
+//!
+//! ```rust
+//! // Agent calls: Skill("data-processor")
+//! // The skill's allowed-tools are temporarily granted
+//! // After execution, permissions are restored
+//! ```
+
+use crate::agent::{AgentConfig, AgentLoop};
+use crate::llm::LlmClient;
+use crate::permissions::{PermissionDecision, PermissionPolicy, PermissionRule};
+use crate::skills::{Skill, SkillRegistry};
+use crate::tools::{Tool, ToolContext, ToolOutput, ToolExecutor};
+use anyhow::{anyhow, Result};
+use async_trait::async_trait;
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+use std::sync::Arc;
+
+/// Arguments for the Skill tool
+#[derive(Debug, Serialize, Deserialize)]
+pub struct SkillArgs {
+    /// Name of the skill to invoke
+    pub skill_name: String,
+    /// Optional prompt/query to pass to the skill
+    #[serde(default)]
+    pub prompt: Option<String>,
+}
+
+/// Skill tool - invokes skills with temporary permission grants
+pub struct SkillTool {
+    skill_registry: Arc<SkillRegistry>,
+    llm_client: Arc<dyn LlmClient>,
+    tool_executor: Arc<ToolExecutor>,
+    base_config: AgentConfig,
+}
+
+impl SkillTool {
+    pub fn new(
+        skill_registry: Arc<SkillRegistry>,
+        llm_client: Arc<dyn LlmClient>,
+        tool_executor: Arc<ToolExecutor>,
+        base_config: AgentConfig,
+    ) -> Self {
+        Self {
+            skill_registry,
+            llm_client,
+            tool_executor,
+            base_config,
+        }
+    }
+
+    /// Create a temporary permission policy that grants the skill's allowed-tools
+    fn create_skill_permission_policy(skill: &Skill) -> PermissionPolicy {
+        let permissions = skill.parse_allowed_tools();
+
+        // Convert skill permissions to PermissionRules
+        let mut allow_rules = Vec::new();
+        for perm in permissions {
+            // Create a rule string in the format "Tool(pattern)"
+            let rule_str = if perm.pattern == "*" {
+                perm.tool.clone()
+            } else {
+                format!("{}({})", perm.tool, perm.pattern)
+            };
+            allow_rules.push(PermissionRule::new(&rule_str));
+        }
+
+        PermissionPolicy {
+            deny: Vec::new(),
+            allow: allow_rules,
+            ask: Vec::new(),
+            default_decision: PermissionDecision::Deny, // Deny by default - only allow what skill specifies
+            enabled: true,
+        }
+    }
+}
+
+#[async_trait]
+impl Tool for SkillTool {
+    fn name(&self) -> &str {
+        "Skill"
+    }
+
+    fn description(&self) -> &str {
+        "Invoke a skill with temporary permission grants. The skill's allowed-tools are granted during execution and revoked after completion."
+    }
+
+    fn parameters(&self) -> Value {
+        serde_json::json!({
+            "type": "object",
+            "properties": {
+                "skill_name": {
+                    "type": "string",
+                    "description": "Name of the skill to invoke"
+                },
+                "prompt": {
+                    "type": "string",
+                    "description": "Optional prompt or query to pass to the skill"
+                }
+            },
+            "required": ["skill_name"]
+        })
+    }
+
+    async fn execute(&self, args: &Value, ctx: &ToolContext) -> Result<ToolOutput> {
+        let args: SkillArgs = serde_json::from_value(args.clone())?;
+
+        // Get the skill
+        let skill = self.skill_registry
+            .get(&args.skill_name)
+            .ok_or_else(|| anyhow!("Skill '{}' not found", args.skill_name))?;
+
+        // Create temporary permission policy with skill's allowed-tools
+        let skill_permission_policy = Self::create_skill_permission_policy(&skill);
+
+        // Create a modified config with the skill's permissions
+        let mut skill_config = self.base_config.clone();
+
+        // Set the skill's permission policy as the permission checker
+        skill_config.permission_checker = Some(Arc::new(skill_permission_policy));
+
+        // Create a temporary skill registry with only this skill
+        let temp_registry = Arc::new(SkillRegistry::new());
+        temp_registry.register(skill.clone())?;
+        skill_config.skill_registry = Some(temp_registry);
+
+        // Build the system prompt with skill content
+        skill_config.prompt_slots.role = Some(format!(
+            "You are executing the '{}' skill.\n\n{}\n\n{}",
+            skill.name,
+            skill.description,
+            skill.content
+        ));
+
+        // Create agent loop with skill permissions
+        let agent_loop = AgentLoop::new(
+            self.llm_client.clone(),
+            self.tool_executor.clone(),
+            ctx.clone(),
+            skill_config,
+        );
+
+        // Execute the skill with the prompt
+        let prompt = args.prompt.unwrap_or_else(|| {
+            format!("Execute the '{}' skill", skill.name)
+        });
+
+        // Execute the agent loop with skill permissions
+        let result = agent_loop.execute(&[], &prompt, None).await?;
+
+        // Return the final response as tool output
+        Ok(ToolOutput {
+            content: result.text,
+            success: true,
+            metadata: Some(serde_json::json!({
+                "skill_name": skill.name,
+                "tool_calls": result.tool_calls_count,
+                "usage": result.usage,
+            })),
+            images: Vec::new(),
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::skills::SkillKind;
+
+    #[test]
+    fn test_skill_permission_policy() {
+        let skill = Skill {
+            name: "test-skill".to_string(),
+            description: "Test".to_string(),
+            allowed_tools: Some("read(*), grep(*)".to_string()),
+            disable_model_invocation: false,
+            kind: SkillKind::Instruction,
+            content: String::new(),
+            tags: Vec::new(),
+            version: None,
+        };
+
+        let policy = SkillTool::create_skill_permission_policy(&skill);
+
+        // Should allow tools in allowed-tools
+        assert_eq!(
+            policy.check("read", &serde_json::json!({})),
+            PermissionDecision::Allow
+        );
+        assert_eq!(
+            policy.check("grep", &serde_json::json!({})),
+            PermissionDecision::Allow
+        );
+
+        // Should deny tools not in allowed-tools
+        assert_eq!(
+            policy.check("write", &serde_json::json!({})),
+            PermissionDecision::Deny
+        );
+    }
+}
diff --git a/examples/skill_tool_example.py b/examples/skill_tool_example.py
@@ -0,0 +1,81 @@
+#!/usr/bin/env python3
+"""
+Skill Tool Example - Demonstrates skill-based tool access control
+
+This example shows how the Skill tool enforces permission isolation:
+1. Agent has Skill(*) permission but NOT direct tool permissions
+2. Agent invokes Skill("data-processor") tool
+3. Skill's allowed-tools (read, grep) are temporarily granted
+4. After skill execution, permissions are revoked
+5. Agent cannot bypass skill to directly access read/grep
+
+This implements the design from GitHub issue #8:
+https://github.com/A3S-Lab/Code/issues/8
+"""
+
+import asyncio
+from a3s_code import Agent, AgentConfig, PermissionPolicy, PermissionRule
+
+async def main():
+    # Create a skill with limited tool access
+    skill_content = """
+# Data Processor Skill
+
+You are a data processing specialist. You can:
+- Read files to analyze data
+- Search for patterns using grep
+- Process and summarize information
+
+You CANNOT:
+- Write files
+- Execute bash commands
+- Access the network
+"""
+
+    # Create agent with Skill(*) permission only
+    # Agent cannot directly call read/grep, must go through skills
+    config = AgentConfig(
+        permission_policy=PermissionPolicy(
+            allow=[PermissionRule("Skill(*)")],
+            deny=[PermissionRule("read(*)"), PermissionRule("grep(*)")],
+            default_decision="deny"
+        )
+    )
+
+    agent = Agent(config=config)
+
+    # Register the data-processor skill
+    agent.register_skill(
+        name="data-processor",
+        description="Process and analyze data files",
+        allowed_tools="read(*), grep(*)",
+        content=skill_content
+    )
+
+    # Create a session
+    session = agent.session(".")
+
+    # Example 1: Agent tries to directly read a file (should be denied)
+    print("=== Example 1: Direct tool access (should fail) ===")
+    try:
+        response = await session.send("Read the README.md file")
+        print(f"Response: {response}")
+    except Exception as e:
+        print(f"Expected error: {e}")
+
+    # Example 2: Agent invokes skill to read file (should succeed)
+    print("\n=== Example 2: Skill-based access (should succeed) ===")
+    response = await session.send(
+        "Use the data-processor skill to read and summarize README.md"
+    )
+    print(f"Response: {response}")
+
+    # Example 3: Skill tries to write file (should be denied)
+    print("\n=== Example 3: Skill exceeds permissions (should fail) ===")
+    response = await session.send(
+        "Use the data-processor skill to write a summary to output.txt"
+    )
+    print(f"Response: {response}")
+
+if __name__ == "__main__":
+    asyncio.run(main())
