distributed-a2a 0.1.26__tar.gz → 0.2.1__tar.gz

{distributed_a2a-0.1.26/distributed_a2a.egg-info → distributed_a2a-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: distributed_a2a
-Version: 0.1.26
+Version: 0.2.1
 Summary: A library for building A2A agents with routing capabilities
 Home-page: https://github.com/Barra-Technologies/distributed-a2a
 Author: Fabian Bell

{distributed_a2a-0.1.26 → distributed_a2a-0.2.1}/distributed_a2a/__init__.py

@@ -1,4 +1,4 @@
-from .client import A2ATimeoutError, RoutingA2AClient
+from .client import A2ATimeoutError, AgentReply, FileRef, RoutingA2AClient
 from .model import (AgentConfig, AgentItem, CardConfig, LLMConfig,
                     RegistryConfig, RegistryItemConfig, RouterConfig,
                     RouterItem, SkillConfig)
@@ -15,6 +15,8 @@ __all__ = [
     "load_router",
     "RoutingA2AClient",
     "A2ATimeoutError",
+    "AgentReply",
+    "FileRef",
     "load_registry",
     "AgentConfig",
     "SkillConfig",

{distributed_a2a-0.1.26 → distributed_a2a-0.2.1}/distributed_a2a/agent.py

@@ -3,7 +3,7 @@ from typing import Any, Literal, cast
 
 from a2a.types import TaskState
 from langchain.agents import create_agent
-from langchain_core.messages import HumanMessage
+from langchain_core.messages import BaseMessage, HumanMessage
 from langchain_core.runnables import RunnableConfig
 from langchain_core.tools import BaseTool
 from langgraph.checkpoint.base import BaseCheckpointSaver
@@ -38,6 +38,21 @@ class StringResponse(AgentResponse):
     response: str = Field(description="The main response to be returned to the user")
 
 
+class AgentInvocation[ResponseT: AgentResponse](BaseModel):
+    """Result of invoking a :class:`StatusAgent`.
+
+    Exposes both the structured response (LLM-visible content the agent decided
+    to return) and the raw message list produced by the underlying LangGraph
+    agent, so callers can mine ``ToolMessage.artifact`` for out-of-band binary
+    payloads (e.g. files returned by MCP tools) without those bytes ever passing
+    through the LLM context.
+    """
+
+    structured: ResponseT
+    messages: list[BaseMessage]
+    model_config = {"arbitrary_types_allowed": True}
+
+
 class StatusAgent[ResponseT: AgentResponse]:
 
     def __init__(self,
@@ -77,7 +92,7 @@ class StatusAgent[ResponseT: AgentResponse]:
 
     async def __call__(self,
                        message: str,
-                       context_id: str | None = None) -> ResponseT:
+                       context_id: str | None = None) -> AgentInvocation[ResponseT]:
         config: RunnableConfig = RunnableConfig(
             configurable={'thread_id': context_id}
         )
@@ -86,4 +101,7 @@ class StatusAgent[ResponseT: AgentResponse]:
             config
         )
         logging.info("agent response: %s", response)
-        return cast(ResponseT, response['structured_response'])
+        return AgentInvocation[ResponseT](
+            structured=cast(ResponseT, response['structured_response']),
+            messages=list(response.get('messages', [])),
+        )

{distributed_a2a-0.1.26 → distributed_a2a-0.2.1}/distributed_a2a/client.py

@@ -1,18 +1,48 @@
 import asyncio
 import json
 import time
+from dataclasses import dataclass, field
 from uuid import uuid4
 
 import httpx
 from a2a.client import (A2ACardResolver, ClientConfig, ClientEvent,
                         ClientFactory, create_text_message_object)
-from a2a.types import (AgentCard, Message, Part, Task, TaskQueryParams,
-                       TaskState, TextPart)
+from a2a.types import (AgentCard, FilePart, FileWithBytes, FileWithUri,
+                       Message, Part, Task, TaskQueryParams, TaskState,
+                       TextPart)
 
 DEFAULT_MAX_POLLS = 50
 DEFAULT_POLL_INTERVAL = 1.0
 
 
+@dataclass
+class FileRef:
+    """A file payload received as part of an A2A agent reply.
+
+    Exactly one of ``bytes_b64`` (for ``FileWithBytes``) or ``uri`` (for
+    ``FileWithUri``) is populated. ``bytes_b64`` is the raw base64 string
+    delivered over the wire by the A2A SDK — the caller is responsible for
+    decoding before forwarding the bytes.
+    """
+
+    name: str
+    mime_type: str
+    bytes_b64: str = ""
+    uri: str | None = None
+
+
+@dataclass
+class AgentReply:
+    """Structured reply from a routing-aware agent.
+
+    Carries the user-visible text (if any) plus zero-or-more files that the
+    agent emitted out-of-band as ``FilePart`` artifacts.
+    """
+
+    text: str | None = None
+    files: list[FileRef] = field(default_factory=list)
+
+
 class A2ATimeoutError(Exception):
     """Raised when polling a remote agent task exceeds max_polls."""
 
@@ -69,8 +99,10 @@ class RemoteAgentConnection:
         response: Task = await self.agent_client.get_task(query_params)
         return response
 
-    async def send_message(self, message_to_send: str, context_id: str,
-                           task_id: None | str = None) -> str | AgentCard | TaskState:
+    async def send_message(self,
+                           message_to_send: str,
+                           context_id: str,
+                           task_id: None | str = None) -> AgentReply | AgentCard | TaskState:
         message: Message = create_text_message_object(content=message_to_send)
         message.message_id = str(uuid4())
         message.context_id = context_id
@@ -112,13 +144,36 @@ class RemoteAgentConnection:
         for artifact in response.artifacts or []:
             match artifact.name, artifact.parts:
                 case 'routing_error', [Part(root=TextPart(text=error_msg)), *_]:
-                    return error_msg
+                    return AgentReply(text=error_msg)
                 case 'rejected', [Part(root=TextPart()), *_]:
                     return TaskState.rejected
                 case 'target_agent', [Part(root=TextPart(text=agent_card_str)), *_]:
                     return AgentCard(**json.loads(agent_card_str))
-                case 'current_result', [Part(root=TextPart(text=result)), *_]:
-                    return result
+
+        text_out: str | None = None
+        files_out: list[FileRef] = []
+        for artifact in response.artifacts or []:
+            for part in artifact.parts or []:
+                root = getattr(part, "root", None)
+                if isinstance(root, TextPart) and artifact.name == "current_result":
+                    text_out = root.text
+                elif isinstance(root, FilePart):
+                    root_file = root.file
+                    if isinstance(root_file, FileWithBytes):
+                        files_out.append(FileRef(
+                            name=root_file.name or artifact.name or "file.bin",
+                            mime_type=root_file.mime_type or "application/octet-stream",
+                            bytes_b64=root_file.bytes,
+                        ))
+                    elif isinstance(root_file, FileWithUri):
+                        files_out.append(FileRef(
+                            name=root_file.name or artifact.name or "file.bin",
+                            mime_type=root_file.mime_type or "application/octet-stream",
+                            uri=root_file.uri,
+                        ))
+
+        if text_out is not None or files_out:
+            return AgentReply(text=text_out, files=files_out)
 
         if task_state == TaskState.rejected:
             return TaskState.rejected
@@ -157,8 +212,10 @@ class RoutingA2AClient:
             await card_resolver.get_agent_card()
         )
 
