From 5d06463b00f6a3f4a375939e1d59b2de600423b9 Mon Sep 17 00:00:00 2001 From: boy Date: Mon, 13 Jul 2026 04:08:02 +0200 Subject: [PATCH 1/7] fix: langfuse v4 SDK compatibility + parent observation lifecycle - 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 --- router/main.py | 128 +++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 119 insertions(+), 9 deletions(-) diff --git a/router/main.py b/router/main.py index 35844949..7f6e205c 100644 --- a/router/main.py +++ b/router/main.py @@ -293,6 +293,47 @@ def get_langfuse(): return _langfuse_client if _langfuse_client is not False else None +def _end_parent_obs(parent_obs, output=None, metadata=None) -> None: + """Safely finalize a Langfuse parent observation (SDK v4: update + end). + + Non-fatal — swallows all exceptions. + """ + if parent_obs is None: + return + try: + if output is not None or metadata is not None: + parent_obs.update(output=output, metadata=metadata) + parent_obs.end() + except Exception: + pass + + +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. + """ + 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 + + +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 + + async def push_aggregate_scores(): """Push aggregate KPIs as Langfuse scores every 5 minutes.""" while True: @@ -1106,7 +1147,7 @@ async def classify_request( if response.status_code != 200: if class_span_obj: try: - class_span_obj.end( + _end_child_span(class_span_obj, output={ "status": response.status_code, "error": "classification_failed", @@ -1143,7 +1184,7 @@ async def classify_request( # Finalize classifier child span if class_span_obj: try: - class_span_obj.end( + _end_child_span(class_span_obj, output={"tier": decision, "raw": raw_result}, metadata={"latency_ms": latency}, ) @@ -1878,6 +1919,21 @@ async def chat_completions(request: Request): client_model = body.get("model", "llm-routing-auto-free") + # Extract session_id and user_id for Langfuse tracing + _trace_session_id = ( + body.get("session_id") + or body.get("session") + or request.headers.get("x-session-id") + ) + if _trace_session_id: + _trace_session_id = str(_trace_session_id) + _trace_user_id = ( + body.get("user") + or request.headers.get("x-user-id") + ) + if _trace_user_id: + _trace_user_id = str(_trace_user_id) + # --- Langfuse parent trace: create early so child spans can reference it --- langfuse_trace_id = None parent_obs = None @@ -1892,6 +1948,13 @@ async def chat_completions(request: Request): name=f"triage-{client_model}", input=last_user_message[:200], level="DEFAULT", + metadata={ + "client_model": client_model, + "environment": os.getenv("ENVIRONMENT", "production"), + "session_id": _trace_session_id or "", + "user_id": _trace_user_id or "", + "tags": [os.getenv("ENVIRONMENT", "production"), "llm-routing"], + }, ) except Exception as e: logger.warning(f"Langfuse trace init failed (non-fatal): {e}") @@ -1922,6 +1985,9 @@ async def chat_completions(request: Request): f"Direct routing: Client requested '{client_model}', skipping classifier" ) else: + # guard: end parent obs before raising + _end_parent_obs(parent_obs, + output={"error": f"Unknown model: {client_model}"}) raise HTTPException( status_code=400, detail=f"Unknown model '{client_model}'. Use 'llm-routing-auto-free' for automatic routing, " @@ -2127,7 +2193,7 @@ async def native_agy_stream_generator(stream_gen, model_name): ) if agy_span_obj: try: - agy_span_obj.end( + _end_child_span(agy_span_obj, output={ "model": model_name, "tokens": token_count, @@ -2139,18 +2205,29 @@ async def native_agy_stream_generator(stream_gen, model_name): ) except Exception: pass + # 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}) + asyncio.create_task(_flush_langfuse_async()) except Exception as stream_err: logger.error( f"Error during native agy stream generation: {type(stream_err).__name__}" ) if agy_span_obj: try: - agy_span_obj.end( + _end_child_span(agy_span_obj, output={"error": type(stream_err).__name__}, metadata={"status": "failed"}, ) except Exception: pass + # End parent trace on stream error + _end_parent_obs(parent_obs, + output={"error": type(stream_err).__name__, + "route": "google_oauth_direct", "stream": True}) raise return StreamingResponse( @@ -2183,7 +2260,7 @@ async def native_agy_stream_generator(stream_gen, model_name): # Finalize agy span if agy_span_obj: try: - agy_span_obj.end( + _end_child_span(agy_span_obj, output={ "model": model_name, "tokens": completion_tokens, @@ -2246,16 +2323,30 @@ async def agy_stream_generator(): "utf-8" ) yield b"data: [DONE]\n\n" + # 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}) + asyncio.create_task(_flush_langfuse_async()) return StreamingResponse( agy_stream_generator(), media_type="text/event-stream" ) else: + # Finalize parent trace for non-streaming agy + _end_parent_obs(parent_obs, + output={"model": model_name, "tier": target_model, + "route": "google_oauth_direct"}, + metadata={"latency_ms": latency_ms, + "completion_tokens": completion_tokens}) + await _flush_langfuse_async() return agy_response except ImportError: if agy_span_obj: try: - agy_span_obj.end( + _end_child_span(agy_span_obj, output={"error": "module_not_available"}, metadata={"status": "skipped"}, ) @@ -2265,7 +2356,7 @@ async def agy_stream_generator(): except Exception as e: if agy_span_obj: try: - agy_span_obj.end( + _end_child_span(agy_span_obj, output={"error": type(e).__name__}, metadata={"status": "failed"}, ) @@ -2446,7 +2537,7 @@ async def stream_generator(): # Finalize LiteLLM span (streaming path) if litellm_span_obj: try: - litellm_span_obj.end( + _end_child_span(litellm_span_obj, output={"model": model_name, "stream": True}, metadata={ "latency_ms": proxy_latency, @@ -2455,8 +2546,19 @@ async def stream_generator(): ) except Exception: pass + # 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}) + asyncio.create_task(_flush_langfuse_async()) except Exception as ex: logger.error(f"Stream error: {ex}") + # 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}) if model_name.startswith("ollama-"): global _ollama_cooldown_until _ollama_cooldown_until = ( @@ -2525,7 +2627,7 @@ async def stream_generator(): # Finalize LiteLLM span (non-streaming path) if litellm_span_obj: try: - litellm_span_obj.end( + _end_child_span(litellm_span_obj, output={ "model": model_name, "tokens": completion_tokens, @@ -2534,6 +2636,14 @@ async def stream_generator(): ) except Exception: pass + # Finalize parent trace (non-streaming path) + _end_parent_obs(parent_obs, + output={"model": model_name, "tier": target_model, + "route": "litellm_fallback"}, + metadata={"latency_ms": proxy_latency, + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens}) + await _flush_langfuse_async() return resp_json else: logger.warning( From 655844db33bac2fae1b6c6959121c34d2e67aa04 Mon Sep 17 00:00:00 2001 From: boy Date: Mon, 13 Jul 2026 04:22:29 +0200 Subject: [PATCH 2/7] =?UTF-8?q?fix:=20address=20all=20PR=20#296=20review?= =?UTF-8?q?=20comments=20=E2=80=94=20Langfuse=20trace=20lifecycle=20harden?= =?UTF-8?q?ing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 : + 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) --- router/main.py | 220 ++++++++++++++++++++++++++----------------------- 1 file changed, 119 insertions(+), 101 deletions(-) diff --git a/router/main.py b/router/main.py index 7f6e205c..0f672be8 100644 --- a/router/main.py +++ b/router/main.py @@ -301,8 +301,13 @@ def _end_parent_obs(parent_obs, output=None, metadata=None) -> None: if parent_obs is None: return try: - if output is not None or metadata is not None: - parent_obs.update(output=output, metadata=metadata) + 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() except Exception: pass @@ -311,13 +316,18 @@ def _end_parent_obs(parent_obs, output=None, metadata=None) -> None: 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. + 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) + 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() except Exception: pass @@ -1145,17 +1155,13 @@ async def classify_request( latency = (time.time() - start_time) * 1000.0 if response.status_code != 200: - if class_span_obj: - try: - _end_child_span(class_span_obj, - output={ - "status": response.status_code, - "error": "classification_failed", - }, - metadata={"latency_ms": latency}, - ) - except Exception: - pass + _end_child_span(class_span_obj, + output={ + "status": response.status_code, + "error": "classification_failed", + }, + metadata={"latency_ms": latency}, + ) logger.error( f"Classification failed with status {response.status_code}: {response.text}" ) @@ -1182,14 +1188,10 @@ async def classify_request( decision = "agent-advanced-core" # Finalize classifier child span - if class_span_obj: - try: - _end_child_span(class_span_obj, - output={"tier": decision, "raw": raw_result}, - metadata={"latency_ms": latency}, - ) - except Exception: - pass + _end_child_span(class_span_obj, + output={"tier": decision, "raw": raw_result}, + metadata={"latency_ms": latency}, + ) # Store in cache triage_cache[normalized_prompt] = (decision, time.time()) @@ -1988,6 +1990,7 @@ async def chat_completions(request: Request): # guard: end parent obs before raising _end_parent_obs(parent_obs, output={"error": f"Unknown model: {client_model}"}) + await _flush_langfuse_async() raise HTTPException( status_code=400, detail=f"Unknown model '{client_model}'. Use 'llm-routing-auto-free' for automatic routing, " @@ -2191,43 +2194,40 @@ async def native_agy_stream_generator(stream_gen, model_name): logger.info( f"✅ native agy stream succeeded: {model_name}, {latency_ms:.0f}ms" ) - if agy_span_obj: - try: - _end_child_span(agy_span_obj, - output={ - "model": model_name, - "tokens": token_count, - }, - metadata={ - "latency_ms": latency_ms, - "tier": target_model, - }, - ) - except Exception: - pass + _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}) - 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) except Exception as stream_err: logger.error( f"Error during native agy stream generation: {type(stream_err).__name__}" ) - if agy_span_obj: - try: - _end_child_span(agy_span_obj, - output={"error": type(stream_err).__name__}, - metadata={"status": "failed"}, - ) - except Exception: - pass + _end_child_span(agy_span_obj, + output={"error": type(stream_err).__name__}, + metadata={"status": "failed"}, + ) # 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) raise return StreamingResponse( @@ -2258,20 +2258,16 @@ async def native_agy_stream_generator(stream_gen, model_name): ) # Finalize agy span - if agy_span_obj: - try: - _end_child_span(agy_span_obj, - output={ - "model": model_name, - "tokens": completion_tokens, - }, - metadata={ - "latency_ms": latency_ms, - "tier": target_model, - }, - ) - except Exception: - pass + _end_child_span(agy_span_obj, + output={ + "model": model_name, + "tokens": completion_tokens, + }, + metadata={ + "latency_ms": latency_ms, + "tier": target_model, + }, + ) if is_stream_requested: # Robust fallback: simulate stream if we requested stream but got buffered response @@ -2329,7 +2325,9 @@ async def agy_stream_generator(): "tier": target_model, "route": "google_oauth_direct"}, metadata={"latency_ms": latency_ms, "completion_tokens": len(content) // 4}) - 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) return StreamingResponse( agy_stream_generator(), media_type="text/event-stream" @@ -2344,24 +2342,16 @@ async def agy_stream_generator(): await _flush_langfuse_async() return agy_response except ImportError: - if agy_span_obj: - try: - _end_child_span(agy_span_obj, - output={"error": "module_not_available"}, - metadata={"status": "skipped"}, - ) - except Exception: - pass + _end_child_span(agy_span_obj, + output={"error": "module_not_available"}, + metadata={"status": "skipped"}, + ) logger.warning("agy_proxy module not available, falling back to LiteLLM") except Exception as e: - if agy_span_obj: - try: - _end_child_span(agy_span_obj, - output={"error": type(e).__name__}, - metadata={"status": "failed"}, - ) - except Exception: - pass + _end_child_span(agy_span_obj, + output={"error": type(e).__name__}, + metadata={"status": "failed"}, + ) logger.error(f"agy proxy failed: {type(e).__name__}, falling back to LiteLLM") if target_model == "llm-routing-agy": @@ -2535,30 +2525,36 @@ async def stream_generator(): route="litellm_fallback", )) # Finalize LiteLLM span (streaming path) - if litellm_span_obj: - try: - _end_child_span(litellm_span_obj, - output={"model": model_name, "stream": True}, - metadata={ - "latency_ms": proxy_latency, - "tokens": completion_chars // 4, - }, - ) - except Exception: - pass + _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}) - 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) except Exception as ex: logger.error(f"Stream error: {ex}") + # End child span before parent on stream error (CodeRabbit: missing finalization) + _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}) + _flush_task = asyncio.create_task(_flush_langfuse_async()) + _background_tasks.add(_flush_task) + _flush_task.add_done_callback(_background_tasks.discard) if model_name.startswith("ollama-"): global _ollama_cooldown_until _ollama_cooldown_until = ( @@ -2585,6 +2581,15 @@ async def stream_generator(): f"LiteLLM stream failed ({r.status_code}): {error_body[:300]}" ) await r.aclose() + # Finalize traces before raising on stream connection failure + _end_child_span(litellm_span_obj, + output={"status": r.status_code, "error": "litellm_stream_failed"}, + metadata={"status": "failed"}, + ) + _end_parent_obs(parent_obs, + output={"error": f"HTTP {r.status_code}", "route": "litellm_fallback", + "stream": True}) + await _flush_langfuse_async() raise HTTPException( status_code=r.status_code, detail="LiteLLM upstream request failed", @@ -2625,17 +2630,13 @@ async def stream_generator(): route="litellm_fallback", )) # Finalize LiteLLM span (non-streaming path) - if litellm_span_obj: - try: - _end_child_span(litellm_span_obj, - output={ - "model": model_name, - "tokens": completion_tokens, - }, - metadata={"latency_ms": proxy_latency}, - ) - except Exception: - pass + _end_child_span(litellm_span_obj, + output={ + "model": model_name, + "tokens": completion_tokens, + }, + metadata={"latency_ms": proxy_latency}, + ) # Finalize parent trace (non-streaming path) _end_parent_obs(parent_obs, output={"model": model_name, "tier": target_model, @@ -2649,6 +2650,15 @@ async def stream_generator(): logger.warning( f"LiteLLM failed ({response.status_code}): {response.text[:300]}" ) + # Finalize traces before raising on non-200 response + _end_child_span(litellm_span_obj, + output={"status": response.status_code, "error": "litellm_upstream_failed"}, + metadata={"status": "failed"}, + ) + _end_parent_obs(parent_obs, + output={"error": f"HTTP {response.status_code}", + "route": "litellm_fallback"}) + await _flush_langfuse_async() raise HTTPException( status_code=response.status_code, detail="LiteLLM upstream request failed", @@ -2657,6 +2667,14 @@ async def stream_generator(): raise except Exception as exc: logger.error(f"httpx call failed: {exc}") + # Finalize traces before raising on proxy exception + _end_child_span(litellm_span_obj, + output={"error": type(exc).__name__}, + metadata={"status": "failed"}, + ) + _end_parent_obs(parent_obs, + output={"error": type(exc).__name__, "route": "litellm_fallback"}) + await _flush_langfuse_async() raise HTTPException( status_code=502, detail="Proxy call failed" ) from exc From 505d2ed6508663f52e1ec79d78dcbefaca4885a0 Mon Sep 17 00:00:00 2001 From: boy Date: Mon, 13 Jul 2026 04:33:41 +0200 Subject: [PATCH 3/7] fix: LLAMA_SERVER_URL via config placeholder + canonical URL (was hardcoded 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 --- .env.dev | 1 + router/config.yaml | 2 ++ router/main.py | 59 +++++++++++++++++++++++++++++++++++++++++++--- 3 files changed, 59 insertions(+), 3 deletions(-) diff --git a/.env.dev b/.env.dev index 10e3f22b..9dc4b1d0 100644 --- a/.env.dev +++ b/.env.dev @@ -28,3 +28,4 @@ ROUTER_IMAGE="localhost/llm-routing-dev:latest" # Shared infrastructure (same llama classifier across dev & prod) LLAMA_CLASSIFIER_URL="https://x570.vendeuvre.lan/llm-routing/llama/v1" +LLAMA_SERVER_URL="https://x570.vendeuvre.lan/llm-routing/llama" diff --git a/router/config.yaml b/router/config.yaml index 4cb81e48..80cb812f 100644 --- a/router/config.yaml +++ b/router/config.yaml @@ -9,6 +9,8 @@ router: api_key: "os.environ/ROUTER_API_KEY" model: "qwen-4b-routing" +llama_server_url: "os.environ/LLAMA_SERVER_URL" + classification_rules: system_prompt: | Classify the coding task complexity. Output ONLY the tier name. diff --git a/router/main.py b/router/main.py index 0f672be8..775f1973 100644 --- a/router/main.py +++ b/router/main.py @@ -28,7 +28,6 @@ from typing import Dict, Optional, Union LITELLM_URL = (os.getenv("LITELLM_ADMIN_URL") or f"http://127.0.0.1:{os.getenv('LITELLM_PORT') or '4000'}").rstrip("/") -LLAMA_SERVER_URL = (os.getenv("LLAMA_SERVER_URL") or "http://127.0.0.1:8080").rstrip("/") LANGFUSE_HOST = (os.getenv("LANGFUSE_HOST") or f"http://127.0.0.1:{os.getenv('LANGFUSE_WEB_PORT') or '3001'}").rstrip("/") GEMINI_OAUTH_CREDS_PATH = "/config/gemini_auth/oauth_creds.json" @@ -135,6 +134,35 @@ def get_classifier_client(): return _classifier_client +_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 + + # Compiled regular expressions for token estimation heuristics WORD_RE = re.compile(r'[a-zA-Z0-9]+') NON_ASCII_RE = re.compile(r'[^\s\x00-\x7F]') @@ -466,6 +494,21 @@ async def push_aggregate_scores(): system_prompt = config.get("classification_rules", {}).get("system_prompt", "") backends = {b["name"]: b for b in config.get("backends", [])} +# --- Resolve llama_server_url from config (os.environ/ pattern) --- +_raw_llama_url = config.get("llama_server_url") or "http://127.0.0.1:8080" +if isinstance(_raw_llama_url, str) and _raw_llama_url.startswith("os.environ/"): + env_var = _raw_llama_url.split("/", 1)[1] + _raw_llama_url = os.environ.get(env_var, "") + if not _raw_llama_url: + if "pytest" in sys.modules: + _raw_llama_url = "http://127.0.0.1:8080" + else: + logger.warning( + 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("/") + # Default colors for tool visualization badges and charts TOOL_COLORS = { "tree": "#34d399", # Green @@ -1063,6 +1106,16 @@ async def check_http_endpoint(url: str) -> bool: return False +async def _check_llama_health() -> bool: + """Check llama-server health using the llama client (verify=False for self-signed TLS).""" + try: + client = get_llama_client() + r = await client.get(f"{LLAMA_SERVER_URL}/health", timeout=3.0) + return r.status_code < 500 + except Exception: + return False + + async def classify_request( prompt: str, bypass_cache: bool = False, langfuse_trace_id: str | None = None ) -> tuple[str, float, bool, str]: @@ -1491,7 +1544,7 @@ async def get_llamacpp_metrics() -> dict: """Fetches live model inventory and slot statistics from the local llama-server.""" result = {"models": [], "slots": [], "build": "unknown"} try: - client = get_http_client() + client = get_llama_client() # Fetch model list r = await client.get(f"{LLAMA_SERVER_URL}/v1/models", timeout=3.0) if r.status_code == 200: @@ -2886,7 +2939,7 @@ async def get_dashboard_data(): asyncio.wait_for(sync_cooldowns_from_valkey(), timeout=2.0), check_tcp_port("127.0.0.1", _valkey_port()), check_http_endpoint(f"http://127.0.0.1:{os.getenv('LITELLM_PORT') or '4000'}/"), - check_http_endpoint(f"{LLAMA_SERVER_URL}/health"), + asyncio.wait_for(_check_llama_health(), timeout=3.0), check_http_endpoint(f"http://127.0.0.1:{os.getenv('LANGFUSE_WEB_PORT') or '3001'}"), get_gemini_oauth_status(), asyncio.wait_for(get_best_free_model(), timeout=5.0), From d763ddb331cbd7305ff730a697d44e96351bf5be Mon Sep 17 00:00:00 2001 From: boy Date: Mon, 13 Jul 2026 04:53:57 +0200 Subject: [PATCH 4/7] fix: handle GeneratorExit in streaming generators + use propagate_attributes 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. --- router/main.py | 131 +++++++++++++++++++++++++++++++++++-------------- 1 file changed, 94 insertions(+), 37 deletions(-) diff --git a/router/main.py b/router/main.py index 775f1973..3bbf6765 100644 --- a/router/main.py +++ b/router/main.py @@ -27,6 +27,11 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator, RootModel from typing import Dict, Optional, Union +try: + from langfuse import propagate_attributes # noqa: F401 +except ImportError: + propagate_attributes = None + LITELLM_URL = (os.getenv("LITELLM_ADMIN_URL") or f"http://127.0.0.1:{os.getenv('LITELLM_PORT') or '4000'}").rstrip("/") LANGFUSE_HOST = (os.getenv("LANGFUSE_HOST") or f"http://127.0.0.1:{os.getenv('LANGFUSE_WEB_PORT') or '3001'}").rstrip("/") @@ -1992,12 +1997,21 @@ async def chat_completions(request: Request): # --- Langfuse parent trace: create early so child spans can reference it --- langfuse_trace_id = None parent_obs = None + _prop_ctx = None lf = get_langfuse() if lf: try: langfuse_trace_id = lf.create_trace_id( seed=f"triage_{stats['total_requests']}" ) + # Propagate session_id/user_id via Langfuse's native session mechanism + if propagate_attributes and (_trace_session_id or _trace_user_id): + _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__() parent_obs = lf.start_observation( trace_context={"trace_id": langfuse_trace_id}, name=f"triage-{client_model}", @@ -2006,15 +2020,18 @@ async def chat_completions(request: Request): metadata={ "client_model": client_model, "environment": os.getenv("ENVIRONMENT", "production"), - "session_id": _trace_session_id or "", - "user_id": _trace_user_id or "", - "tags": [os.getenv("ENVIRONMENT", "production"), "llm-routing"], }, ) except Exception as e: logger.warning(f"Langfuse trace init failed (non-fatal): {e}") langfuse_trace_id = None parent_obs = None + if _prop_ctx: + try: + _prop_ctx.__exit__(None, None, None) + except Exception: + pass + _prop_ctx = None if client_model in AUTO_MODELS or client_model == "llm-routing-ollama": # Full pipeline: classify → route to best tier @@ -2187,6 +2204,7 @@ async def native_agy_stream_generator(stream_gen, model_name): created_time = int(time.time()) chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" token_count = 0 + finalized = False try: async for token in stream_gen: if not token: @@ -2266,6 +2284,7 @@ async def native_agy_stream_generator(stream_gen, model_name): _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__}" @@ -2281,7 +2300,20 @@ async def native_agy_stream_generator(stream_gen, model_name): _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: + _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) return StreamingResponse( native_agy_stream_generator( @@ -2335,9 +2367,29 @@ async def agy_stream_generator(): created_time = int(time.time()) chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" chunk_size = 40 - for i in range(0, len(content), chunk_size): - chunk_text = content[i : i + chunk_size] - chunk_data = { + finalized = False + try: + for i in range(0, len(content), chunk_size): + chunk_text = content[i : i + chunk_size] + chunk_data = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": created_time, + "model": model_name, + "choices": [ + { + "index": 0, + "delta": {"content": chunk_text}, + "finish_reason": None, + } + ], + } + yield f"data: {json.dumps(chunk_data)}\n\n".encode( + "utf-8" + ) + await asyncio.sleep(0.005) + + finish_data = { "id": chunk_id, "object": "chat.completion.chunk", "created": created_time, @@ -2345,42 +2397,33 @@ async def agy_stream_generator(): "choices": [ { "index": 0, - "delta": {"content": chunk_text}, - "finish_reason": None, + "delta": {}, + "finish_reason": "stop", } ], } - yield f"data: {json.dumps(chunk_data)}\n\n".encode( + yield f"data: {json.dumps(finish_data)}\n\n".encode( "utf-8" ) - await asyncio.sleep(0.005) - - finish_data = { - "id": chunk_id, - "object": "chat.completion.chunk", - "created": created_time, - "model": model_name, - "choices": [ - { - "index": 0, - "delta": {}, - "finish_reason": "stop", - } - ], - } - yield f"data: {json.dumps(finish_data)}\n\n".encode( - "utf-8" - ) - yield b"data: [DONE]\n\n" - # 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) + yield b"data: [DONE]\n\n" + # 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) + finalized = True + finally: + if not finalized: + _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) return StreamingResponse( agy_stream_generator(), media_type="text/event-stream" @@ -2536,6 +2579,7 @@ async def stream_generator(): request_tokens = estimate_prompt_tokens(body_to_send) sse_buffer = "" decoder = codecs.getincrementaldecoder("utf-8")() + finalized = False try: async for chunk in r.aiter_bytes(): yield chunk @@ -2594,6 +2638,7 @@ async def stream_generator(): _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 before parent on stream error (CodeRabbit: missing finalization) @@ -2608,6 +2653,7 @@ async def stream_generator(): _flush_task = asyncio.create_task(_flush_langfuse_async()) _background_tasks.add(_flush_task) _flush_task.add_done_callback(_background_tasks.discard) + finalized = True if model_name.startswith("ollama-"): global _ollama_cooldown_until _ollama_cooldown_until = ( @@ -2623,6 +2669,17 @@ async def stream_generator(): f"Failed to save cooldowns to Valkey: {save_err}" ) finally: + if not finalized: + _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() return StreamingResponse( From e832600f21bfd4ae48ae528d40ca9c3584446df1 Mon Sep 17 00:00:00 2001 From: boy Date: Mon, 13 Jul 2026 05:12:36 +0200 Subject: [PATCH 5/7] fix: ensure _prop_ctx context manager is always exited (Gemini Code Assist 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 --- router/main.py | 34 +++++++++++++++++++++++++++++----- 1 file changed, 29 insertions(+), 5 deletions(-) diff --git a/router/main.py b/router/main.py index 3bbf6765..c3fb3d75 100644 --- a/router/main.py +++ b/router/main.py @@ -377,6 +377,20 @@ async def _flush_langfuse_async() -> None: pass +def _close_prop_ctx(prop_ctx): + """Safely exit a propagate_attributes context manager if active. + + Non-fatal — swallows all exceptions. + Returns None after exit for idempotent cleanup. + """ + if prop_ctx is not None: + try: + prop_ctx.__exit__(None, None, None) + except Exception: + pass + return None + + async def push_aggregate_scores(): """Push aggregate KPIs as Langfuse scores every 5 minutes.""" while True: @@ -2027,11 +2041,7 @@ async def chat_completions(request: Request): langfuse_trace_id = None parent_obs = None if _prop_ctx: - try: - _prop_ctx.__exit__(None, None, None) - except Exception: - pass - _prop_ctx = None + _prop_ctx = _close_prop_ctx(_prop_ctx) if client_model in AUTO_MODELS or client_model == "llm-routing-ollama": # Full pipeline: classify → route to best tier @@ -2060,6 +2070,7 @@ async def chat_completions(request: Request): # guard: end parent obs before raising _end_parent_obs(parent_obs, output={"error": f"Unknown model: {client_model}"}) + _close_prop_ctx(_prop_ctx) await _flush_langfuse_async() raise HTTPException( status_code=400, @@ -2281,6 +2292,7 @@ async def native_agy_stream_generator(stream_gen, model_name): "tier": target_model, "route": "google_oauth_direct"}, metadata={"latency_ms": latency_ms, "completion_tokens": token_count}) + _close_prop_ctx(_prop_ctx) _flush_task = asyncio.create_task(_flush_langfuse_async()) _background_tasks.add(_flush_task) _flush_task.add_done_callback(_background_tasks.discard) @@ -2297,6 +2309,7 @@ async def native_agy_stream_generator(stream_gen, model_name): _end_parent_obs(parent_obs, output={"error": type(stream_err).__name__, "route": "google_oauth_direct", "stream": True}) + _close_prop_ctx(_prop_ctx) _flush_task = asyncio.create_task(_flush_langfuse_async()) _background_tasks.add(_flush_task) _flush_task.add_done_callback(_background_tasks.discard) @@ -2311,6 +2324,7 @@ async def native_agy_stream_generator(stream_gen, model_name): _end_parent_obs(parent_obs, output={"error": "cancelled", "route": "google_oauth_direct", "stream": True}) + _close_prop_ctx(_prop_ctx) _flush_task = asyncio.create_task(_flush_langfuse_async()) _background_tasks.add(_flush_task) _flush_task.add_done_callback(_background_tasks.discard) @@ -2412,6 +2426,7 @@ async def agy_stream_generator(): "tier": target_model, "route": "google_oauth_direct"}, metadata={"latency_ms": latency_ms, "completion_tokens": len(content) // 4}) + _close_prop_ctx(_prop_ctx) _flush_task = asyncio.create_task(_flush_langfuse_async()) _background_tasks.add(_flush_task) _flush_task.add_done_callback(_background_tasks.discard) @@ -2421,6 +2436,7 @@ async def agy_stream_generator(): _end_parent_obs(parent_obs, output={"error": "cancelled", "route": "google_oauth_direct", "stream": True}) + _close_prop_ctx(_prop_ctx) _flush_task = asyncio.create_task(_flush_langfuse_async()) _background_tasks.add(_flush_task) _flush_task.add_done_callback(_background_tasks.discard) @@ -2435,6 +2451,7 @@ async def agy_stream_generator(): "route": "google_oauth_direct"}, metadata={"latency_ms": latency_ms, "completion_tokens": completion_tokens}) + _close_prop_ctx(_prop_ctx) await _flush_langfuse_async() return agy_response except ImportError: @@ -2635,6 +2652,7 @@ async def stream_generator(): "tier": target_model, "route": "litellm_fallback"}, metadata={"latency_ms": proxy_latency, "completion_tokens": completion_chars // 4}) + _close_prop_ctx(_prop_ctx) _flush_task = asyncio.create_task(_flush_langfuse_async()) _background_tasks.add(_flush_task) _flush_task.add_done_callback(_background_tasks.discard) @@ -2650,6 +2668,7 @@ async def stream_generator(): _end_parent_obs(parent_obs, output={"error": type(ex).__name__, "route": "litellm_fallback", "stream": True}) + _close_prop_ctx(_prop_ctx) _flush_task = asyncio.create_task(_flush_langfuse_async()) _background_tasks.add(_flush_task) _flush_task.add_done_callback(_background_tasks.discard) @@ -2677,6 +2696,7 @@ async def stream_generator(): _end_parent_obs(parent_obs, output={"error": "cancelled", "route": "litellm_fallback", "stream": True}) + _close_prop_ctx(_prop_ctx) _flush_task = asyncio.create_task(_flush_langfuse_async()) _background_tasks.add(_flush_task) _flush_task.add_done_callback(_background_tasks.discard) @@ -2699,6 +2719,7 @@ async def stream_generator(): _end_parent_obs(parent_obs, output={"error": f"HTTP {r.status_code}", "route": "litellm_fallback", "stream": True}) + _close_prop_ctx(_prop_ctx) await _flush_langfuse_async() raise HTTPException( status_code=r.status_code, @@ -2754,6 +2775,7 @@ async def stream_generator(): metadata={"latency_ms": proxy_latency, "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens}) + _close_prop_ctx(_prop_ctx) await _flush_langfuse_async() return resp_json else: @@ -2768,6 +2790,7 @@ async def stream_generator(): _end_parent_obs(parent_obs, output={"error": f"HTTP {response.status_code}", "route": "litellm_fallback"}) + _close_prop_ctx(_prop_ctx) await _flush_langfuse_async() raise HTTPException( status_code=response.status_code, @@ -2784,6 +2807,7 @@ async def stream_generator(): ) _end_parent_obs(parent_obs, output={"error": type(exc).__name__, "route": "litellm_fallback"}) + _close_prop_ctx(_prop_ctx) await _flush_langfuse_async() raise HTTPException( status_code=502, detail="Proxy call failed" From 07d91c83aec4089e2e8484975987eb536cd4b469 Mon Sep 17 00:00:00 2001 From: boy Date: Mon, 13 Jul 2026 05:22:10 +0200 Subject: [PATCH 6/7] feat: add Langfuse session propagation test to canonical verification - 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 --- .../verify_canonical_endpoints.py | 161 ++++++++++++++++++ 1 file changed, 161 insertions(+) diff --git a/scripts/verification/verify_canonical_endpoints.py b/scripts/verification/verify_canonical_endpoints.py index 23be879f..1c13feeb 100644 --- a/scripts/verification/verify_canonical_endpoints.py +++ b/scripts/verification/verify_canonical_endpoints.py @@ -71,6 +71,8 @@ def _parse(path: Path): "public_base_url": (os.environ.get("PUBLIC_BASE_URL") or env.get("PUBLIC_BASE_URL") or "").rstrip("/"), "minio_s3_port": os.environ.get("MINIO_S3_PORT") or env.get("MINIO_S3_PORT") or "9002", "clickhouse_http_port": os.environ.get("CLICKHOUSE_HTTP_PORT") or env.get("CLICKHOUSE_HTTP_PORT") or "8123", + "langfuse_public_key": os.environ.get("LANGFUSE_PUBLIC_KEY") or env.get("LANGFUSE_PUBLIC_KEY") or "", + "langfuse_secret_key": os.environ.get("LANGFUSE_SECRET_KEY") or env.get("LANGFUSE_SECRET_KEY") or "", } @@ -365,6 +367,164 @@ def test_litellm_direct_chat(cfg: dict) -> tuple[int, int]: return passed, total +def test_langfuse_session_propagation(cfg: dict) -> tuple[int, int]: + """Verify Langfuse session/user propagation and no cross-request leak. + + 1. Send a chat request WITH session_id + user, verify trace has them. + 2. Send a second chat request WITHOUT session info, verify NO leak. + Returns (passed, total). + """ + router_base = f"http://127.0.0.1:{cfg['router_port']}" + lf_host = cfg["langfuse_web_port"] + lf_base = f"http://127.0.0.1:{lf_host}" + pk = cfg["langfuse_public_key"] + sk = cfg["langfuse_secret_key"] + key = cfg["router_api_key"] + + if not pk or not sk: + print( + "\n── Langfuse session propagation — SKIPPED " + "(LANGFUSE_PUBLIC_KEY/SECRET_KEY not set) ──" + ) + return 0, 0 + + auth = (pk, sk) + headers = { + "Authorization": f"Bearer {key}", + "Content-Type": "application/json", + } + session_id = f"e2e-verify-{os.urandom(4).hex()}" + user_id = "verify-script" + passed = total = 0 + + print("\n── Langfuse session propagation ──") + + # --- Step 1: Request WITH session --- + total += 1 + payload = { + "model": "agent-simple-core", + "messages": [{"role": "user", "content": "Say 'test' and nothing else."}], + "max_tokens": 5, + "session_id": session_id, + "user": user_id, + } + try: + r = httpx.post( + f"{router_base}/v1/chat/completions", + json=payload, + headers=headers, + timeout=120, + ) + if r.status_code != 200: + passed += check( + "Session request (200)", False, + f"HTTP {r.status_code}: {r.text[:100]}" + ) + return passed, total + passed += check("Session request (200)", True, session_id[:16]) + except Exception as e: + passed += check("Session request (200)", False, str(e)[:100]) + return passed, total + + # Wait for Langfuse flush (async background task) + time.sleep(2) + + # Query Langfuse for trace with our session + total += 1 + session_trace_id = None + try: + resp = httpx.get( + f"{lf_base}/api/public/traces?page=1&limit=20&orderBy=timestamp.desc", + auth=auth, timeout=10, + ) + if resp.status_code != 200: + passed += check( + "Trace has session", False, + f"Langfuse API HTTP {resp.status_code}" + ) + return passed, total + traces = resp.json().get("data", []) + for t in traces: + if t.get("sessionId") == session_id and t.get("userId") == user_id: + session_trace_id = t["id"] + passed += check( + "Trace has session+user", + True, + f"session={session_id[:12]} user={user_id}", + ) + break + if not session_trace_id: + passed += check( + "Trace has session+user", False, + f"session={session_id[:12]} not found in {len(traces)} recent traces" + ) + except Exception as e: + passed += check("Trace has session+user", False, str(e)[:100]) + + # --- Step 2: Request WITHOUT session (leak test) --- + total += 1 + payload_no_session = { + "model": "agent-simple-core", + "messages": [{"role": "user", "content": "Say 'leaktest' and nothing else."}], + "max_tokens": 5, + } + try: + r2 = httpx.post( + f"{router_base}/v1/chat/completions", + json=payload_no_session, + headers=headers, + timeout=120, + ) + if r2.status_code != 200: + passed += check( + "No-session request (200)", False, + f"HTTP {r2.status_code}: {r2.text[:100]}" + ) + return passed, total + passed += check("No-session request (200)", True, "clean") + except Exception as e: + passed += check("No-session request (200)", False, str(e)[:100]) + return passed, total + + # Wait for flush + time.sleep(2) + + # Verify second request's traces have no session (exclude the known + # session trace from step 1 — it may still be in recent results) + total += 1 + try: + resp2 = httpx.get( + f"{lf_base}/api/public/traces?page=1&limit=10&orderBy=timestamp.desc", + auth=auth, timeout=10, + ) + if resp2.status_code != 200: + passed += check( + "No session leak", False, + f"Langfuse API HTTP {resp2.status_code}" + ) + return passed, total + traces2 = resp2.json().get("data", []) + # Filter out the known session trace from step 1 + leaked = any( + t.get("sessionId") == session_id and t.get("id") != session_trace_id + for t in traces2 + ) + if leaked: + passed += check( + "No session leak", False, + "Previous session leaked into no-session request!" + ) + else: + passed += check( + "No session leak", True, + f"{len(traces2)} recent traces clean (excluded known session trace)" + ) + except Exception as e: + passed += check("No session leak", False, str(e)[:100]) + + return passed, total + + def test_canonical_urls(cfg: dict) -> tuple[int, int, int]: """Verify canonical HTTPS URLs are reachable (if PUBLIC_BASE_URL is set). Returns (passed, total, skipped).""" @@ -460,6 +620,7 @@ def main(): ("Infrastructure", test_infra_health), ("E2E chat (router)", test_e2e_chat), ("LiteLLM direct chat", test_litellm_direct_chat), + ("Langfuse session propagation", test_langfuse_session_propagation), ("Canonical URLs", test_canonical_urls), ]: result = fn(cfg) From a823051f0f3a06c3f339cbe51ef8aabb98ba98b6 Mon Sep 17 00:00:00 2001 From: boy Date: Mon, 13 Jul 2026 05:37:58 +0200 Subject: [PATCH 7/7] =?UTF-8?q?fix:=20address=20Gemini=20review=20?= =?UTF-8?q?=E2=80=94=20contextvar=20task=20isolation,=20excessive=20flush,?= =?UTF-8?q?=20llama=20client=20leak?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- router/main.py | 114 ++++++++++++++++++++++++------------------------- 1 file changed, 57 insertions(+), 57 deletions(-) diff --git a/router/main.py b/router/main.py index c3fb3d75..36958366 100644 --- a/router/main.py +++ b/router/main.py @@ -12,7 +12,7 @@ import yaml import httpx import redis.asyncio as aioredis -from contextlib import asynccontextmanager +from contextlib import asynccontextmanager, nullcontext from dataclasses import dataclass from fastapi import FastAPI, Request, HTTPException, Response @@ -366,17 +366,6 @@ def _end_child_span(span, output=None, metadata=None) -> None: pass -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 - - def _close_prop_ctx(prop_ctx): """Safely exit a propagate_attributes context manager if active. @@ -1084,6 +1073,12 @@ async def lifespan(app: FastAPI): await _classifier_client.aclose() _classifier_client = None + # Close llama client + global _llama_client + if _llama_client is not None: + await _llama_client.aclose() + _llama_client = None + # Close Redis client global _redis_client if _redis_client is not None and _redis_client is not False: @@ -2012,20 +2007,25 @@ async def chat_completions(request: Request): langfuse_trace_id = None parent_obs = None _prop_ctx = None + _is_streaming = body.get("stream", False) lf = get_langfuse() if lf: try: langfuse_trace_id = lf.create_trace_id( seed=f"triage_{stats['total_requests']}" ) - # Propagate session_id/user_id via Langfuse's native session mechanism + # Propagate session_id/user_id via Langfuse's native session mechanism. + # For non-streaming: enter here (same asyncio task, contextvars work). + # For streaming: each generator creates its own context in its own task + # because OpenTelemetry contextvars are task-isolated. if propagate_attributes and (_trace_session_id or _trace_user_id): - _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__() + if not _is_streaming: + _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__() parent_obs = lf.start_observation( trace_context={"trace_id": langfuse_trace_id}, name=f"triage-{client_model}", @@ -2071,7 +2071,6 @@ async def chat_completions(request: Request): _end_parent_obs(parent_obs, output={"error": f"Unknown model: {client_model}"}) _close_prop_ctx(_prop_ctx) - await _flush_langfuse_async() raise HTTPException( status_code=400, detail=f"Unknown model '{client_model}'. Use 'llm-routing-auto-free' for automatic routing, " @@ -2216,6 +2215,16 @@ async def native_agy_stream_generator(stream_gen, model_name): chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" token_count = 0 finalized = False + _native_agy_prop = ( + propagate_attributes( + session_id=_trace_session_id or None, + user_id=_trace_user_id or None, + tags=[os.getenv("ENVIRONMENT", "production"), "llm-routing"], + ) + if propagate_attributes and (_trace_session_id or _trace_user_id) + else nullcontext() + ) + _native_agy_prop.__enter__() try: async for token in stream_gen: if not token: @@ -2292,10 +2301,7 @@ async def native_agy_stream_generator(stream_gen, model_name): "tier": target_model, "route": "google_oauth_direct"}, metadata={"latency_ms": latency_ms, "completion_tokens": token_count}) - _close_prop_ctx(_prop_ctx) - _flush_task = asyncio.create_task(_flush_langfuse_async()) - _background_tasks.add(_flush_task) - _flush_task.add_done_callback(_background_tasks.discard) + _close_prop_ctx(_native_agy_prop) finalized = True except Exception as stream_err: logger.error( @@ -2309,10 +2315,7 @@ async def native_agy_stream_generator(stream_gen, model_name): _end_parent_obs(parent_obs, output={"error": type(stream_err).__name__, "route": "google_oauth_direct", "stream": True}) - _close_prop_ctx(_prop_ctx) - _flush_task = asyncio.create_task(_flush_langfuse_async()) - _background_tasks.add(_flush_task) - _flush_task.add_done_callback(_background_tasks.discard) + _close_prop_ctx(_native_agy_prop) finalized = True raise finally: @@ -2324,10 +2327,7 @@ async def native_agy_stream_generator(stream_gen, model_name): _end_parent_obs(parent_obs, output={"error": "cancelled", "route": "google_oauth_direct", "stream": True}) - _close_prop_ctx(_prop_ctx) - _flush_task = asyncio.create_task(_flush_langfuse_async()) - _background_tasks.add(_flush_task) - _flush_task.add_done_callback(_background_tasks.discard) + _close_prop_ctx(_native_agy_prop) return StreamingResponse( native_agy_stream_generator( @@ -2382,6 +2382,16 @@ async def agy_stream_generator(): chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" chunk_size = 40 finalized = False + _agy_gen_prop = ( + propagate_attributes( + session_id=_trace_session_id or None, + user_id=_trace_user_id or None, + tags=[os.getenv("ENVIRONMENT", "production"), "llm-routing"], + ) + if propagate_attributes and (_trace_session_id or _trace_user_id) + else nullcontext() + ) + _agy_gen_prop.__enter__() try: for i in range(0, len(content), chunk_size): chunk_text = content[i : i + chunk_size] @@ -2426,20 +2436,14 @@ async def agy_stream_generator(): "tier": target_model, "route": "google_oauth_direct"}, metadata={"latency_ms": latency_ms, "completion_tokens": len(content) // 4}) - _close_prop_ctx(_prop_ctx) - _flush_task = asyncio.create_task(_flush_langfuse_async()) - _background_tasks.add(_flush_task) - _flush_task.add_done_callback(_background_tasks.discard) + _close_prop_ctx(_agy_gen_prop) finalized = True finally: if not finalized: _end_parent_obs(parent_obs, output={"error": "cancelled", "route": "google_oauth_direct", "stream": True}) - _close_prop_ctx(_prop_ctx) - _flush_task = asyncio.create_task(_flush_langfuse_async()) - _background_tasks.add(_flush_task) - _flush_task.add_done_callback(_background_tasks.discard) + _close_prop_ctx(_agy_gen_prop) return StreamingResponse( agy_stream_generator(), media_type="text/event-stream" @@ -2452,7 +2456,6 @@ async def agy_stream_generator(): metadata={"latency_ms": latency_ms, "completion_tokens": completion_tokens}) _close_prop_ctx(_prop_ctx) - await _flush_langfuse_async() return agy_response except ImportError: _end_child_span(agy_span_obj, @@ -2597,6 +2600,16 @@ async def stream_generator(): sse_buffer = "" decoder = codecs.getincrementaldecoder("utf-8")() finalized = False + _litellm_gen_prop = ( + propagate_attributes( + session_id=_trace_session_id or None, + user_id=_trace_user_id or None, + tags=[os.getenv("ENVIRONMENT", "production"), "llm-routing"], + ) + if propagate_attributes and (_trace_session_id or _trace_user_id) + else nullcontext() + ) + _litellm_gen_prop.__enter__() try: async for chunk in r.aiter_bytes(): yield chunk @@ -2652,10 +2665,7 @@ async def stream_generator(): "tier": target_model, "route": "litellm_fallback"}, metadata={"latency_ms": proxy_latency, "completion_tokens": completion_chars // 4}) - _close_prop_ctx(_prop_ctx) - _flush_task = asyncio.create_task(_flush_langfuse_async()) - _background_tasks.add(_flush_task) - _flush_task.add_done_callback(_background_tasks.discard) + _close_prop_ctx(_litellm_gen_prop) finalized = True except Exception as ex: logger.error(f"Stream error: {ex}") @@ -2668,10 +2678,7 @@ async def stream_generator(): _end_parent_obs(parent_obs, output={"error": type(ex).__name__, "route": "litellm_fallback", "stream": True}) - _close_prop_ctx(_prop_ctx) - _flush_task = asyncio.create_task(_flush_langfuse_async()) - _background_tasks.add(_flush_task) - _flush_task.add_done_callback(_background_tasks.discard) + _close_prop_ctx(_litellm_gen_prop) finalized = True if model_name.startswith("ollama-"): global _ollama_cooldown_until @@ -2696,10 +2703,7 @@ async def stream_generator(): _end_parent_obs(parent_obs, output={"error": "cancelled", "route": "litellm_fallback", "stream": True}) - _close_prop_ctx(_prop_ctx) - _flush_task = asyncio.create_task(_flush_langfuse_async()) - _background_tasks.add(_flush_task) - _flush_task.add_done_callback(_background_tasks.discard) + _close_prop_ctx(_litellm_gen_prop) await r.aclose() return StreamingResponse( @@ -2720,7 +2724,6 @@ async def stream_generator(): output={"error": f"HTTP {r.status_code}", "route": "litellm_fallback", "stream": True}) _close_prop_ctx(_prop_ctx) - await _flush_langfuse_async() raise HTTPException( status_code=r.status_code, detail="LiteLLM upstream request failed", @@ -2776,7 +2779,6 @@ async def stream_generator(): "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens}) _close_prop_ctx(_prop_ctx) - await _flush_langfuse_async() return resp_json else: logger.warning( @@ -2791,7 +2793,6 @@ async def stream_generator(): output={"error": f"HTTP {response.status_code}", "route": "litellm_fallback"}) _close_prop_ctx(_prop_ctx) - await _flush_langfuse_async() raise HTTPException( status_code=response.status_code, detail="LiteLLM upstream request failed", @@ -2808,7 +2809,6 @@ async def stream_generator(): _end_parent_obs(parent_obs, output={"error": type(exc).__name__, "route": "litellm_fallback"}) _close_prop_ctx(_prop_ctx) - await _flush_langfuse_async() raise HTTPException( status_code=502, detail="Proxy call failed" ) from exc