diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 54b970cdf..400d37fe7 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -213,6 +213,7 @@ async def _process_event_inner(self, event: Event) -> None: 'Consumer[%s]: Request completed', self.active_task._task_id ) self.active_task._request_lock.release() + self.active_task._schedule_idle_release() elif isinstance(event, _RequestStarted): logger.debug( 'Consumer[%s]: Request started', self.active_task._task_id @@ -441,6 +442,7 @@ def __init__( # noqa: PLR0913 push_sender: PushNotificationSender | None = None, on_cleanup: Callable[[ActiveTask], None] | None = None, event_stream: TaskEventStream | None = None, + interrupted_task_timeout: float | None = None, ) -> None: """Initializes the ActiveTask. @@ -454,6 +456,8 @@ def __init__( # noqa: PLR0913 the task from the ActiveTaskRegistry. event_stream: Optional cross-replica stream; applied events are published to it so other replicas can observe them. + interrupted_task_timeout: Optional seconds after which an idle + interrupted task is finished; ``None`` disables it. """ # --- Core Dependencies --- self._agent_executor = agent_executor @@ -466,6 +470,9 @@ def __init__( # noqa: PLR0913 self._push_sender = push_sender self._on_cleanup = on_cleanup self._event_stream = event_stream + self._interrupted_task_timeout = interrupted_task_timeout + # Pending timer that releases a task parked in an interrupted state. + self._idle_release_task: asyncio.Task[None] | None = None # --- Synchronization Primitives --- # `_lock` protects structural lifecycle changes: start(), subscribe() counting, @@ -502,6 +509,11 @@ def task_id(self) -> str: """The ID of the task.""" return self._task_id + @property + def is_finished(self) -> bool: + """Whether the task's active lifecycle has ended.""" + return self._is_finished.is_set() + async def enqueue_request( self, request_context: RequestContext ) -> uuid.UUID: @@ -692,6 +704,7 @@ async def _run_consumer(self) -> None: await EventConsumer(self).run() finally: self._is_finished.set() + self._cancel_idle_release() self._request_queue.shutdown(immediate=True) await self._event_queue_agent.close(immediate=True) @@ -734,6 +747,7 @@ async def subscribe( f'Task {self._task_id} is already completed.' ) self._reference_count += 1 + self._cancel_idle_release() logger.debug( 'Subscribe[%s]: Subscribers count: %d', self._task_id, @@ -814,6 +828,7 @@ async def subscribe( self._reference_count -= 1 # Evaluate if this was the last subscriber on a finished task. await self._maybe_cleanup() + self._schedule_idle_release() async def cancel(self, call_context: ServerCallContext) -> Task: """Cancels the running active task. @@ -940,6 +955,7 @@ async def aclose(self) -> None: # tasks, because their teardown re-acquires it. async with self._lock: self._is_finished.set() + self._cancel_idle_release() background_tasks = [ task for task in (self._producer_task, self._consumer_task) @@ -956,6 +972,45 @@ async def aclose(self) -> None: if isinstance(result, Exception): logger.error('Error during aclose', exc_info=result) + def _schedule_idle_release(self) -> None: + """(Re)starts the timer that releases a parked interrupted task.""" + if self._interrupted_task_timeout is None or self._is_finished.is_set(): + return + self._cancel_idle_release() + self._idle_release_task = asyncio.create_task( + self._release_when_idle(self._interrupted_task_timeout), + name=f'idle-release:{self._task_id}', + ) + + def _cancel_idle_release(self) -> None: + if self._idle_release_task is not None: + self._idle_release_task.cancel() + self._idle_release_task = None + + async def _release_when_idle(self, timeout: float) -> None: + """Finishes the task if it is still parked after `timeout` seconds.""" + await asyncio.sleep(timeout) + # Read before `_lock` so the checks and release below cannot suspend. + task = await self._task_manager.get_task() + async with self._lock: + if ( + task is None + or task.status.state not in INTERRUPTED_TASK_STATES + or self._is_finished.is_set() + # The consumer holds one reference; more means subscribers. + or self._reference_count != 1 + or self._request_lock.locked() + or not self._request_queue.empty() + ): + return + logger.debug( + 'ActiveTask[%s]: Releasing task parked in %s', + self._task_id, + TaskState.Name(task.status.state), + ) + self._is_finished.set() + self._request_queue.shutdown(immediate=True) + async def _maybe_cleanup(self) -> None: """Triggers cleanup if task is finished and has no subscribers. diff --git a/src/a2a/server/agent_execution/active_task_registry.py b/src/a2a/server/agent_execution/active_task_registry.py index cae2a55ce..814c750f9 100644 --- a/src/a2a/server/agent_execution/active_task_registry.py +++ b/src/a2a/server/agent_execution/active_task_registry.py @@ -34,10 +34,12 @@ def __init__( task_store: TaskStore | VersionedTaskStore, push_sender: PushNotificationSender | None = None, event_stream: TaskEventStream | None = None, + interrupted_task_timeout: float | None = None, ): self._agent_executor = agent_executor self._task_store = task_store self._push_sender = push_sender + self._interrupted_task_timeout = interrupted_task_timeout self._event_stream = event_stream self._active_tasks: dict[str, ActiveTask] = {} self._lock = threading.RLock() @@ -57,6 +59,9 @@ async def get_or_create( if self._closed: raise RuntimeError('ActiveTaskRegistry is closed') existing = self._active_tasks.get(task_id) + if existing is not None and existing.is_finished: + # Its removal runs asynchronously, so replace it instead. + existing = None if existing is None: task_manager = TaskManager( task_id=task_id, @@ -73,6 +78,7 @@ async def get_or_create( push_sender=self._push_sender, on_cleanup=self._on_active_task_cleanup, event_stream=self._event_stream, + interrupted_task_timeout=self._interrupted_task_timeout, ) self._active_tasks[task_id] = active_task @@ -111,14 +117,19 @@ async def get_or_create( def _on_active_task_cleanup(self, active_task: ActiveTask) -> None: """Called by ActiveTask when it's finished and has no subscribers.""" logger.debug('Active task %s cleanup scheduled', active_task.task_id) - task = asyncio.create_task(self._remove_task(active_task.task_id)) + task = asyncio.create_task(self._remove_task(active_task)) self._cleanup_tasks.add(task) task.add_done_callback(self._cleanup_tasks.discard) - async def _remove_task(self, task_id: str) -> None: + async def _remove_task(self, active_task: ActiveTask) -> None: + task_id = active_task.task_id with self._lock: - self._active_tasks.pop(task_id, None) - logger.debug('Removed active task for %s from registry', task_id) + # Only remove this instance; it may already have been replaced. + if self._active_tasks.get(task_id) is active_task: + del self._active_tasks[task_id] + logger.debug( + 'Removed active task for %s from registry', task_id + ) async def get(self, task_id: str) -> ActiveTask | None: """Retrieves an existing task.""" diff --git a/src/a2a/server/request_handlers/default_request_handler_v2.py b/src/a2a/server/request_handlers/default_request_handler_v2.py index b135f529f..3496b0d9a 100644 --- a/src/a2a/server/request_handlers/default_request_handler_v2.py +++ b/src/a2a/server/request_handlers/default_request_handler_v2.py @@ -96,6 +96,10 @@ class DefaultRequestHandlerV2(RequestHandler): with `DefaultRequestHandler` but is not used: v2 delegates event streaming to an in-memory `ActiveTaskRegistry`. Passing a non-``None`` value emits a `DeprecationWarning` and logs a warning. + + ``interrupted_task_timeout`` releases the `ActiveTask` of a task idle in + ``input-required`` or ``auth-required`` after that many seconds; a + follow-up recreates it from the store. ``None`` (default) never releases. """ _background_tasks: set[asyncio.Task] @@ -118,6 +122,7 @@ def __init__( # noqa: PLR0913 push_url_validator: Callable[[str], Awaitable[bool]] | None = None, validate_input_modes: bool = False, event_stream: TaskEventStream | None = None, + interrupted_task_timeout: float | None = None, ) -> None: if queue_manager is not None: message = ( @@ -172,6 +177,7 @@ def __init__( # noqa: PLR0913 task_store=self._versioned_store, push_sender=self._push_sender, event_stream=self._event_stream, + interrupted_task_timeout=interrupted_task_timeout, ) self._background_tasks = set() diff --git a/tests/server/request_handlers/test_default_request_handler_v2.py b/tests/server/request_handlers/test_default_request_handler_v2.py index b30fdc695..35799388d 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/tests/server/request_handlers/test_default_request_handler_v2.py @@ -2776,3 +2776,196 @@ async def test_on_message_send_stream_message_response_releases_active_task() -> assert [type(event) for event in events] == [Message] assert await _wait_until_released(handler, baseline) == set() assert handler._active_task_registry._active_tasks == {} + + +class _InterruptingAgent(AgentExecutor): + """Interrupts a new task with `state`; completes it on the follow-up.""" + + def __init__( + self, + state: TaskState = TaskState.TASK_STATE_INPUT_REQUIRED, + complete_on_follow_up: bool = True, + linger: float = 0, + ) -> None: + self.state = state + self.complete_on_follow_up = complete_on_follow_up + self.linger = linger + + async def execute( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + if context.current_task is None: + task = new_task_from_user_message(context.message) + await event_queue.enqueue_event(task) + task_id, context_id = task.id, task.context_id + else: + task_id, context_id = context.task_id, context.context_id + updater = TaskUpdater(event_queue, task_id, context_id) + if context.current_task is not None and self.complete_on_follow_up: + await updater.complete() + return + await updater.update_status(self.state) + # Keeps the request running after the interruption was delivered. + await asyncio.sleep(self.linger) + + async def cancel( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + pass + + +def _interrupting_handler( + agent: AgentExecutor, interrupted_task_timeout: float | None +) -> DefaultRequestHandlerV2: + return DefaultRequestHandlerV2( + agent_executor=agent, + task_store=InMemoryTaskStore(), + agent_card=create_default_agent_card(), + interrupted_task_timeout=interrupted_task_timeout, + ) + + +async def _send( + handler: DefaultRequestHandlerV2, text: str, task: Task | None = None +) -> Task: + message = Message( + message_id=str(uuid.uuid4()), + role=Role.ROLE_USER, + parts=[Part(text=text)], + task_id=task.id if task else '', + context_id=task.context_id if task else '', + ) + result = await handler.on_message_send( + SendMessageRequest(message=message), create_server_call_context() + ) + assert isinstance(result, Task) + return result + + +async def _wait_for_release( + handler: DefaultRequestHandlerV2, task_id: str, timeout: float = 2.0 +) -> bool: + """Waits until the task's ActiveTask has left the registry.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if task_id not in handler._active_task_registry._active_tasks: + return True + await asyncio.sleep(0.01) + return False + + +@pytest.mark.asyncio +async def test_interrupted_task_stays_active_without_timeout() -> None: + """The default keeps today's behavior: no timeout, the ActiveTask waits.""" + handler = _interrupting_handler(_InterruptingAgent(), None) + + task = await _send(handler, 'book a flight') + + assert not await _wait_for_release(handler, task.id, timeout=0.3) + await handler.aclose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + 'state', + [TaskState.TASK_STATE_INPUT_REQUIRED, TaskState.TASK_STATE_AUTH_REQUIRED], +) +async def test_interrupted_task_released_after_timeout( + state: TaskState, +) -> None: + """A parked task's runtime is released, but the stored task is kept.""" + handler = _interrupting_handler(_InterruptingAgent(state), 0.1) + baseline = asyncio.all_tasks() + + task = await _send(handler, 'book a flight') + + assert task.status.state == state + assert await _wait_for_release(handler, task.id) + leftover: set[asyncio.Task] = set() + for _ in range(100): + leftover = asyncio.all_tasks() - baseline - {asyncio.current_task()} + if not leftover: + break + await asyncio.sleep(0.01) + assert leftover == set() + stored = await handler.on_get_task( + GetTaskRequest(id=task.id), create_server_call_context() + ) + assert stored is not None + assert stored.status.state == state + await handler.aclose() + + +@pytest.mark.asyncio +async def test_follow_up_after_release_resumes_from_store() -> None: + handler = _interrupting_handler(_InterruptingAgent(), 0.1) + task = await _send(handler, 'book a flight') + assert await _wait_for_release(handler, task.id) + + result = await _send(handler, 'Friday', task) + + assert result.status.state == TaskState.TASK_STATE_COMPLETED + assert [m.parts[0].text for m in result.history] == [ + 'book a flight', + 'Friday', + ] + await handler.aclose() + + +@pytest.mark.asyncio +async def test_follow_up_resets_interrupted_task_timeout() -> None: + handler = _interrupting_handler( + _InterruptingAgent(complete_on_follow_up=False), 0.5 + ) + task = await _send(handler, 'book a flight') + await asyncio.sleep(0.3) + + await _send(handler, 'Friday', task) + + # 0.6s after the first request, but only 0.3s after the follow-up. + assert not await _wait_for_release(handler, task.id, timeout=0.3) + assert await _wait_for_release(handler, task.id) + await handler.aclose() + + +@pytest.mark.asyncio +async def test_subscriber_keeps_interrupted_task_active() -> None: + handler = _interrupting_handler(_InterruptingAgent(), 0.1) + task = await _send(handler, 'book a flight') + + subscription = handler.on_subscribe_to_task( + SubscribeToTaskRequest(id=task.id), create_server_call_context() + ) + await anext(subscription) + assert not await _wait_for_release(handler, task.id, timeout=0.3) + + await subscription.aclose() + assert await _wait_for_release(handler, task.id) + await handler.aclose() + + +@pytest.mark.asyncio +async def test_running_request_keeps_interrupted_task_active() -> None: + """The timer only counts once the executor has returned.""" + handler = _interrupting_handler(_InterruptingAgent(linger=0.4), 0.1) + task = await _send(handler, 'book a flight') + + assert not await _wait_for_release(handler, task.id, timeout=0.3) + assert await _wait_for_release(handler, task.id) + await handler.aclose() + + +@pytest.mark.asyncio +async def test_follow_up_replaces_finished_registered_active_task() -> None: + """A finished ActiveTask still in the registry must not reject a follow-up.""" + handler = _interrupting_handler(_InterruptingAgent(), None) + task = await _send(handler, 'book a flight') + active_task = handler._active_task_registry._active_tasks[task.id] + # Finish it the way a release does, before its registry cleanup has run. + active_task._is_finished.set() + active_task._request_queue.shutdown(immediate=True) + + result = await _send(handler, 'Friday', task) + + assert result.status.state == TaskState.TASK_STATE_COMPLETED + await handler.aclose()