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