diff --git a/src/a2a/server/agent_execution/active_task_registry.py b/src/a2a/server/agent_execution/active_task_registry.py index ab7d6a11c..bca44162d 100644 --- a/src/a2a/server/agent_execution/active_task_registry.py +++ b/src/a2a/server/agent_execution/active_task_registry.py @@ -69,10 +69,24 @@ async def get_or_create( ) self._active_tasks[task_id] = active_task - await active_task.start( - call_context=call_context, - create_task_if_missing=create_task_if_missing, - ) + try: + await active_task.start( + call_context=call_context, + create_task_if_missing=create_task_if_missing, + ) + except Exception: + # Remove the entry synchronously instead of relying on the + # fire-and-forget _on_active_task_cleanup path scheduled by + # ActiveTask.start() on failure. That cleanup runs in a separate + # task and may not execute before another request checks the + # registry, leaving a zombie entry for a task that never started. + async with self._lock: + self._active_tasks.pop(task_id, None) + logger.debug( + 'Removed failed active task for %s from registry', + task_id, + ) + raise return active_task def _on_active_task_cleanup(self, active_task: ActiveTask) -> None: diff --git a/tests/server/agent_execution/test_active_task_registry.py b/tests/server/agent_execution/test_active_task_registry.py index 16d9c8797..22e127fc8 100644 --- a/tests/server/agent_execution/test_active_task_registry.py +++ b/tests/server/agent_execution/test_active_task_registry.py @@ -105,3 +105,42 @@ async def test_aclose_logs_and_swallows_task_errors(caplog): await registry.aclose() assert 'Error draining active task' in caplog.text + + +@pytest.mark.timeout(5) +@pytest.mark.asyncio +async def test_get_or_create_failed_start_removes_zombie_entry(): + """A failed start() must not leave a zombie registry entry. + + Regression test for zombie-entry cleanup: the entry was inserted before start(), and + the cleanup on failure was fire-and-forget, so a subsequent registry + lookup could still see the never-started task. + """ + from a2a.types.a2a_pb2 import Task, TaskState, TaskStatus + from a2a.utils.errors import InvalidParamsError + + task_store = InMemoryTaskStore() + terminal_task = Task( + id='task-term', + context_id='c1', + status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED), + ) + await task_store.save(terminal_task, ServerCallContext()) + + registry = ActiveTaskRegistry( + agent_executor=_SlowExecutor(), + task_store=task_store, + ) + + with pytest.raises(InvalidParamsError): + await registry.get_or_create( + 'task-term', + call_context=ServerCallContext(), + create_task_if_missing=True, + ) + + # Give any fire-and-forget cleanup tasks a chance to run; the entry must + # already be gone regardless. + await asyncio.sleep(0) + assert await registry.get('task-term') is None + assert 'task-term' not in registry._active_tasks