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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (38) hide show
  1. {distributed_a2a-0.1.26/distributed_a2a.egg-info → distributed_a2a-0.2.2}/PKG-INFO +1 -1
  2. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/__init__.py +3 -1
  3. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/agent.py +21 -3
  4. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/client.py +67 -10
  5. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/executors.py +24 -7
  6. distributed_a2a-0.2.2/distributed_a2a/files.py +68 -0
  7. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2/distributed_a2a.egg-info}/PKG-INFO +1 -1
  8. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a.egg-info/SOURCES.txt +3 -0
  9. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/pyproject.toml +1 -1
  10. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/tests/test_app.py +7 -4
  11. distributed_a2a-0.2.2/tests/test_executor_files.py +216 -0
  12. distributed_a2a-0.2.2/tests/test_files.py +317 -0
  13. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/tests/test_rejection.py +11 -10
  14. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/tests/test_timeout.py +13 -7
  15. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/LICENSE +0 -0
  16. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/MANIFEST.in +0 -0
  17. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/README.md +0 -0
  18. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/config.py +0 -0
  19. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/model.py +0 -0
  20. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/py.typed +0 -0
  21. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/registry.py +0 -0
  22. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/__init__.py +0 -0
  23. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/bootstrap.py +0 -0
  24. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/dynamo_db.py +0 -0
  25. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/in_memory_registry_storage.py +0 -0
  26. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/model.py +0 -0
  27. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/registry_server/storage.py +0 -0
  28. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/router.py +0 -0
  29. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/schemas/agent-schema.json +0 -0
  30. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/schemas/router-agent-schema.json +0 -0
  31. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a/server.py +0 -0
  32. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a.egg-info/dependency_links.txt +0 -0
  33. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a.egg-info/requires.txt +0 -0
  34. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/distributed_a2a.egg-info/top_level.txt +0 -0
  35. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/requirements.txt +0 -0
  36. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/setup.cfg +0 -0
  37. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/setup.py +0 -0
  38. {distributed_a2a-0.1.26 → distributed_a2a-0.2.2}/tests/test_client.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: distributed_a2a
3
- Version: 0.1.26
3
+ Version: 0.2.2
4
4
  Summary: A library for building A2A agents with routing capabilities
5
5
  Home-page: https://github.com/Barra-Technologies/distributed-a2a
6
6
  Author: Fabian Bell
@@ -1,4 +1,4 @@
1
- from .client import A2ATimeoutError, RoutingA2AClient
1
+ from .client import A2ATimeoutError, AgentReply, FileRef, RoutingA2AClient
2
2
  from .model import (AgentConfig, AgentItem, CardConfig, LLMConfig,
3
3
  RegistryConfig, RegistryItemConfig, RouterConfig,
4
4
  RouterItem, SkillConfig)
@@ -15,6 +15,8 @@ __all__ = [
15
15
  "load_router",
16
16
  "RoutingA2AClient",
17
17
  "A2ATimeoutError",
18
+ "AgentReply",
19
+ "FileRef",
18
20
  "load_registry",
19
21
  "AgentConfig",
20
22
  "SkillConfig",
@@ -3,7 +3,7 @@ from typing import Any, Literal, cast
3
3
 
4
4
  from a2a.types import TaskState
5
5
  from langchain.agents import create_agent
6
- from langchain_core.messages import HumanMessage
6
+ from langchain_core.messages import BaseMessage, HumanMessage
7
7
  from langchain_core.runnables import RunnableConfig
8
8
  from langchain_core.tools import BaseTool
9
9
  from langgraph.checkpoint.base import BaseCheckpointSaver
@@ -38,6 +38,21 @@ class StringResponse(AgentResponse):
38
38
  response: str = Field(description="The main response to be returned to the user")
39
39
 
40
40
 
41
+ class AgentInvocation[ResponseT: AgentResponse](BaseModel):
42
+ """Result of invoking a :class:`StatusAgent`.
43
+
44
+ Exposes both the structured response (LLM-visible content the agent decided
45
+ to return) and the raw message list produced by the underlying LangGraph
46
+ agent, so callers can mine ``ToolMessage.artifact`` for out-of-band binary
47
+ payloads (e.g. files returned by MCP tools) without those bytes ever passing
48
+ through the LLM context.
49
+ """
50
+
51
+ structured: ResponseT
52
+ messages: list[BaseMessage]
53
+ model_config = {"arbitrary_types_allowed": True}
54
+
55
+
41
56
  class StatusAgent[ResponseT: AgentResponse]:
42
57
 
43
58
  def __init__(self,
@@ -77,7 +92,7 @@ class StatusAgent[ResponseT: AgentResponse]:
77
92
 
78
93
  async def __call__(self,
79
94
  message: str,
80
- context_id: str | None = None) -> ResponseT:
95
+ context_id: str | None = None) -> AgentInvocation[ResponseT]:
81
96
  config: RunnableConfig = RunnableConfig(
82
97
  configurable={'thread_id': context_id}
83
98
  )
@@ -86,4 +101,7 @@ class StatusAgent[ResponseT: AgentResponse]:
86
101
  config
87
102
  )
88
103
  logging.info("agent response: %s", response)
89
- return cast(ResponseT, response['structured_response'])
104
+ return AgentInvocation[ResponseT](
105
+ structured=cast(ResponseT, response['structured_response']),
106
+ messages=list(response.get('messages', [])),
107
+ )
@@ -1,18 +1,48 @@
1
1
  import asyncio
2
2
  import json
3
3
  import time
4
+ from dataclasses import dataclass, field
4
5
  from uuid import uuid4
5
6
 
6
7
  import httpx
7
8
  from a2a.client import (A2ACardResolver, ClientConfig, ClientEvent,
8
9
  ClientFactory, create_text_message_object)
9
- from a2a.types import (AgentCard, Message, Part, Task, TaskQueryParams,
10
- TaskState, TextPart)
10
+ from a2a.types import (AgentCard, FilePart, FileWithBytes, FileWithUri,
11
+ Message, Part, Task, TaskQueryParams, TaskState,
12
+ TextPart)
11
13
 
12
14
  DEFAULT_MAX_POLLS = 50
13
15
  DEFAULT_POLL_INTERVAL = 1.0
14
16
 
