feat(server): release parked interrupted tasks after a configurable timeout (#1338)

## Description

Second half of #1296, following the direction agreed in
#1296 (comment)
and the maintainer's reply.

A task in `input-required` or `auth-required` keeps its `ActiveTask`
(producer, consumer, two dispatchers, registry entry) until a follow-up
arrives, which may never happen.

### Change

- New `DefaultRequestHandlerV2(..., interrupted_task_timeout: float |
None = None)`. The default `None` keeps today's behavior.
- When the timeout is set, `ActiveTask` starts a timer once a request
finishes or the last subscriber leaves. A new request or subscriber
cancels it. When it fires, the `ActiveTask` is released only if all of
these still hold:
  - the task is in an interrupted state,
  - no request is running or queued,
  - no subscriber is attached.
- The stored task is not modified. A follow-up recreates the
`ActiveTask` from the store.
- `0` releases as soon as the task is idle.

### Registry fix

Removing a finished `ActiveTask` from the registry runs asynchronously.
A follow-up that arrived in between was rejected with `Task … is already
completed.` `get_or_create` now replaces a finished entry, and the
scheduled removal only deletes the instance it was scheduled for. Side
effect without a timeout: a follow-up to a task that just ended now gets
`is in terminal state: …`, the same error as outside that window.

### Tests

Seven new tests in `test_default_request_handler_v2.py`:
- no timeout by default,
- release for input-required and auth-required, with no leftover asyncio
tasks and the stored state unchanged,
- resume from the store after a release,
- reset by a follow-up,
- reset by a subscriber,
- no release while the executor is still running,
- a follow-up to a finished but still registered `ActiveTask` (this one
fails on `main`).

### Checks

- `./scripts/lint.sh`: passed (the 3 `ty` warnings already exist on
`main`).
- `uv run pytest`: 2285 passed. The new tests passed 10/10 repeated runs
and also pass on Python 3.14.
- Not tested: cluster mode (`event_stream`).

Part of #1296

Release-As: 1.3.0

---------

Co-authored-by: Mateusz Iwaniuk <45974414+Iwaniukooo11@users.noreply.github.com>

Author: christophstach

src/a2a/server/agent_execution/active_task.py

@@ -246,6 +246,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
@@ -476,6 +477,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.
 
@@ -489,6 +491,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
@@ -501,6 +505,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,
@@ -537,6 +544,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:
@@ -717,6 +729,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)
 
@@ -737,7 +750,7 @@ async def _run_consumer(self) -> None:
             logger.debug('Consumer[%s]: Finishing', self._task_id)
             await self._maybe_cleanup()
 
-    async def subscribe(
+    async def subscribe(  # noqa: PLR0915
         self,
         *,
         request: RequestContext | None = None,
@@ -753,6 +766,7 @@ async def subscribe(
         """
         logger.debug('Subscribe[%s]: New subscriber', self._task_id)
 
+        self.cancel_idle_release()
         async with self._lock:
             if self._is_finished.is_set():
                 raise UnsupportedOperationError(
@@ -843,6 +857,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.
@@ -968,6 +983,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)
@@ -984,6 +1000,54 @@ 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:
+        """Stops the timer that releases a parked interrupted task, if any."""
+        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."""
+        try:
+            await asyncio.sleep(timeout)
+            # Read before `_lock` so the checks and release below cannot suspend.
+            task = await self._task_manager.get_task()
+        except asyncio.CancelledError:
+            raise
+        except Exception:
+            logger.exception(
+                'ActiveTask[%s]: Idle release check failed', self._task_id
+            )
+            return
+        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.
 

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,12 @@ 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 not None:
+                # A request is about to use it, so it must not be released.
+                existing.cancel_idle_release()
             if existing is None:
                 task_manager = TaskManager(
                     task_id=task_id,
@@ -73,6 +81,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
 
@@ -99,6 +108,7 @@ async def get_or_create(
                 if isinstance(existing_task, StoredTask):
                     existing_task = existing_task.task
                 if not existing_task:
+                    existing.schedule_idle_release()
                     raise TaskNotFoundError
             return existing
 
@@ -111,14 +121,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."""

src/a2a/server/request_handlers/default_request_handler_v2.py

@@ -96,6 +96,11 @@ 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;
+    any other value must be greater than 0.
     """
 
     _background_tasks: set[asyncio.Task]
@@ -118,7 +123,13 @@ 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 (
+            interrupted_task_timeout is not None
+            and interrupted_task_timeout <= 0
+        ):
+            raise ValueError('interrupted_task_timeout must be greater than 0.')
         if queue_manager is not None:
             message = (
                 'A queue_manager was passed to DefaultRequestHandlerV2, but it '
@@ -172,6 +183,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()
 
@@ -260,7 +272,11 @@ async def _cancel_local(
             active_task = await self._active_task_registry.get_or_create(
                 task_id, call_context=context, create_task_if_missing=False
             )
-            result = await active_task.cancel(context)
+            try:
+                result = await active_task.cancel(context)
+            finally:
+                # get_or_create stopped the timer; a failed cancel leaves the task parked.
+                active_task.schedule_idle_release()
         except UnsupportedOperationError as e:
             raise TaskNotCancelableError from e