ship grants, a2a_client, discovery, sandbox SDK + tests

tests/test_sandbox.py

@@ -0,0 +1,205 @@
+from __future__ import annotations
+
+from typing import Sequence
+
+import pytest
+
+from a2a_pack import (
+    A2AAgent,
+    ExecResult,
+    LocalRunContext,
+    NoAuth,
+    RunContext,
+    SandboxClient,
+    SandboxHandle,
+    SandboxSpec,
+    SandboxUnavailable,
+    skill,
+)
+
+
+# ---------------------------------------------------------------------------
+# stub client that records calls — for asserting the surface without needing
+# microsandbox/FUSE/MinIO running.
+# ---------------------------------------------------------------------------
+
+
+class _StubHandle(SandboxHandle):
+    def __init__(self, name: str, exec_log: list[tuple[str, ...]]) -> None:
+        self.name = name
+        self._exec_log = exec_log
+        self.stopped = False
+
+    async def exec(
+        self,
+        cmd: str,
+        args: Sequence[str] | None = None,
+        *,
+        timeout: float | None = None,
+    ) -> ExecResult:
+        a = tuple(args or ())
+        self._exec_log.append(("exec", cmd, *a))
+        return ExecResult(stdout=f"exec:{cmd}:{','.join(a)}", exit_code=0)
+
+    async def shell(
+        self, script: str, *, timeout: float | None = None
+    ) -> ExecResult:
+        self._exec_log.append(("shell", script))
+        return ExecResult(stdout=f"shell:{script}", exit_code=0)
+
+    async def stop(self) -> None:
+        self.stopped = True
+
+    async def kill(self) -> None:
+        self.stopped = True
+
+    async def logs(self, *, tail: int | None = None) -> str:
+        return ""
+
+
+class _StubClient(SandboxClient):
+    def __init__(self) -> None:
+        self.created: list[SandboxSpec] = []
+        self.removed: list[str] = []
+        self._handles: dict[str, _StubHandle] = {}
+        self.exec_log: list[tuple[str, ...]] = []
+
+    async def create(self, spec: SandboxSpec) -> SandboxHandle:
+        self.created.append(spec)
+        h = _StubHandle(spec.name, self.exec_log)
+        self._handles[spec.name] = h
+        return h
+
+    async def get(self, name: str) -> SandboxHandle:
+        return self._handles[name]
+
+    async def list(self) -> list[str]:
+        return list(self._handles)
+
+    async def remove(self, name: str) -> None:
+        self.removed.append(name)
+        self._handles.pop(name, None)
+
+
+# ---------------------------------------------------------------------------
+# context plumbing
+# ---------------------------------------------------------------------------
+
+
+async def test_sandbox_unavailable_when_not_attached():
+    ctx: LocalRunContext[NoAuth] = LocalRunContext(auth=NoAuth())
+    with pytest.raises(SandboxUnavailable):
+        _ = ctx.sandbox
+
+
+async def test_sandbox_attached_returns_client():
+    client = _StubClient()
+    ctx: LocalRunContext[NoAuth] = LocalRunContext(auth=NoAuth(), sandbox=client)
+    assert ctx.sandbox is client
+
+
+# ---------------------------------------------------------------------------
+# agent uses ctx.sandbox via convenience helpers
+# ---------------------------------------------------------------------------
+
+
+class _CoderAgent(A2AAgent):
+    name = "coder"
+    description = "runs code in a sandbox"
+
+    @skill()
+    async def run(self, ctx: RunContext[NoAuth], code: str) -> str:
+        result = await ctx.sandbox.run_python(code)
+        return result.output
+
+
+async def test_agent_uses_run_python_convenience():
+    client = _StubClient()
+    out = await _CoderAgent().local_invoke(
+        "run", sandbox=client, code="print('hello')"
+    )
+    assert out.startswith("exec:python:-c,print('hello')")
+    # Spec was created and torn down (ephemeral)
+    assert len(client.created) == 1
+    assert client.created[0].image == "python:3.11-slim"
+    assert client.removed == [client.created[0].name]
+
+
+async def test_run_shell_one_shot():
+    client = _StubClient()
+
+    class _Agent(A2AAgent):
+        name = "shellbot"
+        description = ""
+
+        @skill()
+        async def go(self, ctx: RunContext[NoAuth]) -> str:
+            r = await ctx.sandbox.run_shell("ls /workspace")
+            return r.stdout
+
+    out = await _Agent().local_invoke("go", sandbox=client)
+    assert "shell:ls /workspace" in out
+    assert client.exec_log == [("shell", "ls /workspace")]
+    # one-shot: created and removed in the same call
+    assert len(client.created) == 1
+    assert len(client.removed) == 1
+
+
+# ---------------------------------------------------------------------------
+# explicit lifecycle (matches microsandbox SDK shape 1:1)
+# ---------------------------------------------------------------------------
+
+
+async def test_explicit_create_and_stop():
+    client = _StubClient()
+
+    class _Agent(A2AAgent):
+        name = "lifecycle"
+        description = ""
+
+        @skill()
+        async def go(self, ctx: RunContext[NoAuth]) -> int:
+            sb = await ctx.sandbox.create(
+                SandboxSpec(name="my-sb", image="python:3.11-slim", workspace="agent-foo")
+            )
+            r1 = await sb.exec("python", ["-c", "print(1)"])
+            r2 = await sb.shell("echo 2")
+            await sb.stop()
+            assert isinstance(r1, ExecResult)
+            assert r2.ok
+            return r1.exit_code
+
+    rc = await _Agent().local_invoke("go", sandbox=client)
+    assert rc == 0
+    assert client.created[0].workspace == "agent-foo"
+
+
+# ---------------------------------------------------------------------------
+# SandboxSpec metadata propagation
+# ---------------------------------------------------------------------------
+
+
+async def test_spec_carries_secrets_and_egress():
+    client = _StubClient()
+
+    class _Agent(A2AAgent):
+        name = "scoped"
+        description = ""
+
+        @skill()
+        async def go(self, ctx: RunContext[NoAuth]) -> str:
+            spec = SandboxSpec(
+                name="net-bound",
+                secrets=("OPENAI_KEY",),
+                egress=("api.openai.com",),
+                labels={"task": "research"},
+            )
+            sb = await ctx.sandbox.create(spec)
+            await sb.stop()
+            return "ok"
+
+    await _Agent().local_invoke("go", sandbox=client)
+    spec = client.created[0]
+    assert spec.secrets == ("OPENAI_KEY",)
+    assert spec.egress == ("api.openai.com",)
+    assert spec.labels == {"task": "research"}