-    async def send_message(self, message: str, context_id: str, depth: int = 0,
-                           rejected_agents: list[str] | None = None) -> str:
+    async def send_message(self, message: str,
+                           context_id: str,
+                           depth: int = 0,
+                           rejected_agents: list[str] | None = None) -> AgentReply:
         if depth > MAX_RECURSION_DEPTH:
             raise Exception("Maximum recursion depth exceeded. This is likely due to an infinite loop in your agent.")
 
@@ -182,7 +239,7 @@ class RoutingA2AClient:
             if rejection_msg not in message:
                 message_to_send = f"{message}\n\n{rejection_msg}"
 
-        agent_response: str | AgentCard | TaskState = await agent_connection.send_message(message_to_send, context_id)
+        agent_response: AgentReply | AgentCard | TaskState = await agent_connection.send_message(message_to_send, context_id)
         if isinstance(agent_response, AgentCard):
             if agent_response.url == self.current_card.url:
                 raise Exception("Agent redirected to itself.")

{distributed_a2a-0.1.26 → distributed_a2a-0.2.1}/distributed_a2a/executors.py

@@ -4,12 +4,15 @@ from typing import Any
 
 from a2a.server.agent_execution import AgentExecutor, RequestContext
 from a2a.server.events import EventQueue
-from a2a.types import (Artifact, TaskArtifactUpdateEvent, TaskState,
-                       TaskStatus, TaskStatusUpdateEvent)
+from a2a.types import (Artifact, FilePart, FileWithBytes, Part,
+                       TaskArtifactUpdateEvent, TaskState, TaskStatus,
+                       TaskStatusUpdateEvent)
 from a2a.utils import new_text_artifact
+from langchain_core.messages import BaseMessage, ToolMessage
 from langchain_core.tools import BaseTool
 from langchain_mcp_adapters.client import MultiServerMCPClient
 from langgraph.checkpoint.base import BaseCheckpointSaver
