[Karthik Thota (karthik-0306)]

Motivation & Context

1. Why is this change required?
   In the single-process BackgroundAgentsProvider, the host persists session state only after an agent run completes. _finalize_task used to drop a finished task from _RuntimeState.in_flight_tasks. If the host then restored a snapshot taken while that task was RUNNING (after a failed run, a rollback, or a state reload), _refresh_task_state could no longer find it and marked it LOST, discarding a result that was already available.

2. What problem does it solve?
   Finished executions stay in memory until they are cleared or the session is released, so a stale RUNNING record can still pick up its real result or error. Because task IDs are reused by continuations and can be reused after a rollback, every execution also gets a persisted identity, so a record can never pick up a different execution's outcome. Separately, release_session(timeout=...) now returns at its deadline even if a task keeps suppressing cancellation.

3. What scenario does it contribute to?
   Single-process background agents used across multi-turn sessions where host persistence timing, rollbacks, or failed turns would otherwise orphan completed tasks, and hosts that rely on a bounded session release during eviction or shutdown.

4. Issue Link:
   Part of [Feature]: BackgroundAgentsProvider cannot tell a task that is running on another process from one whose owner died #8760 (Track 1: single-process retention). Also covers Python: [Bug]: Background session release waits past its finite timeout during cancellation cleanup #9231.

---

Description & Review Guide

• What are the major changes?

  • Retention: _finalize_task no longer removes finished executions. They are removed by background_agents_clear_completed_task or release_session.
  • Execution identity: BackgroundTaskInfo persists an execution_id (a new uuid for every start and continuation). Runtime lookups only match a record's own execution, so a stale snapshot becomes LOST instead of receiving another execution's result or error. State saved before this change loads with no execution_id.
  • Collision-safe IDs: the runtime tracks the highest task ID it has handed out, so a rolled-back snapshot can't reuse an ID that is still tracked.
  • Event-loop ownership: all background tools now run on the event loop, so each tool's check-and-update is atomic with respect to continuation registration.
  • Continue / clear: only the copy of the session that owns the tracked execution can continue it, and clear only releases runtime state that belongs to the record's own execution.
  • Release timeout: _drain_runtime observes cancelled tasks with asyncio.wait(pending, timeout=remaining) instead of asyncio.wait_for(gather(...)), so release returns at the deadline even if a task keeps catching CancelledError. Tasks still running are abandoned with the existing done-callback, which retrieves their eventual exceptions. This was a pre-existing issue, not caused by the retention change.
  • Merged Python: preserve background execution identity across completion waits #9442: its exact-execution wait handling is kept and now checks execution_id. One of its assertions expected finished tasks to be dropped and was updated for the new retention behavior.

• What is the impact of these changes?

  • No more false LOST results after snapshot restores in single-process hosts.
  • release_session(timeout=...) can no longer be held open indefinitely by a task that ignores cancellation.
  • Memory trade-off: finished asyncio.Task references stay in memory until the caller clears the task or the session is released.

• What do you want reviewers to focus on?

  • The execution_id matching in _RuntimeState.get_execution and its use in refresh, wait, continue and clear.
  • The bounded wait in _drain_runtime.
  • The regression tests for stale snapshots, rollback-before-start, event-loop ownership, and repeated cancellation suppression.

---

Validation

• pytest (background agents, harness loop, harness agent): 263 passed, 6 skipped. The new regressions fail on the previous head.
• ruff check / format: clean
• mypy (source and tests), pyrefly, ty (tests): clean
• pyright: no new findings. The only error is the existing tool import, which is also reported without this change.

---

Related Issue

Part of #8760. This PR covers Track 1 only. #8760 should stay open for the Track 2 cross-replica design. Also covers the release-timeout bug reported in #9231.

---

AI Assistance

• No material AI assistance was used.
• This is an AI-assisted contribution. I reviewed, understood, and verified all submitted content and accept responsibility for it.

AI assistance details:
AI assisted with drafting code and regression tests, running validation checks (pytest, ruff, mypy, pyright, pyrefly, ty), and drafting this PR description.

---

Contribution Checklist

• The code builds clean without any errors or warnings
• All unit tests pass, and I have added new tests where possible
• The PR follows the Contribution Guidelines
• This PR links to an agreed issue with no competing open PR, or the Related Issue section documents a trivial-change or repository-automation exception.
• This is not a breaking change.

