fix: close stream generators and cancel embeddings on disconnect#116
Conversation
run_cancellable_stream left the wrapped generator open when the consumer closed the stream early (GeneratorExit), stranding a llama-server --parallel slot or vLLM engine request until GC. Also wraps llama_server's create_embedding in run_cancellable so a disconnected client no longer holds a slot to completion, dedupes _project_embedding_response's usage math into _project_usage, consolidates the duplicated per-loader chat SSE chunk encoder into chat_utils.encode_chat_sse_chunk (switching to model_dump_json for a single serialization pass), and skips the streaming TRACE buffer entirely when TRACE is disabled.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request ensures that the underlying asynchronous generator (work) is properly closed when a streaming generator is closed early by the consumer. It also optimizes chat streaming by avoiding unnecessary string buffering when trace logging is disabled, unifies SSE chunk encoding across inference engines, and wraps embedding requests in a cancellable task. A review comment highlights a potential RuntimeError if work.aclose() is called while the background task driving work.__anext__() is still active, suggesting to cancel and await that task first.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
run_cancellable_stream's finally block called work.aclose() without first cancelling next_item when the driving task is torn down while suspended in asyncio.wait (rather than at the yield) - next_item's __anext__() call still owns work's frame at that point, so aclose() raises RuntimeError: aclose(): asynchronous generator is already running instead of letting CancelledError propagate cleanly.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request enhances the robustness of the asynchronous stream handling in run_cancellable_stream by ensuring that the underlying generator is properly closed on all exit paths and that in-flight tasks are cancelled and awaited during teardown to prevent RuntimeError. It also introduces a shared encode_chat_sse_chunk utility, optimizes trace logging overhead in both llama-server and vLLM inference streams, and adds unit tests to verify early generator closure and task cancellation. There are no review comments, and I have no additional feedback to provide.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
1 participant
run_cancellable_stream left the wrapped generator open when the consumer closed the stream early (GeneratorExit), stranding a llama-server --parallel slot or vLLM engine request until GC. Also wraps llama_server's create_embedding in run_cancellable so a disconnected client no longer holds a slot to completion, dedupes _project_embedding_response's usage math into _project_usage, consolidates the duplicated per-loader chat SSE chunk encoder into chat_utils.encode_chat_sse_chunk (switching to model_dump_json for a single serialization pass), and skips the streaming TRACE buffer entirely when TRACE is disabled.