diff --git a/python/packages/core/agent_framework/_harness/_background_agents.py b/python/packages/core/agent_framework/_harness/_background_agents.py index 631a8bd7705..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 @@ -55,7 +56,13 @@ class BackgroundTaskStatus(str, Enum): @experimental(feature_id=ExperimentalFeature.HARNESS) class BackgroundTaskInfo(SerializationMixin): - """Metadata for a single background task.""" + """Metadata for a single background task. + + ``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() @@ -65,7 +72,8 @@ class BackgroundTaskInfo(SerializationMixin): status: BackgroundTaskStatus result_text: str | None error_text: str | None - __slots__ = ("agent_name", "description", "error_text", "id", "result_text", "status") + execution_id: str | None + __slots__ = ("agent_name", "description", "error_text", "execution_id", "id", "result_text", "status") def __init__( self, @@ -75,6 +83,7 @@ def __init__( status: BackgroundTaskStatus = BackgroundTaskStatus.RUNNING, result_text: str | None = None, error_text: str | None = None, + execution_id: str | None = None, ) -> None: """Initialize a background task info entry.""" self.id = id @@ -83,6 +92,7 @@ def __init__( self.status = status self.result_text = result_text self.error_text = error_text + 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.""" @@ -97,6 +107,8 @@ def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) 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 @@ -109,24 +121,61 @@ 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"), + execution_id=data.get("execution_id"), ) @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 ``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: {}) + 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]]) -> 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.execution_ids[task_id] = execution_id + return execution_id + + 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.execution_ids.pop(task_id, None) + self.background_sessions.pop(task_id, None) # --------------------------------------------------------------------------- @@ -215,9 +264,27 @@ 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.""" + """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. + + Because the retained entry is keyed by task ID, callers must only pass the + 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: + + - :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,9 +296,6 @@ def _finalize_task( else: task_info.status = BackgroundTaskStatus.COMPLETED task_info.result_text = completed_task.result().text - # A stale waiter may observe a task ID already reused by a continuation. - if runtime.in_flight_tasks.get(task_info.id) is completed_task: - runtime.in_flight_tasks.pop(task_info.id, None) def _refresh_task_state( @@ -245,14 +309,17 @@ 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 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 continue if in_flight.done(): - _finalize_task(task_info, in_flight, runtime) + _finalize_task(task_info, in_flight) changed = True if changed: @@ -399,6 +466,7 @@ async def release_session( await asyncio.gather(*completed, return_exceptions=True) finally: runtime.in_flight_tasks.clear() + runtime.execution_ids.clear() runtime.background_sessions.clear() if self._runtime.get(session_id) is runtime: @@ -487,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}" @@ -505,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) @@ -528,12 +597,15 @@ 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 = {t.id: t.execution_id for t in tasks if t.status == BackgroundTaskStatus.RUNNING} for tid in task_ids: - in_flight = runtime.in_flight_tasks.get(tid) - if in_flight is not None: - waitable.append((tid, in_flight)) + if tid in running: + in_flight = runtime.get_execution(tid, running[tid]) + if in_flight is not None: + waitable.append((tid, in_flight)) if not waitable: # Refresh state to catch any that completed. @@ -572,9 +644,10 @@ 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_task is not None: - # Report the waited-on execution without overwriting a newer one. - is_current_task = runtime.in_flight_tasks.get(task_info.id) is completed_task - _finalize_task(task_info, completed_task, runtime) + # Report the waited-on execution without overwriting a newer one: persist only if the + # 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) _save_provider_state(session, provider_state, source_id=source_id) @@ -601,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.""" @@ -614,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.""" @@ -642,13 +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." + # 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} was continued from another copy of this session and cannot be " + "continued here. Start a new task instead." + ) + bg_agent = self._agents[key] async_task = asyncio.create_task(_run_agent(bg_agent.run(text, 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}" + task_info.execution_id = execution_id task_info.status = BackgroundTaskStatus.RUNNING task_info.result_text = None task_info.error_text = None @@ -678,13 +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) - # Clean up runtime references. - runtime.in_flight_tasks.pop(task_id, None) - runtime.background_sessions.pop(task_id, None) + # 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 7e14c074d6f..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. @@ -390,7 +392,11 @@ async def counted_wait(*args: Any, **kwargs: Any) -> Any: agent_release.set() assert await asyncio.gather(*waiters) == ["Task 1 finished with status: completed."] * 2 assert await _invoke_tool(tools["background_agents_get_task_results"], task_id=1) == "shared result" - assert not provider._runtime[session.session_id].in_flight_tasks + # Finished executions are retained until cleared or released. + runtime = provider._runtime[session.session_id] + assert runtime.in_flight_tasks[1].done() + assert await _invoke_tool(tools["background_agents_clear_completed_task"], task_id=1) == "Task 1 cleared." + assert not runtime.in_flight_tasks finally: agent_release.set() await asyncio.gather(*waiters, return_exceptions=True) @@ -1042,3 +1048,692 @@ 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 + assert not runtime.in_flight_tasks + + +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) + + # 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] + + # 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", + ) + + # 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) + + +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) + session = _make_session() + tools = await _get_tools(provider, session) + + # 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", + ) + runtime = provider._get_runtime(session) + await runtime.in_flight_tasks[1] + + # Finalize task 1 so its status is COMPLETED. + tools2 = await _get_tools(provider, session) + await _invoke_tool(tools2["background_agents_get_task_results"], task_id=1) + + # 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() + + continuation_task = runtime.in_flight_tasks[1] + assert not continuation_task.done() + + # 4. Restore snapshot into a new AgentSession. + session2 = AgentSession.from_dict(snapshot) + tools3 = await _get_tools(provider, session2) + + # 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." + + # 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() + + # 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 + + +# --------------------------------------------------------------------------- +# Regression tests: execution identity 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 _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).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]) +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() # first execution, RUNNING + runtime = provider._get_runtime(session) + await runtime.in_flight_tasks[1] + + # 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 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) + 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) + + +@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) + 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() # 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 was continued from another copy of this session") + assert runtime.in_flight_tasks[1] is continuation + assert runtime.execution_ids[1] == continuation_execution_id + + # 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_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) + + 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() + await _invoke_tool(tools["background_agents_continue_task"], task_id=1, text="again") + continuation = runtime.in_flight_tasks[1] + await continuation + + restored = AgentSession.from_dict(completed_snapshot) + restored_tools = await _get_tools(provider, restored) + 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 + + # 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.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)