Skip to content

Commit 6bd7dc2

Browse files
authored
Merge branch 'main' into ae_routes
2 parents 7cef6cb + 2508a9b commit 6bd7dc2

File tree

4 files changed

+59
-4
lines changed

4 files changed

+59
-4
lines changed

src/a2a/client/transports/grpc.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,10 @@ async def get_task(
138138
) -> Task:
139139
"""Retrieves the current state and history of a specific task."""
140140
task = await self.stub.GetTask(
141-
a2a_pb2.GetTaskRequest(name=f'tasks/{request.id}')
141+
a2a_pb2.GetTaskRequest(
142+
name=f'tasks/{request.id}',
143+
history_length=request.history_length,
144+
)
142145
)
143146
return proto_utils.FromProto.task(task)
144147

src/a2a/server/events/event_consumer.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -160,5 +160,5 @@ def agent_task_callback(self, agent_task: asyncio.Task[None]) -> None:
160160
agent_task: The asyncio.Task that completed.
161161
"""
162162
logger.debug('Agent task callback triggered.')
163-
if agent_task.exception() is not None:
163+
if not agent_task.cancelled() and agent_task.done():
164164
self._exception = agent_task.exception()

tests/client/test_grpc_client.py

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -145,11 +145,31 @@ async def test_get_task(
145145
response = await grpc_transport.get_task(params)
146146

147147
mock_grpc_stub.GetTask.assert_awaited_once_with(
148-
a2a_pb2.GetTaskRequest(name=f'tasks/{sample_task.id}')
148+
a2a_pb2.GetTaskRequest(
149+
name=f'tasks/{sample_task.id}', history_length=None
150+
)
149151
)
150152
assert response.id == sample_task.id
151153

152154

155+
@pytest.mark.asyncio
156+
async def test_get_task_with_history(
157+
grpc_transport: GrpcTransport, mock_grpc_stub: AsyncMock, sample_task: Task
158+
):
159+
"""Test retrieving a task with history."""
160+
mock_grpc_stub.GetTask.return_value = proto_utils.ToProto.task(sample_task)
161+
history_len = 10
162+
params = TaskQueryParams(id=sample_task.id, history_length=history_len)
163+
164+
await grpc_transport.get_task(params)
165+
166+
mock_grpc_stub.GetTask.assert_awaited_once_with(
167+
a2a_pb2.GetTaskRequest(
168+
name=f'tasks/{sample_task.id}', history_length=history_len
169+
)
170+
)
171+
172+
153173
@pytest.mark.asyncio
154174
async def test_cancel_task(
155175
grpc_transport: GrpcTransport, mock_grpc_stub: AsyncMock, sample_task: Task

tests/server/events/test_event_consumer.py

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -327,18 +327,22 @@ async def test_consume_all_continues_on_queue_empty_if_not_really_closed(
327327
def test_agent_task_callback_sets_exception(event_consumer: EventConsumer):
328328
"""Test that agent_task_callback sets _exception if the task had one."""
329329
mock_task = MagicMock(spec=asyncio.Task)
330+
mock_task.cancelled.return_value = False
331+
mock_task.done.return_value = True
330332
sample_exception = ValueError('Task failed')
331333
mock_task.exception.return_value = sample_exception
332334

333335
event_consumer.agent_task_callback(mock_task)
334336

335337
assert event_consumer._exception == sample_exception
336-
# mock_task.exception.assert_called_once() # Removing this, as exception() might be called internally by the check
338+
mock_task.exception.assert_called_once()
337339

338340

339341
def test_agent_task_callback_no_exception(event_consumer: EventConsumer):
340342
"""Test that agent_task_callback does nothing if the task has no exception."""
341343
mock_task = MagicMock(spec=asyncio.Task)
344+
mock_task.cancelled.return_value = False
345+
mock_task.done.return_value = True
342346
mock_task.exception.return_value = None # No exception
343347

344348
event_consumer.agent_task_callback(mock_task)
@@ -347,6 +351,34 @@ def test_agent_task_callback_no_exception(event_consumer: EventConsumer):
347351
mock_task.exception.assert_called_once()
348352

349353

354+
def test_agent_task_callback_cancelled_task(event_consumer: EventConsumer):
355+
"""Test that agent_task_callback does nothing if the task has no exception."""
356+
mock_task = MagicMock(spec=asyncio.Task)
357+
mock_task.cancelled.return_value = True
358+
mock_task.done.return_value = True
359+
sample_exception = ValueError('Task still running')
360+
mock_task.exception.return_value = sample_exception
361+
362+
event_consumer.agent_task_callback(mock_task)
363+
364+
assert event_consumer._exception is None # Should remain None
365+
mock_task.exception.assert_not_called()
366+
367+
368+
def test_agent_task_callback_not_done_task(event_consumer: EventConsumer):
369+
"""Test that agent_task_callback does nothing if the task has no exception."""
370+
mock_task = MagicMock(spec=asyncio.Task)
371+
mock_task.cancelled.return_value = False
372+
mock_task.done.return_value = False
373+
sample_exception = ValueError('Task is cancelled')
374+
mock_task.exception.return_value = sample_exception
375+
376+
event_consumer.agent_task_callback(mock_task)
377+
378+
assert event_consumer._exception is None # Should remain None
379+
mock_task.exception.assert_not_called()
380+
381+
350382
@pytest.mark.asyncio
351383
async def test_consume_all_handles_validation_error(
352384
event_consumer: EventConsumer, mock_event_queue: AsyncMock

0 commit comments

Comments
 (0)