15
17
 
18
+ @dataclass
19
+ class FileRef:
20
+ """A file payload received as part of an A2A agent reply.
21
+
22
+ Exactly one of ``bytes_b64`` (for ``FileWithBytes``) or ``uri`` (for
23
+ ``FileWithUri``) is populated. ``bytes_b64`` is the raw base64 string
24
+ delivered over the wire by the A2A SDK — the caller is responsible for
25
+ decoding before forwarding the bytes.
26
+ """
27
+
28
+ name: str
29
+ mime_type: str
30
+ bytes_b64: str = ""
31
+ uri: str | None = None
32
+
33
+
34
+ @dataclass
35
+ class AgentReply:
36
+ """Structured reply from a routing-aware agent.
37
+
38
+ Carries the user-visible text (if any) plus zero-or-more files that the
39
+ agent emitted out-of-band as ``FilePart`` artifacts.
40
+ """
41
+
42
+ text: str | None = None
43
+ files: list[FileRef] = field(default_factory=list)
44
+
45
+
16
46
  class A2ATimeoutError(Exception):
17
47
  """Raised when polling a remote agent task exceeds max_polls."""
18
48
 
@@ -69,8 +99,10 @@ class RemoteAgentConnection:
69
99
  response: Task = await self.agent_client.get_task(query_params)
70
100
  return response
71
101
 
72
- async def send_message(self, message_to_send: str, context_id: str,
73
- task_id: None | str = None) -> str | AgentCard | TaskState:
102
+ async def send_message(self,
103
+ message_to_send: str,
104
+ context_id: str,
105
+ task_id: None | str = None) -> AgentReply | AgentCard | TaskState:
74
106
  message: Message = create_text_message_object(content=message_to_send)
75
107
  message.message_id = str(uuid4())
76
108
  message.context_id = context_id
@@ -112,13 +144,36 @@ class RemoteAgentConnection:
112
144
  for artifact in response.artifacts or []:
113
145
  match artifact.name, artifact.parts:
114
146
  case 'routing_error', [Part(root=TextPart(text=error_msg)), *_]:
115
- return error_msg
147
+ return AgentReply(text=error_msg)
116
148
  case 'rejected', [Part(root=TextPart()), *_]:
117
149
  return TaskState.rejected
118
150
  case 'target_agent', [Part(root=TextPart(text=agent_card_str)), *_]:
119
151
  return AgentCard(**json.loads(agent_card_str))
120
- case 'current_result', [Part(root=TextPart(text=result)), *_]:
121
- return result
152
+
153
+ text_out: str | None = None
154
+ files_out: list[FileRef] = []
155
+ for artifact in response.artifacts or []:
156
+ for part in artifact.parts or []:
157
+ root = getattr(part, "root", None)
158
+ if isinstance(root, TextPart) and artifact.name == "current_result":
159
+ text_out = root.text
160
+ elif isinstance(root, FilePart):
161
+ root_file = root.file
162
+ if isinstance(root_file, FileWithBytes):
163
+ files_out.append(FileRef(
164
+ name=root_file.name or artifact.name or "file.bin",
165
+ mime_type=root_file.mime_type or "application/octet-stream",
166
+ bytes_b64=root_file.bytes,
167
+ ))
168
+ elif isinstance(root_file, FileWithUri):
169
+ files_out.append(FileRef(
170
+ name=root_file.name or artifact.name or "file.bin",
171
+ mime_type=root_file.mime_type or "application/octet-stream",
172
+ uri=root_file.uri,
173
+ ))
174
+
175
+ if text_out is not None or files_out:
176
+ return AgentReply(text=text_out, files=files_out)
122
177
 
123
178
  if task_state == TaskState.rejected:
124
179
  return TaskState.rejected
@@ -157,8 +212,10 @@ class RoutingA2AClient:
157
212
  await card_resolver.get_agent_card()
158
213
  )
159
214
 
160
- async def send_message(self, message: str, context_id: str, depth: int = 0,
161
- rejected_agents: list[str] | None = None) -> str:
215
+ async def send_message(self, message: str,
216
+ context_id: str,
217
+ depth: int = 0,
218
+ rejected_agents: list[str] | None = None) -> AgentReply:
162
219
  if depth > MAX_RECURSION_DEPTH:
163
220
  raise Exception("Maximum recursion depth exceeded. This is likely due to an infinite loop in your agent.")
164
221
 
@@ -182,7 +239,7 @@ class RoutingA2AClient:
182
239
  if rejection_msg not in message:
183
240
  message_to_send = f"{message}\n\n{rejection_msg}"
184
241
 
185
- agent_response: str | AgentCard | TaskState = await agent_connection.send_message(message_to_send, context_id)
242
+ agent_response: AgentReply | AgentCard | TaskState = await agent_connection.send_message(message_to_send, context_id)
186
243
  if isinstance(agent_response, AgentCard):
187
244
  if agent_response.url == self.current_card.url:
188
245
  raise Exception("Agent redirected to itself.")
@@ -4,8 +4,8 @@ from typing import Any
4
4
 
5
5
  from a2a.server.agent_execution import AgentExecutor, RequestContext
6
6
  from a2a.server.events import EventQueue
7
- from a2a.types import (Artifact, TaskArtifactUpdateEvent, TaskState,
8
- TaskStatus, TaskStatusUpdateEvent)
7
+ from a2a.types import (Artifact, FilePart, Part, TaskArtifactUpdateEvent,
8
+ TaskState, TaskStatus, TaskStatusUpdateEvent)
9
9
  from a2a.utils import new_text_artifact
10
10
  from langchain_core.tools import BaseTool
11
11
  from langchain_mcp_adapters.client import MultiServerMCPClient
@@ -13,6 +13,7 @@ from langgraph.checkpoint.base import BaseCheckpointSaver
13
13
 
14
14
  from .agent import RoutingResponse, StatusAgent, StringResponse
15
15
  from .config import settings
16
+ from .files import extract_file_parts
16
17
  from .model import AgentConfig, RouterConfig
17
18
  from .registry import AgentRegistryLookupClient, McpRegistryLookup
18
19
 
@@ -116,21 +117,36 @@ class RoutingAgentExecutor(AgentExecutor):
116
117
  task_id=context.task_id
117
118
  ))
