fix: GeneratorExit handling in streaming generators + Langfuse native sessions

[sheepdestroyer]

What

Addresses all review comments on PR #297 and fixes additional issues discovered during review.

Changes

GeneratorExit handling (Gemini Code Assist review)

All three streaming generators now handle early client disconnect (GeneratorExit, a BaseException that bypasses except Exception):

• native_agy_stream_generator: added finalized flag + finally block
• agy_stream_generator (simulated): wrapped in try/finally with finalized flag
• stream_generator (LiteLLM fallback): extended existing finally to handle non-finalized cancellation

Langfuse native sessions (propagate_attributes)

• Replaced metadata-only session_id/user_id/tags with Langfuse's native propagate_attributes() context manager
• propagate_attributes(session_id=..., user_id=..., tags=...) is called BEFORE start_observation() so all child observations inherit the session context
• Metadata now only carries client_model and environment (dimensions, not session attributes)

Verification

• 193/193 tests pass (101 top-level + 92 router)
• 24/24 canonical endpoints pass (dev deployment)
• Langfuse session propagation confirmed: sessionId=e2e-pr297-session and userId=pr297-test visible on child observations via Langfuse API
• Syntax check clean

Closes #297

Summary by CodeRabbit

• Bug Fixes

  • Improved reliability for streaming and non-streaming responses when errors or cancellations occur.
  • Ensured request monitoring data is finalized consistently, including during failed requests.
  • Improved handling of unknown or unsupported model selections with clearer failure behavior.

• Improvements

  • Added more reliable connectivity and health checking for the Llama service.
  • Improved request metadata tracking for more complete observability in dashboards.

[sourcery-ai]

Sorry @sheepdestroyer, you have reached your weekly rate limit of 500000 diff characters.

Please try again later or upgrade to continue using Sourcery

[coderabbitai]

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 51515582-e29f-4980-8588-32b29f6be9cc

📥 Commits

Reviewing files that changed from the base of the PR and between 9560f48 and e832600.

📒 Files selected for processing (3)

• .env.dev
• router/config.yaml
• router/main.py

---

📝 Walkthrough

Walkthrough

The router adds environment-driven Llama server connectivity, a dedicated HTTP client and health checks, plus Langfuse v4-compatible trace propagation, finalization, and flushing across AGY and LiteLLM success, error, streaming, and cancellation paths.

Changes

Llama routing and Langfuse lifecycle

Layer / File(s)	Summary
Llama configuration and health connectivity .env.dev, router/config.yaml, router/main.py	Adds LLAMA_SERVER_URL, resolves it through configuration, introduces a dedicated Llama client, routes metrics through it, and updates dashboard health checks.
Langfuse finalization and trace initialization router/main.py	Adds safe parent/child observation helpers, propagation context handling, request identity metadata, asynchronous flushing, and unknown-model cleanup.
AGY trace finalization paths router/main.py	Finalizes native and simulated AGY traces across streaming, non-streaming, success, error, and cancellation paths.
LiteLLM proxy trace finalization paths router/main.py	Finalizes LiteLLM traces exactly once across streaming, non-streaming, non-200, exception, and cancellation paths.

Estimated code review effort: 4 (Complex) | ~60 minutes

Possibly related PRs

• sheepdestroyer/LLM-Routing#47: Updates overlapping Llama server URL configuration and llama.cpp routing call sites.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch pr-285

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[gemini-code-assist]

Code Review

This pull request introduces robust Langfuse tracing integration, including session and user ID propagation, helper functions for safely finalizing parent observations and child spans, and asynchronous event flushing. It also refactors the llama-server configuration to resolve URLs dynamically and introduces a dedicated get_llama_client singleton with configurable TLS verification. The feedback highlights a critical context leak in router/main.py where the context manager _prop_ctx is entered manually but never exited on successful execution paths or during certain exceptions.

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.

[gemini-code-assist]

The context manager _prop_ctx (returned by propagate_attributes) is entered manually using _prop_ctx.__enter__() but is never exited on the successful execution path of the chat_completions endpoint, nor is it exited when raising an HTTPException (e.g., for unknown models). This causes a context leak where the session_id, user_id, and tags of one request can persist and affect subsequent requests processed by the same worker thread or asyncio task context.

To fix this cleanly without having to indent the entire 900-line function, you can extract the core logic of chat_completions into a helper function (e.g., _chat_completions_impl) and wrap the call in a with block:

if propagate_attributes and (_trace_session_id or _trace_user_id):
    with propagate_attributes(
        session_id=_trace_session_id or None,
        user_id=_trace_user_id or None,
        tags=[os.getenv("ENVIRONMENT", "production"), "llm-routing"],
    ):
        return await _chat_completions_impl(request, body, ...)
else:
    return await _chat_completions_impl(request, body, ...)

[sheepdestroyer]

Closing in favor of a fresh PR with Gemini Code Assist review fix applied (propagate_attributes context leak). See new PR.