diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index b0154c8d6..718db7500 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -574,6 +574,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)