Skip to content
13 changes: 13 additions & 0 deletions src/a2a/server/agent_execution/active_task.py
Original file line number Diff line number Diff line change
Expand Up @@ -587,6 +587,19 @@ async def _run_consumer(self) -> None:
self._is_finished.set()
self._request_queue.shutdown(immediate=True)
await self._event_queue_agent.close(immediate=True)

if self._producer_task and not self._producer_task.done():
try:
await self._producer_task
except asyncio.CancelledError:
pass
except Exception as e: # noqa: BLE001
logger.debug(
'Consumer[%s]: Awaited producer_task raised %r',
self._task_id,
e,
)

async with self._lock:
self._reference_count -= 1
logger.debug('Consumer[%s]: Finishing', self._task_id)
Expand Down
84 changes: 84 additions & 0 deletions tests/server/agent_execution/test_active_task.py
Original file line number Diff line number Diff line change
Expand Up @@ -895,3 +895,87 @@ async def execute_mock(req, q):
assert len(events) == 0

await active_task.cancel(request_context)


@pytest.mark.timeout(5)
@pytest.mark.asyncio
async def test_producer_awaited_on_normal_completion():
"""Verify producer is awaited before cleanup to prevent GC warnings.

Regression test for #1121: when the consumer finishes first (normal
completion path), the producer must be awaited before
_maybe_cleanup() releases the ActiveTask reference, otherwise asyncio
logs "Task was destroyed but it is pending!".
"""
agent_executor = Mock()
task_manager = Mock()
cleanup_called = False

def on_cleanup(_task: ActiveTask) -> None:
nonlocal cleanup_called
cleanup_called = True

active_task = ActiveTask(
agent_executor=agent_executor,
task_id='test-task-id',
task_manager=task_manager,
on_cleanup=on_cleanup,
)

execute_started = asyncio.Event()
execute_barrier = asyncio.Event()

async def execute_mock(req, q):
execute_started.set()
await execute_barrier.wait()

agent_executor.execute = AsyncMock(side_effect=execute_mock)
agent_executor.cancel = AsyncMock()
task_manager.get_task = AsyncMock(
return_value=Task(
id='test-task-id',
status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
)
)
task_manager.save_task_event = AsyncMock()
task_manager.ensure_task_id = AsyncMock(
return_value=Task(
id='test-task-id',
status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
)
)
task_manager.process = AsyncMock(side_effect=lambda x: x)

request_context = Mock(spec=RequestContext)
request_context.call_context = ServerCallContext()
request_context.context_id = 'test-context'
request_context.message = None

await active_task.enqueue_request(request_context)
await active_task.start(
call_context=ServerCallContext(), create_task_if_missing=True
)

await execute_started.wait()

if active_task._consumer_task:
active_task._consumer_task.cancel()

# Release the producer so it can finish its work and loop back
# to get(), where it will receive QueueShutDown.
execute_barrier.set()

try:
await active_task._consumer_task
except asyncio.CancelledError:
pass

await active_task._is_finished.wait()

assert active_task._producer_task is not None
assert active_task._producer_task.done(), (
'Producer task should be done after consumer teardown'
)
assert cleanup_called, (
'on_cleanup should be called after producer is drained'
)
Loading