fix: refresh idle ActiveTask snapshots on reuse

src/a2a/server/agent_execution/active_task.py

@@ -391,6 +391,8 @@ def __init__(
 
         # `_request_lock` protects parallel request processing.
         self._request_lock = asyncio.Lock()
+        # `_snapshot_lock` serializes cache reads with request-boundary refreshes.
+        self._snapshot_lock = asyncio.Lock()
 
         # _task_created is set when initial version of task is stored in DB.
         self._task_created = asyncio.Event()
@@ -419,6 +421,14 @@ def task_id(self) -> str:
         """The ID of the task."""
         return self._task_id
 
+    @staticmethod
+    def _raise_if_task_terminal(task: Task) -> None:
+        """Rejects operations that would restart a terminal task."""
+        if task.status.state in TERMINAL_TASK_STATES:
+            raise InvalidParamsError(
+                message=f'Task {task.id} is in terminal state: {task.status.state}'
+            )
+
     async def enqueue_request(
         self, request_context: RequestContext
     ) -> uuid.UUID:
@@ -427,6 +437,24 @@ async def enqueue_request(
         await self._request_queue.put((request_context, request_id))
         return request_id
 
+    async def refresh_task_if_idle(
+        self, call_context: ServerCallContext
+    ) -> None:
+        """Drops an idle task snapshot before a new request boundary.
+
+        An ``ActiveTask`` remains registered while a task waits for
+        human-in-the-loop input,
+        so another process may persist a newer task snapshot in the meantime.
+        The request lock, rather than the subscriber count, defines whether
+        execution is idle. A subscriber may detach before background artifact
+        streaming finishes, while passive subscribers may remain after the
+        previous request is fully persisted.
+        """
+        async with self._snapshot_lock, self._lock:
+            if not self._request_lock.locked():
+                self._task_manager._call_context = call_context
+                self._task_manager._current_task = None
+
     async def start(
         self,
         call_context: ServerCallContext,
@@ -468,10 +496,7 @@ async def start(
 
                 if task:
                     self._task_created.set()
-                    if task.status.state in TERMINAL_TASK_STATES:
-                        raise InvalidParamsError(
-                            message=f'Task {task.id} is in terminal state: {task.status.state}'
-                        )
+                    self._raise_if_task_terminal(task)
                 elif not create_task_if_missing:
                     raise TaskNotFoundError
 
@@ -510,20 +535,34 @@ async def _run_producer(self) -> None:
         """
         logger.debug('Producer[%s]: Started', self._task_id)
         request_context = None
+        task_missing_at_boundary = False
         try:
             while True:
                 (
                     request_context,
                     request_id,
                 ) = await self._request_queue.get()
                 await self._request_lock.acquire()
-                # TODO: Should we create task manager every time?
-                self._task_manager._call_context = request_context.call_context
-
-                request_context.current_task = (
-                    await self._task_manager.get_task()
-                )
+                # Order task loading after any idle refresh that began before
+                # this request acquired `_request_lock`. Later refreshes observe
+                # the held request lock and leave the streaming snapshot intact.
+                async with self._snapshot_lock:
+                    self._task_manager._call_context = (
+                        request_context.call_context
+                    )
+                    # This is the request boundary for queued requests that
+                    # were discovered while the previous request was active.
+                    self._task_manager._current_task = None
+                    request_context.current_task = (
+                        await self._task_manager.get_task()
+                    )
 
+                if (
+                    request_context.current_task is None
+                    and self._task_created.is_set()
+                ):
+                    task_missing_at_boundary = True
+                    raise TaskNotFoundError(f'Task {self._task_id} not found')
                 logger.debug(
                     'Producer[%s]: Executing agent task %s',
                     self._task_id,
@@ -562,7 +601,7 @@ async def _run_producer(self) -> None:
             )
             # Persist the failure directly instead of relying on the closing
             # event queue to carry a final status update.
-            if request_context:
+            if request_context and not task_missing_at_boundary:
                 task = await self._task_manager.ensure_task_id(
                     self._task_id,
                     request_context.context_id or '',
@@ -643,6 +682,7 @@ async def subscribe(
                     self._task_id,
                 )
                 task = await self.get_task()
+                self._raise_if_task_terminal(task)
                 yield task
 
             while True:
@@ -715,9 +755,13 @@ async def cancel(self, call_context: ServerCallContext) -> Task:
         logger.debug('Cancel[%s]: Cancelling task', self._task_id)
 
         # TODO: Conflicts with call_context on the pending request.
-        self._task_manager._call_context = call_context
-
-        task = await self._task_manager.get_task()
+        async with self._snapshot_lock:
+            self._task_manager._call_context = call_context
+            task = await self._task_manager.get_task()
+        if task is None and self._task_created.is_set():
+            raise TaskNotFoundError(f'Task {self._task_id} not found')
+        if task is not None and task.status.state in TERMINAL_TASK_STATES:
+            return task
         request_context = RequestContext(
             call_context=call_context,
             task_id=self._task_id,
@@ -843,7 +887,8 @@ async def _mark_task_as_failed(self, exception: Exception) -> Task | None:
     async def get_task(self) -> Task:
         """Get task from db."""
         await self._task_created.wait()
-        task = await self._task_manager.get_task()
+        async with self._snapshot_lock:
+            task = await self._task_manager.get_task()
         if not task:
             raise RuntimeError('Task should have been created')
         return task

src/a2a/server/agent_execution/active_task_registry.py

@@ -46,28 +46,38 @@ async def get_or_create(
         initial_message: Message | None = None,
     ) -> ActiveTask:
         """Retrieves an existing ActiveTask or creates a new one."""
-        async with self._lock:
-            if self._closed:
-                raise RuntimeError('ActiveTaskRegistry is closed')
-            if task_id in self._active_tasks:
-                return self._active_tasks[task_id]
-
-            task_manager = TaskManager(
-                task_id=task_id,
-                context_id=context_id,
-                task_store=self._task_store,
-                initial_message=initial_message,
-                context=call_context,
-            )
-
-            active_task = ActiveTask(
-                agent_executor=self._agent_executor,
-                task_id=task_id,
-                task_manager=task_manager,
-                push_sender=self._push_sender,
-                on_cleanup=self._on_active_task_cleanup,
-            )
-            self._active_tasks[task_id] = active_task
+        while True:
+            async with self._lock:
+                if self._closed:
+                    raise RuntimeError('ActiveTaskRegistry is closed')
+                active_task = self._active_tasks.get(task_id)
+                if active_task is None:
+                    task_manager = TaskManager(
+                        task_id=task_id,
+                        context_id=context_id,
+                        task_store=self._task_store,
+                        initial_message=initial_message,
+                        context=call_context,
+                    )
+
+                    active_task = ActiveTask(
+                        agent_executor=self._agent_executor,
+                        task_id=task_id,
+                        task_manager=task_manager,
+                        push_sender=self._push_sender,
+                        on_cleanup=self._on_active_task_cleanup,
+                    )
+                    self._active_tasks[task_id] = active_task
+                    break
+
+            # A refresh can wait behind task-store I/O, so do not hold the
+            # global registry lock while synchronizing this individual task.
+            await active_task.refresh_task_if_idle(call_context)
+            async with self._lock:
+                if self._closed:
+                    raise RuntimeError('ActiveTaskRegistry is closed')
+                if self._active_tasks.get(task_id) is active_task:
+                    return active_task
 
         await active_task.start(
             call_context=call_context,

src/a2a/server/request_handlers/default_request_handler_v2.py

@@ -164,6 +164,13 @@ async def on_cancel_task(  # noqa: D102
         context: ServerCallContext,
     ) -> Task | None:
         task_id = params.id
+        task = await self.task_store.get(task_id, context)
+        if not task:
+            raise TaskNotFoundError(f'Task {task_id} not found')
+        if task.status.state in TERMINAL_TASK_STATES:
+            raise TaskNotCancelableError(
+                message=f'Task cannot be canceled - current state: {task.status.state}'
+            )
 
         try:
             active_task = await self._active_task_registry.get_or_create(
@@ -203,6 +210,10 @@ async def _setup_active_task(
             task = await self.task_store.get(original_task_id, call_context)
             if not task:
                 raise TaskNotFoundError(f'Task {original_task_id} not found')
+            if task.status.state in TERMINAL_TASK_STATES:
+                raise InvalidParamsError(
+                    message=f'Task {task.id} is in terminal state: {task.status.state}'
+                )
 
         # Build context to resolve or generate missing IDs
         request_context = await self._request_context_builder.build(

tests/server/agent_execution/test_active_task.py

@@ -27,6 +27,13 @@
 logger = logging.getLogger(__name__)
 
 
+def _working_task() -> Task:
+    return Task(
+        id='test-task-id',
+        status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
+    )
+
+
 class TestActiveTask:
     """Tests for the ActiveTask class."""
 
@@ -105,10 +112,9 @@ async def execute_mock(req, q):
         agent_executor.execute = AsyncMock(side_effect=execute_mock)
         agent_executor.cancel = AsyncMock()
         task_manager.get_task.side_effect = [
-            Task(
-                id='test-task-id',
-                status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
-            )
+            _working_task(),
+            _working_task(),
+            _working_task(),
         ] + [
             Task(
                 id='test-task-id',
@@ -153,10 +159,8 @@ async def execute_mock(req, q):
 
         agent_executor.execute = AsyncMock(side_effect=execute_mock)
         task_manager.get_task.side_effect = [
-            Task(
-                id='test-task-id',
-                status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
-            )
+            _working_task(),
+            _working_task(),
         ] + [task_obj] * 10
 
         await active_task.start(
@@ -200,10 +204,8 @@ async def execute_mock(req, q):
 
         agent_executor.execute = AsyncMock(side_effect=execute_mock)
         task_manager.get_task.side_effect = [
-            Task(
-                id='test-task-id',
-                status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
-            )
+            _working_task(),
+            _working_task(),
         ] + [task_obj] * 10
 
         await active_task.start(
@@ -266,10 +268,8 @@ async def execute_mock(req, q):
 
         agent_executor.execute = AsyncMock(side_effect=execute_mock)
         task_manager.get_task.side_effect = [
-            Task(
-                id='test-task-id',
-                status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
-            )
+            _working_task(),
+            _working_task(),
         ] + [task_obj] * 10
 
         await active_task.start(
@@ -380,10 +380,8 @@ async def execute_mock(req, q):
 
         agent_executor.execute = AsyncMock(side_effect=execute_mock)
         task_manager.get_task.side_effect = [
-            Task(
-                id='test-task-id',
-                status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
-            )
+            _working_task(),
+            _working_task(),
         ] + [task_obj] * 10
 
         await active_task.start(
@@ -517,10 +515,8 @@ async def execute_mock(req, q):
 
         agent_executor.execute = AsyncMock(side_effect=execute_mock)
         task_manager.get_task.side_effect = [
-            Task(
-                id='test-task-id',
-                status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
-            )
+            _working_task(),
+            _working_task(),
         ] + [task_obj] * 10
 
         await active_task.start(
@@ -662,10 +658,8 @@ async def execute_mock(req, q):
 
         agent_executor.execute = AsyncMock(side_effect=execute_mock)
         task_manager.get_task.side_effect = [
-            Task(
-                id='test-task-id',
-                status=TaskStatus(state=TaskState.TASK_STATE_WORKING),
-            )
+            _working_task(),
+            _working_task(),
         ] + [task_obj] * 10
 
         await active_task.start(