feat: add sandbox support to AgentSkills plugin (3/N)

src/strands/vended_plugins/skills/agent_skills.py

@@ -3,11 +3,20 @@
 This module provides the AgentSkills class that extends the Plugin base class
 to add Agent Skills support. The plugin registers a tool for activating
 skills, and injects skill metadata into the system prompt.
+
+Sandbox skill sources (e.g., ``"sandbox:///home/skills"``) are loaded
+asynchronously during ``init_agent()`` using the agent's sandbox instance,
+following the same deferred-loading pattern as the TypeScript SDK.
+
+The skill catalog is maintained **per-agent** using a ``WeakKeyDictionary``,
+so a single plugin instance can be safely shared across multiple agents
+(even with different sandboxes) without skill cross-contamination.
 """
 
 from __future__ import annotations
 
 import logging
+import weakref
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, TypeAlias
 from xml.sax.saxutils import escape
@@ -27,9 +36,10 @@
 _DEFAULT_STATE_KEY = "agent_skills"
 _RESOURCE_DIRS = ("scripts", "references", "assets")
 _DEFAULT_MAX_RESOURCE_FILES = 20
+_SANDBOX_PREFIX = "sandbox://"
 
 SkillSource: TypeAlias = str | Path | Skill
-"""A single skill source: path string, Path object, or Skill instance."""
+"""A single skill source: path string, Path object, Skill instance, or ``"sandbox:///path"`` string."""
 
 SkillSources: TypeAlias = SkillSource | list[SkillSource]
 """One or more skill sources."""
@@ -52,7 +62,16 @@ class AgentSkills(Plugin):
     3. Session persistence of active skill state via ``agent.state``
 
     Skills can be provided as filesystem paths (to individual skill directories or
-    parent directories containing multiple skills) or as pre-built ``Skill`` instances.
+    parent directories containing multiple skills), ``"sandbox:///path"`` URIs that
+    load skills from the agent's sandbox at init time, or pre-built ``Skill`` instances.
+
+    Sandbox sources are loaded asynchronously during ``init_agent()`` using the
+    agent's sandbox, following the same deferred-loading pattern as the TypeScript SDK.
+
+    The skill catalog is stored **per-agent** so that a single plugin instance
+    can be safely attached to multiple agents with different sandboxes.
+    Synchronous sources (filesystem, URL, Skill instances) are shared as the
+    base set; sandbox-resolved skills are added per-agent on top of that base.
 
     Example:
         ```python
@@ -62,11 +81,21 @@ class AgentSkills(Plugin):
         # Load from filesystem
         plugin = AgentSkills(skills=["./skills/pdf-processing", "./skills/"])
 
+        # Load from agent's sandbox (resolved at agent init)
+        plugin = AgentSkills(skills=["sandbox:///home/skills"])
+
+        # Mix local and sandbox sources
+        plugin = AgentSkills(skills=["./local-skills/", "sandbox:///home/skills"])
+
         # Or provide Skill instances directly
         skill = Skill(name="my-skill", description="A custom skill", instructions="Do the thing")
         plugin = AgentSkills(skills=[skill])
 
         agent = Agent(plugins=[plugin])
+
+        # Safe to share across multiple agents with different sandboxes
+        agent_a = Agent(sandbox=sandbox_a, plugins=[plugin])
+        agent_b = Agent(sandbox=sandbox_b, plugins=[plugin])
         ```
     """
 
@@ -81,34 +110,148 @@ def __init__(
     ) -> None:
         """Initialize the AgentSkills plugin.
 
+        Synchronous sources (filesystem paths, ``Skill`` instances, HTTPS URLs)
+        are resolved immediately into the base skill set. Sandbox sources
+        (``"sandbox:///path"``) are stored as pending and resolved per-agent
+        during ``init_agent()`` when each agent's sandbox is available.
+
         Args:
             skills: One or more skill sources. Can be a single value or a list. Each element can be:
 
                 - A ``str`` or ``Path`` to a skill directory (containing SKILL.md)
                 - A ``str`` or ``Path`` to a parent directory (containing skill subdirectories)
                 - A ``Skill`` dataclass instance
                 - An ``https://`` URL pointing directly to raw SKILL.md content
+                - A ``"sandbox:///path"`` URI to load from the agent's sandbox at init time
             state_key: Key used to store plugin state in ``agent.state``.
             max_resource_files: Maximum number of resource files to list in skill responses.
             strict: If True, raise on skill validation issues. If False (default), warn and load anyway.
         """
         self._strict = strict
-        self._skills: dict[str, Skill] = self._resolve_skills(_normalize_sources(skills))
         self._state_key = state_key
         self._max_resource_files = max_resource_files
