fix(sdk): agent block due to incorrect queue close (#1260)

Signed-off-by: Radek Ježek <radek.jezek@ibm.com>

Author: jezekra1

agents/chat/uv.lock

@@ -6,7 +6,7 @@ readme = "README.md"
 authors = [{ name = "IBM Corp." }]
 requires-python = ">=3.11"
 dependencies = [
-    "a2a-sdk==0.3.5",
+    "a2a-sdk==0.3.7",
     "objprint>=0.3.0",
     "uvicorn>=0.35.0",
     "asyncclick>=8.1.8",

agents/form/uv.lock

@@ -6,9 +6,11 @@
 from asyncio import CancelledError
 from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator
 from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
-from typing import NamedTuple, TypeAlias, cast
+from datetime import datetime, timedelta
+from typing import NamedTuple, TypeAlias, TypedDict, cast
 
 import janus
+from a2a.client import create_text_message_object
 from a2a.server.agent_execution import AgentExecutor, RequestContext
 from a2a.server.events import EventQueue, QueueManager
 from a2a.server.tasks import TaskUpdater
@@ -40,7 +42,7 @@
 from beeai_sdk.server.dependencies import extract_dependencies
 from beeai_sdk.server.logging import logger
 from beeai_sdk.server.store.context_store import ContextStore
-from beeai_sdk.server.utils import cancel_task
+from beeai_sdk.server.utils import cancel_task, close_queue
 
 AgentFunction: TypeAlias = Callable[[], AsyncGenerator[RunYield, RunYieldResume]]
 AgentFunctionFactory: TypeAlias = Callable[
@@ -278,15 +280,25 @@ async def agent_generator():
     return decorator
 
 
+class RunningTask(TypedDict):
+    task: asyncio.Task
+    last_invocation: datetime
+
+
 class Executor(AgentExecutor):
     def __init__(
-        self, execute_fn: AgentFunctionFactory, queue_manager: QueueManager, context_store: ContextStore
+        self,
+        execute_fn: AgentFunctionFactory,
+        queue_manager: QueueManager,
+        context_store: ContextStore,
+        task_timeout: timedelta,
     ) -> None:
         self._agent_executor_span = execute_fn
         self._queue_manager = queue_manager
-        self._running_tasks: dict[str, asyncio.Task] = {}
+        self._running_tasks: dict[str, RunningTask] = {}
         self._cancel_queues: dict[str, EventQueue] = {}
         self._context_store = context_store
+        self._task_timeout = task_timeout
 
     async def _watch_for_cancellation(self, task_id: str, task: asyncio.Task) -> None:
         cancel_queue = await self._queue_manager.create_or_tap(f"_cancel_{task_id}")
@@ -333,7 +345,11 @@ def with_context(message: Message | None = None) -> Message | None:
                 value: RunYieldResume = None
                 opened_artifacts: set[str] = set()
                 while True:
+                    # update invocation time
+                    self._running_tasks[task_updater.task_id]["last_invocation"] = datetime.now()
+
                     yielded_value = await agent_generator_fn.asend(value)
+
                     match yielded_value:
                         case str(text):
                             await task_updater.update_status(
@@ -438,10 +454,16 @@ def with_context(message: Message | None = None) -> Message | None:
         except Exception as ex:
             logger.error("Error when executing agent", exc_info=ex)
             await task_updater.failed(task_updater.new_agent_message(parts=[Part(root=TextPart(text=str(ex)))]))
-        finally:
-            await self._queue_manager.close(f"_event_{task_updater.task_id}")
-            await self._queue_manager.close(f"_resume_{task_updater.task_id}")
+        finally:  # cleanup
             await cancel_task(cancellation_task)
+            is_cancelling = bool(current_task.cancelling())
+            try:
+                async with asyncio.timeout(10):  # grace period to read all events from queue
+                    await close_queue(self._queue_manager, f"_event_{context.task_id}", immediate=is_cancelling)
+                    await close_queue(self._queue_manager, f"_resume_{context.task_id}", immediate=is_cancelling)
+            except (TimeoutError, CancelledError):
+                await close_queue(self._queue_manager, f"_event_{context.task_id}", immediate=True)
+                await close_queue(self._queue_manager, f"_resume_{context.task_id}", immediate=True)
 
     async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
         assert context.message  # this is only executed in the context of SendMessage request
@@ -472,10 +494,12 @@ async def execute(self, context: RequestContext, event_queue: EventQueue) -> Non
                     resume_queue=resume_queue,
                 )
 
-                self._running_tasks[context.task_id] = asyncio.create_task(run_generator)
-                self._running_tasks[context.task_id].add_done_callback(
-                    lambda _: self._running_tasks.pop(context.task_id)  # pyright: ignore [reportArgumentType]
+                self._running_tasks[context.task_id] = RunningTask(
+                    task=asyncio.create_task(run_generator), last_invocation=datetime.now()
                 )
+                asyncio.create_task(
+                    self._schedule_run_cleanup(task_id=context.task_id, task_timeout=self._task_timeout)
+                ).add_done_callback(lambda _: ...)
 
             while True:
                 # Forward messages to local event queue
@@ -490,18 +514,39 @@ async def execute(self, context: RequestContext, event_queue: EventQueue) -> Non
             # When a streaming request is canceled, this executor is canceled first meaning that "cancellation" event
             # passed from the agent's long_running_event_queue is not forwarded. Instead of shielding this function,
             # we report the cancellation explicitly
+            await self._cancel_task(context.task_id)
             local_updater = TaskUpdater(event_queue, task_id=context.task_id, context_id=context.context_id)
             await local_updater.cancel()
         except Exception as ex:
             logger.error("Error executing agent", exc_info=ex)
             local_updater = TaskUpdater(event_queue, task_id=context.task_id, context_id=context.context_id)
             await local_updater.failed(local_updater.new_agent_message(parts=[Part(root=TextPart(text=str(ex)))]))
 
+    async def _cancel_task(self, task_id: str):
+        if queue := self._cancel_queues.get(task_id):
+            await queue.enqueue_event(create_text_message_object(content="canceled"))
+
+    async def _schedule_run_cleanup(self, task_id: str, task_timeout: timedelta):
+        task = self._running_tasks.get(task_id)
+        assert task
+
+        try:
+            while not task["task"].done():
+                await asyncio.sleep(5)
+                if not task["task"].done() and task["last_invocation"] + task_timeout < datetime.now():
+                    # Task might be stuck waiting for queue events to be processed
+                    logger.warning(f"Task {task_id} did not finish in {task_timeout}")
+                    await self._cancel_task(task_id)
+                    break
+        except Exception as ex:
+            logger.error("Error when cleaning up task", exc_info=ex)
+        finally:
+            self._running_tasks.pop(task_id)
+
     async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
         if not context.task_id or not context.context_id:
             raise ValueError("Task ID and context ID must be set to cancel a task")
         try:
-            if context.current_task and (queue := self._cancel_queues.get(context.task_id)):
-                await queue.enqueue_event(context.current_task)
+            await self._cancel_task(task_id=context.task_id)
         finally:
             await TaskUpdater(event_queue, task_id=context.task_id, context_id=context.context_id).cancel()

agents/rag/uv.lock

@@ -2,6 +2,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
+from datetime import timedelta
+
 from a2a.server.agent_execution import RequestContextBuilder
 from a2a.server.apps.jsonrpc import A2AFastAPIApplication
 from a2a.server.apps.rest import A2ARESTFastAPIApplication
@@ -29,13 +31,14 @@ def create_app(
     lifespan: Lifespan[AppType] | None = None,
     dependencies: list[Depends] | None = None,  # pyright: ignore [reportGeneralTypeIssues]
     override_interfaces: bool = True,
+    task_timeout: timedelta = timedelta(minutes=10),
     **kwargs,
 ) -> FastAPI:
     queue_manager = queue_manager or InMemoryQueueManager()
     task_store = task_store or InMemoryTaskStore()
     context_store = context_store or InMemoryContextStore()
     http_handler = DefaultRequestHandler(
-        agent_executor=Executor(agent.execute, queue_manager, context_store=context_store),
+        agent_executor=Executor(agent.execute, queue_manager, context_store=context_store, task_timeout=task_timeout),
         task_store=task_store,
         queue_manager=queue_manager,
         push_config_store=push_config_store,

apps/beeai-cli/uv.lock

@@ -8,6 +8,7 @@
 from collections.abc import AsyncGenerator, Awaitable, Callable
 from configparser import RawConfigParser
 from contextlib import asynccontextmanager, nullcontext, suppress
+from datetime import timedelta
 from ssl import CERT_NONE
 from typing import IO, Any, Literal
 from urllib.parse import urljoin
@@ -65,6 +66,7 @@ async def serve(
         task_store: TaskStore | None = None,
         context_store: ContextStore | None = None,
         queue_manager: QueueManager | None = None,
+        task_timeout: timedelta = timedelta(minutes=10),
         push_config_store: PushNotificationConfigStore | None = None,
         push_sender: PushNotificationSender | None = None,
         request_context_builder: RequestContextBuilder | None = None,
@@ -162,6 +164,7 @@ async def _lifespan_fn(app: FastAPI) -> AsyncGenerator[None, None]:
             queue_manager=queue_manager,
             push_config_store=push_config_store,
             push_sender=push_sender,
+            task_timeout=task_timeout,
             request_context_builder=request_context_builder,
         )