Repository navigation
Python: fix(core): keep finished background tasks until cleared or released - #9387
Karthik Thota (karthik-0306) wants to merge 6 commits into
Conversation
…leased Part of microsoft#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
GitHub Actions
Inactive
GitHub Actions
Inactive
agent-framework-automation
Bot
added
the
python
There was a problem hiding this comment.
🟡 Changes recommended
Reused task IDs can associate stale snapshots with results from a later continuation.
1 open finding
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.
…execution attribution (microsoft#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.
GitHub Actions
Inactive
Eduard van Valkenburg (eavanvalkenburg)
left a comment
Eduard van Valkenburg (eavanvalkenburg)
left a comment
There was a problem hiding this comment.
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.
GitHub Actions
Inactive
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)
left a comment
Eduard van Valkenburg (eavanvalkenburg)
left a comment
There was a problem hiding this comment.
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.
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.
GitHub Actions
Inactive
|
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. |
…hed-tasks-until-cleared # Conflicts: # python/packages/core/agent_framework/_harness/_background_agents.py
GitHub Actions
Inactive
|
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. |
| raise RuntimeError("Session runtime is closed; cannot start background task.") | ||
|
|
||
| self.in_flight_tasks[task_id] = task | ||
| self.task_generations[task_id] = generation |
There was a problem hiding this comment.
Karthik Thota (@karthik-0306), initial starts still call track_task() with the default generation of 1 and obtain their ID from the snapshotted next_task_id. Restoring a snapshot from before a task was started therefore lets a different task overwrite the same ID/generation. Restoring the original task's RUNNING snapshot then accepts the replacement task's result as its own. If the original execution is still running, its runtime handle is also overwritten. Please make initial-start ID/execution allocation collision-safe against the existing runtime, and add a rollback-before-start regression.
There was a problem hiding this comment.
Good catch. Switched from per-task generations to a uuid execution_id per execution, and the runtime now tracks the highest task ID it handed out so a rolled-back snapshot can't reuse one. Added rollback-before-start tests, including one across a released runtime.
| # 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: |
There was a problem hiding this comment.
Karthik Thota (@karthik-0306), the generation check and forget_task() are not atomic with continuation startup. This synchronous clear tool runs in a worker thread, whereas continue_task explicitly runs on the event loop. A clear from a completed snapshot can read generation 1, then a continuation can register generation 2 before the clear resumes and deletes its task handle and child session. The continuation remains running, and release_session(cancel_running=True) can no longer cancel it. Please serialize this check-and-delete with continuation registration, for example by giving clear the same event-loop ownership, and cover the concurrent interleaving.
There was a problem hiding this comment.
You're right, clear was running in a worker thread. All the background tools now run on the event loop, so clear's check and delete can't interleave with a continuation. Clear also only drops runtime state when it belongs to the record's own execution. Added a test that each tool runs on the loop thread.
Eduard van Valkenburg (eavanvalkenburg)
left a comment
Eduard van Valkenburg (eavanvalkenburg)
left a comment
There was a problem hiding this comment.
Karthik Thota (@karthik-0306) Please address the inline generation-consistency race before merge.
| return None | ||
| return self.task_generations.get(task_id, 1) | ||
|
|
||
| def get_execution(self, task_id: int, generation: int) -> asyncio.Task[AgentResponse[Any]] | None: |
There was a problem hiding this comment.
Karthik Thota (@karthik-0306) Generation validation and task retrieval are separate operations here. A generation-1 result reader can pass current_generation(), then a continuation can register and finalize generation 2 before in_flight_tasks.get(), causing generation 2’s result to be accepted and persisted under generation 1. Please give the result/list tools’ refresh-and-save transaction the same event-loop ownership as start/continue, or otherwise synchronize it with continuation registration, and add a regression that advances and finalizes a continuation between validation and retrieval.
There was a problem hiding this comment.
Same root cause, same fix. Results and list now run on the event loop too, so a continuation can't sneak in between the check and the lookup. The loop-thread test covers it, and there's a test checking a stale reader only ever sees its own result or LOST.
GitHub Actions
Inactive
|
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. |
Motivation & Context
Why is this change required?
In the single-process
BackgroundAgentsProvider, session state is persisted by the host only after an agent run completes. If the host restores a session snapshot taken while a task wasRUNNING(e.g. after a run failure, turn rollback, or host state reload), any completed or failed task that had already been finalized in memory had itsasyncio.Taskpurged from_RuntimeState.in_flight_tasks. On subsequent turn invocations,_refresh_task_statecould not find the task inin_flight_tasksand misclassified it asLOST, discarding a result or error that had already completed successfully.What problem does it solve?
It ensures that completed/failed task instances remain available in memory so that a stale persisted
RUNNINGrecord can reconcile against the finished task and return the actual result or exception rather than reporting"Task state was lost".What scenario does it contribute to?
Single-process background agents executing across multi-turn sessions where host persistence timing, rollbacks, or intermediate turn failures could otherwise orphan completed tasks.
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 in-flight task retention).
Description & Review Guide
What are the major changes?
_background_agents.py: Removedruntime.in_flight_tasks.pop(task_info.id, None)from_finalize_task. Tasks are now retained inin_flight_tasksuntil explicitly removed viabackground_agents_clear_completed_taskor session teardown viaBackgroundAgentsProvider.release_session.test_harness_background_agents.py: Added 5 deterministic regression tests (usingawait taskinstead ofasyncio.sleep):"Task state was lost".clear_completed_taskexplicitly removes the task fromruntime.in_flight_tasks.release_sessionclears all in-flight tasks and releases the runtime state.What is the impact of these changes?
LOSTtask states caused by session persistence desync in single-process scenarios.asyncio.Taskreferences remain in memory for the duration of the session until cleared by the caller or released with the session.What do you want reviewers to focus on?
_finalize_taskvsclear_completed_task/release_session.test_harness_background_agents.py.Validation
52 passed, 0 failedin 0.77sRelated Issue
Part of #8760 (Does not close #8760; addresses Track 1 single-process gap while Track 2 cross-process ADR remains open).
AI Assistance
AI assistance details:
AI assisted with drafting regression tests, formatting docstrings, running validation checks (pytest, ruff, pyright, mypy), and drafting this PR description.
Contribution Checklist