+
+        # Split sources into sync and deferred (sandbox) groups
+        all_sources = _normalize_sources(skills)
+        sync_sources: list[SkillSource] = []
+        self._sandbox_sources: list[str] = []
+
+        for source in all_sources:
+            if isinstance(source, str) and source.startswith(_SANDBOX_PREFIX):
+                self._sandbox_sources.append(source[len(_SANDBOX_PREFIX) :])
+            else:
+                sync_sources.append(source)
+
+        # Resolve synchronous sources immediately — shared across all agents
+        self._base_skills: dict[str, Skill] = self._resolve_skills(sync_sources)
+
+        # Per-agent skill catalogs (base + sandbox-resolved skills)
+        # Uses WeakKeyDictionary so entries are cleaned up when agents are GC'd
+        self._agent_skills: weakref.WeakKeyDictionary[Agent, dict[str, Skill]] = weakref.WeakKeyDictionary()
+
         super().__init__()
 
-    def init_agent(self, agent: Agent) -> None:
+    def _get_skills(self, agent: Agent | None = None) -> dict[str, Skill]:
+        """Get the skill catalog for a specific agent.
+
+        Returns the per-agent catalog if the agent has been initialized,
+        otherwise falls back to the base (sync-resolved) skills.
+
+        Args:
+            agent: The agent to get skills for. If None, returns base skills.
+
+        Returns:
+            Dict mapping skill names to Skill instances.
+        """
+        if agent is not None and agent in self._agent_skills:
+            return self._agent_skills[agent]
+        if agent is not None and agent not in self._agent_skills and self._sandbox_sources:
+            logger.warning(
+                "agent has not been initialized yet — sandbox skills are not available; returning base skills only"
+            )
+        return self._base_skills
+
+    async def init_agent(self, agent: Agent) -> None:
         """Initialize the plugin with an agent instance.
 
-        Decorated hooks and tools are auto-registered by the plugin registry.
+        Creates a per-agent skill catalog by copying the base skills and
+        then resolving any sandbox sources using the agent's sandbox. This
+        ensures each agent gets its own skill set without cross-contamination.
 
         Args:
             agent: The agent instance to extend with skills support.
         """
-        if not self._skills:
+        # Start with a COPY of base skills for this agent
+        agent_skills = dict(self._base_skills)
+
+        # Resolve deferred sandbox sources from THIS agent's sandbox
+        if self._sandbox_sources:
+            sandbox_skills = await self._resolve_sandbox_skills(agent)
+            agent_skills.update(sandbox_skills)
+
+        # Store per-agent catalog
+        self._agent_skills[agent] = agent_skills
+
+        if not agent_skills:
             logger.warning("no skills were loaded, the agent will have no skills available")
-        logger.debug("skill_count=<%d> | skills plugin initialized", len(self._skills))
+        logger.debug("skill_count=<%d> | skills plugin initialized for agent", len(agent_skills))
+
+    async def _resolve_sandbox_skills(self, agent: Agent) -> dict[str, Skill]:
+        """Resolve sandbox skill sources using the agent's sandbox.
+
+        Each sandbox source path is treated as either a single skill directory
+        (if it contains SKILL.md) or a parent directory containing skill
+        subdirectories.
+
+        Args:
+            agent: The agent whose sandbox to load skills from.
+
+        Returns:
+            Dict mapping skill names to Skill instances loaded from the sandbox.
+        """
+        resolved: dict[str, Skill] = {}
+
+        sandbox = getattr(agent, "sandbox", None)
+        if sandbox is None:
+            logger.warning(
+                "agent has no sandbox configured — skipping %d sandbox skill source(s)",
+                len(self._sandbox_sources),
+            )
+            return resolved
+
+        for sandbox_path in self._sandbox_sources:
+            logger.debug("sandbox_path=<%s> | resolving sandbox skill source", sandbox_path)
+
+            try:
+                # First try loading as a single skill
+                skill = await Skill.from_sandbox(sandbox, sandbox_path, strict=self._strict)
+                if skill.name in resolved:
+                    logger.warning("name=<%s> | duplicate skill name, overwriting previous skill", skill.name)
+                resolved[skill.name] = skill
+                logger.debug(
+                    "sandbox_path=<%s>, name=<%s> | loaded single skill from sandbox",
+                    sandbox_path,
+                    skill.name,
+                )
+            except FileNotFoundError:
+                # Not a single skill — try as parent directory
+                try:
+                    skills = await Skill.from_sandbox_directory(sandbox, sandbox_path, strict=self._strict)
+                    for skill in skills:
+                        if skill.name in resolved:
+                            logger.warning("name=<%s> | duplicate skill name, overwriting previous skill", skill.name)
+                        resolved[skill.name] = skill
+                    logger.debug(
+                        "sandbox_path=<%s>, count=<%d> | loaded skills from sandbox directory",
+                        sandbox_path,
+                        len(skills),
+                    )
+                except Exception as e:
+                    logger.warning("sandbox_path=<%s> | failed to load sandbox skills: %s", sandbox_path, e)
+            except Exception as e:
+                logger.warning("sandbox_path=<%s> | failed to load sandbox skill: %s", sandbox_path, e)
+
+        return resolved
 
     @tool(context=True)
     def skills(self, skill_name: str, tool_context: ToolContext) -> str:  # noqa: D417
