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