ship grants, a2a_client, discovery, sandbox SDK + tests
This commit is contained in:
robert
261
tests/test_multi_agent.py
Normal file
261
tests/test_multi_agent.py
Normal file
@@ -0,0 +1,261 @@
|
||||
"""End-to-end tests for the agent-to-agent + grant handoff seam."""
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from a2a_pack import (
|
||||
A2AAgent,
|
||||
DiscoveredAgent,
|
||||
FileType,
|
||||
Grant,
|
||||
GrantInvalid,
|
||||
InMemoryA2AClient,
|
||||
InMemoryDiscovery,
|
||||
LocalRunContext,
|
||||
LocalWorkspaceClient,
|
||||
NoAuth,
|
||||
RunContext,
|
||||
WorkspaceAccess,
|
||||
WorkspaceMode,
|
||||
mint_grant,
|
||||
skill,
|
||||
verify_grant,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# grant tokens
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_mint_and_verify_round_trip():
|
||||
grant, token = mint_grant(
|
||||
issuer="main", audience="graph", bucket="user-42-files"
|
||||
)
|
||||
out = verify_grant(token)
|
||||
assert out.grant_id == grant.grant_id
|
||||
assert out.bucket == "user-42-files"
|
||||
assert out.audience == "graph"
|
||||
assert out.expires_at > int(time.time())
|
||||
|
||||
|
||||
def test_tampered_grant_rejected():
|
||||
_, token = mint_grant(issuer="main", audience="graph", bucket="b")
|
||||
payload, sig = token.rsplit(".", 1)
|
||||
forged = payload + "x." + sig
|
||||
with pytest.raises(GrantInvalid):
|
||||
verify_grant(forged)
|
||||
|
||||
|
||||
def test_expired_grant_rejected():
|
||||
_, token = mint_grant(
|
||||
issuer="main",
|
||||
audience="graph",
|
||||
bucket="b",
|
||||
ttl_seconds=-1, # already expired
|
||||
)
|
||||
with pytest.raises(GrantInvalid, match="expired"):
|
||||
verify_grant(token)
|
||||
|
||||
|
||||
def test_signature_mismatch_rejected():
|
||||
_, token = mint_grant(
|
||||
issuer="main", audience="graph", bucket="b", secret=b"key-A"
|
||||
)
|
||||
with pytest.raises(GrantInvalid, match="signature"):
|
||||
verify_grant(token, secret=b"key-B")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# workspace.delegate() mints a grant; receiving side verifies it
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_workspace_delegate_mints_valid_grant():
|
||||
files = {"data/sales.xlsx": b"...", "secrets/.env": b"DO_NOT"}
|
||||
ws = LocalWorkspaceClient(
|
||||
files=files,
|
||||
access=WorkspaceAccess.dynamic(
|
||||
max_files=5,
|
||||
allowed_modes=(WorkspaceMode.READ_ONLY,),
|
||||
deny_patterns=("secrets/**",),
|
||||
),
|
||||
bucket="user-42-files",
|
||||
issuer="main-agent",
|
||||
)
|
||||
token = await ws.delegate(
|
||||
audience="graph-agent",
|
||||
allow_patterns=("*.xlsx",),
|
||||
deny_patterns=("secrets/**",),
|
||||
outputs_prefix="charts/",
|
||||
ttl_seconds=300,
|
||||
)
|
||||
grant = verify_grant(token)
|
||||
assert grant.bucket == "user-42-files"
|
||||
assert grant.audience == "graph-agent"
|
||||
assert "*.xlsx" in grant.allow_patterns
|
||||
assert "secrets/**" in grant.deny_patterns
|
||||
assert grant.outputs_prefix == "charts/"
|
||||
assert grant.mode is WorkspaceMode.READ_ONLY
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Two-agent demo: main calls graph in-process, hands a grant
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _GraphAgent(A2AAgent):
|
||||
name = "graph-agent"
|
||||
description = "Generates a chart from spreadsheet files"
|
||||
tools_used = ("matplotlib",)
|
||||
workspace_access = WorkspaceAccess.dynamic(
|
||||
max_files=8,
|
||||
allowed_modes=(WorkspaceMode.READ_ONLY,),
|
||||
deny_patterns=("secrets/**",),
|
||||
)
|
||||
|
||||
@skill(description="Render a chart", tags=["visualization", "spreadsheet"])
|
||||
async def generate_dashboard(
|
||||
self, ctx: RunContext[NoAuth], prompt: str
|
||||
) -> dict:
|
||||
# Workspace was attached by the runtime from the caller's grant.
|
||||
ws = ctx.workspace
|
||||
return {
|
||||
"prompt": prompt,
|
||||
"bucket": getattr(ws, "bucket", None),
|
||||
"deny_patterns": list(getattr(ws, "_access").deny_patterns),
|
||||
}
|
||||
|
||||
|
||||
class _MainAgent(A2AAgent):
|
||||
name = "main-agent"
|
||||
description = "Orchestrates user files via discovered agents"
|
||||
|
||||
@skill(description="Find a viz agent and delegate the chart")
|
||||
async def make_chart(self, ctx: RunContext[NoAuth], prompt: str) -> dict:
|
||||
hits = await ctx.discover.find_agents(tags=["visualization"])
|
||||
assert hits, "no graph agent in registry"
|
||||
graph = hits[0]
|
||||
|
||||
token = await ctx.workspace.delegate(
|
||||
audience=graph.name,
|
||||
allow_patterns=("*.xlsx",),
|
||||
deny_patterns=("secrets/**",),
|
||||
outputs_prefix="charts/",
|
||||
ttl_seconds=300,
|
||||
)
|
||||
result = await ctx.call(
|
||||
graph.name, "generate_dashboard", args={"prompt": prompt}, grant=token
|
||||
)
|
||||
return {"called": graph.name, "out": result.result, "grant_id": result.grant_id}
|
||||
|
||||
|
||||
def _build_a2a_router(
|
||||
agents: dict[str, A2AAgent], caller_workspace: LocalWorkspaceClient
|
||||
):
|
||||
"""In-memory router. Builds a callee-side ctx whose workspace is bounded
|
||||
by the inbound grant — same shape the HTTP server adapter does."""
|
||||
|
||||
def factory(agent: A2AAgent, grant_token: str | None):
|
||||
ws = None
|
||||
if grant_token is not None:
|
||||
grant = verify_grant(grant_token)
|
||||
ws = LocalWorkspaceClient(
|
||||
files={
|
||||
p: b
|
||||
for p, b in caller_workspace._files.items()
|
||||
if not any(p.startswith(d.rstrip("*").rstrip("/")) for d in grant.deny_patterns if d)
|
||||
},
|
||||
access=WorkspaceAccess.dynamic(
|
||||
max_files=64,
|
||||
allowed_modes=(WorkspaceMode.READ_ONLY,),
|
||||
deny_patterns=tuple(grant.deny_patterns),
|
||||
),
|
||||
bucket=grant.bucket,
|
||||
issuer=grant.audience,
|
||||
)
|
||||
return LocalRunContext(auth=NoAuth(), workspace=ws)
|
||||
|
||||
return InMemoryA2AClient(agents=agents, ctx_factory=factory)
|
||||
|
||||
|
||||
async def test_main_agent_discovers_and_delegates_to_graph_agent():
|
||||
main = _MainAgent()
|
||||
graph = _GraphAgent()
|
||||
|
||||
user_workspace = LocalWorkspaceClient(
|
||||
files={
|
||||
"sales.xlsx": b"q1,q2,q3\n10,20,30\n",
|
||||
"secrets/.env": b"NEVER",
|
||||
},
|
||||
access=WorkspaceAccess.dynamic(
|
||||
max_files=10,
|
||||
allowed_modes=(
|
||||
WorkspaceMode.READ_ONLY,
|
||||
WorkspaceMode.READ_WRITE_OVERLAY,
|
||||
),
|
||||
deny_patterns=("secrets/**",),
|
||||
),
|
||||
bucket="user-42-files",
|
||||
issuer="user-42",
|
||||
)
|
||||
|
||||
discovery = InMemoryDiscovery(
|
||||
{graph.name: DiscoveredAgent(name=graph.name, url=None, card=graph.card())}
|
||||
)
|
||||
router = _build_a2a_router({graph.name: graph}, caller_workspace=user_workspace)
|
||||
|
||||
out = await main.local_invoke(
|
||||
"make_chart",
|
||||
workspace=user_workspace,
|
||||
a2a=router,
|
||||
discover=discovery,
|
||||
prompt="weekly burn rate",
|
||||
)
|
||||
|
||||
assert out["called"] == "graph-agent"
|
||||
assert out["out"]["bucket"] == "user-42-files"
|
||||
# Callee saw the deny patterns we minted in the grant
|
||||
assert "secrets/**" in out["out"]["deny_patterns"]
|
||||
# And the call was audit-tagged
|
||||
assert out["grant_id"]
|
||||
|
||||
|
||||
async def test_no_grant_means_no_workspace_for_callee():
|
||||
"""If the main agent didn't delegate, the callee can't touch any workspace."""
|
||||
from a2a_pack import SkillInvocationError
|
||||
|
||||
class _Greedy(A2AAgent):
|
||||
name = "greedy"
|
||||
description = ""
|
||||
|
||||
@skill()
|
||||
async def steal(self, ctx: RunContext[NoAuth]) -> str:
|
||||
return str(ctx.workspace.bucket) # type: ignore[attr-defined]
|
||||
|
||||
class _Caller(A2AAgent):
|
||||
name = "caller"
|
||||
description = ""
|
||||
|
||||
@skill()
|
||||
async def go(self, ctx: RunContext[NoAuth]) -> str:
|
||||
r = await ctx.call("greedy", "steal", args={}, grant=None)
|
||||
return str(r.result)
|
||||
|
||||
router = InMemoryA2AClient(
|
||||
agents={"greedy": _Greedy()},
|
||||
ctx_factory=lambda agent, grant: LocalRunContext(auth=NoAuth()),
|
||||
)
|
||||
discovery = InMemoryDiscovery({})
|
||||
with pytest.raises(SkillInvocationError) as ei:
|
||||
await _Caller().local_invoke("go", a2a=router, discover=discovery)
|
||||
# Trace back to the PermissionError: callee accessed ctx.workspace with
|
||||
# nothing bound, so the runtime denied it.
|
||||
chain = []
|
||||
err: BaseException | None = ei.value
|
||||
while err is not None:
|
||||
chain.append(err)
|
||||
err = err.__cause__
|
||||
assert any(isinstance(e, PermissionError) for e in chain)
|
||||
205
tests/test_sandbox.py
Normal file
205
tests/test_sandbox.py
Normal file
@@ -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"}
|
||||
Reference in New Issue
Block a user