[Copilot]

🟡 Changes recommended

Reused task IDs can associate stale snapshots with results from a later continuation.

1 open finding

• Retained task ID can apply results to the wrong execution · New

What changed in this PR

Retains finished background tasks to reconcile stale single-process session snapshots.

Changes:

• Stops removing tasks during finalization.
• Adds regression and lifecycle cleanup tests.

File	Description
_background_agents.py	Retains finalized tasks and documents cleanup lifecycle.
test_harness_background_agents.py	Tests snapshot recovery, failures, cancellation, clearing, and release.

🧠 Review effort: Balanced

---

💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.

[Eduard van Valkenburg (eavanvalkenburg)]

Karthik Thota (@karthik-0306) background_agents_clear_completed_task can clear a newer running continuation when the caller restores an older completed generation of the same task id. The restored generation passes the non-running check, but runtime.clear_task(task_id) removes every (task_id, generation) entry and the shared background session. A newer generation can therefore keep running untracked; later refreshes mark it lost, and session release can no longer cancel or drain it.

Please make clear generation-aware so it removes only the completed generation, or explicitly reject/cancel newer running generations before removing their runtime handles. Add a regression that restores generation 1 while generation 2 is still running and then attempts to clear generation 1.

[Karthik Thota (karthik-0306)]

Karthik Thota (Karthik Thota (@karthik-0306)) background_agents_clear_completed_task can clear a newer running continuation when the caller restores an older completed generation of the same task id. The restored generation passes the non-running check, but runtime.clear_task(task_id) removes every (task_id, generation) entry and the shared background session. A newer generation can therefore keep running untracked; later refreshes mark it lost, and session release can no longer cancel or drain it.

Please make clear generation-aware so it removes only the completed generation, or explicitly reject/cancel newer running generations before removing their runtime handles. Add a regression that restores generation 1 while generation 2 is still running and then attempts to clear generation 1.

Addressed in 6b90a3d. Generation tracking is removed. background_agents_clear_completed_task now drops the runtime handle and background session only when the local task is done or absent. If a continuation is still running, its handle stays tracked, so release_session(cancel_running=True) can still cancel and drain it. The persisted record is still removed because the caller asked to clear it. Regression added: restore the completed snapshot while the continuation runs, clear, then confirm the continuation is still tracked and release cancels it.

[Eduard van Valkenburg (eavanvalkenburg)]

Karthik Thota (@karthik-0306) The final head again keys retained runtime tasks only by task_id. background_agents_continue_task() overwrites that entry for a newer execution, while _refresh_task_state() reconciles any persisted RUNNING record against whichever task currently occupies the ID. A restored stale snapshot can therefore receive the result or error of a later continuation. Please preserve a serialized execution identity/generation (or otherwise prove the retained task is the same execution) and add a stale-snapshot-plus-continuation regression.

[Karthik Thota (karthik-0306)]

Eduard van Valkenburg (@eavanvalkenburg) Restored an execution identity. BackgroundTaskInfo now persists a generation (defaults to 1 for older state), and it goes up on every continuation. Runtime records which generation each tracked run belongs to, and refresh, wait and finalize only match a record's own generation. So a restored stale RUNNING snapshot becomes LOST instead of receiving a later continuation's result or error. Continuing from a stale snapshot is rejected while a newer run is still going, and otherwise moves past every generation seen. Clearing only releases runtime state when no newer generation exists. Added regression tests: restore a gen-1 RUNNING snapshot after gen 2 completes or fails, stale wait, and stale continue/clear.

[Karthik Thota (karthik-0306)]

Eduard van Valkenburg (@eavanvalkenburg) . Merged main to resolve the conflict with #9442. That PR's exact-execution wait handling is kept and now also checks the persisted generation. _finalize_task still keeps finished runs until clear/release, which is the core of this PR, so I updated the one #9442 assertion that expected them to be dropped.

Since your last review, 2ac9ae0 adds the persisted per-task generation you asked for, with stale-snapshot + continuation regression tests. Could you take another look, and approve the CI workflows if they need it? Happy to adjust anything.

[Eduard van Valkenburg (eavanvalkenburg)]