@@ -120,13 +263,15 @@ def skills(self, skill_name: str, tool_context: ToolContext) -> str:  # noqa: D4
         Args:
             skill_name: Name of the skill to activate.
         """
+        agent_skills = self._get_skills(tool_context.agent)
+
         if not skill_name:
-            available = ", ".join(self._skills)
+            available = ", ".join(agent_skills)
             return f"Error: skill_name is required. Available skills: {available}"
 
-        found = self._skills.get(skill_name)
+        found = agent_skills.get(skill_name)
         if found is None:
-            available = ", ".join(self._skills)
+            available = ", ".join(agent_skills)
             return f"Skill '{skill_name}' not found. Available skills: {available}"
 
         logger.debug("skill_name=<%s> | skill activated", skill_name)
@@ -154,7 +299,7 @@ def _on_before_invocation(self, event: BeforeInvocationEvent) -> None:
         state_data = agent.state.get(self._state_key)
         last_injected_xml = state_data.get("last_injected_xml") if isinstance(state_data, dict) else None
 
-        skills_xml = self._generate_skills_xml()
+        skills_xml = self._generate_skills_xml(agent)
         content = agent.system_prompt_content
 
         if content is not None:
@@ -183,13 +328,20 @@ def _on_before_invocation(self, event: BeforeInvocationEvent) -> None:
             self._set_state_field(agent, "last_injected_xml", new_injected_xml)
             agent.system_prompt = new_prompt
 
-    def get_available_skills(self) -> list[Skill]:
+    def get_available_skills(self, agent: Agent | None = None) -> list[Skill]:
         """Get the list of available skills.
 
+        When called with an agent, returns that agent's full skill catalog
+        (base + sandbox-resolved). Without an agent, returns only the base
+        (sync-resolved) skills.
+
+        Args:
+            agent: Optional agent to get per-agent skills for.
+
         Returns:
-            A copy of the current skills list.
+            A copy of the skills list.
         """
-        return list(self._skills.values())
+        return list(self._get_skills(agent).values())
 
     def set_available_skills(self, skills: SkillSources) -> None:
         """Set the available skills, replacing any existing ones.
@@ -199,14 +351,31 @@ def set_available_skills(self, skills: SkillSources) -> None:
         parent directory containing skill subdirectories, or an ``https://``
         URL pointing directly to raw SKILL.md content.
 
-        Note: this does not persist state or deactivate skills on any agent.
-        Active skill state is managed per-agent and will be reconciled on the
-        next tool call or invocation.
+        Note: Sandbox sources (``"sandbox:///path"``) are NOT supported in
+        ``set_available_skills`` because no agent context is available.
+        Use ``"sandbox:..."`` sources in the constructor instead.
+
+        Note: this replaces the base skill set and clears all per-agent
+        caches so they will be rebuilt on the next ``init_agent``.
 
         Args:
             skills: One or more skill sources to resolve and set.
+
+        Raises:
+            ValueError: If a sandbox source is passed (not supported here).
         """
-        self._skills = self._resolve_skills(_normalize_sources(skills))
+        sources = _normalize_sources(skills)
+        for source in sources:
+            if isinstance(source, str) and source.startswith(_SANDBOX_PREFIX):
+                raise ValueError(
+                    f"Sandbox sources ('{source}') are not supported in set_available_skills(). "
+                    "Use sandbox sources in the AgentSkills constructor instead."
+                )
+
+        self._base_skills = self._resolve_skills(sources)
+        self._sandbox_sources = []
+        # Clear per-agent caches so they pick up new base skills
+        self._agent_skills.clear()
 
     def _format_skill_response(self, skill: Skill) -> str:
         """Format the tool response when a skill is activated.
@@ -274,22 +443,27 @@ def _list_skill_resources(self, skill_path: Path) -> list[str]:
 
         return files
 
-    def _generate_skills_xml(self) -> str:
+    def _generate_skills_xml(self, agent: Agent | None = None) -> str:
         """Generate the XML block listing available skills for the system prompt.
 
         When no skills are loaded, returns a block indicating no skills are available.
         Otherwise includes a ``<location>`` element for skills loaded from the filesystem,
         following the AgentSkills.io integration spec.
 
+        Args:
+            agent: Optional agent to generate per-agent skills XML for.
+
         Returns:
             XML-formatted string with skill metadata.
         """
-        if not self._skills:
+        skills = self._get_skills(agent)
+
+        if not skills:
             return "<available_skills>\nNo skills are currently available.\n</available_skills>"
 
         lines: list[str] = ["<available_skills>"]
 
-        for skill in self._skills.values():
+        for skill in skills.values():
             lines.append("<skill>")
             lines.append(f"<name>{escape(skill.name)}</name>")
             lines.append(f"<description>{escape(skill.description)}</description>")
@@ -307,6 +481,9 @@ def _resolve_skills(self, sources: list[SkillSource]) -> dict[str, Skill]:
         a path to a parent directory containing multiple skills, or an
         HTTPS URL pointing to a SKILL.md file.
 
+        Note: Sandbox sources should be resolved separately via
+        ``_resolve_sandbox_skills()``.
+
         Args:
             sources: List of skill sources to resolve.