From 62e53521f6d6444d441900a364c81ee9db34a565 Mon Sep 17 00:00:00 2001 From: Long Bui Date: Thu, 3 Sep 2026 15:49:06 +0700 Subject: [PATCH] fix(server): drop stale ActiveTask snapshot when reusing an idle task A reused ActiveTask keeps its TaskManager._current_task snapshot across non-terminal interrupts (input_required/auth_required) because cleanup only runs on terminal states. The per-request get_task() in _run_producer intends to refresh the task each request, but get_task() short-circuits on the cached snapshot, so the refresh is a no-op. With more than one replica sharing a TaskStore, a resume routed back to the pod holding the stale snapshot silently overwrites artifacts, history and status that another replica persisted while the task was parked. Drop the cached snapshot when the registry reuses an idle ActiveTask (reference_count <= 1, no subscriber stream in flight), forcing the next request to re-read the store at the request boundary. A still-streaming task keeps its snapshot, so an open artifact is not lost mid-stream. --- src/a2a/server/agent_execution/active_task.py | 11 +++ .../agent_execution/active_task_registry.py | 5 ++ src/a2a/server/tasks/task_manager.py | 4 + .../agent_execution/test_active_task.py | 22 +++++ .../test_active_task_registry.py | 88 +++++++++++++++++++ tests/server/tasks/test_task_manager.py | 22 +++++ 6 files changed, 152 insertions(+) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 720d18bff..c8c554fa8 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -419,6 +419,17 @@ def task_id(self) -> str: """The ID of the task.""" return self._task_id + def refresh_on_reuse(self) -> None: + """Drops the cached task snapshot when the task is idle. + + Forces the next request to re-read the store instead of resuming from a + stale pre-interrupt snapshot. Skipped while a subscriber stream is in + flight (`_reference_count > 1`), since re-reading mid-stream would lose + the open artifact. + """ + if self._reference_count <= 1: + self._task_manager.invalidate_cached_task() + async def enqueue_request( self, request_context: RequestContext ) -> uuid.UUID: diff --git a/src/a2a/server/agent_execution/active_task_registry.py b/src/a2a/server/agent_execution/active_task_registry.py index 15e9c6350..0595b6cef 100644 --- a/src/a2a/server/agent_execution/active_task_registry.py +++ b/src/a2a/server/agent_execution/active_task_registry.py @@ -52,6 +52,11 @@ 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: + # Drop the reused task's stale snapshot so the next request + # re-reads the store instead of overwriting another replica's + # writes. No-op while a subscriber stream is in flight. + existing.refresh_on_reuse() if existing is None: task_manager = TaskManager( task_id=task_id, diff --git a/src/a2a/server/tasks/task_manager.py b/src/a2a/server/tasks/task_manager.py index c9dfc879f..38402483d 100644 --- a/src/a2a/server/tasks/task_manager.py +++ b/src/a2a/server/tasks/task_manager.py @@ -155,6 +155,10 @@ async def get_task(self) -> Task | None: logger.debug('Task %s not found.', self.task_id) return self._current_task + def invalidate_cached_task(self) -> None: + """Drops the cached snapshot so the next `get_task` re-reads the store.""" + self._current_task = None + async def save_task_event( self, event: Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent ) -> Task | None: diff --git a/tests/server/agent_execution/test_active_task.py b/tests/server/agent_execution/test_active_task.py index 5a06f4d34..81c86bb6a 100644 --- a/tests/server/agent_execution/test_active_task.py +++ b/tests/server/agent_execution/test_active_task.py @@ -72,6 +72,28 @@ async def active_task( push_sender=push_sender, ) + @pytest.mark.asyncio + async def test_refresh_on_reuse_drops_snapshot_when_idle( + self, active_task: ActiveTask, task_manager: Mock + ) -> None: + """An idle reused task (reference_count <= 1) invalidates its snapshot.""" + active_task._reference_count = 1 + + active_task.refresh_on_reuse() + + task_manager.invalidate_cached_task.assert_called_once_with() + + @pytest.mark.asyncio + async def test_refresh_on_reuse_keeps_snapshot_when_streaming( + self, active_task: ActiveTask, task_manager: Mock + ) -> None: + """A task with an in-flight subscriber keeps its snapshot.""" + active_task._reference_count = 2 + + active_task.refresh_on_reuse() + + task_manager.invalidate_cached_task.assert_not_called() + @pytest.mark.asyncio async def test_active_task_already_started( self, active_task: ActiveTask, request_context: Mock diff --git a/tests/server/agent_execution/test_active_task_registry.py b/tests/server/agent_execution/test_active_task_registry.py index 7f0e3e46b..e2ef7d6d3 100644 --- a/tests/server/agent_execution/test_active_task_registry.py +++ b/tests/server/agent_execution/test_active_task_registry.py @@ -170,3 +170,91 @@ async def test_get_or_create_cache_hit_is_owner_scoped(): assert again is active await registry.aclose() + + +@pytest.mark.timeout(5) +@pytest.mark.asyncio +async def test_reused_idle_task_drops_stale_snapshot(): + """Issue #1188: reusing an idle ActiveTask after a non-terminal interrupt + must drop its cached TaskManager snapshot, so the per-request get_task() + in _run_producer re-reads the store instead of resuming from a + pre-interrupt snapshot that would overwrite state another replica wrote. + """ + store = InMemoryTaskStore() + registry = ActiveTaskRegistry( + agent_executor=_SlowExecutor(), task_store=store + ) + ctx = _ctx('alice') + + await store.save( + Task( + id='task-1', + status=TaskStatus(state=TaskState.TASK_STATE_INPUT_REQUIRED), + ), + ctx, + ) + active = await registry.get_or_create( + 'task-1', call_context=ctx, create_task_if_missing=True + ) + + # Simulate the producer having cached a pre-interrupt snapshot, then the + # task going idle (its previous request's subscriber has detached). + stale = Task( + id='task-1', + status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED), + ) + active._task_manager._current_task = stale + assert active._reference_count == 1 # idle: no in-flight subscriber + + reused = await registry.get_or_create( + 'task-1', call_context=ctx, create_task_if_missing=False + ) + + assert reused is active + assert active._task_manager._current_task is None + + await registry.aclose() + + +@pytest.mark.timeout(5) +@pytest.mark.asyncio +async def test_reused_streaming_task_keeps_snapshot(): + """Issue #1188 guard: a reused ActiveTask with a subscriber stream still in + flight (reference_count > 1) must KEEP its snapshot. Re-reading the store + mid-stream would drop the open artifact and the next append=True chunk + would fail with InvalidAgentResponseError. + """ + store = InMemoryTaskStore() + registry = ActiveTaskRegistry( + agent_executor=_SlowExecutor(), task_store=store + ) + ctx = _ctx('alice') + + await store.save( + Task( + id='task-1', + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ), + ctx, + ) + active = await registry.get_or_create( + 'task-1', call_context=ctx, create_task_if_missing=True + ) + + snapshot = Task( + id='task-1', + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ) + active._task_manager._current_task = snapshot + # Simulate an in-flight subscriber tailing the current stream. + active._reference_count = 2 + + reused = await registry.get_or_create( + 'task-1', call_context=ctx, create_task_if_missing=False + ) + + assert reused is active + assert active._task_manager._current_task is snapshot + + active._reference_count = 1 + await registry.aclose() diff --git a/tests/server/tasks/test_task_manager.py b/tests/server/tasks/test_task_manager.py index 7daa3ab3d..f743a6062 100644 --- a/tests/server/tasks/test_task_manager.py +++ b/tests/server/tasks/test_task_manager.py @@ -111,6 +111,28 @@ async def test_get_task_nonexistent( mock_task_store.get.assert_called_once_with(MINIMAL_TASK_ID, TEST_CONTEXT) +@pytest.mark.asyncio +async def test_invalidate_cached_task_forces_store_reread( + task_manager: TaskManager, mock_task_store: AsyncMock +) -> None: + """After invalidation, get_task re-reads the store instead of returning the + cached snapshot (issue #1188).""" + stale = create_minimal_task() + fresh = create_minimal_task() + fresh.status.state = TaskState.TASK_STATE_INPUT_REQUIRED + mock_task_store.get.return_value = fresh + + # Prime the cache; a second get_task without invalidation stays cached. + task_manager._current_task = stale + assert await task_manager.get_task() is stale + mock_task_store.get.assert_not_called() + + task_manager.invalidate_cached_task() + + assert await task_manager.get_task() is fresh + mock_task_store.get.assert_called_once_with(MINIMAL_TASK_ID, TEST_CONTEXT) + + @pytest.mark.asyncio async def test_save_task_event_new_task( task_manager: TaskManager, mock_task_store: AsyncMock