Karthik Thota (@karthik-0306) Please address the inline generation-consistency race before merge.

[Karthik Thota (karthik-0306)]

Eduard van Valkenburg (@eavanvalkenburg) addressed all three. The root causes were the per-task generation counter and the read/clear tools running off the event loop, so I fixed those directly instead of patching each race. The new tests fail on the previous commit. Let me know if you'd like anything changed.

[Eduard van Valkenburg (eavanvalkenburg)]

Karthik Thota (@karthik-0306) The current Related Issue wording says Does not close #8760, but GitHub still records #8760 as a closing issue because the closing keyword remains adjacent to the reference. Since the broader Track 2 work is intentionally unfinished, please reword this without a closing keyword so merging this PR cannot auto-close #8760. The code review itself is clean; this metadata correction is the remaining approval blocker.

[Karthik Thota (karthik-0306)]

Eduard van Valkenburg (@eavanvalkenburg) fixed, it now says "Part of #8760" with no closing keyword, and the description matches the current code.

For Track 2, I'm planning to start with an ADR in docs/decisions (task ownership, outcomes and cancellation across replicas), and could cover Jeffin's case from #9531 too. Once that's agreed I'd open a fresh PR and close #9040. Does that sound right? Should someone from the .NET side join on the ADR?

Also, .NET's FinalizeTask has the same issue. I'd like to send a separate PR for it once this one merges, if that works for you.

[Eduard van Valkenburg (eavanvalkenburg)]

Karthik Thota (@karthik-0306) There is a separate timeout problem in _drain_runtime() (python/packages/core/agent_framework/_harness/_background_agents.py:516-519): asyncio.wait_for(gather(...), timeout=...) waits for cancellation to finish after its deadline. A child that repeatedly catches CancelledError can therefore keep release_session(timeout=...) pending indefinitely, retaining the closed runtime and blocking host eviction/shutdown. A child waiting on an exit gate inside a loop that catches cancellation reproduced this at both the PR head and base; release only completed after opening the gate. The existing timeout test suppresses just the first cancellation, so it misses this case. Please bound observation of the cancelled tasks with asyncio.wait(..., timeout=remaining) (and retrieve eventual exceptions from unfinished tasks), with a regression that suppresses repeated cancellation. This is an affected pre-existing cleanup defect, not a regression introduced by the retention change.

[Karthik Thota (karthik-0306)]

Eduard van Valkenburg (@eavanvalkenburg) done. _drain_runtime now uses asyncio.wait(pending, timeout=remaining), so release returns at the deadline even if a task keeps catching cancellation. Tasks still running are abandoned with the existing done-callback, so their later exceptions are retrieved. Added a regression with a task that suppresses every cancellation; it fails on the previous commit and passes now.

This also fixes #9231, so cc westey (@westey-m) and Charan Rathore (@charan-rathore) , who were looking at it.

[Eduard van Valkenburg (eavanvalkenburg)]

Karthik Thota (@karthik-0306) There is an exception-observation regression in the new _drain_runtime() timeout path at python/packages/core/agent_framework/_harness/_background_agents.py:522-529. asyncio.wait() does not retrieve exceptions, but its completed set is discarded and only not_done receives exception-retrieval callbacks. With two children—one raising during cancellation and another exceeding the release deadline—the failed child is dropped when release_session() clears the runtime, and its failure reaches the event loop as “Task exception was never retrieved.” I reproduced this at the current head; the previous gather(..., return_exceptions=True) consumed the same failure.

Please retrieve the completed tasks' exceptions before returning on timeout, retain the callbacks for unfinished tasks, and add a mixed-outcome regression.

[Karthik Thota (karthik-0306)]

Eduard van Valkenburg (@eavanvalkenburg) good catch, asyncio.wait doesn't consume exceptions and I was only adding the callback to unfinished tasks. Fixed, and I went through the other ways release can stop waiting:

• On timeout, exceptions of all finished tracked tasks are retrieved, and unfinished ones keep the callback.
• If release_session itself is cancelled mid-wait, the same handling runs before the cancellation propagates.
• Abandoned tasks are now strongly referenced until they finish, so they can't be garbage collected while pending.

Added regressions for the mixed-outcome case, a cancelled release, and an abandoned task that only awaits an unreferenced future. Each fails on the previous commit and passes now.