From 2514167932dad4961091357671edd43f0b2949cf Mon Sep 17 00:00:00 2001
From: karthik-0306 <karthikthota7777@gmail.com>
Date: Wed, 14 Oct 2026 21:24:35 +0530
Subject: [PATCH 1/4] 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 <karthikthota7777@gmail.com>
Date: Wed, 14 Oct 2026 22:03:48 +0530
Subject: [PATCH 2/4] 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 <karthikthota7777@gmail.com>
Date: Thu, 15 Oct 2026 13:49:41 +0530
Subject: [PATCH 3/4] 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 <karthikthota7777@gmail.com>
Date: Fri, 16 Oct 2026 15:45:39 +0530
Subject: [PATCH 4/4] 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)