ship grants, a2a_client, discovery, sandbox SDK + tests

tests/test_multi_agent.py

@@ -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)