+from mcp.types import BlobResourceContents, EmbeddedResource, ImageContent
 
 from .agent import RoutingResponse, StatusAgent, StringResponse
 from .config import settings
@@ -19,6 +22,32 @@ from .registry import AgentRegistryLookupClient, McpRegistryLookup
 logger = logging.getLogger(__name__)
 
 
+def _extract_file_parts(messages: list[BaseMessage]) -> list[tuple[str, FilePart]]:
+    out: list[tuple[str, FilePart]] = []
+    for message in messages:
+        if not isinstance(message, ToolMessage):
+            continue
+        artifact = getattr(message, "artifact", None)
+        if not artifact:
+            continue
+        for block in artifact:
+            if isinstance(block, EmbeddedResource) and isinstance(block.resource, BlobResourceContents):
+                uri_str = str(block.resource.uri)
+                name = uri_str.rsplit("/", 1)[-1] or "file.bin"
+                out.append((name, FilePart(file=FileWithBytes(
+                    name=name,
+                    mime_type=block.resource.mimeType or "application/octet-stream",
+                    bytes=block.resource.blob,
+                ))))
+            elif isinstance(block, ImageContent):
+                out.append(("image", FilePart(file=FileWithBytes(
+                    name="image",
+                    mime_type=block.mimeType,
+                    bytes=block.data,
+                ))))
+    return out
+
+
 class RoutingFailed(Exception):
     def __init__(self, message: str) -> None:
         super().__init__(message)
@@ -116,21 +145,36 @@ class RoutingAgentExecutor(AgentExecutor):
                 task_id=context.task_id
             ))
             await self.reinitialize_agent_with_tools()
-            agent_response: StringResponse = await self.agent(message=context.get_user_input(),
-                                                              context_id=context.context_id)
+            invocation = await self.agent(message=context.get_user_input(),
+                                          context_id=context.context_id)
+            agent_response: StringResponse = invocation.structured
 
             artifact: Artifact
+            file_parts: list[tuple[str, FilePart]] = []
             if agent_response.status == TaskState.rejected:
                 artifact = await _route_request_to_matching_agent(self.routing_agent, self.agent_registry, context)
             else:
                 logger.info(f"Request with id {context.context_id} was successfully processed by agent.")
+                file_parts = _extract_file_parts(invocation.messages)
                 artifact = new_text_artifact(
                     name='current_result',
                     description='Result of request to agent.',
                     text=f"*{self.agent_config.agent.card.name}*: {agent_response.response}"
                 )
 
