a2a/tests/test_sandbox.py

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"}