molecule-core/workspace/builtin_tools/a2a_tools.py
Hongming Wang 65b531acf6 fix(workspace): tag self-originated A2A POSTs with X-Workspace-ID
Workspace runtime fired four classes of A2A request to the platform
without the X-Workspace-ID header that identifies the source
workspace: heartbeat self-messages, initial_prompt, idle-loop fires,
and peer-to-peer A2A from runtime tools. The platform's a2a_receive
logger keys source_id off that header — without it, every such row
was written with source_id=NULL, which the canvas's My Chat tab
filters as ?source=canvas (i.e. "user typed this") and rendered the
internal triggers as if the human user had sent them. The
"Delegation results are ready..." heartbeat trigger was visible to
end users in the chat history; delegate_task A2A calls between agents
were misclassified the same way.

Centralise the header construction in a new platform_auth helper
self_source_headers(workspace_id) that returns auth_headers() PLUS
{X-Workspace-ID: <id>}. Apply it to:

  - heartbeat.py self-message (refactored from inline header dict)
  - main.py initial_prompt POST
  - main.py idle_prompt POST
  - a2a_client.py send_a2a_message (peer A2A from runtime)
  - builtin_tools/a2a_tools.py delegate_task (was missing ALL headers)

Tests:
  - test_heartbeat.py asserts the X-Workspace-ID header is set on
    the self-message POST.
  - test_a2a_tools_module.py asserts the same on delegate_task POSTs;
    FakeClient.post mocks updated to accept the headers kwarg.

Production effect lands the moment workspace containers are rebuilt
with this code; existing rows in activity_logs keep their NULL
source_id (legacy data). The canvas-side filter (#follow-up)
covers the historical-rows case until backfill.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-24 19:54:43 -07:00

91 lines
3.4 KiB
Python

"""A2A communication tools — framework-agnostic delegation and peer discovery.
These are plain async functions that any adapter can wrap in its native tool format.
The LangChain @tool versions are in tools/delegation.py.
"""
import os
import uuid
import httpx
PLATFORM_URL = os.environ.get("PLATFORM_URL", "http://host.docker.internal:8080")
WORKSPACE_ID = os.environ.get("WORKSPACE_ID", "")
async def list_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")
if resp.status_code == 200:
return resp.json()
return []
except Exception:
return []
async def delegate_task(workspace_id: str, task: str) -> str:
"""Send a task to a peer workspace via A2A and return the response text."""
async with httpx.AsyncClient(timeout=120.0) as client:
# Discover target URL
try:
resp = await client.get(
f"{PLATFORM_URL}/registry/discover/{workspace_id}",
headers={"X-Workspace-ID": WORKSPACE_ID},
)
if resp.status_code != 200:
return f"Error: cannot reach workspace {workspace_id} (status {resp.status_code})"
target_url = resp.json().get("url", "")
if not target_url:
return f"Error: workspace {workspace_id} has no URL"
except Exception as e:
return f"Error discovering workspace: {e}"
# Send A2A message. X-Workspace-ID identifies us as the source —
# without it the platform's a2a_receive logger writes
# source_id=NULL and the recipient's My Chat tab renders the
# delegation as if a human user typed it. Same hazard fixed
# in heartbeat.py / a2a_client.py / main.py initial+idle flows.
try:
a2a_resp = await client.post(
target_url,
headers={"X-Workspace-ID": WORKSPACE_ID},
json={
"jsonrpc": "2.0",
"id": str(uuid.uuid4()),
"method": "message/send",
"params": {
"message": {
"role": "user",
"messageId": str(uuid.uuid4()),
"parts": [{"kind": "text", "text": task}],
},
},
},
)
data = a2a_resp.json()
if "result" in data:
parts = data["result"].get("parts", [])
return parts[0].get("text", "(no text)") if parts else str(data["result"])
elif "error" in data:
return f"Error: {data['error'].get('message', str(data['error']))}"
return str(data)
except Exception as e:
return f"Error sending A2A message: {e}"
async def get_peers_summary() -> str:
"""Return a formatted string of available peers for system prompts."""
peers = await list_peers()
if not peers:
return "No peers available."
lines = []
for p in peers:
name = p.get("name", "Unknown")
pid = p.get("id", "")
role = p.get("role", "")
status = p.get("status", "")
lines.append(f"- {name} (ID: {pid}) — {role} [{status}]")
return "Available peers:\n" + "\n".join(lines)