118
119
  await self.reinitialize_agent_with_tools()
119
- agent_response: StringResponse = await self.agent(message=context.get_user_input(),
120
- context_id=context.context_id)
120
+ invocation = await self.agent(message=context.get_user_input(),
121
+ context_id=context.context_id)
122
+ agent_response: StringResponse = invocation.structured
121
123
 
122
124
  artifact: Artifact
125
+ file_parts: list[tuple[str, FilePart]] = []
123
126
  if agent_response.status == TaskState.rejected:
124
127
  artifact = await _route_request_to_matching_agent(self.routing_agent, self.agent_registry, context)
125
128
  else:
126
129
  logger.info(f"Request with id {context.context_id} was successfully processed by agent.")
130
+ file_parts = extract_file_parts(invocation.messages)
127
131
  artifact = new_text_artifact(
128
132
  name='current_result',
129
133
  description='Result of request to agent.',
130
134
  text=f"*{self.agent_config.agent.card.name}*: {agent_response.response}"
131
135
  )
132
136
 
133
- # publish actual result
137
+ for name, file_part in file_parts:
138
+ await event_queue.enqueue_event(TaskArtifactUpdateEvent(
139
+ append=False,
140
+ last_chunk=False,
141
+ context_id=context.context_id,
142
+ task_id=context.task_id,
143
+ artifact=Artifact(
144
+ artifact_id=name,
145
+ name=name,
146
+ parts=[Part(root=file_part)],
147
+ ),
148
+ ))
149
+
134
150
  await event_queue.enqueue_event(TaskArtifactUpdateEvent(
135
151
  append=False,
136
152
  context_id=context.context_id,
@@ -297,8 +313,9 @@ class RoutingExecutor(AgentExecutor):
297
313
  async def _route_request_to_matching_agent(routing_agent: StatusAgent[RoutingResponse],
298
314
  agent_registry: AgentRegistryLookupClient,
299
315
  context: RequestContext) -> Artifact:
300
- routing_agent_response: RoutingResponse = await routing_agent(message=context.get_user_input(),
301
- context_id=context.context_id)
316
+ invocation = await routing_agent(message=context.get_user_input(),
317
+ context_id=context.context_id)
318
+ routing_agent_response: RoutingResponse = invocation.structured
302
319
  agent_name: str | None = routing_agent_response.agent_name
303
320
  logger.info(f"routing response received: {routing_agent_response}")
304
321
  if agent_name is None:
@@ -0,0 +1,68 @@
1
+ import json
2
+ import mimetypes
3
+ from typing import Any
4
+
5
+ from a2a.types import FilePart, FileWithBytes
6
+ from langchain_core.messages import BaseMessage, ToolMessage
7
+
8
+ _LANGCHAIN_BINARY_BLOCK_TYPES: dict[str, str] = {
9
+ "file": "attachment",
10
+ "image": "image",
11
+ }
12
+
13
+
14
+ def _filename_from_text_block(block: dict[str, Any]) -> str | None:
15
+ text = block.get("text")
16
+ if not isinstance(text, str):
17
+ return None
18
+ try:
19
+ payload = json.loads(text)
20
+ except (ValueError, TypeError):
21
+ return None
22
+ if isinstance(payload, dict):
23
+ name = payload.get("filename")
24
+ if isinstance(name, str) and name:
25
+ return name
26
+ return None
27
+
28
+
29
+ def extract_file_parts(messages: list[BaseMessage]) -> list[tuple[str, FilePart]]:
30
+ parts: list[tuple[str, FilePart]] = []
31
+ for message in messages:
32
+ if not (isinstance(message, ToolMessage)
33
+ and isinstance(message.content, list)):
34
+ continue
35
+ pending_name: str | None = None
36
+ counters: dict[str, int] = {"file": 0, "image": 0}
37
+ for block in message.content:
38
+ if not isinstance(block, dict):
39
+ continue
40
+ block_type = block.get("type")
41
+ if not isinstance(block_type, str):
42
+ continue
43
+ if block_type == "text":
44
+ hint = _filename_from_text_block(block)
45
+ if hint:
46
+ pending_name = hint
47
+ continue
48
+ kind = _LANGCHAIN_BINARY_BLOCK_TYPES.get(block_type)
49
+ if kind is None:
50
+ continue
51
+ b64 = block.get("base64")
52
+ if not isinstance(b64, str) or not b64:
53
+ continue
54
+ mime_type = block.get("mime_type") or "application/octet-stream"
55
+ if pending_name is not None:
56
+ name = pending_name
57
+ pending_name = None
58
+ else:
59
+ index = counters[block_type]
60
+ counters[block_type] = index + 1
61
+ guessed_ext = mimetypes.guess_extension(mime_type)
62
+ ext = f"-{guessed_ext}" if guessed_ext is not None else ""
63
+ suffix = f"-{index}" if index > 0 else ""
64
+ name = f"{kind}{suffix}{ext}"
65
+ parts.append((name, FilePart(file=FileWithBytes(
66
+ name=name, mime_type=mime_type, bytes=b64,
67
+ ))))
68
+ return parts
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: distributed_a2a
3
- Version: 0.1.26
3
+ Version: 0.2.2
4
4
  Summary: A library for building A2A agents with routing capabilities
5
5
  Home-page: https://github.com/Barra-Technologies/distributed-a2a
6
6
  Author: Fabian Bell
@@ -9,6 +9,7 @@ distributed_a2a/agent.py
9
9
  distributed_a2a/client.py
10
10
  distributed_a2a/config.py
11
11
  distributed_a2a/executors.py
12
+ distributed_a2a/files.py
12
13
  distributed_a2a/model.py
13
14
  distributed_a2a/py.typed
14
15
  distributed_a2a/registry.py
@@ -29,5 +30,7 @@ distributed_a2a/schemas/agent-schema.json
29
30
  distributed_a2a/schemas/router-agent-schema.json
30
31
  tests/test_app.py
31
32
  tests/test_client.py
33
+ tests/test_executor_files.py
34
+ tests/test_files.py
32
35
  tests/test_rejection.py
33
36
  tests/test_timeout.py
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "distributed_a2a"
7
- version = "0.1.26"
7
+ version = "0.2.2"
8
8
  description = "A library for building A2A agents with routing capabilities"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.14"
@@ -59,10 +59,12 @@ async def test_app_completed_path(fake_registry_server: str, fake_completed_llm:
59
59
  with FakeAgent(fake_registry_server, fake_completed_llm, "test-agent") as agent:
60
60
  # When
61
61
  client = RoutingA2AClient(initial_url=f"http://127.0.0.1:{agent.app_port}/{agent.name}")
62
- response = await client.send_message(message="Hello", context_id="test-context")
62
+ reply = await client.send_message(message="Hello", context_id="test-context")
63
63
 
64
64
  # Then: Check the response
65
- assert "This is a mock response from the fake OpenAI server." in response
65
+ assert reply.text is not None
66
+ assert "This is a mock response from the fake OpenAI server." in reply.text
67
+ assert reply.files == []
66
68
 
67
69
 
68
70
  @pytest.mark.asyncio
@@ -75,7 +77,8 @@ async def test_app_redirect_path(fake_registry_server: str, fake_completed_llm:
75
77
  client = RoutingA2AClient(initial_url=f"http://127.0.0.1:{first_agent.app_port}")
76
78
 
77
79
  # When
78
- response = await client.send_message(message="Hello", context_id="test-context")
80
+ reply = await client.send_message(message="Hello", context_id="test-context")
79
81
 
80
82
  # Then
81
- assert FINAL_RESPONSE in response
83
+ assert reply.text is not None
84
+ assert FINAL_RESPONSE in reply.text
@@ -0,0 +1,216 @@
1
+ from __future__ import annotations
2
+
3
+ import base64
4
+ from types import SimpleNamespace
5
+ from typing import Any
6
+
7
+ import pytest
8
+ from a2a.server.agent_execution import RequestContext
9
+ from a2a.server.events import EventQueue
10
+ from a2a.types import FilePart, FileWithBytes
11
+ from a2a.types import Message as A2AMessage
12
+ from a2a.types import (MessageSendParams, Part, Role, TaskArtifactUpdateEvent,
13
+ TaskState, TaskStatusUpdateEvent, TextPart)
14
+ from langchain_core.messages import BaseMessage, HumanMessage, ToolMessage
15
+
16
+ from distributed_a2a.agent import AgentInvocation, StringResponse
17
+ from distributed_a2a.executors import RoutingAgentExecutor
18
+
19
+
20
+ class _StubStatusAgent:
21
+ def __init__(self, response: StringResponse, messages: list[BaseMessage]) -> None:
22
+ self._response = response
23
+ self._messages = messages
24
+
25
+ async def __call__(self, message: str, context_id: str | None = None) -> AgentInvocation[StringResponse]:
26
+ return AgentInvocation[StringResponse](
27
+ structured=self._response,
28
+ messages=self._messages,
29
+ )
30
+
31
+
32
+ def _make_request_context() -> RequestContext:
33
+ msg = A2AMessage(
34
+ message_id="m-1",
35
+ context_id="ctx-1",
36
+ task_id="task-1",
37
+ role=Role.user,
38
+ parts=[Part(root=TextPart(text="render a CV please"))],
39
+ )
40
+ return RequestContext(
41
+ request=MessageSendParams(message=msg),
42
+ task_id="task-1",
43
+ context_id="ctx-1",
44
+ )
45
+
46
+
47
+ async def _drain_queue(queue: EventQueue) -> list[Any]:
48
+ events: list[Any] = []
49
+ while True:
50
+ try:
51
+ evt = await queue.dequeue_event(no_wait=True)
52
+ except Exception:
53
+ break
54
+ events.append(evt)
55
+ queue.task_done()
56
+ return events
57
+
58
+
59
+ @pytest.mark.asyncio
60
+ async def test_executor_emits_no_file_event_when_no_artifacts() -> None:
61
+ executor = RoutingAgentExecutor.__new__(RoutingAgentExecutor)
62
+ executor.agent_config = SimpleNamespace( # type: ignore[assignment]
63
+ agent=SimpleNamespace(card=SimpleNamespace(name="plain-agent")),
64
+ )
65
+ executor.agent = _StubStatusAgent( # type: ignore[assignment]
66
+ StringResponse(status=TaskState.completed, response="No files here."),
67
+ [HumanMessage(content="hi")],
68
+ )
69
+
70
+ async def _noop_reinit() -> None:
71
+ return None
72
+
73
+ executor.reinitialize_agent_with_tools = _noop_reinit # type: ignore[method-assign]
74
+
75
+ ctx = _make_request_context()
76
+ queue = EventQueue()
77
+ await executor.execute(ctx, queue)
78
+ events = await _drain_queue(queue)
79
+
80
+ artifact_events = [e for e in events if isinstance(e, TaskArtifactUpdateEvent)]
81
+ assert len(artifact_events) == 1
82
+ assert artifact_events[0].artifact.name == "current_result"
83
+ assert artifact_events[0].last_chunk is True
84
+
85
+
86
+ @pytest.mark.asyncio
87
+ async def test_executor_emits_file_part_from_langchain_content_block() -> None:
88
+ docx_b64 = base64.b64encode(b"PK\x03\x04 fake docx").decode("ascii")
89
+ summary_json = (
90
+ '{"filename": "cv-bob.docx", '
91
+ '"mime_type": "application/vnd.openxmlformats-officedocument.'
92
+ 'wordprocessingml.document"}'
93
+ )
94
+ tool_msg = ToolMessage(
95
+ content=[
96
+ {"type": "text", "text": summary_json, "id": "lc_text_1"},
97
+ {"type": "file",
98
+ "base64": docx_b64,
99
+ "mime_type": (
100
+ "application/vnd.openxmlformats-officedocument."
101
+ "wordprocessingml.document"
102
+ ),
103
+ "id": "lc_file_1"},
104
+ ],
105
+ tool_call_id="call-cv",
106
+ artifact={"structured_content": {"filename": "cv-bob.docx"}},
107
+ )
108
+
109
+ executor = RoutingAgentExecutor.__new__(RoutingAgentExecutor)
110
+ executor.agent_config = SimpleNamespace( # type: ignore[assignment]
111
+ agent=SimpleNamespace(card=SimpleNamespace(name="cv-agent")),
112
+ )
113
+ executor.agent = _StubStatusAgent( # type: ignore[assignment]
114
+ StringResponse(status=TaskState.completed,
115
+ response="Here is your CV."),
116
+ [HumanMessage(content="render a CV please"), tool_msg],
117
+ )
118
+
119
+ async def _noop_reinit() -> None:
120
+ return None
121
+
122
+ executor.reinitialize_agent_with_tools = _noop_reinit # type: ignore[method-assign]
123
+
124
+ ctx = _make_request_context()
125
+ queue = EventQueue()
126
+ await executor.execute(ctx, queue)
127
+ events = await _drain_queue(queue)
128
+
129
+ artifact_events = [e for e in events if isinstance(e, TaskArtifactUpdateEvent)]
130
+ assert len(artifact_events) == 2, (
131
+ "Expected the executor to emit a file artifact followed by the text "
132
+ "artifact when the tool response uses the production LangChain "
133
+ "content-block shape."
134
+ )
135
+ file_event, text_event = artifact_events
136
+ assert file_event.last_chunk is False
137
+ assert file_event.artifact.name == "cv-bob.docx"
138
+ file_part = file_event.artifact.parts[0].root
139
+ assert isinstance(file_part, FilePart)
140
+ assert isinstance(file_part.file, FileWithBytes)
141
+ assert file_part.file.name == "cv-bob.docx"
142
+ assert file_part.file.bytes == docx_b64
143
+
144
+ assert text_event.last_chunk is True
145
+ assert text_event.artifact.name == "current_result"
146
+
147
+
148
+ @pytest.mark.asyncio
149
+ async def test_executor_emits_one_file_event_per_file_block() -> None:
150
+ b64_a = base64.b64encode(b"aaa docx bytes").decode("ascii")
151
+ b64_b = base64.b64encode(b"bbb docx bytes").decode("ascii")
152
+ docx_mime = (
153
+ "application/vnd.openxmlformats-officedocument."
154
+ "wordprocessingml.document"
155
+ )
156
+ tool_msg = ToolMessage(
157
+ content=[
158
+ {"type": "text", "text": '{"filename": "cv-a.docx"}', "id": "t1"},
159
+ {"type": "file", "base64": b64_a, "mime_type": docx_mime, "id": "f1"},
160
+ {"type": "text", "text": '{"filename": "cv-b.docx"}', "id": "t2"},
161
+ {"type": "file", "base64": b64_b, "mime_type": docx_mime, "id": "f2"},
162
+ ],
163
+ tool_call_id="call-multi",
164
+ )
165
+
166
+ executor = RoutingAgentExecutor.__new__(RoutingAgentExecutor)
167
+ executor.agent_config = SimpleNamespace( # type: ignore[assignment]
168
+ agent=SimpleNamespace(card=SimpleNamespace(name="cv-agent")),
169
+ )
170
+ executor.agent = _StubStatusAgent( # type: ignore[assignment]
171
+ StringResponse(status=TaskState.completed,
172
+ response="Here are your two CVs."),
173
+ [HumanMessage(content="render two CVs"), tool_msg],
174
+ )
175
+
176
+ async def _noop_reinit() -> None:
177
+ return None
178
+
179
+ executor.reinitialize_agent_with_tools = _noop_reinit # type: ignore[method-assign]
180
+
181
+ ctx = _make_request_context()
182
+ queue = EventQueue()
183
+ await executor.execute(ctx, queue)
184
+ events = await _drain_queue(queue)
185
+
186
+ artifact_events = [e for e in events if isinstance(e, TaskArtifactUpdateEvent)]
187
+ assert len(artifact_events) == 3, (
188
+ "Expected 2 file artifacts (last_chunk=False) + 1 text artifact "
189
+ f"(last_chunk=True). Got: "
190
+ f"{[(e.artifact.name, e.last_chunk) for e in artifact_events]}"
191
+ )
192
+ file_a_event, file_b_event, text_event = artifact_events
193
+
194
+ assert file_a_event.last_chunk is False
195
+ assert file_a_event.artifact.name == "cv-a.docx"
196
+ file_a_part = file_a_event.artifact.parts[0].root
197
+ assert isinstance(file_a_part, FilePart)
198
+ assert isinstance(file_a_part.file, FileWithBytes)
199
+ assert file_a_part.file.name == "cv-a.docx"
200
+ assert file_a_part.file.bytes == b64_a
201
+
202
+ assert file_b_event.last_chunk is False
203
+ assert file_b_event.artifact.name == "cv-b.docx"
204
+ file_b_part = file_b_event.artifact.parts[0].root
205
+ assert isinstance(file_b_part, FilePart)
206
+ assert isinstance(file_b_part.file, FileWithBytes)
207
+ assert file_b_part.file.name == "cv-b.docx"
208
+ assert file_b_part.file.bytes == b64_b
209
+
210
+ assert text_event.last_chunk is True
211
+ assert text_event.artifact.name == "current_result"
212
+
213
+ final_status = [e for e in events
214
+ if isinstance(e, TaskStatusUpdateEvent) and e.final]
215
+ assert len(final_status) == 1
216
+ assert final_status[0].status.state == TaskState.completed
@@ -0,0 +1,317 @@
1
+ from __future__ import annotations
2
+
3
+ import base64
4
+ from collections.abc import AsyncGenerator
5
+
6
+ import httpx
7
+ import pytest
8
+ from a2a.types import (AgentCapabilities, AgentCard, Artifact, FilePart,
9
+ FileWithBytes, FileWithUri, Message, Part, Task,
10
+ TaskState, TaskStatus, TextPart)
11
+ from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
12
+
13
+ from distributed_a2a.client import AgentReply, RemoteAgentConnection
14
+ from distributed_a2a.files import extract_file_parts
15
+
16
+
17
+ def _b64(payload: bytes) -> str:
18
+ return base64.b64encode(payload).decode("ascii")
19
+
20
+
21
+ def test_extract_file_parts_ignores_non_tool_messages_and_string_content() -> None:
22
+ messages = [
23
+ HumanMessage(content="hi"),
24
+ AIMessage(content="hello"),
25
+ ToolMessage(content="just text", tool_call_id="c1"),
26
+ ]
27
+ assert extract_file_parts(messages) == []
28
+
29
+
30
+ def test_extract_file_parts_reads_langchain_file_content_block() -> None:
31
+ docx_b64 = _b64(b"PK\x03\x04 real docx bytes")
32
+ summary = (
33
+ '{"filename": "cv-alice.docx", '
34
+ '"mime_type": "application/vnd.openxmlformats-officedocument.'
35
+ 'wordprocessingml.document", "size_bytes": 4321}'
36
+ )
37
+ tool_msg = ToolMessage(
38
+ content=[
39
+ {"type": "text", "text": summary, "id": "lc_text_1"},
40
+ {"type": "file",
41
+ "base64": docx_b64,
42
+ "mime_type": (
43
+ "application/vnd.openxmlformats-officedocument."
44
+ "wordprocessingml.document"
45
+ ),
46
+ "id": "lc_file_1"},
47
+ ],
48
+ tool_call_id="call-cv",
49
+ artifact={"structured_content": {"filename": "cv-alice.docx"}},
50
+ )
51
+
52
+ parts = extract_file_parts([HumanMessage(content="hi"), tool_msg])
53
+
54
+ assert len(parts) == 1
55
+ name, file_part = parts[0]
56
+ assert name == "cv-alice.docx"
57
+ assert isinstance(file_part, FilePart)
58
+ assert isinstance(file_part.file, FileWithBytes)
59
+ assert file_part.file.name == "cv-alice.docx"
60
+ assert file_part.file.mime_type == (
61
+ "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
62
+ )
63
+ assert file_part.file.bytes == docx_b64
64
+
65
+
66
+ def test_extract_file_parts_reads_langchain_image_content_block() -> None:
67
+ png_b64 = _b64(b"\x89PNG\r\n\x1a\n fake")
68
+ tool_msg = ToolMessage(
69
+ content=[
70
+ {"type": "image", "base64": png_b64,
71
+ "mime_type": "image/png", "id": "lc_img_1"},
72
+ ],
73
+ tool_call_id="call-img",
74
+ )
75
+
76
+ parts = extract_file_parts([tool_msg])
77
+
78
+ assert len(parts) == 1
79
+ name, file_part = parts[0]
80
+ assert name.startswith("image")
81
+ assert name.endswith(".png")
82
+ assert isinstance(file_part.file, FileWithBytes)
83
+ assert file_part.file.mime_type == "image/png"
84
+ assert file_part.file.bytes == png_b64
85
+
86
+
87
+ def test_extract_file_parts_synthesises_name_without_filename_hint() -> None:
88
+ docx_b64 = _b64(b"PK\x03\x04 bytes")
89
+ tool_msg = ToolMessage(
90
+ content=[
91
+ {"type": "text", "text": "no json here", "id": "lc_text_1"},
92
+ {"type": "file",
93
+ "base64": docx_b64,
94
+ "mime_type": (
95
+ "application/vnd.openxmlformats-officedocument."
96
+ "wordprocessingml.document"
97
+ ),
98
+ "id": "lc_file_1"},
99
+ ],
100
+ tool_call_id="call-cv",
101
+ )
102
+
103
+ parts = extract_file_parts([tool_msg])
104
+
105
+ assert len(parts) == 1
106
+ name, file_part = parts[0]
107
+ assert name.startswith("attachment")
108
+ assert name.endswith(".docx")
109
+ assert isinstance(file_part.file, FileWithBytes)
110
+ assert file_part.file.bytes == docx_b64
111
+
112
+
113
+ def test_extract_file_parts_ignores_url_only_file_block() -> None:
114
+ tool_msg = ToolMessage(
115
+ content=[
116
+ {"type": "file",
117
+ "url": "https://example.com/report.pdf",
118
+ "mime_type": "application/pdf",
119
+ "id": "lc_file_1"},
120
+ ],
121
+ tool_call_id="call-url-only",
122
+ )
123
+ assert extract_file_parts([tool_msg]) == []
124
+
125
+
126
+ def test_extract_file_parts_matches_multiple_filenames_by_order() -> None:
127
+ """Two [text-summary, file] pairs in the same ToolMessage — each
128
+ text-summary's ``filename`` is consumed by the immediately following
129
+ file block."""
130
+ b64_a = _b64(b"aaa")
131
+ b64_b = _b64(b"bbb")
132
+ tool_msg = ToolMessage(
133
+ content=[
134
+ {"type": "text", "text": '{"filename": "a.bin"}',
135
+ "id": "lc_text_1"},
136
+ {"type": "file", "base64": b64_a,
137
+ "mime_type": "application/octet-stream", "id": "lc_file_1"},
138
+ {"type": "text", "text": '{"filename": "b.bin"}',
139
+ "id": "lc_text_2"},
140
+ {"type": "file", "base64": b64_b,
141
+ "mime_type": "application/octet-stream", "id": "lc_file_2"},
142
+ ],
143
+ tool_call_id="call-multi",
144
+ )
145
+
146
+ parts = extract_file_parts([tool_msg])
147
+
148
+ assert [name for name, _ in parts] == ["a.bin", "b.bin"]
149
+ assert parts[0][1].file.bytes == b64_a # type: ignore[union-attr]
150
+ assert parts[1][1].file.bytes == b64_b # type: ignore[union-attr]
151
+
152
+
153
+ class _StubAgentClient:
154
+ def __init__(self, task: Task):
155
+ self._task = task
156
+
157
+ async def send_message(self, _message: Message) -> AsyncGenerator[tuple[Task, None], None]:
158
+ yield self._task, None
159
+
160
+
161
+ class _StubClientFactory:
162
+ def __init__(self, _config: object, task: Task):
163
+ self._task = task
164
+
165
+ def create(self, _agent_card: AgentCard) -> _StubAgentClient:
166
+ return _StubAgentClient(self._task)
167
+
168
+
169
+ def _patch_client_factory(monkeypatch: pytest.MonkeyPatch, task: Task) -> None:
170
+ def _ctor(config: object) -> _StubClientFactory:
171
+ return _StubClientFactory(config, task)
172
+ monkeypatch.setattr("distributed_a2a.client.ClientFactory", _ctor)
173
+
174
+
175
+ def _agent_card() -> AgentCard:
176
+ return AgentCard(
177
+ name="stub",
178
+ description="stub",
179
+ url="http://127.0.0.1:0",
180
+ version="1.0.0",
181
+ default_input_modes=["text"],
182
+ default_output_modes=["text"],
183
+ capabilities=AgentCapabilities(streaming=False, push_notifications=False),
184
+ skills=[],
185
+ preferred_transport="JSONRPC",
186
+ )
187
+
188
+
189
+ @pytest.mark.asyncio
190
+ async def test_remote_agent_connection_returns_text_and_files_in_agent_reply(
191
+ monkeypatch: pytest.MonkeyPatch) -> None:
192
+ docx_b64 = _b64(b"PK\x03\x04 docx-bytes")
193
+ completed_task = Task(
194
+ id="task-files",
195
+ context_id="ctx-files",
196
+ status=TaskStatus(state=TaskState.completed),
197
+ artifacts=[
198
+ Artifact(
199
+ artifact_id="cv-foo.docx",
200
+ name="cv-foo.docx",
201
+ parts=[Part(root=FilePart(file=FileWithBytes(
202
+ name="cv-foo.docx",
203
+ mime_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
204
+ bytes=docx_b64,
205
+ )))],
206
+ ),
207
+ Artifact(
208
+ artifact_id="current_result",
209
+ name="current_result",
210
+ parts=[Part(root=TextPart(text="Here is your CV."))],
211
+ ),
212
+ ],
213
+ )
214
+ factory_task = completed_task
215
+ _patch_client_factory(monkeypatch, factory_task)
216
+
217
+ async with httpx.AsyncClient() as http_client:
218
+ conn = RemoteAgentConnection(_agent_card(), http_client)
219
+ reply = await conn.send_message("hello", "ctx-files")
220
+
221
+ assert isinstance(reply, AgentReply)
222
+ assert reply.text == "Here is your CV."
223
+ assert len(reply.files) == 1
224
+ f = reply.files[0]
225
+ assert f.name == "cv-foo.docx"
226
+ assert f.mime_type == (
227
+ "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
228
+ )
229
+ assert f.bytes_b64 == docx_b64
230
+ assert f.uri is None
231
+
232
+
233
+ @pytest.mark.asyncio
234
+ async def test_remote_agent_connection_handles_file_with_uri(
235
+ monkeypatch: pytest.MonkeyPatch) -> None:
236
+ completed_task = Task(
237
+ id="task-uri",
238
+ context_id="ctx-uri",
239
+ status=TaskStatus(state=TaskState.completed),
240
+ artifacts=[
241
+ Artifact(
242
+ artifact_id="report",
243
+ name="report",
244
+ parts=[Part(root=FilePart(file=FileWithUri(
245
+ name="report.pdf",
246
+ mime_type="application/pdf",
247
+ uri="https://example.com/report.pdf",
248
+ )))],
249
+ ),
250
+ ],
251
+ )
252
+ _patch_client_factory(monkeypatch, completed_task)
253
+
254
+ async with httpx.AsyncClient() as http_client:
255
+ conn = RemoteAgentConnection(_agent_card(), http_client)
256
+ reply = await conn.send_message("hello", "ctx-uri")
257
+
258
+ assert isinstance(reply, AgentReply)
259
+ assert reply.text is None
260
+ assert len(reply.files) == 1
261
+ f = reply.files[0]
262
+ assert f.name == "report.pdf"
263
+ assert f.mime_type == "application/pdf"
264
+ assert f.bytes_b64 == ""
265
+ assert f.uri == "https://example.com/report.pdf"
266
+
267
+
268
+ @pytest.mark.asyncio
269
+ async def test_remote_agent_connection_gathers_multiple_files_in_order(
270
+ monkeypatch: pytest.MonkeyPatch) -> None:
271
+ b64_a = _b64(b"aaa")
272
+ b64_b = _b64(b"bbb")
273
+ completed_task = Task(
274
+ id="task-multi",
275
+ context_id="ctx-multi",
276
+ status=TaskStatus(state=TaskState.completed),
277
+ artifacts=[
278
+ Artifact(
279
+ artifact_id="a.bin",
280
+ name="a.bin",
281
+ parts=[Part(root=FilePart(file=FileWithBytes(
282
+ name="a.bin",
283
+ mime_type="application/octet-stream",
284
+ bytes=b64_a,
285
+ )))],
286
+ ),
287
+ Artifact(
288
+ artifact_id="b.bin",
289
+ name="b.bin",
290
+ parts=[Part(root=FilePart(file=FileWithBytes(
291
+ name="b.bin",
292
+ mime_type="application/octet-stream",
293
+ bytes=b64_b,
294
+ )))],
295
+ ),
296
+ Artifact(
297
+ artifact_id="current_result",
298
+ name="current_result",
299
+ parts=[Part(root=TextPart(text="Two files ready."))],
300
+ ),
301
+ ],
302
+ )
303
+ _patch_client_factory(monkeypatch, completed_task)
304
+
305
+ async with httpx.AsyncClient() as http_client:
306
+ conn = RemoteAgentConnection(_agent_card(), http_client)
307
+ reply = await conn.send_message("hello", "ctx-multi")
308
+
309
+ assert isinstance(reply, AgentReply)
310
+ assert reply.text == "Two files ready."
311
+ assert [f.name for f in reply.files] == ["a.bin", "b.bin"]
312
+ assert [f.mime_type for f in reply.files] == [
313
+ "application/octet-stream", "application/octet-stream",
314
+ ]
315
+ assert [f.bytes_b64 for f in reply.files] == [b64_a, b64_b]
316
+ assert all(f.uri is None for f in reply.files)
317
+
@@ -1,7 +1,7 @@
1
1
  import pytest
2
2
  from a2a.types import AgentCapabilities, AgentCard, TaskState
3
3
 
4
- from distributed_a2a.client import RoutingA2AClient
4
+ from distributed_a2a.client import AgentReply, RoutingA2AClient
5
5
 
6
6
 
7
7
  def build_card(name: str, url: str) -> AgentCard:
@@ -38,7 +38,7 @@ async def test_rejection_triggers_automated_rerouting(monkeypatch: pytest.Monkey
38
38
  def __init__(self, agent_card: AgentCard, _client: object, **_kwargs: object) -> None:
39
39
  self.agent_card = agent_card
40
40
 
41
- async def send_message(self, message_to_send: str, _context_id: str) -> str | AgentCard | TaskState:
41
+ async def send_message(self, message_to_send: str, _context_id: str) -> AgentReply | AgentCard | TaskState:
42
42
  if self.agent_card.name == "Router":
43
43
  router_messages.append(message_to_send)
44
44
  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
49
49
  return TaskState.rejected
50
50
 
51
51
  if self.agent_card.name == "success-agent":
52
- return "final answer"
52
+ return AgentReply(text="final answer")
53
53
 
54
54
  raise AssertionError(f"Unexpected agent {self.agent_card.name}")
55
55
 
@@ -57,7 +57,8 @@ async def test_rejection_triggers_automated_rerouting(monkeypatch: pytest.Monkey
57
57
 
58
58
  result = await client.send_message("Hello", context_id="ctx-1")
59
59
 
60
- assert result == "final answer"
60
+ assert isinstance(result, AgentReply)
61
+ assert result.text == "final answer"
61
62
  assert len(router_messages) == 2
62
63
  assert "Please exclude the following agents from routing" not in router_messages[0]
63
64
  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
88
89
  def __init__(self, agent_card: AgentCard, _client: object, **_kwargs: object) -> None:
89
90
  self.agent_card = agent_card
90
91
 
91
- async def send_message(self, message_to_send: str, _context_id: str) -> str | AgentCard | TaskState:
92
+ async def send_message(self, message_to_send: str, _context_id: str) -> AgentReply | AgentCard | TaskState:
92
93
  idx = next_step()
93
94
  match idx:
94
95
  case 0:
@@ -104,7 +105,7 @@ async def test_rejected_agents_reset_between_calls(monkeypatch: pytest.MonkeyPat
104
105
  return success_card
105
106
  case 3:
106
107
  assert self.agent_card.name == "success-agent"
107
- return "first response"
108
+ return AgentReply(text="first response")
108
109
  case 4:
109
110
  assert self.agent_card.name == "success-agent"
110
111
  return TaskState.rejected
@@ -115,7 +116,7 @@ async def test_rejected_agents_reset_between_calls(monkeypatch: pytest.MonkeyPat
115
116
  return rejecting_card
116
117
  case 6:
117
118
  assert self.agent_card.name == "rejecting-agent"
118
- return "second response"
119
+ return AgentReply(text="second response")
119
120
  case _:
120
121
  raise AssertionError(f"Unexpected step {idx}")
121
122
 
@@ -124,8 +125,8 @@ async def test_rejected_agents_reset_between_calls(monkeypatch: pytest.MonkeyPat
124
125
  first = await client.send_message("First", context_id="ctx-1")
125
126
  second = await client.send_message("Second", context_id="ctx-2")
126
127
 
127
- assert first == "first response"
128
- assert second == "second response"
128
+ assert isinstance(first, AgentReply) and first.text == "first response"
129
+ assert isinstance(second, AgentReply) and second.text == "second response"
129
130
 
130
131
 
131
132
  @pytest.mark.asyncio
@@ -147,7 +148,7 @@ async def test_fails_when_router_returns_already_rejected_agent(monkeypatch: pyt
147
148
  def __init__(self, agent_card: AgentCard, _client: object, **_kwargs: object) -> None:
148
149
  self.agent_card = agent_card
149
150
 
150
- async def send_message(self, message_to_send: str, _context_id: str) -> str | AgentCard | TaskState:
151
+ async def send_message(self, message_to_send: str, _context_id: str) -> AgentReply | AgentCard | TaskState:
151
152
  idx = step["idx"]
152
153
  step["idx"] += 1
153
154
  if idx == 0:
@@ -8,8 +8,8 @@ import pytest
8
8
  from a2a.types import (AgentCapabilities, AgentCard, Artifact, Message, Part,
9
9
  Role, Task, TaskState, TaskStatus, TextPart)
10
10
 
11
- from distributed_a2a.client import (A2ATimeoutError, RemoteAgentConnection,
12
- RoutingA2AClient)
11
+ from distributed_a2a.client import (A2ATimeoutError, AgentReply,
12
+ RemoteAgentConnection, RoutingA2AClient)
13
13
 
14
14
 
15
15
  def _agent_card() -> AgentCard:
@@ -124,7 +124,9 @@ async def test_send_message_honors_completion_on_final_allowed_poll(
124
124
  )
125
125
  result = await connection.send_message("hi", "ctx-x")
126
126
 
127
- assert result == "all done"
127
+ assert isinstance(result, AgentReply)
128
+ assert result.text == "all done"
129
+ assert result.files == []
128
130
  assert scripted.get_task_calls == 3
129
131
  assert scripted._get_states == []
130
132
 
@@ -156,7 +158,9 @@ async def test_send_message_returns_immediately_if_already_completed(
156
158
  )
157
159
  result = await connection.send_message("hi", "ctx-x")
158
160
 
159
- assert result == "all done"
161
+ assert isinstance(result, AgentReply)
162
+ assert result.text == "all done"
163
+ assert result.files == []
160
164
 
161
165
 
162
166
  def test_a2a_timeout_error_message_includes_diagnostics() -> None:
@@ -182,8 +186,8 @@ async def test_routing_a2a_client_propagates_poll_kwargs(
182
186
  captured["poll_interval"] = poll_interval
183
187
  self.agent_card = agent_card
184
188
 
185
- async def send_message(self, _message: str, _context_id: str) -> str:
186
- return "done"
189
+ async def send_message(self, _message: str, _context_id: str) -> AgentReply:
190
+ return AgentReply(text="done")
187
191
 
188
192
  monkeypatch.setattr("distributed_a2a.client.RemoteAgentConnection", _Capturing)
189
193
 
@@ -197,5 +201,7 @@ async def test_routing_a2a_client_propagates_poll_kwargs(
197
201
  monkeypatch.setattr(client, "fetch_initial_card", _fetch)
198
202
 
199
203
  result = await client.send_message("hi", context_id="ctx-1")
200
- assert result == "done"
204
+ assert isinstance(result, AgentReply)
205
+ assert result.text == "done"
206
+ assert result.files == []
201
207
  assert captured == {"max_polls": 7, "poll_interval": 0.25}