Skip to content

fix: keep strong references to PromptStreamConsumer background tasks#821

Open
kratos0718 wants to merge 1 commit into
future-agi:mainfrom
kratos0718:fix/prompt-stream-background-tasks
Open

fix: keep strong references to PromptStreamConsumer background tasks#821
kratos0718 wants to merge 1 commit into
future-agi:mainfrom
kratos0718:fix/prompt-stream-background-tasks

Conversation

@kratos0718

Copy link
Copy Markdown

What

PromptStreamConsumer launched its streaming work with bare asyncio.create_task(...) and threw away the returned task in three places (execute_template_async, execute_improve_prompt_async, execute_generate_prompt_async).

The event loop only keeps a weak reference to a task, so a fire-and-forget task can be garbage-collected before it finishes — silently dropping a user's prompt execution with no result and no error. Per the CPython docs: "Save a reference to the result of this function, to avoid a task disappearing mid-execution."

Closes #819.

How

  • Added a self._background_tasks set and a small _spawn() helper that holds a strong reference to each task and removes it via add_done_callback once it completes (no leak).
  • Routed the three create_task call sites through _spawn().
  • Cancel any still-pending tasks in disconnect() so they don't outlive the connection.

Notes

  • No behavior change for the happy path — same coroutines, same scheduling; this just prevents the GC race.
  • Found while reviewing the async paths with a static analyzer (codehound). It flagged one more spot — model_hub/utils/kb_helpers.py uses a deprecated asyncio.get_event_loop() plus a fire-and-forget ensure_future — happy to send a separate small PR for that to keep this one focused.

Testing

  • python -m py_compile clean; static analyzer reports the file clean after the change.

The consumer launched execute/improve/generate work with bare
asyncio.create_task() and discarded the result. The event loop only
holds a weak reference to a task, so a fire-and-forget task can be
garbage-collected mid-run, silently dropping a user's prompt execution.

Track spawned tasks in a set via a _spawn() helper that keeps a strong
reference until completion, and cancel any pending tasks on disconnect.

Closes future-agi#819
@abhijaisrivastava15 abhijaisrivastava15 added the bug Something isn't working label Jun 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

backend bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Background tasks in PromptStreamConsumer can be garbage-collected mid-run (silent prompt-execution drops)

2 participants