diff --git a/workspace/executor_helpers.py b/workspace/executor_helpers.py index ad98ab33..68497d50 100644 --- a/workspace/executor_helpers.py +++ b/workspace/executor_helpers.py @@ -933,3 +933,60 @@ def collect_outbound_files(reply_text: str) -> list[dict[str, str]]: if staged is not None: out.append(staged) return out + + +def new_response_message( + context: Any, + text: str = "", + files: list[dict[str, str]] | None = None, +) -> Any: + """Build an A2A v1 protobuf response Message with task/context correlation. + + Adapter executors should use this instead of ``a2a.helpers.new_text_message`` + (which omits ``task_id`` / ``context_id``) so the platform's a2a proxy can + reliably correlate the response to the originating task. Mirrors the shape + used by ``workspace/a2a_executor.py``'s own response construction so all + runtime paths produce the same Message envelope. + + Args: + context: The ``RequestContext`` from the inbound A2A request. Reads + ``context.task_id`` and ``context.context_id``; both fall back to + fresh UUIDs when ``None`` (RequestContextBuilder always sets them + in production; the fallback exists for unit tests). + text: Response text. Empty string omits the text Part — useful when + replying with files only. + files: Optional list of ``{"path": ..., "name": ..., "mime_type": ...}`` + dicts (e.g. the output of :func:`collect_outbound_files`). Each + becomes a Part with ``url="workspace:"``, ``filename``, and + ``media_type`` set. + + Returns: + A v1 protobuf ``a2a.types.Message`` ready to pass to + ``event_queue.enqueue_event(...)``. + + Why this exists: a2a-sdk v1 replaced the v0 Pydantic discriminated-union + types (``Part(root=TextPart(...))`` / ``Part(root=FilePart(file= + FileWithUri(...)))``) with a flat protobuf Part struct. Templates that + were written against v0 + then auto-renamed have shipped without + ``task_id``/``context_id`` correlation; this helper centralizes the + canonical pattern. + """ + # Lazy import: a2a.types is provided by a2a-sdk which is a runtime + # dependency every adapter image already has. Importing here keeps the + # module load path lean for callers that don't construct messages. + from a2a.types import Message, Part, Role + + parts: list = [Part(text=text)] if text else [] + for f in files or []: + parts.append(Part( + url="workspace:" + f["path"], + filename=f["name"], + media_type=f["mime_type"], + )) + return Message( + message_id=_uuid.uuid4().hex, + role=Role.ROLE_AGENT, + parts=parts, + task_id=getattr(context, "task_id", None) or _uuid.uuid4().hex, + context_id=getattr(context, "context_id", None) or _uuid.uuid4().hex, + ) diff --git a/workspace/tests/conftest.py b/workspace/tests/conftest.py index 066cc21b..1aacd9a1 100644 --- a/workspace/tests/conftest.py +++ b/workspace/tests/conftest.py @@ -65,15 +65,42 @@ def _make_a2a_mocks(): tasks_mod.TaskUpdater = TaskUpdater - # a2a.types needs Part stub for artifact construction (v1: Part takes text= directly, no TextPart) + # a2a.types needs stubs for Part, Message, Role. + # v1 Part: flat protobuf with optional text/url/filename/media_type/raw/data fields. + # v1 Message: has message_id, role, parts, task_id, context_id, etc. + # Stubs preserve all kwargs so tests can assert on any field. types_mod = ModuleType("a2a.types") class Part: - """Stub for A2A Part (v1: takes text= kwarg directly).""" + """Stub for A2A Part (v1: flat protobuf with optional fields).""" def __init__(self, text=None, root=None, **kwargs): self.text = text + # Preserve every other kwarg as an attribute so tests can + # assert on Part(url=..., filename=..., media_type=...). + for k, v in kwargs.items(): + setattr(self, k, v) + + class Message: + """Stub for A2A Message (v1: protobuf with snake_case fields).""" + def __init__(self, message_id="", role=0, parts=None, task_id="", + context_id="", **kwargs): + self.message_id = message_id + self.role = role + self.parts = list(parts) if parts is not None else [] + self.task_id = task_id + self.context_id = context_id + for k, v in kwargs.items(): + setattr(self, k, v) + + class _RoleEnum: + """Stub for A2A Role enum (v1 protobuf: ROLE_UNSPECIFIED=0, ROLE_USER=1, ROLE_AGENT=2).""" + ROLE_UNSPECIFIED = 0 + ROLE_USER = 1 + ROLE_AGENT = 2 types_mod.Part = Part + types_mod.Message = Message + types_mod.Role = _RoleEnum # a2a.helpers (v1: moved from a2a.utils, renamed new_agent_text_message # → new_text_message). Mock both names — production code only calls diff --git a/workspace/tests/test_executor_helpers.py b/workspace/tests/test_executor_helpers.py index 884d9245..195d8dda 100644 --- a/workspace/tests/test_executor_helpers.py +++ b/workspace/tests/test_executor_helpers.py @@ -937,3 +937,89 @@ def test_collect_outbound_files_deduplicates(tmp_path, monkeypatch): reply = f"Wrote {report}. Again at {report}." out = collect_outbound_files(reply) assert len(out) == 1 + + +# ============================================================================ +# new_response_message — A2A v1 protobuf Message envelope with task/context +# correlation. Replaces ad-hoc per-template Message construction so every +# adapter response threads task_id/context_id back to the platform. +# ============================================================================ + + +def test_new_response_message_text_only(): + """Text-only response sets one text Part; role=ROLE_AGENT; + task_id/context_id passed through from context.""" + from executor_helpers import new_response_message + from a2a.types import Role + + ctx = SimpleNamespace(task_id="task-abc", context_id="ctx-xyz") + msg = new_response_message(ctx, "hello world") + + assert msg.role == Role.ROLE_AGENT + assert msg.task_id == "task-abc" + assert msg.context_id == "ctx-xyz" + assert len(msg.parts) == 1 + assert msg.parts[0].text == "hello world" + # message_id should be a 32-char hex (uuid4().hex) + assert len(msg.message_id) == 32 + + +def test_new_response_message_with_files(): + """Files become file Parts with workspace: URI scheme, filename, + media_type. Text Part comes first when text is non-empty.""" + from executor_helpers import new_response_message + + ctx = SimpleNamespace(task_id="t", context_id="c") + files = [ + {"path": "/workspace/.molecule/chat-uploads/a.png", "name": "a.png", "mime_type": "image/png"}, + {"path": "/workspace/.molecule/chat-uploads/b.txt", "name": "b.txt", "mime_type": "text/plain"}, + ] + msg = new_response_message(ctx, "see attachments", files=files) + + assert len(msg.parts) == 3 # 1 text + 2 file parts + assert msg.parts[0].text == "see attachments" + assert msg.parts[1].url == "workspace:/workspace/.molecule/chat-uploads/a.png" + assert msg.parts[1].filename == "a.png" + assert msg.parts[1].media_type == "image/png" + assert msg.parts[2].url == "workspace:/workspace/.molecule/chat-uploads/b.txt" + + +def test_new_response_message_files_only_no_text(): + """Empty text omits the text Part — useful when replying with files only.""" + from executor_helpers import new_response_message + + ctx = SimpleNamespace(task_id="t", context_id="c") + files = [{"path": "/x.txt", "name": "x.txt", "mime_type": "text/plain"}] + msg = new_response_message(ctx, "", files=files) + + assert len(msg.parts) == 1 + assert msg.parts[0].url == "workspace:/x.txt" + + +def test_new_response_message_falls_back_when_context_ids_unset(): + """RequestContextBuilder always populates task_id/context_id in + production, but unit tests + edge cases may have None. Helper falls + back to fresh UUIDs so the resulting Message is still well-formed.""" + from executor_helpers import new_response_message + + ctx = SimpleNamespace(task_id=None, context_id=None) + msg = new_response_message(ctx, "hi") + + # Both should be 32-char hex UUIDs (fallback path) + assert len(msg.task_id) == 32 + assert len(msg.context_id) == 32 + # And they should be DIFFERENT (not accidentally the same uuid) + assert msg.task_id != msg.context_id + + +def test_new_response_message_handles_missing_attrs(): + """getattr with default — context object lacking task_id/context_id + attributes entirely (not just None) still works.""" + from executor_helpers import new_response_message + + class BareContext: + pass + + msg = new_response_message(BareContext(), "hi") + assert len(msg.task_id) == 32 # fallback uuid + assert len(msg.context_id) == 32