-            # publish actual result
+            for name, file_part in file_parts:
+                await event_queue.enqueue_event(TaskArtifactUpdateEvent(
+                    append=False,
+                    last_chunk=False,
+                    context_id=context.context_id,
+                    task_id=context.task_id,
+                    artifact=Artifact(
+                        artifact_id=name,
+                        name=name,
+                        parts=[Part(root=file_part)],
+                    ),
+                ))
+
             await event_queue.enqueue_event(TaskArtifactUpdateEvent(
                 append=False,
                 context_id=context.context_id,
@@ -297,8 +341,9 @@ class RoutingExecutor(AgentExecutor):
 async def _route_request_to_matching_agent(routing_agent: StatusAgent[RoutingResponse],
                                            agent_registry: AgentRegistryLookupClient,
                                            context: RequestContext) -> Artifact:
-    routing_agent_response: RoutingResponse = await routing_agent(message=context.get_user_input(),
-                                                                  context_id=context.context_id)
+    invocation = await routing_agent(message=context.get_user_input(),
+                                     context_id=context.context_id)
+    routing_agent_response: RoutingResponse = invocation.structured
     agent_name: str | None = routing_agent_response.agent_name
     logger.info(f"routing response received: {routing_agent_response}")
     if agent_name is None:

{distributed_a2a-0.1.26 → distributed_a2a-0.2.1/distributed_a2a.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: distributed_a2a
-Version: 0.1.26
+Version: 0.2.1
 Summary: A library for building A2A agents with routing capabilities
 Home-page: https://github.com/Barra-Technologies/distributed-a2a
 Author: Fabian Bell

{distributed_a2a-0.1.26 → distributed_a2a-0.2.1}/distributed_a2a.egg-info/SOURCES.txt

@@ -29,5 +29,7 @@ distributed_a2a/schemas/agent-schema.json
 distributed_a2a/schemas/router-agent-schema.json
 tests/test_app.py
 tests/test_client.py
+tests/test_executor_files.py
+tests/test_files.py
 tests/test_rejection.py
 tests/test_timeout.py

{distributed_a2a-0.1.26 → distributed_a2a-0.2.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "distributed_a2a"
-version = "0.1.26"
+version = "0.2.1"
 description = "A library for building A2A agents with routing capabilities"
 readme = "README.md"
 requires-python = ">=3.14"

{distributed_a2a-0.1.26 → distributed_a2a-0.2.1}/tests/test_app.py

@@ -59,10 +59,12 @@ async def test_app_completed_path(fake_registry_server: str, fake_completed_llm:
     with FakeAgent(fake_registry_server, fake_completed_llm, "test-agent") as agent:
         # When
         client = RoutingA2AClient(initial_url=f"http://127.0.0.1:{agent.app_port}/{agent.name}")
-        response = await client.send_message(message="Hello", context_id="test-context")
+        reply = await client.send_message(message="Hello", context_id="test-context")
 
         # Then: Check the response
-        assert "This is a mock response from the fake OpenAI server." in response
+        assert reply.text is not None
+        assert "This is a mock response from the fake OpenAI server." in reply.text
+        assert reply.files == []
 
 
 @pytest.mark.asyncio
@@ -75,7 +77,8 @@ async def test_app_redirect_path(fake_registry_server: str, fake_completed_llm:
                 client = RoutingA2AClient(initial_url=f"http://127.0.0.1:{first_agent.app_port}")
 
                 # When
-                response = await client.send_message(message="Hello", context_id="test-context")
+                reply = await client.send_message(message="Hello", context_id="test-context")
 
                 # Then
-                assert FINAL_RESPONSE in response
+                assert reply.text is not None
+                assert FINAL_RESPONSE in reply.text

distributed_a2a-0.2.1/tests/test_executor_files.py

@@ -0,0 +1,182 @@
+from __future__ import annotations
+
+import base64
+from typing import Any
+
+import pytest
+from a2a.server.agent_execution import RequestContext
+from a2a.server.events import EventQueue
+from a2a.types import FilePart, FileWithBytes
+from a2a.types import Message as A2AMessage
+from a2a.types import (MessageSendParams, Part, Role, TaskArtifactUpdateEvent,
+                       TaskState, TaskStatusUpdateEvent, TextPart)
+from langchain_core.messages import BaseMessage, HumanMessage, ToolMessage
+from mcp.types import BlobResourceContents, EmbeddedResource
+from pydantic import AnyUrl
+
+from distributed_a2a.agent import AgentInvocation, StringResponse
+from distributed_a2a.executors import RoutingAgentExecutor
+
+
+class _StubStatusAgent:
+    """Minimal stand-in for ``StatusAgent`` used in the executor."""
+
+    def __init__(self, response: StringResponse, messages: list[BaseMessage]) -> None:
+        self._response = response
+        self._messages = messages
+
+    async def __call__(self, message: str, context_id: str | None = None) -> AgentInvocation[StringResponse]:
+        return AgentInvocation[StringResponse](
+            structured=self._response,
+            messages=self._messages,
+        )
+
+
+def _make_request_context() -> RequestContext:
+    msg = A2AMessage(
+        message_id="m-1",
+        context_id="ctx-1",
+        task_id="task-1",
+        role=Role.user,
+        parts=[Part(root=TextPart(text="render a CV please"))],
+    )
+    return RequestContext(
+        request=MessageSendParams(message=msg),
+        task_id="task-1",
+        context_id="ctx-1",
+    )
+
+
+def _make_tool_message_with_docx(b64: str) -> ToolMessage:
+    return ToolMessage(
+        content='{"filename": "cv-foo.docx"}',
+        tool_call_id="call-1",
+        artifact=[
+            EmbeddedResource(
+                type="resource",
+                resource=BlobResourceContents(
+                    uri=AnyUrl("cv://cv-foo.docx"),
+                    mimeType=(
+                        "application/vnd.openxmlformats-officedocument."
+                        "wordprocessingml.document"
+                    ),
+                    blob=b64,
+                ),
+            ),
+        ],
+    )
+
+
+async def _drain_queue(queue: EventQueue) -> list[Any]:
+    """Dequeue every event currently in ``queue`` without blocking.
+
+    Uses ``no_wait=True`` and stops on the first empty/closed exception. We
+    intentionally avoid ``await queue.close()`` because the graceful close
+    waits for ``queue.join()`` — which requires the consumer to call
+    ``task_done()`` for every enqueued item — and would otherwise deadlock
+    the test once we've finished dequeuing.
+    """
+    events: list[Any] = []
+    while True:
+        try:
+            evt = await queue.dequeue_event(no_wait=True)
+        except Exception:
+            break
+        events.append(evt)
+        queue.task_done()
+    return events
+
+
+@pytest.mark.asyncio
+async def test_executor_emits_file_part_event_for_embedded_resource(
+        monkeypatch: pytest.MonkeyPatch) -> None:
+    docx_b64 = base64.b64encode(b"PK\x03\x04 fake docx").decode("ascii")
+    tool_msg = _make_tool_message_with_docx(docx_b64)
+
+    # Build the executor without invoking its real __init__ (which requires
+    # API keys + registry network calls). We only need ``execute`` to use:
+    #   - self.agent (stubbed)
+    #   - self.reinitialize_agent_with_tools (stubbed no-op)
+    #   - self.agent_config.agent.card.name (stubbed via SimpleNamespace)
+    from types import SimpleNamespace
+
+    executor = RoutingAgentExecutor.__new__(RoutingAgentExecutor)
+    executor.agent_config = SimpleNamespace(  # type: ignore[assignment]
+        agent=SimpleNamespace(card=SimpleNamespace(name="cv-agent")),
+    )
+    executor.agent = _StubStatusAgent(  # type: ignore[assignment]
+        StringResponse(status=TaskState.completed,
+                       response="Here is your CV."),
+        [HumanMessage(content="render a CV please"), tool_msg],
+    )
+
+    async def _noop_reinit() -> None:
+        return None
+
+    executor.reinitialize_agent_with_tools = _noop_reinit  # type: ignore[method-assign]
+
+    ctx = _make_request_context()
+    queue = EventQueue()
+    await executor.execute(ctx, queue)
+    events = await _drain_queue(queue)
+
+    # Expected sequence: working status, file artifact, text artifact, final status.
+    artifact_events = [e for e in events if isinstance(e, TaskArtifactUpdateEvent)]
+    status_events = [e for e in events if isinstance(e, TaskStatusUpdateEvent)]
+
+    assert len(artifact_events) == 2, (
+        f"Expected 2 artifact events (file + text), got: "
+        f"{[(e.artifact.name, e.last_chunk) for e in artifact_events]}"
+    )
+
+    file_event, text_event = artifact_events
+    # 1) file artifact carries a FilePart, is NOT the final chunk
+    assert file_event.last_chunk is False
+    assert file_event.artifact.name == "cv-foo.docx"
+    assert len(file_event.artifact.parts) == 1
+    file_part = file_event.artifact.parts[0].root
+    assert isinstance(file_part, FilePart)
+    assert isinstance(file_part.file, FileWithBytes)
+    assert file_part.file.name == "cv-foo.docx"
+    assert file_part.file.bytes == docx_b64
+
+    # 2) text artifact carries the LLM-visible summary, IS the final chunk
+    assert text_event.last_chunk is True
+    assert text_event.artifact.name == "current_result"
+    text_part = text_event.artifact.parts[0].root
+    assert isinstance(text_part, TextPart)
+    assert text_part.text == "*cv-agent*: Here is your CV."
+
+    # 3) Final status event is `completed`
+    final_status = [e for e in status_events if e.final]
+    assert len(final_status) == 1
+    assert final_status[0].status.state == TaskState.completed
+
+
+@pytest.mark.asyncio
+async def test_executor_emits_no_file_event_when_no_artifacts() -> None:
+    from types import SimpleNamespace
+
+    executor = RoutingAgentExecutor.__new__(RoutingAgentExecutor)
+    executor.agent_config = SimpleNamespace(  # type: ignore[assignment]
+        agent=SimpleNamespace(card=SimpleNamespace(name="plain-agent")),
+    )
+    executor.agent = _StubStatusAgent(  # type: ignore[assignment]
+        StringResponse(status=TaskState.completed, response="No files here."),
+        [HumanMessage(content="hi")],
+    )
+
+    async def _noop_reinit() -> None:
+        return None
+
+    executor.reinitialize_agent_with_tools = _noop_reinit  # type: ignore[method-assign]
+
+    ctx = _make_request_context()
+    queue = EventQueue()
+    await executor.execute(ctx, queue)
+    events = await _drain_queue(queue)
+
+    artifact_events = [e for e in events if isinstance(e, TaskArtifactUpdateEvent)]
+    assert len(artifact_events) == 1
+    assert artifact_events[0].artifact.name == "current_result"
+    assert artifact_events[0].last_chunk is True

distributed_a2a-0.2.1/tests/test_files.py

@@ -0,0 +1,237 @@
+from __future__ import annotations
+
+import base64
+from collections.abc import AsyncGenerator
+
+import httpx
+import pytest
+from a2a.types import (AgentCapabilities, AgentCard, Artifact, FilePart,
+                       FileWithBytes, FileWithUri, Message, Part, Task,
+                       TaskState, TaskStatus, TextPart)
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+from mcp.types import (BlobResourceContents, EmbeddedResource, ImageContent,
+                       TextResourceContents)
+from pydantic import AnyUrl
+
+from distributed_a2a.client import AgentReply, RemoteAgentConnection
+from distributed_a2a.executors import _extract_file_parts
+
+
+def _b64(payload: bytes) -> str:
+    return base64.b64encode(payload).decode("ascii")
+
+
+def test_extract_file_parts_picks_up_embedded_resource() -> None:
+    docx_b64 = _b64(b"PK\x03\x04 fake docx bytes")
+    tool_msg = ToolMessage(
+        content="{\"filename\": \"cv-foo.docx\"}",
+        tool_call_id="call-1",
+        artifact=[
+            EmbeddedResource(
+                type="resource",
+                resource=BlobResourceContents(
+                    uri=AnyUrl("cv://cv-foo.docx"),
+                    mimeType="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+                    blob=docx_b64,
+                ),
+            ),
+        ],
+    )
+
+    parts = _extract_file_parts([HumanMessage(content="hi"), tool_msg])
+
+    assert len(parts) == 1
+    name, file_part = parts[0]
+    assert name == "cv-foo.docx"
+    assert isinstance(file_part, FilePart)
+    assert isinstance(file_part.file, FileWithBytes)
+    assert file_part.file.name == "cv-foo.docx"
+    assert file_part.file.mime_type == (
+        "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+    )
+    assert file_part.file.bytes == docx_b64
+
+
+def test_extract_file_parts_handles_image_content() -> None:
+    png_b64 = _b64(b"\x89PNG\r\n\x1a\n fake")
+    tool_msg = ToolMessage(
+        content="",
+        tool_call_id="call-img",
+        artifact=[
+            ImageContent(type="image", data=png_b64, mimeType="image/png"),
+        ],
+    )
+
+    parts = _extract_file_parts([tool_msg])
+
+    assert len(parts) == 1
+    name, file_part = parts[0]
+    assert name == "image"
+    assert isinstance(file_part.file, FileWithBytes)
+    assert file_part.file.mime_type == "image/png"
+    assert file_part.file.bytes == png_b64
+
+
+def test_extract_file_parts_ignores_non_tool_messages_and_empty_artifacts() -> None:
+    messages = [
+        HumanMessage(content="hi"),
+        AIMessage(content="hello"),
+        ToolMessage(content="just text", tool_call_id="c1"),
+        ToolMessage(content="empty artifact", tool_call_id="c2", artifact=[]),
+    ]
+    assert _extract_file_parts(messages) == []
+
+
+def test_extract_file_parts_skips_text_resource_contents() -> None:
+    tool_msg = ToolMessage(
+        content="",
+        tool_call_id="call-x",
+        artifact=[
+            EmbeddedResource(
+                type="resource",
+                resource=TextResourceContents(
+                    uri=AnyUrl("text://note.txt"),
+                    mimeType="text/plain",
+                    text="some inline text",
+                ),
+            ),
+        ],
+    )
+    assert _extract_file_parts([tool_msg]) == []
+
+
+def test_extract_file_parts_falls_back_to_octet_stream_when_mime_missing() -> None:
+    tool_msg = ToolMessage(
+        content="",
+        tool_call_id="call-x",
+        artifact=[
+            EmbeddedResource(
+                type="resource",
+                resource=BlobResourceContents(
+                    uri=AnyUrl("file:///tmp/x.bin"),
+                    blob=_b64(b"x"),
+                ),
+            ),
+        ],
+    )
+    parts = _extract_file_parts([tool_msg])
+    assert len(parts) == 1
+    name, file_part = parts[0]
+    assert name == "x.bin"
+    assert isinstance(file_part.file, FileWithBytes)
+    assert file_part.file.mime_type == "application/octet-stream"
+
+
+class _StubAgentClient:
+    def __init__(self, task: Task):
+        self._task = task
+
+    async def send_message(self, _message: Message) -> AsyncGenerator[tuple[Task, None], None]:
+        yield self._task, None
+
+
+class _StubClientFactory:
+    def __init__(self, _config: object, task: Task):
+        self._task = task
+
+    def create(self, _agent_card: AgentCard) -> _StubAgentClient:
+        return _StubAgentClient(self._task)
+
+
+def _patch_client_factory(monkeypatch: pytest.MonkeyPatch, task: Task) -> None:
+    def _ctor(config: object) -> _StubClientFactory:
+        return _StubClientFactory(config, task)
+    monkeypatch.setattr("distributed_a2a.client.ClientFactory", _ctor)
+
+
+def _agent_card() -> AgentCard:
+    return AgentCard(
+        name="stub",
+        description="stub",
+        url="http://127.0.0.1:0",
+        version="1.0.0",
+        default_input_modes=["text"],
+        default_output_modes=["text"],
+        capabilities=AgentCapabilities(streaming=False, push_notifications=False),
+        skills=[],
+        preferred_transport="JSONRPC",
+    )
+
+
+@pytest.mark.asyncio
+async def test_remote_agent_connection_returns_text_and_files_in_agent_reply(
+        monkeypatch: pytest.MonkeyPatch) -> None:
+    docx_b64 = _b64(b"PK\x03\x04 docx-bytes")
+    completed_task = Task(
+        id="task-files",
+        context_id="ctx-files",
+        status=TaskStatus(state=TaskState.completed),
+        artifacts=[
+            Artifact(
+                artifact_id="cv-foo.docx",
+                name="cv-foo.docx",
+                parts=[Part(root=FilePart(file=FileWithBytes(
+                    name="cv-foo.docx",
+                    mime_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+                    bytes=docx_b64,
+                )))],
+            ),
+            Artifact(
+                artifact_id="current_result",
+                name="current_result",
+                parts=[Part(root=TextPart(text="Here is your CV."))],
+            ),
+        ],
+    )
+    factory_task = completed_task
+    _patch_client_factory(monkeypatch, factory_task)
+
+    async with httpx.AsyncClient() as http_client:
+        conn = RemoteAgentConnection(_agent_card(), http_client)
+        reply = await conn.send_message("hello", "ctx-files")
+
+    assert isinstance(reply, AgentReply)
+    assert reply.text == "Here is your CV."
+    assert len(reply.files) == 1
+    f = reply.files[0]
+    assert f.name == "cv-foo.docx"
+    assert f.mime_type == (
+        "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+    )
+    assert f.bytes_b64 == docx_b64
+    assert f.uri is None
+
+
+@pytest.mark.asyncio
+async def test_remote_agent_connection_handles_file_with_uri(
+        monkeypatch: pytest.MonkeyPatch) -> None:
+    completed_task = Task(
+        id="task-uri",
+        context_id="ctx-uri",
+        status=TaskStatus(state=TaskState.completed),
+        artifacts=[
+            Artifact(
+                artifact_id="report",
+                name="report",
+                parts=[Part(root=FilePart(file=FileWithUri(
+                    name="report.pdf",
+                    mime_type="application/pdf",
+                    uri="https://example.com/report.pdf",
+                )))],
+            ),
+        ],
+    )
+    _patch_client_factory(monkeypatch, completed_task)
+
+    async with httpx.AsyncClient() as http_client:
+        conn = RemoteAgentConnection(_agent_card(), http_client)
+        reply = await conn.send_message("hello", "ctx-uri")
+
+    assert isinstance(reply, AgentReply)
+    assert reply.text is None
+    assert len(reply.files) == 1
+    f = reply.files[0]
+    assert f.name == "report.pdf"
+    assert f.mime_type == "application/pdf"
+    assert f.bytes_b64 == ""
+    assert f.uri == "https://example.com/report.pdf"

{distributed_a2a-0.1.26 → distributed_a2a-0.2.1}/tests/test_rejection.py

@@ -1,7 +1,7 @@
 import pytest
 from a2a.types import AgentCapabilities, AgentCard, TaskState
 
-from distributed_a2a.client import RoutingA2AClient
+from distributed_a2a.client import AgentReply, RoutingA2AClient
 
 
 def build_card(name: str, url: str) -> AgentCard:
@@ -38,7 +38,7 @@ async def test_rejection_triggers_automated_rerouting(monkeypatch: pytest.Monkey
         def __init__(self, agent_card: AgentCard, _client: object, **_kwargs: object) -> None:
             self.agent_card = agent_card
 
-        async def send_message(self, message_to_send: str, _context_id: str) -> str | AgentCard | TaskState:
+        async def send_message(self, message_to_send: str, _context_id: str) -> AgentReply | AgentCard | TaskState:
             if self.agent_card.name == "Router":
                 router_messages.append(message_to_send)
                 if "Please exclude the following agents from routing: rejecting-agent" in message_to_send:
@@ -49,7 +49,7 @@ async def test_rejection_triggers_automated_rerouting(monkeypatch: pytest.Monkey
                 return TaskState.rejected
 
             if self.agent_card.name == "success-agent":
-                return "final answer"
+                return AgentReply(text="final answer")
 
             raise AssertionError(f"Unexpected agent {self.agent_card.name}")
 
@@ -57,7 +57,8 @@ async def test_rejection_triggers_automated_rerouting(monkeypatch: pytest.Monkey
 
     result = await client.send_message("Hello", context_id="ctx-1")
 
-    assert result == "final answer"
+    assert isinstance(result, AgentReply)
+    assert result.text == "final answer"
     assert len(router_messages) == 2
     assert "Please exclude the following agents from routing" not in router_messages[0]
     assert "Please exclude the following agents from routing: rejecting-agent" in router_messages[1]
@@ -88,7 +89,7 @@ async def test_rejected_agents_reset_between_calls(monkeypatch: pytest.MonkeyPat
         def __init__(self, agent_card: AgentCard, _client: object, **_kwargs: object) -> None:
             self.agent_card = agent_card
 
-        async def send_message(self, message_to_send: str, _context_id: str) -> str | AgentCard | TaskState:
+        async def send_message(self, message_to_send: str, _context_id: str) -> AgentReply | AgentCard | TaskState:
             idx = next_step()
             match idx:
                 case 0:
@@ -104,7 +105,7 @@ async def test_rejected_agents_reset_between_calls(monkeypatch: pytest.MonkeyPat
                     return success_card
                 case 3:
                     assert self.agent_card.name == "success-agent"
-                    return "first response"
+                    return AgentReply(text="first response")
                 case 4:
                     assert self.agent_card.name == "success-agent"
                     return TaskState.rejected
@@ -115,7 +116,7 @@ async def test_rejected_agents_reset_between_calls(monkeypatch: pytest.MonkeyPat
                     return rejecting_card
                 case 6:
                     assert self.agent_card.name == "rejecting-agent"
-                    return "second response"
+                    return AgentReply(text="second response")
                 case _:
                     raise AssertionError(f"Unexpected step {idx}")
 
@@ -124,8 +125,8 @@ async def test_rejected_agents_reset_between_calls(monkeypatch: pytest.MonkeyPat
     first = await client.send_message("First", context_id="ctx-1")
     second = await client.send_message("Second", context_id="ctx-2")
 
-    assert first == "first response"
-    assert second == "second response"
+    assert isinstance(first, AgentReply) and first.text == "first response"
+    assert isinstance(second, AgentReply) and second.text == "second response"
 
 
 @pytest.mark.asyncio
@@ -147,7 +148,7 @@ async def test_fails_when_router_returns_already_rejected_agent(monkeypatch: pyt
         def __init__(self, agent_card: AgentCard, _client: object, **_kwargs: object) -> None:
             self.agent_card = agent_card
 
-        async def send_message(self, message_to_send: str, _context_id: str) -> str | AgentCard | TaskState:
+        async def send_message(self, message_to_send: str, _context_id: str) -> AgentReply | AgentCard | TaskState:
             idx = step["idx"]
             step["idx"] += 1
             if idx == 0:

{distributed_a2a-0.1.26 → distributed_a2a-0.2.1}/tests/test_timeout.py

@@ -8,8 +8,8 @@ import pytest
 from a2a.types import (AgentCapabilities, AgentCard, Artifact, Message, Part,
                        Role, Task, TaskState, TaskStatus, TextPart)
 
-from distributed_a2a.client import (A2ATimeoutError, RemoteAgentConnection,
-                                    RoutingA2AClient)
+from distributed_a2a.client import (A2ATimeoutError, AgentReply,
+                                    RemoteAgentConnection, RoutingA2AClient)
 
 
 def _agent_card() -> AgentCard:
@@ -124,7 +124,9 @@ async def test_send_message_honors_completion_on_final_allowed_poll(
         )
         result = await connection.send_message("hi", "ctx-x")
 
-    assert result == "all done"
+    assert isinstance(result, AgentReply)
+    assert result.text == "all done"
+    assert result.files == []
     assert scripted.get_task_calls == 3
     assert scripted._get_states == []
 
@@ -156,7 +158,9 @@ async def test_send_message_returns_immediately_if_already_completed(
         )
         result = await connection.send_message("hi", "ctx-x")
 
-    assert result == "all done"
+    assert isinstance(result, AgentReply)
+    assert result.text == "all done"
+    assert result.files == []
 
 
 def test_a2a_timeout_error_message_includes_diagnostics() -> None:
@@ -182,8 +186,8 @@ async def test_routing_a2a_client_propagates_poll_kwargs(
             captured["poll_interval"] = poll_interval
             self.agent_card = agent_card
 
-        async def send_message(self, _message: str, _context_id: str) -> str:
-            return "done"
+        async def send_message(self, _message: str, _context_id: str) -> AgentReply:
+            return AgentReply(text="done")
 
     monkeypatch.setattr("distributed_a2a.client.RemoteAgentConnection", _Capturing)
 
@@ -197,5 +201,7 @@ async def test_routing_a2a_client_propagates_poll_kwargs(
     monkeypatch.setattr(client, "fetch_initial_card", _fetch)
 
     result = await client.send_message("hi", context_id="ctx-1")
-    assert result == "done"
+    assert isinstance(result, AgentReply)
+    assert result.text == "done"
+    assert result.files == []
     assert captured == {"max_polls": 7, "poll_interval": 0.25}