molecule-ai/molecule-core
35
0
Fork 2
Code Issues 16 Pull Requests 2 Actions 49 Packages Projects Releases Wiki Activity

Merge pull request #46 from Molecule-AI/fix/a2a-client-auth-headers

Browse Source 
fix(security): complete Phase 30.6 auth headers in a2a_client — fixes post-deploy break in get_peers
This commit is contained in:
Hongming Wang 2026-04-14 07:18:16 -07:00 committed by GitHub
commit cfe1912997
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 21 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
workspace-template/a2a_client.py
View File

@ -31,7 +31,7 @@ async def discover_peer(target_id: str) -> dict | None:
try:
resp = await client.get(
f"{PLATFORM_URL}/registry/discover/{target_id}",
headers={"X-Workspace-ID": WORKSPACE_ID},
headers={"X-Workspace-ID": WORKSPACE_ID, **auth_headers()},
)
if resp.status_code == 200:
return resp.json()
@ -84,7 +84,10 @@ async def get_peers() -> list[dict]:
"""Get this workspace's peers from the platform registry."""
async with httpx.AsyncClient(timeout=10.0) as client:
try:
resp = await client.get(f"{PLATFORM_URL}/registry/{WORKSPACE_ID}/peers")
resp = await client.get(
f"{PLATFORM_URL}/registry/{WORKSPACE_ID}/peers",
headers={"X-Workspace-ID": WORKSPACE_ID, **auth_headers()},
)
if resp.status_code == 200:
return resp.json()
return []

22
workspace-template/tests/test_a2a_client.py
View File

@ -119,14 +119,11 @@ class TestDiscoverPeer:
await a2a_client.discover_peer("ws-xyz")
mock_client.get.assert_called_once()
call_args = mock_client.get.call_args
url = call_args[0][0] if call_args[0] else call_args[1].get("url") or call_args[0][0]
# The first positional arg is the URL
positional_url = mock_client.get.call_args.args[0]
assert "ws-xyz" in positional_url
assert mock_client.get.call_args.kwargs.get("headers") == {
"X-Workspace-ID": a2a_client.WORKSPACE_ID
}
# X-Workspace-ID must be present; bearer token also merged in when available
headers_sent = mock_client.get.call_args.kwargs.get("headers", {})
assert headers_sent.get("X-Workspace-ID") == a2a_client.WORKSPACE_ID
# ---------------------------------------------------------------------------
@ -312,6 +309,19 @@ class TestGetPeers:
url = mock_client.get.call_args.args[0]
assert "peers" in url
async def test_request_sends_workspace_id_header(self):
"""GET /registry/:id/peers must send X-Workspace-ID header (Phase 30.6)."""
import a2a_client
resp = _make_response(200, [])
mock_client = _make_mock_client(get_resp=resp)
with patch("a2a_client.httpx.AsyncClient", return_value=mock_client):
await a2a_client.get_peers()
headers_sent = mock_client.get.call_args.kwargs.get("headers", {})
assert headers_sent.get("X-Workspace-ID") == a2a_client.WORKSPACE_ID
# ---------------------------------------------------------------------------
# get_workspace_info