From 2514167932dad4961091357671edd43f0b2949cf Mon Sep 17 00:00:00 2001 From: karthik-0306 Date: Wed, 14 Oct 2026 21:24:35 +0530 Subject: [PATCH 1/5] Python: fix(core): keep finished background tasks until cleared or released Part of #8760 ### Description Keep finished asyncio.Task objects in _RuntimeState.in_flight_tasks inside _finalize_task. They remain in memory until explicitly removed by background_agents_clear_completed_task or torn down via BackgroundAgentsProvider.release_session. ### Bug & Repro In the single-process implementation, _finalize_task extracted results and immediately removed the task from in_flight_tasks. However, host session state is persisted only after the agent run completes. If the host restores a session snapshot taken while a task was RUNNING (e.g. after a run failure, turn rollback, or host desync), subsequent calls to _refresh_task_state found no entry in in_flight_tasks and marked the task as LOST, silently discarding a completed or failed result. ### Trade-off Finished asyncio.Task objects remain referenced in in_flight_tasks until explicitly cleared via clear_completed_task or released via release_session, instead of being discarded immediately at finalization. ### Validation - pytest: 52 passed, 0 failed in 0.77s - ruff check: All checks passed - ruff format: 2 files formatted clean - pyright: 0 errors, 0 warnings, 0 informations - mypy: Success: no issues found --- .../_harness/_background_agents.py | 18 +- .../core/test_harness_background_agents.py | 175 ++++++++++++++++++ 2 files changed, 191 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_harness/_background_agents.py b/python/packages/core/agent_framework/_harness/_background_agents.py index 6cd70a34037..f28c059511a 100644 --- a/python/packages/core/agent_framework/_harness/_background_agents.py +++ b/python/packages/core/agent_framework/_harness/_background_agents.py @@ -217,7 +217,20 @@ def _finalize_task( completed_task: asyncio.Task[AgentResponse[Any]], runtime: _RuntimeState, ) -> None: - """Extract results from a completed asyncio task and update task info.""" + """Extract results from a completed asyncio task and update task info. + + The finished :class:`asyncio.Task` is intentionally **not** removed from + ``runtime.in_flight_tasks`` here. Session state is persisted by the host + only after the agent run completes; if the run fails, is stopped early, or + the host restores an older session snapshot, the persisted record will still + say ``RUNNING`` while this in-memory task no longer exists. On the next + call to :func:`_refresh_task_state` a missing entry would be misidentified + as ``LOST``, discarding a result that is already available. + + Tasks are removed from ``in_flight_tasks`` by two callers: + - :func:`background_agents_clear_completed_task` — explicit caller-driven cleanup. + - :meth:`BackgroundAgentsProvider.release_session` — bulk cleanup on session teardown. + """ if completed_task.cancelled(): task_info.status = BackgroundTaskStatus.FAILED task_info.error_text = "Task was canceled." @@ -229,7 +242,8 @@ def _finalize_task( else: task_info.status = BackgroundTaskStatus.COMPLETED task_info.result_text = completed_task.result().text - runtime.in_flight_tasks.pop(task_info.id, None) + # NOTE: do NOT pop task_info.id from runtime.in_flight_tasks here. + # See docstring above for the reasoning. def _refresh_task_state( diff --git a/python/packages/core/tests/core/test_harness_background_agents.py b/python/packages/core/tests/core/test_harness_background_agents.py index 88e776a0c81..1e8f38610de 100644 --- a/python/packages/core/tests/core/test_harness_background_agents.py +++ b/python/packages/core/tests/core/test_harness_background_agents.py @@ -920,3 +920,178 @@ async def fake_drain(*args: Any, **kwargs: Any) -> None: assert drain_called is False assert session.session_id not in provider._runtime + + +# --------------------------------------------------------------------------- +# Regression tests: finished tasks must survive until cleared or released +# (Part of #8760 - single-process snapshot/state desync fix) +# --------------------------------------------------------------------------- + + +async def test_snapshot_restore_returns_real_result_not_lost() -> None: + """A session snapshot taken while a task is RUNNING must still return the + real result after the task finishes and the snapshot is restored. + + Repro for the single-process gap: _finalize_task previously removed the + finished asyncio.Task from in_flight_tasks. If the host persists session + state only after the run completes and a failure (or snapshot restoration) + occurs before that, the next refresh would find no local task and mark it + LOST, silently discarding the result. + """ + provider = _make_provider(_FakeAgent("Worker", response_text="hello from worker")) + session = _make_session() + tools = await _get_tools(provider, session) + + # Start a task and immediately snapshot the session state (task is RUNNING). + await _invoke_tool( + tools["background_agents_start_task"], + agent_name="Worker", + input="do work", + description="regression task", + ) + snapshot = session.to_dict() # captured while task is still RUNNING + + # Let the task complete deterministically (no sleep). + runtime = provider._get_runtime(session) + task = runtime.in_flight_tasks[1] + await task # wait for the asyncio.Task to finish + + # Finalize by refreshing (simulates what before_run / get_task_results does). + tools2 = await _get_tools(provider, session) + result_before_restore = await _invoke_tool(tools2["background_agents_get_task_results"], task_id=1) + assert result_before_restore == "hello from worker" + + # Now restore the pre-finalization snapshot into the same provider. + # This simulates the host restoring an older session state (e.g. after a + # failed persist or a multi-turn rollback). + session2 = AgentSession.from_dict(snapshot) + tools3 = await _get_tools(provider, session2) + + # The task is still in in_flight_tasks (finished but not yet cleared). + # get_task_results must return the real result, not "Task state was lost". + result_after_restore = await _invoke_tool(tools3["background_agents_get_task_results"], task_id=1) + assert result_after_restore == "hello from worker", f"Expected real result but got: {result_after_restore!r}" + + await provider.release_session(session) + await provider.release_session(session2) + + +async def test_snapshot_restore_preserves_error_text_for_failed_task() -> None: + """Same as above but for a task that raises; error_text must survive restore.""" + provider = _make_provider(_FakeAgent("Failer", should_fail=True)) + session = _make_session() + tools = await _get_tools(provider, session) + + await _invoke_tool( + tools["background_agents_start_task"], + agent_name="Failer", + input="trigger failure", + description="failure regression", + ) + snapshot = session.to_dict() + + runtime = provider._get_runtime(session) + task = runtime.in_flight_tasks[1] + with suppress(Exception): + await task + + # Finalize via a refresh so task_info gets FAILED status. + tools2 = await _get_tools(provider, session) + await _invoke_tool(tools2["background_agents_get_task_results"], task_id=1) + + # Restore snapshot and confirm the failure is still readable. + session2 = AgentSession.from_dict(snapshot) + tools3 = await _get_tools(provider, session2) + result = await _invoke_tool(tools3["background_agents_get_task_results"], task_id=1) + assert "Agent execution failed" in result, f"Unexpected result: {result!r}" + + await provider.release_session(session) + await provider.release_session(session2) + + +async def test_snapshot_restore_preserves_error_text_for_cancelled_task() -> None: + """Cancelled task error_text must survive snapshot restoration.""" + provider = _make_provider(_FakeAgent("Slow", delay=60.0)) + session = _make_session() + tools = await _get_tools(provider, session) + + await _invoke_tool( + tools["background_agents_start_task"], + agent_name="Slow", + input="start slow task", + description="cancellation regression", + ) + snapshot = session.to_dict() + + runtime = provider._get_runtime(session) + task = runtime.in_flight_tasks[1] + task.cancel() + with suppress(asyncio.CancelledError): + await task + + # Finalize via a refresh. + tools2 = await _get_tools(provider, session) + await _invoke_tool(tools2["background_agents_get_task_results"], task_id=1) + + # Restore snapshot and confirm cancellation error survives. + session2 = AgentSession.from_dict(snapshot) + tools3 = await _get_tools(provider, session2) + result = await _invoke_tool(tools3["background_agents_get_task_results"], task_id=1) + assert "canceled" in result.lower(), f"Unexpected result: {result!r}" + + await provider.release_session(session) + await provider.release_session(session2) + + +async def test_clear_completed_task_removes_entry_from_in_flight_tasks() -> None: + """clear_completed_task must remove the finished task from in_flight_tasks.""" + provider = _make_provider(_FakeAgent("Worker", response_text="done")) + session = _make_session() + tools = await _get_tools(provider, session) + + await _invoke_tool( + tools["background_agents_start_task"], + agent_name="Worker", + input="work", + description="clear test", + ) + + runtime = provider._get_runtime(session) + await runtime.in_flight_tasks[1] + + # Finalize and then clear. + tools2 = await _get_tools(provider, session) + await _invoke_tool(tools2["background_agents_get_task_results"], task_id=1) + clear_result = await _invoke_tool(tools2["background_agents_clear_completed_task"], task_id=1) + + assert clear_result == "Task 1 cleared." + assert 1 not in runtime.in_flight_tasks, "Task should be removed by clear_completed_task" + + await provider.release_session(session) + + +async def test_release_session_clears_in_flight_tasks() -> None: + """release_session must clear all in_flight_tasks including finished ones.""" + provider = _make_provider(_FakeAgent("Worker", response_text="done")) + session = _make_session() + tools = await _get_tools(provider, session) + + await _invoke_tool( + tools["background_agents_start_task"], + agent_name="Worker", + input="work", + description="release test", + ) + + runtime = provider._get_runtime(session) + await runtime.in_flight_tasks[1] + + # Finalize (task is now finished but still in in_flight_tasks). + tools2 = await _get_tools(provider, session) + await _invoke_tool(tools2["background_agents_get_task_results"], task_id=1) + + assert 1 in runtime.in_flight_tasks, "Finished task should still be present before release" + + await provider.release_session(session) + + assert session.session_id not in provider._runtime, "Runtime should be cleaned up after release" From 6a12a170627912b4539dc1f14a57d4d9ef074dfc Mon Sep 17 00:00:00 2001 From: karthik-0306 Date: Wed, 14 Oct 2026 22:03:48 +0530 Subject: [PATCH 2/5] Python: Record generation token on background tasks to prevent wrong execution attribution (#8760) Address review finding on reusable task IDs in continuation flows: - Track a monotonic generation token on BackgroundTaskInfo and serialize it in to_dict/from_dict. - Dual-key _InFlightTasks with (task_id, generation) so each execution of a continued task is stored discretely. - _refresh_task_state reconciles RUNNING tasks against their specific generation, preventing a restored pre-continuation snapshot from erroneously receiving the result of a later continuation. - Clear and release paths clean up all executions associated with the task ID. - Add regression test test_reused_task_id_does_not_associate_stale_snapshot_with_continuation. --- .../_harness/_background_agents.py | 149 ++++++++++++++++-- .../core/test_harness_background_agents.py | 82 +++++++++- 2 files changed, 213 insertions(+), 18 deletions(-) diff --git a/python/packages/core/agent_framework/_harness/_background_agents.py b/python/packages/core/agent_framework/_harness/_background_agents.py index f28c059511a..ff0b3aed370 100644 --- a/python/packages/core/agent_framework/_harness/_background_agents.py +++ b/python/packages/core/agent_framework/_harness/_background_agents.py @@ -65,7 +65,8 @@ class BackgroundTaskInfo(SerializationMixin): status: BackgroundTaskStatus result_text: str | None error_text: str | None - __slots__ = ("agent_name", "description", "error_text", "id", "result_text", "status") + generation: int + __slots__ = ("agent_name", "description", "error_text", "generation", "id", "result_text", "status") def __init__( self, @@ -75,6 +76,7 @@ def __init__( status: BackgroundTaskStatus = BackgroundTaskStatus.RUNNING, result_text: str | None = None, error_text: str | None = None, + generation: int = 1, ) -> None: """Initialize a background task info entry.""" self.id = id @@ -83,6 +85,7 @@ def __init__( self.status = status self.result_text = result_text self.error_text = error_text + self.generation = generation def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: """Serialize for session state persistence.""" @@ -92,6 +95,7 @@ def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) "agent_name": self.agent_name, "description": self.description, "status": self.status.value, + "generation": self.generation, } if not exclude_none or self.result_text is not None: data["result_text"] = self.result_text @@ -109,24 +113,125 @@ def from_dict(cls, data: MutableMapping[str, Any], **kwargs: Any) -> BackgroundT status=BackgroundTaskStatus(data["status"]), result_text=data.get("result_text"), error_text=data.get("error_text"), + generation=data.get("generation", 1), ) +_TaskKey = tuple[int, int] | int + + +class _InFlightTasks(dict[_TaskKey, asyncio.Task[AgentResponse[Any]]]): + """Non-serializable container for tracked in-flight tasks. + + Supports dual keying: + - Composite key ``(task_id, generation)``: isolates discrete continuation + executions under the same reusable task ID so older session snapshots + reconcile with their own generation rather than a later continuation. + - Integer key ``task_id``: direct fallback lookup and backwards-compatible + access for callers and existing unit tests. + """ + + def __getitem__(self, key: _TaskKey) -> asyncio.Task[AgentResponse[Any]]: + if super().__contains__(key): + return super().__getitem__(key) + if isinstance(key, int): + matches: list[tuple[int, asyncio.Task[AgentResponse[Any]]]] = [] + for k, v in self.items(): + if isinstance(k, tuple) and k[0] == key: + matches.append((k[1], v)) + if matches: + matches.sort(key=lambda item: item[0], reverse=True) + return matches[0][1] + raise KeyError(key) + + def get(self, key: _TaskKey, default: Any = None, /) -> Any: + try: + return self[key] + except KeyError: + return default + + def get_exact(self, key: _TaskKey) -> asyncio.Task[AgentResponse[Any]] | None: + """Retrieve task for an exact key without composite fallback.""" + if super().__contains__(key): + return super().__getitem__(key) + return None + + def __contains__(self, key: object) -> bool: + if super().__contains__(key): + return True + if isinstance(key, int): + for k in self: + if isinstance(k, tuple) and k[0] == key: + return True + return False + + def pop(self, key: _TaskKey, *args: Any) -> Any: + if super().__contains__(key): + val = super().pop(key) + if isinstance(key, int): + composites: list[tuple[int, int]] = [ + k for k in list(self.keys()) if isinstance(k, tuple) and k[0] == key + ] + for ck in composites: + super().pop(ck, None) + return val + if isinstance(key, int): + composites_to_pop: list[tuple[int, int]] = [ + k for k in list(self.keys()) if isinstance(k, tuple) and k[0] == key + ] + if composites_to_pop: + composites_to_pop.sort(key=lambda k: k[1]) + popped_val: asyncio.Task[AgentResponse[Any]] | None = None + for ck in composites_to_pop: + popped_val = super().pop(ck) + return popped_val + if args: + return args[0] + raise KeyError(key) + + @dataclass class _RuntimeState: """Non-serializable per-session runtime state for background tasks.""" - in_flight_tasks: dict[int, asyncio.Task[AgentResponse[Any]]] = field(default_factory=lambda: {}) + in_flight_tasks: _InFlightTasks = field(default_factory=_InFlightTasks) background_sessions: dict[int, AgentSession] = field(default_factory=lambda: {}) closed: bool = False - def track_task(self, task_id: int, task: asyncio.Task[AgentResponse[Any]]) -> None: + def track_task( + self, + task_id: int, + task: asyncio.Task[AgentResponse[Any]], + *, + generation: int = 1, + ) -> None: """Track a background task if this runtime is still open.""" if self.closed: task.cancel() raise RuntimeError("Session runtime is closed; cannot start background task.") - self.in_flight_tasks[task_id] = task + self.in_flight_tasks[task_id, generation] = task + + def get_task( + self, + task_id: int, + generation: int | None = None, + ) -> asyncio.Task[AgentResponse[Any]] | None: + """Retrieve a task for a given task ID and optional generation.""" + if generation is not None: + task = self.in_flight_tasks.get_exact((task_id, generation)) + if task is not None: + return task + for k in self.in_flight_tasks: + if isinstance(k, tuple) and k[0] == task_id: + return None + return self.in_flight_tasks.get_exact(task_id) + return self.in_flight_tasks.get(task_id) + + def clear_task(self, task_id: int) -> None: + """Remove all executions and background sessions associated with a task ID.""" + self.in_flight_tasks.pop(task_id, None) + self.background_sessions.pop(task_id, None) # --------------------------------------------------------------------------- @@ -227,6 +332,11 @@ def _finalize_task( call to :func:`_refresh_task_state` a missing entry would be misidentified as ``LOST``, discarding a result that is already available. + Furthermore, each execution records a monotonic ``generation`` in both + ``task_info`` and ``runtime.in_flight_tasks[(task_id, generation)]`` so + that restoring an earlier snapshot of a continued task will reconcile + with that snapshot's execution rather than a later continuation. + Tasks are removed from ``in_flight_tasks`` by two callers: - :func:`background_agents_clear_completed_task` — explicit caller-driven cleanup. - :meth:`BackgroundAgentsProvider.release_session` — bulk cleanup on session teardown. @@ -242,7 +352,7 @@ def _finalize_task( else: task_info.status = BackgroundTaskStatus.COMPLETED task_info.result_text = completed_task.result().text - # NOTE: do NOT pop task_info.id from runtime.in_flight_tasks here. + # NOTE: do NOT pop task from runtime.in_flight_tasks here. # See docstring above for the reasoning. @@ -257,7 +367,7 @@ def _refresh_task_state( if task_info.status != BackgroundTaskStatus.RUNNING: continue - in_flight = runtime.in_flight_tasks.get(task_info.id) + in_flight = runtime.get_task(task_info.id, task_info.generation) if in_flight is None: task_info.status = BackgroundTaskStatus.LOST changed = True @@ -503,9 +613,16 @@ def background_agents_start_task(agent_name: str, input: str, description: str) sub_session = bg_agent.create_session() + task_info = BackgroundTaskInfo( + id=task_id, + agent_name=agent_name, + description=description, + generation=1, + ) + async_task = asyncio.create_task(_run_agent(bg_agent.run(input, session=sub_session))) try: - runtime.track_task(task_id, async_task) + runtime.track_task(task_id, async_task, generation=task_info.generation) except RuntimeError as exc: return f"Error: {exc}" @@ -513,11 +630,6 @@ def background_agents_start_task(agent_name: str, input: str, description: str) provider_state["next_task_id"] = task_id + 1 - task_info = BackgroundTaskInfo( - id=task_id, - agent_name=agent_name, - description=description, - ) tasks = _get_tasks(provider_state) tasks.append(task_info) _save_tasks(provider_state, tasks) @@ -542,8 +654,11 @@ async def background_agents_wait_for_first_completion(task_ids: list[int]) -> st # Collect in-flight tasks matching the requested IDs. waitable: list[tuple[int, asyncio.Task[AgentResponse[Any]]]] = [] + tasks = _get_tasks(provider_state) for tid in task_ids: - in_flight = runtime.in_flight_tasks.get(tid) + t_info = next((t for t in tasks if t.id == tid and t.status == BackgroundTaskStatus.RUNNING), None) + gen = t_info.generation if t_info is not None else None + in_flight = runtime.get_task(tid, gen) if in_flight is not None: waitable.append((tid, in_flight)) @@ -582,7 +697,7 @@ async def background_agents_wait_for_first_completion(task_ids: list[int]) -> st tasks = _get_tasks(provider_state) task_info = next((t for t in tasks if t.id == completed_id), None) if task_info is not None and completed_id is not None: - completed_task = runtime.in_flight_tasks.get(completed_id) + completed_task = runtime.get_task(completed_id, task_info.generation) if completed_task is not None: _finalize_task(task_info, completed_task, runtime) _save_tasks(provider_state, tasks) @@ -653,9 +768,10 @@ def background_agents_continue_task(task_id: int, text: str) -> str: bg_agent = self._agents[key] + task_info.generation += 1 async_task = asyncio.create_task(_run_agent(bg_agent.run(text, session=sub_session))) try: - runtime.track_task(task_id, async_task) + runtime.track_task(task_id, async_task, generation=task_info.generation) except RuntimeError as exc: return f"Error: {exc}" task_info.status = BackgroundTaskStatus.RUNNING @@ -688,8 +804,7 @@ def background_agents_clear_completed_task(task_id: int) -> str: _save_tasks(provider_state, tasks) # Clean up runtime references. - runtime.in_flight_tasks.pop(task_id, None) - runtime.background_sessions.pop(task_id, None) + runtime.clear_task(task_id) _save_provider_state(session, provider_state, source_id=source_id) return f"Task {task_id} cleared." diff --git a/python/packages/core/tests/core/test_harness_background_agents.py b/python/packages/core/tests/core/test_harness_background_agents.py index 1e8f38610de..bd655e678de 100644 --- a/python/packages/core/tests/core/test_harness_background_agents.py +++ b/python/packages/core/tests/core/test_harness_background_agents.py @@ -1094,4 +1094,84 @@ async def test_release_session_clears_in_flight_tasks() -> None: await provider.release_session(session) - assert session.session_id not in provider._runtime, "Runtime should be cleaned up after release" + assert session.session_id not in provider._runtime + + +async def test_reused_task_id_does_not_associate_stale_snapshot_with_continuation() -> None: + """A stale snapshot must reconcile with its own execution, not a later continuation. + + Regression test for reused task IDs: + If task 1 is snapshotted while RUNNING, finishes, is continued (starting a + new execution under task ID 1), and the second execution finishes, restoring + the original snapshot must reconcile with the first execution's result, NOT + the continuation's result. + """ + + class _ContinuationAgent: + def __init__(self, name: str = "Worker") -> None: + self.name = name + self.description = "Worker that tracks runs" + self.call_count = 0 + + def create_session(self, *, session_id: str | None = None) -> AgentSession: + return AgentSession(session_id=session_id) + + def get_session(self, service_session_id: Any, *, session_id: str | None = None) -> AgentSession: + return AgentSession(service_session_id=service_session_id, session_id=session_id) + + async def run(self, messages: Any = None, **kwargs: Any) -> AgentResponse[Any]: + del messages, kwargs + self.call_count += 1 + return AgentResponse( + messages=[Message(role="assistant", contents=[f"result generation {self.call_count}"])] + ) + + agent = _ContinuationAgent("Worker") + provider = _make_provider(agent) # type: ignore[arg-type] + session = _make_session() + tools = await _get_tools(provider, session) + + # 1. Start task 1 on Worker (generation 1). + await _invoke_tool( + tools["background_agents_start_task"], + agent_name="Worker", + input="first run", + description="task 1", + ) + # Immediately capture snapshot while task 1 is RUNNING (generation 1). + snapshot_gen1 = session.to_dict() + + runtime = provider._get_runtime(session) + task_1 = runtime.in_flight_tasks[(1, 1)] + await task_1 + + # 2. Finalize first run via refresh. + tools2 = await _get_tools(provider, session) + res1 = await _invoke_tool(tools2["background_agents_get_task_results"], task_id=1) + assert res1 == "result generation 1" + + # 3. Continue task 1 with new input (generation 2). + continue_msg = await _invoke_tool(tools2["background_agents_continue_task"], task_id=1, text="second run") + assert "continued" in continue_msg.lower() + + task_2 = runtime.in_flight_tasks[(1, 2)] + await task_2 + + # Finalize continuation via refresh. + tools3 = await _get_tools(provider, session) + res2 = await _invoke_tool(tools3["background_agents_get_task_results"], task_id=1) + assert res2 == "result generation 2" + + # 4. Now restore the generation 1 snapshot into the same provider. + session_restored = AgentSession.from_dict(snapshot_gen1) + tools_restored = await _get_tools(provider, session_restored) + + # Refreshing on the restored snapshot must reconcile with generation 1 task, + # NOT with generation 2 task. + res_restored = await _invoke_tool(tools_restored["background_agents_get_task_results"], task_id=1) + assert res_restored == "result generation 1", ( + f"Expected generation 1 result but got continuation result: {res_restored!r}" + ) + + await provider.release_session(session) + await provider.release_session(session_restored) From 6b90a3d4a206f5c028088b31f618115fbd324836 Mon Sep 17 00:00:00 2001 From: karthik-0306 Date: Thu, 15 Oct 2026 13:49:41 +0530 Subject: [PATCH 3/5] Python: Simplify background task retention without generation tracking (#8760) --- .../_harness/_background_agents.py | 159 +++--------------- .../core/test_harness_background_agents.py | 137 +++++++++------ 2 files changed, 110 insertions(+), 186 deletions(-) diff --git a/python/packages/core/agent_framework/_harness/_background_agents.py b/python/packages/core/agent_framework/_harness/_background_agents.py index ff0b3aed370..859472bdd76 100644 --- a/python/packages/core/agent_framework/_harness/_background_agents.py +++ b/python/packages/core/agent_framework/_harness/_background_agents.py @@ -65,8 +65,7 @@ class BackgroundTaskInfo(SerializationMixin): status: BackgroundTaskStatus result_text: str | None error_text: str | None - generation: int - __slots__ = ("agent_name", "description", "error_text", "generation", "id", "result_text", "status") + __slots__ = ("agent_name", "description", "error_text", "id", "result_text", "status") def __init__( self, @@ -76,7 +75,6 @@ def __init__( status: BackgroundTaskStatus = BackgroundTaskStatus.RUNNING, result_text: str | None = None, error_text: str | None = None, - generation: int = 1, ) -> None: """Initialize a background task info entry.""" self.id = id @@ -85,7 +83,6 @@ def __init__( self.status = status self.result_text = result_text self.error_text = error_text - self.generation = generation def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: """Serialize for session state persistence.""" @@ -95,7 +92,6 @@ def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) "agent_name": self.agent_name, "description": self.description, "status": self.status.value, - "generation": self.generation, } if not exclude_none or self.result_text is not None: data["result_text"] = self.result_text @@ -113,125 +109,24 @@ def from_dict(cls, data: MutableMapping[str, Any], **kwargs: Any) -> BackgroundT status=BackgroundTaskStatus(data["status"]), result_text=data.get("result_text"), error_text=data.get("error_text"), - generation=data.get("generation", 1), ) -_TaskKey = tuple[int, int] | int - - -class _InFlightTasks(dict[_TaskKey, asyncio.Task[AgentResponse[Any]]]): - """Non-serializable container for tracked in-flight tasks. - - Supports dual keying: - - Composite key ``(task_id, generation)``: isolates discrete continuation - executions under the same reusable task ID so older session snapshots - reconcile with their own generation rather than a later continuation. - - Integer key ``task_id``: direct fallback lookup and backwards-compatible - access for callers and existing unit tests. - """ - - def __getitem__(self, key: _TaskKey) -> asyncio.Task[AgentResponse[Any]]: - if super().__contains__(key): - return super().__getitem__(key) - if isinstance(key, int): - matches: list[tuple[int, asyncio.Task[AgentResponse[Any]]]] = [] - for k, v in self.items(): - if isinstance(k, tuple) and k[0] == key: - matches.append((k[1], v)) - if matches: - matches.sort(key=lambda item: item[0], reverse=True) - return matches[0][1] - raise KeyError(key) - - def get(self, key: _TaskKey, default: Any = None, /) -> Any: - try: - return self[key] - except KeyError: - return default - - def get_exact(self, key: _TaskKey) -> asyncio.Task[AgentResponse[Any]] | None: - """Retrieve task for an exact key without composite fallback.""" - if super().__contains__(key): - return super().__getitem__(key) - return None - - def __contains__(self, key: object) -> bool: - if super().__contains__(key): - return True - if isinstance(key, int): - for k in self: - if isinstance(k, tuple) and k[0] == key: - return True - return False - - def pop(self, key: _TaskKey, *args: Any) -> Any: - if super().__contains__(key): - val = super().pop(key) - if isinstance(key, int): - composites: list[tuple[int, int]] = [ - k for k in list(self.keys()) if isinstance(k, tuple) and k[0] == key - ] - for ck in composites: - super().pop(ck, None) - return val - if isinstance(key, int): - composites_to_pop: list[tuple[int, int]] = [ - k for k in list(self.keys()) if isinstance(k, tuple) and k[0] == key - ] - if composites_to_pop: - composites_to_pop.sort(key=lambda k: k[1]) - popped_val: asyncio.Task[AgentResponse[Any]] | None = None - for ck in composites_to_pop: - popped_val = super().pop(ck) - return popped_val - if args: - return args[0] - raise KeyError(key) - - @dataclass class _RuntimeState: """Non-serializable per-session runtime state for background tasks.""" - in_flight_tasks: _InFlightTasks = field(default_factory=_InFlightTasks) + in_flight_tasks: dict[int, asyncio.Task[AgentResponse[Any]]] = field(default_factory=lambda: {}) background_sessions: dict[int, AgentSession] = field(default_factory=lambda: {}) closed: bool = False - def track_task( - self, - task_id: int, - task: asyncio.Task[AgentResponse[Any]], - *, - generation: int = 1, - ) -> None: + def track_task(self, task_id: int, task: asyncio.Task[AgentResponse[Any]]) -> None: """Track a background task if this runtime is still open.""" if self.closed: task.cancel() raise RuntimeError("Session runtime is closed; cannot start background task.") - self.in_flight_tasks[task_id, generation] = task - - def get_task( - self, - task_id: int, - generation: int | None = None, - ) -> asyncio.Task[AgentResponse[Any]] | None: - """Retrieve a task for a given task ID and optional generation.""" - if generation is not None: - task = self.in_flight_tasks.get_exact((task_id, generation)) - if task is not None: - return task - for k in self.in_flight_tasks: - if isinstance(k, tuple) and k[0] == task_id: - return None - return self.in_flight_tasks.get_exact(task_id) - return self.in_flight_tasks.get(task_id) - - def clear_task(self, task_id: int) -> None: - """Remove all executions and background sessions associated with a task ID.""" - self.in_flight_tasks.pop(task_id, None) - self.background_sessions.pop(task_id, None) + self.in_flight_tasks[task_id] = task # --------------------------------------------------------------------------- @@ -332,12 +227,8 @@ def _finalize_task( call to :func:`_refresh_task_state` a missing entry would be misidentified as ``LOST``, discarding a result that is already available. - Furthermore, each execution records a monotonic ``generation`` in both - ``task_info`` and ``runtime.in_flight_tasks[(task_id, generation)]`` so - that restoring an earlier snapshot of a continued task will reconcile - with that snapshot's execution rather than a later continuation. - Tasks are removed from ``in_flight_tasks`` by two callers: + - :func:`background_agents_clear_completed_task` — explicit caller-driven cleanup. - :meth:`BackgroundAgentsProvider.release_session` — bulk cleanup on session teardown. """ @@ -352,8 +243,6 @@ def _finalize_task( else: task_info.status = BackgroundTaskStatus.COMPLETED task_info.result_text = completed_task.result().text - # NOTE: do NOT pop task from runtime.in_flight_tasks here. - # See docstring above for the reasoning. def _refresh_task_state( @@ -367,7 +256,7 @@ def _refresh_task_state( if task_info.status != BackgroundTaskStatus.RUNNING: continue - in_flight = runtime.get_task(task_info.id, task_info.generation) + in_flight = runtime.in_flight_tasks.get(task_info.id) if in_flight is None: task_info.status = BackgroundTaskStatus.LOST changed = True @@ -613,16 +502,9 @@ def background_agents_start_task(agent_name: str, input: str, description: str) sub_session = bg_agent.create_session() - task_info = BackgroundTaskInfo( - id=task_id, - agent_name=agent_name, - description=description, - generation=1, - ) - async_task = asyncio.create_task(_run_agent(bg_agent.run(input, session=sub_session))) try: - runtime.track_task(task_id, async_task, generation=task_info.generation) + runtime.track_task(task_id, async_task) except RuntimeError as exc: return f"Error: {exc}" @@ -630,6 +512,11 @@ def background_agents_start_task(agent_name: str, input: str, description: str) provider_state["next_task_id"] = task_id + 1 + task_info = BackgroundTaskInfo( + id=task_id, + agent_name=agent_name, + description=description, + ) tasks = _get_tasks(provider_state) tasks.append(task_info) _save_tasks(provider_state, tasks) @@ -655,12 +542,12 @@ async def background_agents_wait_for_first_completion(task_ids: list[int]) -> st # Collect in-flight tasks matching the requested IDs. waitable: list[tuple[int, asyncio.Task[AgentResponse[Any]]]] = [] tasks = _get_tasks(provider_state) + running_ids = {t.id for t in tasks if t.status == BackgroundTaskStatus.RUNNING} for tid in task_ids: - t_info = next((t for t in tasks if t.id == tid and t.status == BackgroundTaskStatus.RUNNING), None) - gen = t_info.generation if t_info is not None else None - in_flight = runtime.get_task(tid, gen) - if in_flight is not None: - waitable.append((tid, in_flight)) + if tid in running_ids: + in_flight = runtime.in_flight_tasks.get(tid) + if in_flight is not None: + waitable.append((tid, in_flight)) if not waitable: # Refresh state to catch any that completed. @@ -697,7 +584,7 @@ async def background_agents_wait_for_first_completion(task_ids: list[int]) -> st tasks = _get_tasks(provider_state) task_info = next((t for t in tasks if t.id == completed_id), None) if task_info is not None and completed_id is not None: - completed_task = runtime.get_task(completed_id, task_info.generation) + completed_task = runtime.in_flight_tasks.get(completed_id) if completed_task is not None: _finalize_task(task_info, completed_task, runtime) _save_tasks(provider_state, tasks) @@ -768,10 +655,9 @@ def background_agents_continue_task(task_id: int, text: str) -> str: bg_agent = self._agents[key] - task_info.generation += 1 async_task = asyncio.create_task(_run_agent(bg_agent.run(text, session=sub_session))) try: - runtime.track_task(task_id, async_task, generation=task_info.generation) + runtime.track_task(task_id, async_task) except RuntimeError as exc: return f"Error: {exc}" task_info.status = BackgroundTaskStatus.RUNNING @@ -803,8 +689,11 @@ def background_agents_clear_completed_task(task_id: int) -> str: tasks = [t for t in tasks if t.id != task_id] _save_tasks(provider_state, tasks) - # Clean up runtime references. - runtime.clear_task(task_id) + # Clean up runtime references only if the local task is done or absent. + in_flight = runtime.in_flight_tasks.get(task_id) + if in_flight is None or in_flight.done(): + runtime.in_flight_tasks.pop(task_id, None) + runtime.background_sessions.pop(task_id, None) _save_provider_state(session, provider_state, source_id=source_id) return f"Task {task_id} cleared." diff --git a/python/packages/core/tests/core/test_harness_background_agents.py b/python/packages/core/tests/core/test_harness_background_agents.py index bd655e678de..3b87a5644b5 100644 --- a/python/packages/core/tests/core/test_harness_background_agents.py +++ b/python/packages/core/tests/core/test_harness_background_agents.py @@ -1095,83 +1095,118 @@ async def test_release_session_clears_in_flight_tasks() -> None: await provider.release_session(session) assert session.session_id not in provider._runtime + assert not runtime.in_flight_tasks -async def test_reused_task_id_does_not_associate_stale_snapshot_with_continuation() -> None: - """A stale snapshot must reconcile with its own execution, not a later continuation. +async def test_wait_blocks_for_running_task_and_ignores_finished() -> None: + """wait_for_first_completion ignores finished retained tasks and waits for running ones.""" + provider = _make_provider( + _FakeAgent("Fast", response_text="fast done"), + _FakeAgent("Slow", response_text="slow done", delay=0.02), + ) + session = _make_session() + tools = await _get_tools(provider, session) - Regression test for reused task IDs: - If task 1 is snapshotted while RUNNING, finishes, is continued (starting a - new execution under task ID 1), and the second execution finishes, restoring - the original snapshot must reconcile with the first execution's result, NOT - the continuation's result. - """ + # 1. Start and finish task 1. + await _invoke_tool( + tools["background_agents_start_task"], + agent_name="Fast", + input="work", + description="task 1", + ) + runtime = provider._get_runtime(session) + await runtime.in_flight_tasks[1] - class _ContinuationAgent: - def __init__(self, name: str = "Worker") -> None: - self.name = name - self.description = "Worker that tracks runs" - self.call_count = 0 + # Finalize task 1 so its status is COMPLETED and its task is retained. + tools2 = await _get_tools(provider, session) + await _invoke_tool(tools2["background_agents_get_task_results"], task_id=1) + + # 2. Start task 2 which takes longer. + await _invoke_tool( + tools2["background_agents_start_task"], + agent_name="Slow", + input="slow work", + description="task 2", + ) - def create_session(self, *, session_id: str | None = None) -> AgentSession: - return AgentSession(session_id=session_id) + # 3. wait([1, 2]) must block for running task 2 and NOT return finished task 1. + wait_result = await _invoke_tool( + tools2["background_agents_wait_for_first_completion"], + task_ids=[1, 2], + ) + assert "Task 2 finished with status: completed." in wait_result + assert "Task 1" not in wait_result + + # 4. wait([1]) on finished task returns "not running; current status" text. + finished_wait_result = await _invoke_tool( + tools2["background_agents_wait_for_first_completion"], + task_ids=[1], + ) + assert "Task 1 is not running; current status: completed." in finished_wait_result + + await provider.release_session(session) - def get_session(self, service_session_id: Any, *, session_id: str | None = None) -> AgentSession: - return AgentSession(service_session_id=service_session_id, session_id=session_id) - async def run(self, messages: Any = None, **kwargs: Any) -> AgentResponse[Any]: - del messages, kwargs - self.call_count += 1 - return AgentResponse( - messages=[Message(role="assistant", contents=[f"result generation {self.call_count}"])] - ) +async def test_clear_completed_task_does_not_clear_in_flight_continuation() -> None: + """Restoring an older completed snapshot and clearing task 1 must not clear a running continuation.""" + + class _ContinuationAgent(_FakeAgent): + def __init__(self, name: str = "Worker") -> None: + super().__init__(name, response_text="done") + self.release = asyncio.Event() + self._runs = 0 + + async def run( + self, messages: Any = None, *, stream: bool = False, session: Any = None, **kwargs: Any + ) -> AgentResponse[Any]: + self._runs += 1 + if self._runs > 1: + await self.release.wait() + return await super().run(messages, stream=stream, session=session, **kwargs) agent = _ContinuationAgent("Worker") - provider = _make_provider(agent) # type: ignore[arg-type] + provider = _make_provider(agent) session = _make_session() tools = await _get_tools(provider, session) - # 1. Start task 1 on Worker (generation 1). + # 1. Start task 1 and let it complete. await _invoke_tool( tools["background_agents_start_task"], agent_name="Worker", input="first run", description="task 1", ) - # Immediately capture snapshot while task 1 is RUNNING (generation 1). - snapshot_gen1 = session.to_dict() - runtime = provider._get_runtime(session) - task_1 = runtime.in_flight_tasks[(1, 1)] - await task_1 + await runtime.in_flight_tasks[1] - # 2. Finalize first run via refresh. + # Finalize task 1 so its status is COMPLETED. tools2 = await _get_tools(provider, session) - res1 = await _invoke_tool(tools2["background_agents_get_task_results"], task_id=1) - assert res1 == "result generation 1" + await _invoke_tool(tools2["background_agents_get_task_results"], task_id=1) - # 3. Continue task 1 with new input (generation 2). + # 2. Snapshot the completed state. + snapshot = session.to_dict() + + # 3. continue_task with slow agent. continue_msg = await _invoke_tool(tools2["background_agents_continue_task"], task_id=1, text="second run") assert "continued" in continue_msg.lower() - task_2 = runtime.in_flight_tasks[(1, 2)] - await task_2 + continuation_task = runtime.in_flight_tasks[1] + assert not continuation_task.done() - # Finalize continuation via refresh. - tools3 = await _get_tools(provider, session) - res2 = await _invoke_tool(tools3["background_agents_get_task_results"], task_id=1) - assert res2 == "result generation 2" + # 4. Restore snapshot into a new AgentSession. + session2 = AgentSession.from_dict(snapshot) + tools3 = await _get_tools(provider, session2) - # 4. Now restore the generation 1 snapshot into the same provider. - session_restored = AgentSession.from_dict(snapshot_gen1) - tools_restored = await _get_tools(provider, session_restored) + # 5. Call clear_completed_task(1) on restored session. + clear_msg = await _invoke_tool(tools3["background_agents_clear_completed_task"], task_id=1) + assert clear_msg == "Task 1 cleared." - # Refreshing on the restored snapshot must reconcile with generation 1 task, - # NOT with generation 2 task. - res_restored = await _invoke_tool(tools_restored["background_agents_get_task_results"], task_id=1) - assert res_restored == "result generation 1", ( - f"Expected generation 1 result but got continuation result: {res_restored!r}" - ) + # 6. Continuation is still in in_flight_tasks. + assert 1 in runtime.in_flight_tasks + assert runtime.in_flight_tasks[1] is continuation_task + assert not continuation_task.done() - await provider.release_session(session) - await provider.release_session(session_restored) + # 7. release_session(cancel_running=True) cancels it. + await provider.release_session(session, cancel_running=True) + assert continuation_task.cancelled() + assert not runtime.in_flight_tasks From 2ac9ae05bfa8c0b19a64537ca44ebf3e3b48e5e2 Mon Sep 17 00:00:00 2001 From: karthik-0306 Date: Fri, 16 Oct 2026 15:45:39 +0530 Subject: [PATCH 4/5] Python: Track execution generation on background tasks (#8760) Persist a per-task generation so a restored older session snapshot is never reconciled against a later continuation that reused the same task ID. Continue/clear are generation-aware. --- .../_harness/_background_agents.py | 96 ++++++++-- .../core/test_harness_background_agents.py | 168 ++++++++++++++++++ 2 files changed, 245 insertions(+), 19 deletions(-) diff --git a/python/packages/core/agent_framework/_harness/_background_agents.py b/python/packages/core/agent_framework/_harness/_background_agents.py index 859472bdd76..3da182b5f06 100644 --- a/python/packages/core/agent_framework/_harness/_background_agents.py +++ b/python/packages/core/agent_framework/_harness/_background_agents.py @@ -55,7 +55,12 @@ class BackgroundTaskStatus(str, Enum): @experimental(feature_id=ExperimentalFeature.HARNESS) class BackgroundTaskInfo(SerializationMixin): - """Metadata for a single background task.""" + """Metadata for a single background task. + + ``generation`` identifies which execution of the task this record describes. It starts at 1 + and increases every time the task is continued, so a restored older snapshot of the session + is never reconciled against a later execution that reused the same task ID. + """ DEFAULT_EXCLUDE: ClassVar[set[str]] = set() @@ -65,7 +70,8 @@ class BackgroundTaskInfo(SerializationMixin): status: BackgroundTaskStatus result_text: str | None error_text: str | None - __slots__ = ("agent_name", "description", "error_text", "id", "result_text", "status") + generation: int + __slots__ = ("agent_name", "description", "error_text", "generation", "id", "result_text", "status") def __init__( self, @@ -75,6 +81,7 @@ def __init__( status: BackgroundTaskStatus = BackgroundTaskStatus.RUNNING, result_text: str | None = None, error_text: str | None = None, + generation: int = 1, ) -> None: """Initialize a background task info entry.""" self.id = id @@ -83,6 +90,7 @@ def __init__( self.status = status self.result_text = result_text self.error_text = error_text + self.generation = generation def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: """Serialize for session state persistence.""" @@ -92,6 +100,7 @@ def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) "agent_name": self.agent_name, "description": self.description, "status": self.status.value, + "generation": self.generation, } if not exclude_none or self.result_text is not None: data["result_text"] = self.result_text @@ -109,24 +118,51 @@ def from_dict(cls, data: MutableMapping[str, Any], **kwargs: Any) -> BackgroundT status=BackgroundTaskStatus(data["status"]), result_text=data.get("result_text"), error_text=data.get("error_text"), + generation=data.get("generation", 1), ) @dataclass class _RuntimeState: - """Non-serializable per-session runtime state for background tasks.""" + """Non-serializable per-session runtime state for background tasks. + + ``in_flight_tasks`` holds the latest execution started for each task ID and + ``task_generations`` records which generation that execution belongs to. Lookups made on + behalf of a persisted :class:`BackgroundTaskInfo` go through :meth:`get_execution` so that a + record only ever sees the execution it was created for. + """ in_flight_tasks: dict[int, asyncio.Task[AgentResponse[Any]]] = field(default_factory=lambda: {}) + task_generations: dict[int, int] = field(default_factory=lambda: {}) background_sessions: dict[int, AgentSession] = field(default_factory=lambda: {}) closed: bool = False - def track_task(self, task_id: int, task: asyncio.Task[AgentResponse[Any]]) -> None: + def track_task(self, task_id: int, task: asyncio.Task[AgentResponse[Any]], *, generation: int = 1) -> None: """Track a background task if this runtime is still open.""" if self.closed: task.cancel() raise RuntimeError("Session runtime is closed; cannot start background task.") self.in_flight_tasks[task_id] = task + self.task_generations[task_id] = generation + + def current_generation(self, task_id: int) -> int | None: + """Return the generation of the tracked execution for a task ID, if any.""" + if task_id not in self.in_flight_tasks: + return None + return self.task_generations.get(task_id, 1) + + def get_execution(self, task_id: int, generation: int) -> asyncio.Task[AgentResponse[Any]] | None: + """Return the tracked execution for a task ID only if it belongs to the given generation.""" + if self.current_generation(task_id) != generation: + return None + return self.in_flight_tasks.get(task_id) + + def forget_task(self, task_id: int) -> None: + """Drop all runtime references for a task ID.""" + self.in_flight_tasks.pop(task_id, None) + self.task_generations.pop(task_id, None) + self.background_sessions.pop(task_id, None) # --------------------------------------------------------------------------- @@ -215,7 +251,6 @@ def _save_tasks(state: dict[str, Any], tasks: list[BackgroundTaskInfo]) -> None: def _finalize_task( task_info: BackgroundTaskInfo, completed_task: asyncio.Task[AgentResponse[Any]], - runtime: _RuntimeState, ) -> None: """Extract results from a completed asyncio task and update task info. @@ -227,6 +262,11 @@ def _finalize_task( call to :func:`_refresh_task_state` a missing entry would be misidentified as ``LOST``, discarding a result that is already available. + Because the retained entry is keyed by task ID, callers must only pass the + execution matching ``task_info.generation`` (see + :meth:`_RuntimeState.get_execution`); a later continuation of the same task + ID must never be applied to an older record. + Tasks are removed from ``in_flight_tasks`` by two callers: - :func:`background_agents_clear_completed_task` — explicit caller-driven cleanup. @@ -256,14 +296,16 @@ def _refresh_task_state( if task_info.status != BackgroundTaskStatus.RUNNING: continue - in_flight = runtime.in_flight_tasks.get(task_info.id) + # A missing execution, or one from a different generation (the task was continued from + # another snapshot of this session), means this record's execution is no longer tracked. + in_flight = runtime.get_execution(task_info.id, task_info.generation) if in_flight is None: task_info.status = BackgroundTaskStatus.LOST changed = True continue if in_flight.done(): - _finalize_task(task_info, in_flight, runtime) + _finalize_task(task_info, in_flight) changed = True if changed: @@ -410,6 +452,7 @@ async def release_session( await asyncio.gather(*completed, return_exceptions=True) finally: runtime.in_flight_tasks.clear() + runtime.task_generations.clear() runtime.background_sessions.clear() if self._runtime.get(session_id) is runtime: @@ -539,13 +582,14 @@ async def background_agents_wait_for_first_completion(task_ids: list[int]) -> st if not task_ids: return "Error: No task IDs provided." - # Collect in-flight tasks matching the requested IDs. + # Collect the executions that the requested running records refer to. waitable: list[tuple[int, asyncio.Task[AgentResponse[Any]]]] = [] tasks = _get_tasks(provider_state) - running_ids = {t.id for t in tasks if t.status == BackgroundTaskStatus.RUNNING} + running = {t.id: t.generation for t in tasks if t.status == BackgroundTaskStatus.RUNNING} for tid in task_ids: - if tid in running_ids: - in_flight = runtime.in_flight_tasks.get(tid) + generation = running.get(tid) + if generation is not None: + in_flight = runtime.get_execution(tid, generation) if in_flight is not None: waitable.append((tid, in_flight)) @@ -584,9 +628,9 @@ async def background_agents_wait_for_first_completion(task_ids: list[int]) -> st tasks = _get_tasks(provider_state) task_info = next((t for t in tasks if t.id == completed_id), None) if task_info is not None and completed_id is not None: - completed_task = runtime.in_flight_tasks.get(completed_id) + completed_task = runtime.get_execution(completed_id, task_info.generation) if completed_task is not None: - _finalize_task(task_info, completed_task, runtime) + _finalize_task(task_info, completed_task) _save_tasks(provider_state, tasks) _save_provider_state(session, provider_state, source_id=source_id) @@ -653,13 +697,26 @@ def background_agents_continue_task(task_id: int, text: str) -> str: if sub_session is None: return f"Error: Session for task {task_id} is no longer available." + # This record is not running, so a pending tracked execution must be a newer one started + # from another snapshot of this session. Running both against one child session is unsafe. + current = runtime.in_flight_tasks.get(task_id) + if current is not None and not current.done(): + return ( + f"Error: Task {task_id} has a newer execution that is still running. " + "Wait for it to complete before continuing." + ) + bg_agent = self._agents[key] + # Always move past any generation this runtime has seen so the new execution can't be + # mistaken for one referenced by another snapshot. + generation = max(task_info.generation, runtime.current_generation(task_id) or 0) + 1 async_task = asyncio.create_task(_run_agent(bg_agent.run(text, session=sub_session))) try: - runtime.track_task(task_id, async_task) + runtime.track_task(task_id, async_task, generation=generation) except RuntimeError as exc: return f"Error: {exc}" + task_info.generation = generation task_info.status = BackgroundTaskStatus.RUNNING task_info.result_text = None task_info.error_text = None @@ -689,11 +746,12 @@ def background_agents_clear_completed_task(task_id: int) -> str: tasks = [t for t in tasks if t.id != task_id] _save_tasks(provider_state, tasks) - # Clean up runtime references only if the local task is done or absent. - in_flight = runtime.in_flight_tasks.get(task_id) - if in_flight is None or in_flight.done(): - runtime.in_flight_tasks.pop(task_id, None) - runtime.background_sessions.pop(task_id, None) + # Release runtime references unless a newer execution of this task ID exists (started from + # another snapshot of this session); that execution must stay tracked so it can still be + # retrieved, cancelled, or drained by release_session. + current_generation = runtime.current_generation(task_id) + if current_generation is None or current_generation <= task_info.generation: + runtime.forget_task(task_id) _save_provider_state(session, provider_state, source_id=source_id) return f"Task {task_id} cleared." diff --git a/python/packages/core/tests/core/test_harness_background_agents.py b/python/packages/core/tests/core/test_harness_background_agents.py index 3b87a5644b5..a21ab0c473e 100644 --- a/python/packages/core/tests/core/test_harness_background_agents.py +++ b/python/packages/core/tests/core/test_harness_background_agents.py @@ -1210,3 +1210,171 @@ async def run( await provider.release_session(session, cancel_running=True) assert continuation_task.cancelled() assert not runtime.in_flight_tasks + + +# --------------------------------------------------------------------------- +# Regression tests: execution identity (generation) for reused task IDs +# --------------------------------------------------------------------------- + + +class _ScriptedAgent(_FakeAgent): + """Agent whose N-th run returns ``responses[N]`` and optionally waits on a gate or fails.""" + + def __init__(self, name: str, responses: list[str], *, fail_runs: set[int] | None = None) -> None: + super().__init__(name) + self.responses = responses + self.fail_runs = fail_runs or set() + self.gates: dict[int, asyncio.Event] = {} + self.runs = 0 + + async def run( + self, messages: Any = None, *, stream: bool = False, session: Any = None, **kwargs: Any + ) -> AgentResponse[Any]: + index = self.runs + self.runs += 1 + gate = self.gates.get(index) + if gate is not None: + await gate.wait() + if index in self.fail_runs: + raise RuntimeError(f"run {index} failed") + return AgentResponse(messages=[Message(role="assistant", contents=[self.responses[index]])]) + + +def test_task_info_generation_round_trip_and_legacy_default() -> None: + """generation is persisted, and state written before it existed loads as generation 1.""" + info = BackgroundTaskInfo(id=1, agent_name="A", description="d", generation=3) + assert BackgroundTaskInfo.from_dict(info.to_dict()).generation == 3 + + legacy = {"id": 1, "agent_name": "A", "description": "d", "status": "running"} + assert BackgroundTaskInfo.from_dict(legacy).generation == 1 + + +@pytest.mark.parametrize("continuation_fails", [False, True]) +async def test_restored_running_snapshot_does_not_receive_continuation_outcome(continuation_fails: bool) -> None: + """A stale RUNNING snapshot must not be finalized with a later continuation's result or error.""" + agent = _ScriptedAgent("Worker", ["first result", "second result"], fail_runs={1} if continuation_fails else None) + provider = _make_provider(agent) + session = _make_session() + tools = await _get_tools(provider, session) + + await _invoke_tool(tools["background_agents_start_task"], agent_name="Worker", input="go", description="t") + stale_snapshot = session.to_dict() # generation 1, RUNNING + runtime = provider._get_runtime(session) + await runtime.in_flight_tasks[1] + + # The live session retrieves generation 1 and continues the task (generation 2). + assert await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) == "first result" + await _invoke_tool(tools["background_agents_continue_task"], task_id=1, text="again") + continuation = runtime.in_flight_tasks[1] + with suppress(Exception): + await continuation + + # The host restores the snapshot taken while generation 1 was running. + restored = AgentSession.from_dict(stale_snapshot) + restored_tools = await _get_tools(provider, restored) + stale_result = await _invoke_tool(restored_tools["background_agents_get_task_results"], task_id=1) + assert stale_result == "Task state was lost (reference unavailable)." + + # The live session still sees its own continuation outcome, and it is still tracked. + live_result = await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) + assert live_result == ("Task failed: run 1 failed" if continuation_fails else "second result") + assert runtime.in_flight_tasks[1] is continuation + + await provider.release_session(session) + + +async def test_restored_running_snapshot_does_not_wait_on_continuation() -> None: + """Waiting from a stale snapshot must not block on (or finalize) a newer running continuation.""" + agent = _ScriptedAgent("Worker", ["first result", "second result"]) + agent.gates[1] = asyncio.Event() + provider = _make_provider(agent) + session = _make_session() + tools = await _get_tools(provider, session) + + await _invoke_tool(tools["background_agents_start_task"], agent_name="Worker", input="go", description="t") + stale_snapshot = session.to_dict() + runtime = provider._get_runtime(session) + await runtime.in_flight_tasks[1] + await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) + await _invoke_tool(tools["background_agents_continue_task"], task_id=1, text="again") + continuation = runtime.in_flight_tasks[1] + + restored = AgentSession.from_dict(stale_snapshot) + restored_tools = await _get_tools(provider, restored) + wait_result = await _invoke_tool(restored_tools["background_agents_wait_for_first_completion"], task_ids=[1]) + assert wait_result == "Task 1 is not running; current status: lost." + assert runtime.in_flight_tasks[1] is continuation + assert not continuation.done() + + agent.gates[1].set() + wait_result = await _invoke_tool(tools["background_agents_wait_for_first_completion"], task_ids=[1]) + assert wait_result == "Task 1 finished with status: completed." + assert await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) == "second result" + + await provider.release_session(session) + + +async def test_continue_from_stale_snapshot_rejected_while_newer_execution_running() -> None: + """Continuing an old completed snapshot must not overwrite a newer running continuation.""" + agent = _ScriptedAgent("Worker", ["first result", "second result", "third result"]) + agent.gates[1] = asyncio.Event() + provider = _make_provider(agent) + session = _make_session() + tools = await _get_tools(provider, session) + + await _invoke_tool(tools["background_agents_start_task"], agent_name="Worker", input="go", description="t") + runtime = provider._get_runtime(session) + await runtime.in_flight_tasks[1] + await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) + completed_snapshot = session.to_dict() # generation 1, COMPLETED + await _invoke_tool(tools["background_agents_continue_task"], task_id=1, text="again") + continuation = runtime.in_flight_tasks[1] + + restored = AgentSession.from_dict(completed_snapshot) + restored_tools = await _get_tools(provider, restored) + result = await _invoke_tool(restored_tools["background_agents_continue_task"], task_id=1, text="stale") + assert result.startswith("Error: Task 1 has a newer execution that is still running.") + assert runtime.in_flight_tasks[1] is continuation + assert runtime.task_generations[1] == 2 + + await provider.release_session(session, cancel_running=True) + assert continuation.cancelled() + + +async def test_continue_from_stale_snapshot_starts_newer_generation() -> None: + """Continuing from an older snapshot after a newer execution finished must use a fresh generation.""" + agent = _ScriptedAgent("Worker", ["first result", "second result", "third result"]) + provider = _make_provider(agent) + session = _make_session() + tools = await _get_tools(provider, session) + + await _invoke_tool(tools["background_agents_start_task"], agent_name="Worker", input="go", description="t") + runtime = provider._get_runtime(session) + await runtime.in_flight_tasks[1] + await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) + completed_snapshot = session.to_dict() # generation 1, COMPLETED + await _invoke_tool(tools["background_agents_continue_task"], task_id=1, text="again") + await runtime.in_flight_tasks[1] + assert await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) == "second result" + + restored = AgentSession.from_dict(completed_snapshot) + restored_tools = await _get_tools(provider, restored) + await _invoke_tool(restored_tools["background_agents_continue_task"], task_id=1, text="from snapshot") + third = runtime.in_flight_tasks[1] + assert runtime.task_generations[1] == 3 + assert restored.state["background_agents"]["tasks"][0]["generation"] == 3 + await third + assert await _invoke_tool(restored_tools["background_agents_get_task_results"], task_id=1) == "third result" + + # The live session (generation 2) clears its record without dropping generation 3. + assert await _invoke_tool(tools["background_agents_clear_completed_task"], task_id=1) == "Task 1 cleared." + assert runtime.in_flight_tasks[1] is third + assert 1 in runtime.background_sessions + + # Clearing the record that owns the latest generation releases the runtime references. + assert await _invoke_tool(restored_tools["background_agents_clear_completed_task"], task_id=1) == "Task 1 cleared." + assert 1 not in runtime.in_flight_tasks + assert 1 not in runtime.task_generations + assert 1 not in runtime.background_sessions + + await provider.release_session(session) From 1a3fadf9375564effbc2ceabb9a51aff07b73b8d Mon Sep 17 00:00:00 2001 From: karthik-0306 Date: Mon, 26 Oct 2026 22:19:07 +0530 Subject: [PATCH 5/5] Python: Use unique execution IDs and run background tools on the event loop (#8760) --- .../_harness/_background_agents.py | 133 ++++---- .../core/test_harness_background_agents.py | 303 ++++++++++++++++-- 2 files changed, 343 insertions(+), 93 deletions(-) diff --git a/python/packages/core/agent_framework/_harness/_background_agents.py b/python/packages/core/agent_framework/_harness/_background_agents.py index fb93292aec3..6692b0fb622 100644 --- a/python/packages/core/agent_framework/_harness/_background_agents.py +++ b/python/packages/core/agent_framework/_harness/_background_agents.py @@ -11,6 +11,7 @@ import asyncio import logging +import uuid from collections.abc import Awaitable, MutableMapping, Sequence from dataclasses import dataclass, field from enum import Enum @@ -57,9 +58,10 @@ class BackgroundTaskStatus(str, Enum): class BackgroundTaskInfo(SerializationMixin): """Metadata for a single background task. - ``generation`` identifies which execution of the task this record describes. It starts at 1 - and increases every time the task is continued, so a restored older snapshot of the session - is never reconciled against a later execution that reused the same task ID. + ``execution_id`` identifies the execution (initial start or continuation) that this record + describes. A new, globally unique ID is assigned to every execution, so a record restored from + an older session snapshot is never reconciled against a different execution, even when that + execution reuses the same task ID. """ DEFAULT_EXCLUDE: ClassVar[set[str]] = set() @@ -70,8 +72,8 @@ class BackgroundTaskInfo(SerializationMixin): status: BackgroundTaskStatus result_text: str | None error_text: str | None - generation: int - __slots__ = ("agent_name", "description", "error_text", "generation", "id", "result_text", "status") + execution_id: str | None + __slots__ = ("agent_name", "description", "error_text", "execution_id", "id", "result_text", "status") def __init__( self, @@ -81,7 +83,7 @@ def __init__( status: BackgroundTaskStatus = BackgroundTaskStatus.RUNNING, result_text: str | None = None, error_text: str | None = None, - generation: int = 1, + execution_id: str | None = None, ) -> None: """Initialize a background task info entry.""" self.id = id @@ -90,7 +92,7 @@ def __init__( self.status = status self.result_text = result_text self.error_text = error_text - self.generation = generation + self.execution_id = execution_id def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: """Serialize for session state persistence.""" @@ -100,12 +102,13 @@ def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) "agent_name": self.agent_name, "description": self.description, "status": self.status.value, - "generation": self.generation, } if not exclude_none or self.result_text is not None: data["result_text"] = self.result_text if not exclude_none or self.error_text is not None: data["error_text"] = self.error_text + if not exclude_none or self.execution_id is not None: + data["execution_id"] = self.execution_id return data @classmethod @@ -118,7 +121,7 @@ def from_dict(cls, data: MutableMapping[str, Any], **kwargs: Any) -> BackgroundT status=BackgroundTaskStatus(data["status"]), result_text=data.get("result_text"), error_text=data.get("error_text"), - generation=data.get("generation", 1), + execution_id=data.get("execution_id"), ) @@ -126,42 +129,52 @@ def from_dict(cls, data: MutableMapping[str, Any], **kwargs: Any) -> BackgroundT class _RuntimeState: """Non-serializable per-session runtime state for background tasks. - ``in_flight_tasks`` holds the latest execution started for each task ID and - ``task_generations`` records which generation that execution belongs to. Lookups made on - behalf of a persisted :class:`BackgroundTaskInfo` go through :meth:`get_execution` so that a - record only ever sees the execution it was created for. + ``in_flight_tasks`` holds the latest execution started for each task ID and ``execution_ids`` + records which execution that is. Lookups made on behalf of a persisted + :class:`BackgroundTaskInfo` go through :meth:`get_execution`, so a record only ever sees the + execution it was created for. + + All reads and writes happen on the event loop thread (every background tool is invoked there), + so each tool's check-and-update sequence is atomic with respect to the others. """ in_flight_tasks: dict[int, asyncio.Task[AgentResponse[Any]]] = field(default_factory=lambda: {}) - task_generations: dict[int, int] = field(default_factory=lambda: {}) + execution_ids: dict[int, str] = field(default_factory=lambda: {}) background_sessions: dict[int, AgentSession] = field(default_factory=lambda: {}) + next_task_id: int = 1 closed: bool = False - def track_task(self, task_id: int, task: asyncio.Task[AgentResponse[Any]], *, generation: int = 1) -> None: - """Track a background task if this runtime is still open.""" + def allocate_task_id(self, persisted_next_task_id: int) -> int: + """Return a task ID that was never handed out by this runtime. + + A restored older session snapshot carries a stale ``next_task_id``; reusing it would let a + new task overwrite the runtime entry of a task started after that snapshot was taken. + """ + task_id = max(persisted_next_task_id, self.next_task_id) + self.next_task_id = task_id + 1 + return task_id + + def track_task(self, task_id: int, task: asyncio.Task[AgentResponse[Any]]) -> str: + """Track a new execution for a task ID if this runtime is still open and return its execution ID.""" if self.closed: task.cancel() raise RuntimeError("Session runtime is closed; cannot start background task.") + execution_id = uuid.uuid4().hex self.in_flight_tasks[task_id] = task - self.task_generations[task_id] = generation + self.execution_ids[task_id] = execution_id + return execution_id - def current_generation(self, task_id: int) -> int | None: - """Return the generation of the tracked execution for a task ID, if any.""" - if task_id not in self.in_flight_tasks: - return None - return self.task_generations.get(task_id, 1) - - def get_execution(self, task_id: int, generation: int) -> asyncio.Task[AgentResponse[Any]] | None: - """Return the tracked execution for a task ID only if it belongs to the given generation.""" - if self.current_generation(task_id) != generation: + def get_execution(self, task_id: int, execution_id: str | None) -> asyncio.Task[AgentResponse[Any]] | None: + """Return the tracked execution for a task ID only if it is the given execution.""" + if execution_id is None or self.execution_ids.get(task_id) != execution_id: return None return self.in_flight_tasks.get(task_id) def forget_task(self, task_id: int) -> None: """Drop all runtime references for a task ID.""" self.in_flight_tasks.pop(task_id, None) - self.task_generations.pop(task_id, None) + self.execution_ids.pop(task_id, None) self.background_sessions.pop(task_id, None) @@ -263,9 +276,9 @@ def _finalize_task( as ``LOST``, discarding a result that is already available. Because the retained entry is keyed by task ID, callers must only pass the - execution matching ``task_info.generation`` (see - :meth:`_RuntimeState.get_execution`); a later continuation of the same task - ID must never be applied to an older record. + execution matching ``task_info.execution_id`` (see + :meth:`_RuntimeState.get_execution`); a different execution under the same + task ID must never be applied to the record. Tasks are removed from ``in_flight_tasks`` by two callers: @@ -296,9 +309,10 @@ def _refresh_task_state( if task_info.status != BackgroundTaskStatus.RUNNING: continue - # A missing execution, or one from a different generation (the task was continued from - # another snapshot of this session), means this record's execution is no longer tracked. - in_flight = runtime.get_execution(task_info.id, task_info.generation) + # A missing execution, or a different one under the same task ID (for example, the task was + # continued from another snapshot of this session), means this record's execution is no + # longer tracked. + in_flight = runtime.get_execution(task_info.id, task_info.execution_id) if in_flight is None: task_info.status = BackgroundTaskStatus.LOST changed = True @@ -452,7 +466,7 @@ async def release_session( await asyncio.gather(*completed, return_exceptions=True) finally: runtime.in_flight_tasks.clear() - runtime.task_generations.clear() + runtime.execution_ids.clear() runtime.background_sessions.clear() if self._runtime.get(session_id) is runtime: @@ -541,13 +555,13 @@ def background_agents_start_task(agent_name: str, input: str, description: str) return f"Error: No background agent found with name '{agent_name}'. Available agents: {available}" bg_agent = self._agents[key] - task_id = provider_state.get("next_task_id", 1) + task_id = runtime.allocate_task_id(provider_state.get("next_task_id", 1)) sub_session = bg_agent.create_session() async_task = asyncio.create_task(_run_agent(bg_agent.run(input, session=sub_session))) try: - runtime.track_task(task_id, async_task) + execution_id = runtime.track_task(task_id, async_task) except RuntimeError as exc: return f"Error: {exc}" @@ -559,6 +573,7 @@ def background_agents_start_task(agent_name: str, input: str, description: str) id=task_id, agent_name=agent_name, description=description, + execution_id=execution_id, ) tasks = _get_tasks(provider_state) tasks.append(task_info) @@ -585,11 +600,10 @@ async def background_agents_wait_for_first_completion(task_ids: list[int]) -> st # Collect the executions that the requested running records refer to. waitable: list[tuple[int, asyncio.Task[AgentResponse[Any]]]] = [] tasks = _get_tasks(provider_state) - running = {t.id: t.generation for t in tasks if t.status == BackgroundTaskStatus.RUNNING} + running = {t.id: t.execution_id for t in tasks if t.status == BackgroundTaskStatus.RUNNING} for tid in task_ids: - generation = running.get(tid) - if generation is not None: - in_flight = runtime.get_execution(tid, generation) + if tid in running: + in_flight = runtime.get_execution(tid, running[tid]) if in_flight is not None: waitable.append((tid, in_flight)) @@ -631,8 +645,8 @@ async def background_agents_wait_for_first_completion(task_ids: list[int]) -> st task_info = next((t for t in tasks if t.id == completed_id), None) if task_info is not None and completed_task is not None: # Report the waited-on execution without overwriting a newer one: persist only if the - # record's generation still refers to the execution this waiter observed. - is_current_task = runtime.get_execution(task_info.id, task_info.generation) is completed_task + # record still refers to the execution this waiter observed. + is_current_task = runtime.get_execution(task_info.id, task_info.execution_id) is completed_task _finalize_task(task_info, completed_task) if is_current_task: _save_tasks(provider_state, tasks) @@ -660,6 +674,8 @@ def background_agents_get_task_results(task_id: int) -> str: return f"Task {task_id} is still running." return f"Task {task_id} has status: {task_info.status.value}." + background_agents_get_task_results._invoke_sync_on_event_loop = True # pyright: ignore[reportPrivateUsage] + @tool(name="background_agents_get_all_tasks", approval_mode="never_require") def background_agents_get_all_tasks() -> str: """List all background tasks with their IDs, statuses, agent names, and descriptions.""" @@ -673,6 +689,8 @@ def background_agents_get_all_tasks() -> str: lines.append(f"- Task {t.id} [{t.status.value}] ({t.agent_name}): {t.description}") return "\n".join(lines) + background_agents_get_all_tasks._invoke_sync_on_event_loop = True # pyright: ignore[reportPrivateUsage] + @tool(name="background_agents_continue_task", approval_mode="never_require") def background_agents_continue_task(task_id: int, text: str) -> str: """Send follow-up input to a completed or failed task to resume its work.""" @@ -701,26 +719,23 @@ def background_agents_continue_task(task_id: int, text: str) -> str: if sub_session is None: return f"Error: Session for task {task_id} is no longer available." - # This record is not running, so a pending tracked execution must be a newer one started - # from another snapshot of this session. Running both against one child session is unsafe. - current = runtime.in_flight_tasks.get(task_id) - if current is not None and not current.done(): + # Only the copy of the session that owns the tracked execution may continue it. Otherwise + # the task was already continued from another snapshot of this session: continuing again + # would run concurrently with, or discard the unread result of, that other execution. + if runtime.execution_ids.get(task_id) != task_info.execution_id: return ( - f"Error: Task {task_id} has a newer execution that is still running. " - "Wait for it to complete before continuing." + f"Error: Task {task_id} was continued from another copy of this session and cannot be " + "continued here. Start a new task instead." ) bg_agent = self._agents[key] - # Always move past any generation this runtime has seen so the new execution can't be - # mistaken for one referenced by another snapshot. - generation = max(task_info.generation, runtime.current_generation(task_id) or 0) + 1 async_task = asyncio.create_task(_run_agent(bg_agent.run(text, session=sub_session))) try: - runtime.track_task(task_id, async_task, generation=generation) + execution_id = runtime.track_task(task_id, async_task) except RuntimeError as exc: return f"Error: {exc}" - task_info.generation = generation + task_info.execution_id = execution_id task_info.status = BackgroundTaskStatus.RUNNING task_info.result_text = None task_info.error_text = None @@ -750,16 +765,18 @@ def background_agents_clear_completed_task(task_id: int) -> str: tasks = [t for t in tasks if t.id != task_id] _save_tasks(provider_state, tasks) - # Release runtime references unless a newer execution of this task ID exists (started from - # another snapshot of this session); that execution must stay tracked so it can still be - # retrieved, cancelled, or drained by release_session. - current_generation = runtime.current_generation(task_id) - if current_generation is None or current_generation <= task_info.generation: + # Release runtime references only when they belong to this record's execution. A different + # execution under this task ID (started from another snapshot of this session) must stay + # tracked so it can still be retrieved, cancelled, or drained by release_session. + tracked_execution_id = runtime.execution_ids.get(task_id) + if tracked_execution_id is None or tracked_execution_id == task_info.execution_id: runtime.forget_task(task_id) _save_provider_state(session, provider_state, source_id=source_id) return f"Task {task_id} cleared." + background_agents_clear_completed_task._invoke_sync_on_event_loop = True # pyright: ignore[reportPrivateUsage] + # Inject instructions and current task status. context.extend_instructions(self.source_id, [self._instructions]) context.extend_tools( diff --git a/python/packages/core/tests/core/test_harness_background_agents.py b/python/packages/core/tests/core/test_harness_background_agents.py index 1904b5100f8..421a4c08f9f 100644 --- a/python/packages/core/tests/core/test_harness_background_agents.py +++ b/python/packages/core/tests/core/test_harness_background_agents.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import threading from contextlib import suppress from typing import Any @@ -17,6 +18,7 @@ Message, ServiceSessionId, ) +from agent_framework._harness import _background_agents from agent_framework._sessions import SessionContext # Suppress "coroutine was never awaited" warnings from task cancellation in tests. @@ -1339,7 +1341,7 @@ async def run( # --------------------------------------------------------------------------- -# Regression tests: execution identity (generation) for reused task IDs +# Regression tests: execution identity for reused task IDs # --------------------------------------------------------------------------- @@ -1366,13 +1368,39 @@ async def run( return AgentResponse(messages=[Message(role="assistant", contents=[self.responses[index]])]) -def test_task_info_generation_round_trip_and_legacy_default() -> None: - """generation is persisted, and state written before it existed loads as generation 1.""" - info = BackgroundTaskInfo(id=1, agent_name="A", description="d", generation=3) - assert BackgroundTaskInfo.from_dict(info.to_dict()).generation == 3 +def _task_record(session: AgentSession, task_id: int) -> dict[str, Any]: + return next(t for t in session.state["background_agents"]["tasks"] if t["id"] == task_id) + + +def test_task_info_execution_id_round_trip_and_legacy_state() -> None: + """execution_id is persisted, and state written before it existed loads without one.""" + info = BackgroundTaskInfo(id=1, agent_name="A", description="d", execution_id="abc") + assert BackgroundTaskInfo.from_dict(info.to_dict()).execution_id == "abc" legacy = {"id": 1, "agent_name": "A", "description": "d", "status": "running"} - assert BackgroundTaskInfo.from_dict(legacy).generation == 1 + assert BackgroundTaskInfo.from_dict(legacy).execution_id is None + + +async def test_every_execution_gets_a_distinct_execution_id() -> None: + """Starts and continuations each record a new execution ID on the task.""" + agent = _ScriptedAgent("Worker", ["first result", "second result"]) + provider = _make_provider(agent) + session = _make_session() + tools = await _get_tools(provider, session) + + await _invoke_tool(tools["background_agents_start_task"], agent_name="Worker", input="go", description="t") + runtime = provider._get_runtime(session) + first_execution_id = _task_record(session, 1)["execution_id"] + assert runtime.execution_ids[1] == first_execution_id + await runtime.in_flight_tasks[1] + await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) + + await _invoke_tool(tools["background_agents_continue_task"], task_id=1, text="again") + second_execution_id = _task_record(session, 1)["execution_id"] + assert second_execution_id != first_execution_id + assert runtime.execution_ids[1] == second_execution_id + + await provider.release_session(session) @pytest.mark.parametrize("continuation_fails", [False, True]) @@ -1384,18 +1412,18 @@ async def test_restored_running_snapshot_does_not_receive_continuation_outcome(c tools = await _get_tools(provider, session) await _invoke_tool(tools["background_agents_start_task"], agent_name="Worker", input="go", description="t") - stale_snapshot = session.to_dict() # generation 1, RUNNING + stale_snapshot = session.to_dict() # first execution, RUNNING runtime = provider._get_runtime(session) await runtime.in_flight_tasks[1] - # The live session retrieves generation 1 and continues the task (generation 2). + # The live session retrieves the first execution and continues the task. assert await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) == "first result" await _invoke_tool(tools["background_agents_continue_task"], task_id=1, text="again") continuation = runtime.in_flight_tasks[1] with suppress(Exception): await continuation - # The host restores the snapshot taken while generation 1 was running. + # The host restores the snapshot taken while the first execution was running. restored = AgentSession.from_dict(stale_snapshot) restored_tools = await _get_tools(provider, restored) stale_result = await _invoke_tool(restored_tools["background_agents_get_task_results"], task_id=1) @@ -1440,8 +1468,9 @@ async def test_restored_running_snapshot_does_not_wait_on_continuation() -> None await provider.release_session(session) -async def test_continue_from_stale_snapshot_rejected_while_newer_execution_running() -> None: - """Continuing an old completed snapshot must not overwrite a newer running continuation.""" +@pytest.mark.parametrize("other_execution_state", ["running", "completed"]) +async def test_continue_from_stale_snapshot_is_rejected(other_execution_state: str) -> None: + """A task continued from one copy of the session cannot be continued again from an older copy.""" agent = _ScriptedAgent("Worker", ["first result", "second result", "third result"]) agent.gates[1] = asyncio.Event() provider = _make_provider(agent) @@ -1452,24 +1481,32 @@ async def test_continue_from_stale_snapshot_rejected_while_newer_execution_runni runtime = provider._get_runtime(session) await runtime.in_flight_tasks[1] await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) - completed_snapshot = session.to_dict() # generation 1, COMPLETED + completed_snapshot = session.to_dict() # first execution, COMPLETED await _invoke_tool(tools["background_agents_continue_task"], task_id=1, text="again") continuation = runtime.in_flight_tasks[1] + continuation_execution_id = runtime.execution_ids[1] + if other_execution_state == "completed": + agent.gates[1].set() + await continuation restored = AgentSession.from_dict(completed_snapshot) restored_tools = await _get_tools(provider, restored) result = await _invoke_tool(restored_tools["background_agents_continue_task"], task_id=1, text="stale") - assert result.startswith("Error: Task 1 has a newer execution that is still running.") + assert result.startswith("Error: Task 1 was continued from another copy of this session") assert runtime.in_flight_tasks[1] is continuation - assert runtime.task_generations[1] == 2 + assert runtime.execution_ids[1] == continuation_execution_id - await provider.release_session(session, cancel_running=True) - assert continuation.cancelled() + # The other copy's continuation, including an unread result, is untouched. + agent.gates[1].set() + await continuation + assert await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) == "second result" + + await provider.release_session(session) -async def test_continue_from_stale_snapshot_starts_newer_generation() -> None: - """Continuing from an older snapshot after a newer execution finished must use a fresh generation.""" - agent = _ScriptedAgent("Worker", ["first result", "second result", "third result"]) +async def test_clear_from_stale_snapshot_keeps_other_execution() -> None: + """Clearing an older copy's record keeps a finished execution that another copy still owns.""" + agent = _ScriptedAgent("Worker", ["first result", "second result"]) provider = _make_provider(agent) session = _make_session() tools = await _get_tools(provider, session) @@ -1478,29 +1515,225 @@ async def test_continue_from_stale_snapshot_starts_newer_generation() -> None: runtime = provider._get_runtime(session) await runtime.in_flight_tasks[1] await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) - completed_snapshot = session.to_dict() # generation 1, COMPLETED + completed_snapshot = session.to_dict() await _invoke_tool(tools["background_agents_continue_task"], task_id=1, text="again") - await runtime.in_flight_tasks[1] - assert await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) == "second result" + continuation = runtime.in_flight_tasks[1] + await continuation restored = AgentSession.from_dict(completed_snapshot) restored_tools = await _get_tools(provider, restored) - await _invoke_tool(restored_tools["background_agents_continue_task"], task_id=1, text="from snapshot") - third = runtime.in_flight_tasks[1] - assert runtime.task_generations[1] == 3 - assert restored.state["background_agents"]["tasks"][0]["generation"] == 3 - await third - assert await _invoke_tool(restored_tools["background_agents_get_task_results"], task_id=1) == "third result" - - # The live session (generation 2) clears its record without dropping generation 3. - assert await _invoke_tool(tools["background_agents_clear_completed_task"], task_id=1) == "Task 1 cleared." - assert runtime.in_flight_tasks[1] is third + assert await _invoke_tool(restored_tools["background_agents_clear_completed_task"], task_id=1) == "Task 1 cleared." + assert runtime.in_flight_tasks[1] is continuation assert 1 in runtime.background_sessions - # Clearing the record that owns the latest generation releases the runtime references. - assert await _invoke_tool(restored_tools["background_agents_clear_completed_task"], task_id=1) == "Task 1 cleared." + # The owning copy can still read and then clear it, which releases the runtime references. + assert await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) == "second result" + assert await _invoke_tool(tools["background_agents_clear_completed_task"], task_id=1) == "Task 1 cleared." assert 1 not in runtime.in_flight_tasks - assert 1 not in runtime.task_generations + assert 1 not in runtime.execution_ids assert 1 not in runtime.background_sessions await provider.release_session(session) + + +async def test_start_after_rollback_does_not_reuse_tracked_task_id() -> None: + """Restoring a snapshot from before a task started must not let a new task take over its ID.""" + agent = _ScriptedAgent("Worker", ["original result", "replacement result"]) + agent.gates[0] = asyncio.Event() + provider = _make_provider(agent) + session = _make_session() + tools = await _get_tools(provider, session) + before_start = session.to_dict() # next_task_id == 1 + + await _invoke_tool(tools["background_agents_start_task"], agent_name="Worker", input="go", description="original") + after_start = session.to_dict() # original task RUNNING + runtime = provider._get_runtime(session) + original = runtime.in_flight_tasks[1] + + # Roll back to before the original task started and start a different task. + rolled_back = AgentSession.from_dict(before_start) + rolled_back_tools = await _get_tools(provider, rolled_back) + result = await _invoke_tool( + rolled_back_tools["background_agents_start_task"], agent_name="Worker", input="go", description="replacement" + ) + assert result == "Background task 2 started on agent 'Worker'." + assert runtime.in_flight_tasks[1] is original + await runtime.in_flight_tasks[2] + + # Restoring the original task's RUNNING snapshot must not accept the replacement's result. + restored = AgentSession.from_dict(after_start) + restored_tools = await _get_tools(provider, restored) + assert ( + await _invoke_tool(restored_tools["background_agents_get_task_results"], task_id=1) + == "Task 1 is still running." + ) + agent.gates[0].set() + await original + assert await _invoke_tool(restored_tools["background_agents_get_task_results"], task_id=1) == "original result" + + await provider.release_session(session) + + +async def test_start_after_rollback_and_clear_does_not_reuse_task_id() -> None: + """A task ID stays reserved for the runtime's lifetime, even after its task was cleared.""" + provider = _make_provider(_ScriptedAgent("Worker", ["first result", "second result"])) + session = _make_session() + tools = await _get_tools(provider, session) + before_start = session.to_dict() + + await _invoke_tool(tools["background_agents_start_task"], agent_name="Worker", input="go", description="t1") + runtime = provider._get_runtime(session) + await runtime.in_flight_tasks[1] + await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) + assert await _invoke_tool(tools["background_agents_clear_completed_task"], task_id=1) == "Task 1 cleared." + + rolled_back = AgentSession.from_dict(before_start) + rolled_back_tools = await _get_tools(provider, rolled_back) + result = await _invoke_tool( + rolled_back_tools["background_agents_start_task"], agent_name="Worker", input="go", description="t2" + ) + assert result == "Background task 2 started on agent 'Worker'." + + await provider.release_session(session) + + +async def test_snapshot_from_released_runtime_never_matches_a_new_execution() -> None: + """Execution IDs are globally unique, so a new runtime's execution never matches an old record.""" + agent = _ScriptedAgent("Worker", ["original result", "replacement result"]) + agent.gates[0] = asyncio.Event() + provider = _make_provider(agent) + session = _make_session() + tools = await _get_tools(provider, session) + before_start = session.to_dict() + + await _invoke_tool(tools["background_agents_start_task"], agent_name="Worker", input="go", description="original") + after_start = session.to_dict() + await asyncio.sleep(0) # let the original execution start (it then blocks on its gate) + assert agent.runs == 1 + await provider.release_session(session, cancel_running=True) + + # A fresh runtime hands out task ID 1 again for a different task. + rolled_back = AgentSession.from_dict(before_start) + rolled_back_tools = await _get_tools(provider, rolled_back) + result = await _invoke_tool( + rolled_back_tools["background_agents_start_task"], agent_name="Worker", input="go", description="replacement" + ) + assert result == "Background task 1 started on agent 'Worker'." + await provider._get_runtime(rolled_back).in_flight_tasks[1] + + restored = AgentSession.from_dict(after_start) + restored_tools = await _get_tools(provider, restored) + stale_result = await _invoke_tool(restored_tools["background_agents_get_task_results"], task_id=1) + assert stale_result == "Task state was lost (reference unavailable)." + + await provider.release_session(rolled_back) + + +@pytest.mark.parametrize( + ("tool_name", "arguments"), + [ + ("background_agents_start_task", {"agent_name": "Worker", "input": "go", "description": "t2"}), + ("background_agents_wait_for_first_completion", {"task_ids": [1]}), + ("background_agents_get_task_results", {"task_id": 1}), + ("background_agents_get_all_tasks", {}), + ("background_agents_continue_task", {"task_id": 1, "text": "again"}), + ("background_agents_clear_completed_task", {"task_id": 1}), + ], +) +async def test_background_tools_run_on_the_event_loop_thread( + monkeypatch: pytest.MonkeyPatch, tool_name: str, arguments: dict[str, Any] +) -> None: + """Every tool touches runtime state on the event loop, never from a worker thread. + + Continuation registration runs on the event loop; running the other tools there too makes each + tool's check-and-update atomic with respect to it. + """ + provider = _make_provider(_ScriptedAgent("Worker", ["first result", "second result", "third result"])) + session = _make_session() + tools = await _get_tools(provider, session) + await _invoke_tool(tools["background_agents_start_task"], agent_name="Worker", input="go", description="t1") + runtime = provider._get_runtime(session) + await runtime.in_flight_tasks[1] + if tool_name != "background_agents_wait_for_first_completion": + await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) + + loop_thread = threading.get_ident() + seen_threads: set[int] = set() + real_get_tasks = _background_agents._get_tasks + + def recording_get_tasks(state: dict[str, Any]) -> list[BackgroundTaskInfo]: + seen_threads.add(threading.get_ident()) + return real_get_tasks(state) + + monkeypatch.setattr(_background_agents, "_get_tasks", recording_get_tasks) + await _invoke_tool(tools[tool_name], **arguments) + + assert seen_threads == {loop_thread} + await provider.release_session(session) + + +@pytest.mark.parametrize("continue_first", [True, False]) +async def test_concurrent_stale_clear_and_continuation_keep_continuation_tracked(continue_first: bool) -> None: + """A stale clear racing a continuation must never strand the continuation untracked.""" + agent = _ScriptedAgent("Worker", ["first result", "second result"]) + agent.gates[1] = asyncio.Event() + provider = _make_provider(agent) + session = _make_session() + tools = await _get_tools(provider, session) + + await _invoke_tool(tools["background_agents_start_task"], agent_name="Worker", input="go", description="t") + runtime = provider._get_runtime(session) + await runtime.in_flight_tasks[1] + await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) + restored = AgentSession.from_dict(session.to_dict()) # another copy of the same COMPLETED execution + restored_tools = await _get_tools(provider, restored) + + continue_call = _invoke_tool(tools["background_agents_continue_task"], task_id=1, text="again") + clear_call = _invoke_tool(restored_tools["background_agents_clear_completed_task"], task_id=1) + calls = [continue_call, clear_call] if continue_first else [clear_call, continue_call] + results = await asyncio.gather(*calls) + continue_result = results[0] if continue_first else results[1] + + if continue_result == "Task 1 continued with new input.": + # The continuation must still be tracked so release can cancel it. + continuation = runtime.in_flight_tasks[1] + assert not continuation.done() + await provider.release_session(session, cancel_running=True) + assert continuation.cancelled() + else: + # The clear won; the continuation was refused instead of starting untracked. + assert continue_result == "Error: Session for task 1 is no longer available." + assert 1 not in runtime.in_flight_tasks + assert agent.runs == 1 + await provider.release_session(session) + + +@pytest.mark.parametrize("read_first", [True, False]) +async def test_concurrent_stale_result_read_and_continuation_never_crosses_executions(read_first: bool) -> None: + """A stale reader racing a continuation must never accept the continuation's result as its own.""" + agent = _ScriptedAgent("Worker", ["first result", "second result"]) + provider = _make_provider(agent) + session = _make_session() + tools = await _get_tools(provider, session) + + await _invoke_tool(tools["background_agents_start_task"], agent_name="Worker", input="go", description="t") + stale_snapshot = session.to_dict() # first execution, RUNNING + runtime = provider._get_runtime(session) + await runtime.in_flight_tasks[1] + await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) + restored = AgentSession.from_dict(stale_snapshot) + restored_tools = await _get_tools(provider, restored) + + read_call = _invoke_tool(restored_tools["background_agents_get_task_results"], task_id=1) + continue_call = _invoke_tool(tools["background_agents_continue_task"], task_id=1, text="again") + calls = [read_call, continue_call] if read_first else [continue_call, read_call] + results = await asyncio.gather(*calls) + stale_result = results[0] if read_first else results[1] + assert stale_result in {"first result", "Task state was lost (reference unavailable)."} + + # Let the continuation finish; the stale record still never reports it. + await runtime.in_flight_tasks[1] + assert await _invoke_tool(restored_tools["background_agents_get_task_results"], task_id=1) == stale_result + assert await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) == "second result" + + await provider.release_session(session)