Comprehensive Fixes: Worktree Leak Cleanups, LiteLLM Authentication, and Code Review Feedback#11
Comprehensive Fixes: Worktree Leak Cleanups, LiteLLM Authentication, and Code Review Feedback#11sheepdestroyer wants to merge 34 commits into
Conversation
Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com>
Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com>
Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com>
Co-authored-by: sheepdestroyer <1377479+sheepdestroyer@users.noreply.github.com>
- Move _last_stats_save inside try block (only advance on successful write) - Add save_persisted_stats(force=True) + timeline flush in lifespan shutdown - Remove redundant import time inside save_persisted_stats - Fix misleading comment about independent throttle timers - Use time.monotonic() instead of time.time() for throttle timestamps - Add logger.warning to bare except in timeline write
- Add logging to bare except Exception: pass in shutdown timeline flush - Move asyncio.create_task(push_aggregate_scores()) before yield so it runs during app lifetime - Use atomic file writes (temp file + os.replace) in save_persisted_stats - Use atomic file writes in shutdown timeline flush and record_tool_usage timeline write - Restructure lifespan to single yield point (remove else: yield; return pattern) - Capture time.monotonic() once in record_tool_usage and reuse for guard + assignment
- Extract shared _atomic_write_json_sync / _atomic_write_json_async helpers - Make save_persisted_stats async, using _atomic_write_json_async with run_in_executor + deep-copy to prevent concurrent modification - Update all 4 async call sites (lifespan, classify_request x2, chat_completions) to await save_persisted_stats() - In record_tool_usage (sync), fire-and-forget both stats and timeline writes via loop.run_in_executor with deep-copied data; fall back to sync write when no event loop is running (early startup) - Replace duplicated atomic write logic in lifespan shutdown with shared helper
…jor and empty ignore blocks
…ner build, and docstrings
…sync, native streaming)
- router/main.py: wrap lifespan yield in try/finally so stats flush always runs; add _background_tasks set to prevent fire-and-forget task GC (RUF006) - router/static/visualizer.html: add keyboard accessibility (tabindex, role, onkeydown) to prompt list items; switch annotation keys from array index to stable djb2 prompt hash with backward-compatible index fallback - scripts/benchmark_classifier.py: safe tier field lookup (tier/llm_tier/clf_tier), guard per_tier keying to skip ERROR/unknown labels, safe choices[0] access - scripts/extract_complex.py: use forward iteration for first user message, add system-note filtering ([System: / [Note:]) - scripts/extract_gapfill.py: same as extract_complex.py - scripts/retry_errors.py: safe status.get() to prevent AttributeError, schema-aware ERROR detection (tier + clf_tier), only increment fixed on success - scripts/reclassify_all.py: safe choices[0] fallback to prevent IndexError - test_circuit_breaker.py: replace == True/False with idiomatic assert (14 sites) - verify_breaker.py: replace == True/False with idiomatic assert (5 sites)
…dation, scripts, and test suite cleanups
…le null content safely
…d fix SAST/lint comments
…oldown for llm-routing-ollama on rate limit
LiteLLM Community Edition's deployment cooldown is unreliable for single-deployment model groups and fallback-target groups. The previous approach (multi-deployment replicas, allowed_fails_policy) did not reliably cool down llm-routing-ollama, causing crashloops when Ollama was rate-limited. New approach: the triage router manages Ollama cooldowns internally. When Ollama fails (429/502/503), the router activates a 5-minute cooldown (configurable via OLLAMA_COOLDOWN_SECONDS env var). During cooldown, all Ollama requests are immediately rejected: - Auto modes: silently fall back to the free tier - Direct/fallback mode: return 429 so LiteLLM skips to openrouter-auto Changes: - router/main.py: Add _ollama_cooldown_until state, cooldown check before Ollama proxy call, and activation on failure. Add Prometheus metrics (ollama_cooldown_active, ollama_cooldown_remaining_seconds). - litellm/config.yaml: Remove test-fallback-model, remove duplicate llm-routing-ollama deployment (127.0.0.2), restore Ollama api_base to production (api.ollama.com), remove enterprise-only allowed_fails_policy. - README.md: Update sequence diagrams, Fallback and Cooldown Behavior section, and metrics table to document router-side cooldown.
…ent verification scripts
…ing in sed, expected model asserts in tests, synced metrics documentation
Reviewer's GuideImplements router-side Ollama cooldown and gated routing, cleans up path and credential leaks in deployment/scripts, wires LiteLLM master key from env instead of hardcoding, hardens httpx client lifecycle, and refreshes routing/telemetry docs and verification scripts. Sequence diagram for Ollama gated routing and router-side cooldownsequenceDiagram
actor Client
participant LiteLLM as LiteLLM_Gateway
participant Router as Triage_Router
participant Ollama as Ollama_API
participant OpenRouter as OpenRouter_auto
%% First request: llm-routing-ollama hits Ollama and fails
Client->>LiteLLM: POST /v1/chat/completions (model=llm-routing-ollama)
LiteLLM->>Router: POST /v1/chat/completions (model=llm-routing-ollama)
Router->>Router: classify_request
Router->>Router: map to ollama-deepseek-v4-pro / -flash
Router->>LiteLLM: execute_proxy(ollama-deepseek-v4-*)
LiteLLM->>Ollama: /api/chat (ollama_chat provider)
Ollama-->>LiteLLM: 429 / 5xx error
LiteLLM-->>Router: HTTP error
Router->>Router: activate _ollama_cooldown_until
Router-->>LiteLLM: HTTP 429 (Ollama cooled down)
LiteLLM-->>Client: HTTP 429
%% Second request: LiteLLM falls through to openrouter-auto
Client->>LiteLLM: POST /v1/chat/completions (model=agent-advanced-core)
LiteLLM->>Router: POST /v1/chat/completions (model=agent-advanced-core)
Router->>Router: should_try_ollama = True
Router->>Router: [now < _ollama_cooldown_until?]
alt Cooldown_active
Router-->>LiteLLM: HTTP 429 (Ollama backend cooled down)
LiteLLM->>OpenRouter: openrouter-auto fallback chain
OpenRouter-->>LiteLLM: completion
LiteLLM-->>Client: completion
end
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
|
Warning Review limit reached
More reviews will be available in 20 minutes and 12 seconds. Learn how PR review limits work. Your organization has used up its prepaid credits, and credit purchases are no longer available. Enable the review add-on in the billing tab to keep reviews running — you're only billed for reviews past your plan's rate limits ($0.25/file). ⌛ How to resolve this issue?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 credits. 🚦 How do rate limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan refill rate. 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, the refill rate gradually slows as usage increases. The highest same-day bursts are limited more strictly. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (10)
📝 WalkthroughWalkthroughThis PR introduces a router-side Ollama cooldown mechanism with 5-minute duration, persisted to Valkey/Redis across the triage router, circuit breaker, and agy_proxy. It refactors LiteLLM fallback chains, centralizes LiteLLM forwarding in ChangesOllama Cooldown & Valkey State Management
Portability, Deployment Hardening & Scripts Maintenance
Sequence Diagram(s)sequenceDiagram
participant Client
participant Triage as Triage Router (main.py)
participant Valkey
participant AGYProxy as agy_proxy.py
participant LiteLLM
participant Ollama
Client->>Triage: POST /v1/chat/completions (llm-routing-ollama or auto)
Triage->>Valkey: sync_cooldowns_from_valkey()
Valkey-->>Triage: _ollama_cooldown_until, circuit breaker state
alt Ollama cooldown active & direct mode
Triage-->>Client: HTTP 429
note over Triage,LiteLLM: LiteLLM falls through to openrouter-auto
else Ollama cooldown active & auto mode
Triage->>LiteLLM: forward to classified free-tier (silent fallback)
LiteLLM-->>Client: response
else No cooldown
Triage->>AGYProxy: try_agy_proxy (if AGY tier)
AGYProxy->>Valkey: sync_cooldowns_from_valkey()
AGYProxy->>Ollama: streaming or non-streaming request
alt Success
AGYProxy->>Valkey: save_cooldowns_to_valkey()
Ollama-->>Client: response
else Quota exhausted / failure
AGYProxy->>Valkey: save_cooldowns_to_valkey()
Triage->>Triage: activate _ollama_cooldown_until (+300s)
Triage->>Valkey: save_cooldowns_to_valkey()
Triage->>LiteLLM: execute_proxy() fallback
LiteLLM-->>Client: response
end
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
There was a problem hiding this comment.
Hey - I've left some high level feedback:
- In
start-stack.sh, you now substitute the LiteLLM master key intopod.yamlbut never validate thatLITELLM_MASTER_KEYis set and non-empty; consider failing fast with a clear error if the variable is missing to avoid deploying a gateway with an invalid auth key. - The new verification scripts handle authentication inconsistently (some parse
LITELLM_MASTER_KEYfrom.env, whileverify_ollama_routing.pyhardcodesgateway-pass); it would be more robust to centralize env parsing and always read the key from the same source to prevent drift between test scripts and production configuration.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In `start-stack.sh`, you now substitute the LiteLLM master key into `pod.yaml` but never validate that `LITELLM_MASTER_KEY` is set and non-empty; consider failing fast with a clear error if the variable is missing to avoid deploying a gateway with an invalid auth key.
- The new verification scripts handle authentication inconsistently (some parse `LITELLM_MASTER_KEY` from `.env`, while `verify_ollama_routing.py` hardcodes `gateway-pass`); it would be more robust to centralize env parsing and always read the key from the same source to prevent drift between test scripts and production configuration.Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
There was a problem hiding this comment.
Code Review
This pull request introduces a router-side Ollama cooldown mechanism and gated routing modes to prevent cascading fallback loops, along with path dynamicization, new verification scripts, and general code cleanups. The reviewer feedback highlights several opportunities to improve robustness and performance: storing the cooldown state in Redis (Valkey) to support multi-worker environments, reusing a shared httpx.AsyncClient to prevent socket exhaustion, refining prompt token estimation to avoid inflating metrics, propagating non-transient HTTP errors instead of masking them as 429s, and offloading synchronous file reads to a thread pool to prevent blocking the asyncio event loop.
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.
| _ollama_cooldown_until: float = 0.0 # monotonic timestamp when cooldown expires | ||
| OLLAMA_COOLDOWN_SECONDS: int = int(os.getenv("OLLAMA_COOLDOWN_SECONDS", "300")) # 5 min default |
There was a problem hiding this comment.
Using a global in-memory variable (_ollama_cooldown_until) to track the cooldown state is not robust in multi-worker environments (such as Uvicorn with multiple workers), as each worker process will maintain its own independent cooldown state. Since Valkey (Redis) is already part of the stack and running on port 6379, consider storing the cooldown state in Redis (e.g., using a key like ollama_cooldown with a TTL of 300 seconds). This ensures consistent cooldown behavior across all worker processes.
There was a problem hiding this comment.
Each free models and endpoints like agy & ollama should have cooldown values stored global, and updated opportunistically.
| client = None | ||
| should_close_client = True | ||
| try: | ||
| client = httpx.AsyncClient(timeout=3600.0) |
There was a problem hiding this comment.
Creating a new httpx.AsyncClient on every request is inefficient because it defeats connection pooling and can lead to socket exhaustion under high load. Consider reusing a single shared httpx.AsyncClient instance across requests by storing it in FastAPI's request.app.state during lifespan or on first access.
| client = None | |
| should_close_client = True | |
| try: | |
| client = httpx.AsyncClient(timeout=3600.0) | |
| if not hasattr(request.app.state, "http_client"): | |
| request.app.state.http_client = httpx.AsyncClient(timeout=3600.0) | |
| client = request.app.state.http_client | |
| should_close_client = False | |
| try: |
| finally: | ||
| await r.aclose() | ||
| await client.aclose() |
There was a problem hiding this comment.
| async def stream_generator(): | ||
| """Asynchronous generator that yields streaming chunks from LiteLLM completions response and logs usage stats on completion.""" | ||
| completion_chars = 0 | ||
| request_tokens = len(json.dumps(body_to_send)) // 4 |
There was a problem hiding this comment.
Estimating prompt tokens using len(json.dumps(body_to_send)) // 4 includes all JSON formatting, keys, and API parameters (like model, temperature, etc.), which artificially inflates the token metrics. Consider estimating based on the actual text content of the messages (e.g., sum(len(m.get('content') or '') for m in messages) // 4), similar to the implementation in native_agy_stream_generator.
| else: | ||
| error_body = await r.aread() if r else b"" | ||
| logger.warning(f"LiteLLM stream failed ({r.status_code}): {error_body[:300]}") | ||
| await r.aclose(); await client.aclose() | ||
| raise HTTPException(status_code=502, detail=f"LiteLLM failed: {r.status_code}") | ||
| else: | ||
| logger.info(f"Proxying to LiteLLM as model={model_name}") | ||
| response = await client.post(f"{backend_api_base}/chat/completions", json=body_to_send, headers=headers) | ||
| await client.aclose() | ||
| if response.status_code == 200: | ||
| proxy_latency = (time.time() - proxy_start) * 1000.0 | ||
| stats["total_proxy_time_ms"] += proxy_latency | ||
| stats["avg_proxy_latency_ms"] = stats["total_proxy_time_ms"] / stats["total_requests"] | ||
| resp_json = response.json() | ||
| usage = resp_json.get("usage", {}) | ||
| prompt_tokens = usage.get("prompt_tokens", len(json.dumps(body_to_send)) // 4) | ||
| completion_tokens = usage.get("completion_tokens", len(json.dumps(resp_json)) // 4) | ||
| record_tool_usage(active_tool, prompt_tokens, completion_tokens, model_name, proxy_latency, route="litellm_fallback") | ||
| # Finalize LiteLLM span (non-streaming path) | ||
| if litellm_span_obj: | ||
| try: | ||
| litellm_span_obj.end( | ||
| output={"model": model_name, "tokens": completion_tokens}, | ||
| metadata={"latency_ms": proxy_latency}, | ||
| ) | ||
| except Exception: | ||
| pass | ||
| return resp_json | ||
| # Direct/fallback llm-routing-ollama request: return 429 to | ||
| # signal LiteLLM to skip this model group in the fallback chain | ||
| logger.error(f"Ollama proxy failed ({e.detail}) for direct/fallback request, returning 429") | ||
| raise HTTPException( | ||
| status_code=429, | ||
| detail="Ollama backend rate limited/unavailable" | ||
| ) from e |
There was a problem hiding this comment.
Converting all HTTPException errors to a 429 for direct/fallback requests can mask critical non-transient errors (such as 401 Unauthorized or 400 Bad Request due to misconfiguration), making troubleshooting extremely difficult. Consider only raising a 429 if the original error status code indicates a transient failure (e.g., 429, 502, 503), and propagating other status codes as-is.
| with open(ann_path, "r", encoding="utf-8") as f: | ||
| existing = json.load(f) |
There was a problem hiding this comment.
Performing synchronous file reads (open and json.load) inside an async function blocks the asyncio event loop, which can degrade performance and increase latency for concurrent requests. Consider offloading the file read to a thread pool executor using asyncio.get_running_loop().run_in_executor, similar to how _atomic_write_json_async is implemented.
…fecycles, and handle transient exception status codes
There was a problem hiding this comment.
Actionable comments posted: 10
🤖 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.
Inline comments:
In `@README.md`:
- Around line 202-203: The routing documentation for the
`llm-routing-auto-agy-ollama` routing mode is inaccurate. The current wording
states that Ollama chaining only occurs for advanced/reasoning tasks, but the
implementation also enables Ollama chaining for `agent-complex-core` (which maps
to `ollama-deepseek-v4-flash`). Update the documentation descriptions for
`llm-routing-auto-agy-ollama` across all three sections mentioned (lines
202-203, 259-260, and 774-775) to accurately reflect that Ollama is enabled for
both advanced/reasoning tasks AND `agent-complex-core` tasks, ensuring the
wording aligns with the actual runtime behavior and prevents operator confusion.
In `@router/agy_proxy.py`:
- Around line 205-210: The code contains multiple instances of bare `except
Exception: pass` blocks that silently suppress Valkey sync/persist operation
failures, including the sync_cooldowns_from_valkey call and similar patterns at
the other mentioned locations. Replace all five instances of these exception
handlers (at the sync_cooldowns_from_valkey call and the other persist/sync
operations at lines 332-336, 345-349, 404-408, and 442-446) with proper error
logging instead of silent suppression. Each exception handler should log the
error details using the appropriate logger to ensure visibility into persistence
failures and maintain consistent state across worker processes.
- Around line 292-295: The client.send() call on line 294 is passing a timeout
parameter that httpx.AsyncClient.send() does not accept, which will cause a
TypeError at runtime. Replace the build_request() and send() pattern with a
direct client.request("POST", url, json=payload, stream=True,
timeout=tier_timeout + 5.0) call instead, which properly supports the timeout
argument. Alternatively, if build_request() and send() must be used, set the
timeout via the request object's extensions before passing it to send(), or rely
on the client's default timeout configuration.
In `@router/main.py`:
- Around line 1663-1691: The HTTPException responses at lines 1666 and 1691 are
returning raw LiteLLM error body details to API callers, which can leak internal
diagnostics and provider information. Replace the detail field in both
HTTPException calls with generic, client-friendly error messages instead of
including the raw error_body or response.text. Keep the detailed error
information only in the existing logger.warning calls which already log the full
error bodies for debugging purposes internally.
- Around line 28-37: The issue is that when Valkey client initialization fails
in the exception handler, setting _redis_client to False permanently prevents
any future retry attempts, causing a single transient failure to disable state
sharing for the entire process lifetime. To fix this, remove the line that sets
_redis_client = False after the warning log, and instead let the variable remain
None so that subsequent calls to the initialization function will attempt to
reconnect to Valkey again rather than permanently giving up. This allows the
function to retry initialization on transient failures while maintaining the
same fallback-to-None behavior for the caller.
In `@scripts/verification/mock_rate_limit_server.py`:
- Around line 5-7: The `format` parameter in the `log_message` method shadows
Python's builtin `format` function, which causes linting failures. Rename the
`format` parameter to a descriptive alternative name like `msg` or
`message_format`, then update all references to this parameter within the method
body (including the format string operation `format%args`) to use the new
parameter name instead.
In `@scripts/verification/verify_direct_ollama_cooldown.py`:
- Around line 28-32: The metric parsing logic in the function that extracts
triage_requests_total is directly converting the string value to int, which
fails when Prometheus outputs floating point counters like "1.0". To fix this,
convert the metric value to float first before converting it to int in the
return statement where line.split()[1] is processed, ensuring that both integer
and decimal values from Prometheus counters are handled correctly and the
function returns the actual count instead of falling back to 0.
In `@scripts/verification/verify_ollama_cooldown.py`:
- Around line 29-34: The metrics value parsing in the Prometheus line handler is
brittle because Prometheus exposition format can return float values like `2.0`,
which cannot be directly converted to int and will raise a ValueError, causing
the function to return 0 and break cooldown assertion logic. Fix the value
extraction by parsing the split line value as a float first before converting to
int, changing `int(line.split()[1])` to `int(float(line.split()[1]))` to
properly handle Prometheus float exposition format.
In `@scripts/verification/verify_ollama_routing.py`:
- Around line 50-52: The exception handling in the except Exception as e block
is swallowing errors and allowing the script to continue execution, which
permits false-positive verification results when routing validation actually
fails. Instead of just printing the error message and continuing, the exception
should be re-raised or the script should exit with a failure status (using raise
or sys.exit with a non-zero code) to ensure the verification actually fails when
requests encounter exceptions, preventing the script from exiting successfully
when validation was never completed.
In `@start-stack.sh`:
- Line 311: The LITELLM_MASTER_KEY is being exposed in process command-line
arguments when the shell expands $ESCAPED_LITELLM_MASTER_KEY into the sed
command on line 311, line 317, and line 343. Instead of embedding the expanded
key directly in the sed command arguments where it can be inspected via process
listing, refactor the key substitution to occur within a renderer process or
environment. Consider using a heredoc, reading from stdin, or passing the key
through environment variables to the rendering tool (podman in this case) rather
than as part of the visible sed command arguments. Apply this same pattern fix
to all three occurrences of the LITELLM_MASTER_KEY substitution.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 6f0d3a80-f51d-467a-a212-5fcac1eca318
📒 Files selected for processing (23)
README.mdhello.pylitellm/config.yamlrouter/Containerfilerouter/agy_proxy.pyrouter/circuit_breaker.pyrouter/free_models_roster.jsonrouter/main.pyscripts/README.mdscripts/backup.shscripts/extract_gapfill.pyscripts/extract_prompts.pyscripts/reclassify_all.pyscripts/retry_errors.pyscripts/verification/mock_rate_limit_server.pyscripts/verification/verify_direct_ollama_cooldown.pyscripts/verification/verify_ollama_cooldown.pyscripts/verification/verify_ollama_routing.pystart-stack.shsync_gemini_token.pytest_antigravity.pytest_classifier_accuracy.pytest_goose.py
💤 Files with no reviewable changes (3)
- hello.py
- test_goose.py
- scripts/reclassify_all.py
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces router-side Ollama cooldown management to prevent cascading fallback loops, updates fallback configurations in LiteLLM, and adds several verification scripts. The reviewer identified critical issues in router/main.py where list-based multimodal or structured message payloads could cause crashes during classification, proxying, or fingerprinting. Additionally, the reviewer recommended using asyncio.gather for concurrent Valkey operations, reusing estimate_prompt_tokens for consistent token estimation, correcting token calculations in the streaming path, and ensuring that verification scripts exit with non-zero status codes on failure to prevent silent CI/CD passes.
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.
| last_user_message = "" | ||
| for msg in reversed(messages): | ||
| if msg.get("role") == "user": | ||
| last_user_message = msg.get("content", "") | ||
| last_user_message = msg.get("content") or "" | ||
| break |
There was a problem hiding this comment.
If the last user message is multimodal or structured (where content is a list of dicts), last_user_message will be a list. Passing this list to classify_request will cause a crash with AttributeError: 'list' object has no attribute 'strip' when it tries to normalize the prompt. You should safely extract the text string from the list if content is not a string.
last_user_message = ""
for msg in reversed(messages):
if msg.get("role") == "user":
content = msg.get("content") or ""
if isinstance(content, list):
content = "".join(block.get("text", "") for block in content if isinstance(block, dict) and block.get("type") == "text")
last_user_message = content
break| last_prompt = "" | ||
| for msg in reversed(messages): | ||
| if msg.get("role") == "user": | ||
| last_prompt = msg.get("content", "") | ||
| last_prompt = msg.get("content") or "" | ||
| break |
There was a problem hiding this comment.
Similar to last_user_message, if the user message has a list-based content (multimodal/structured), last_prompt will be a list. This can cause issues or crashes downstream in try_agy_proxy or when logging. You should safely extract the text string from the list.
| last_prompt = "" | |
| for msg in reversed(messages): | |
| if msg.get("role") == "user": | |
| last_prompt = msg.get("content", "") | |
| last_prompt = msg.get("content") or "" | |
| break | |
| last_prompt = "" | |
| for msg in reversed(messages): | |
| if msg.get("role") == "user": | |
| content = msg.get("content") or "" | |
| if isinstance(content, list): | |
| content = "".join(block.get("text", "") for block in content if isinstance(block, dict) and block.get("type") == "text") | |
| last_prompt = content | |
| break |
| fingerprint_parts = [] | ||
| for msg in messages[:4]: | ||
| c = msg.get("content", "") or "" | ||
| c = msg.get("content") or "" | ||
| if c: | ||
| fingerprint_parts.append(c[:200]) | ||
| fingerprint = "|".join(fingerprint_parts) |
There was a problem hiding this comment.
If content is a list (which is common for multimodal or structured tool-calling payloads), c will be a list. Attempting to slice it (c[:200]) and append it to fingerprint_parts will result in a list of dicts being appended. Subsequently, "|".join(fingerprint_parts) will raise a TypeError: sequence item 0: expected str instance, list found and crash the request. You should safely extract the text content if it is a list before constructing the fingerprint.
| fingerprint_parts = [] | |
| for msg in messages[:4]: | |
| c = msg.get("content", "") or "" | |
| c = msg.get("content") or "" | |
| if c: | |
| fingerprint_parts.append(c[:200]) | |
| fingerprint = "|".join(fingerprint_parts) | |
| fingerprint_parts = [] | |
| for msg in messages[:4]: | |
| c = msg.get("content") or "" | |
| if isinstance(c, list): | |
| c = "".join(block.get("text", "") for block in c if isinstance(block, dict) and block.get("type") == "text") | |
| if isinstance(c, str) and c: | |
| fingerprint_parts.append(c[:200]) | |
| fingerprint = "|".join(fingerprint_parts) |
| async def sync_from_valkey(self, redis_client) -> None: | ||
| """Synchronize both sub-breakers from Valkey.""" | ||
| await self.google.sync_from_valkey(redis_client) | ||
| await self.vendor.sync_from_valkey(redis_client) | ||
|
|
||
| async def save_to_valkey(self, redis_client) -> None: | ||
| """Persist both sub-breakers to Valkey.""" | ||
| await self.google.save_to_valkey(redis_client) | ||
| await self.vendor.save_to_valkey(redis_client) |
There was a problem hiding this comment.
The synchronization and persistence of the Google and Vendor sub-breakers are currently performed sequentially. Since these involve asynchronous network I/O operations to Valkey/Redis, they can be executed concurrently using asyncio.gather to reduce latency and improve throughput.
| async def sync_from_valkey(self, redis_client) -> None: | |
| """Synchronize both sub-breakers from Valkey.""" | |
| await self.google.sync_from_valkey(redis_client) | |
| await self.vendor.sync_from_valkey(redis_client) | |
| async def save_to_valkey(self, redis_client) -> None: | |
| """Persist both sub-breakers to Valkey.""" | |
| await self.google.save_to_valkey(redis_client) | |
| await self.vendor.save_to_valkey(redis_client) | |
| async def sync_from_valkey(self, redis_client) -> None: | |
| """Synchronize both sub-breakers from Valkey.""" | |
| import asyncio | |
| await asyncio.gather( | |
| self.google.sync_from_valkey(redis_client), | |
| self.vendor.sync_from_valkey(redis_client) | |
| ) | |
| async def save_to_valkey(self, redis_client) -> None: | |
| """Persist both sub-breakers to Valkey.""" | |
| import asyncio | |
| await asyncio.gather( | |
| self.google.save_to_valkey(redis_client), | |
| self.vendor.save_to_valkey(redis_client) | |
| ) |
| # Approximate prompt tokens based on messages characters | ||
| prompt_chars = sum(len(m.get("content", "")) for m in messages) | ||
| prompt_chars = sum(len(m.get("content") or "") for m in messages) | ||
| approx_prompt_tokens = max(1, prompt_chars // 4) |
There was a problem hiding this comment.
Instead of using a fragile inline character-based approximation that can fail or underestimate tokens if content is a list (e.g., in multimodal or tool-calling payloads), you should reuse the existing estimate_prompt_tokens(body) function defined at the top of the file for consistency and robustness.
| # Approximate prompt tokens based on messages characters | |
| prompt_chars = sum(len(m.get("content", "")) for m in messages) | |
| prompt_chars = sum(len(m.get("content") or "") for m in messages) | |
| approx_prompt_tokens = max(1, prompt_chars // 4) | |
| approx_prompt_tokens = estimate_prompt_tokens(body) |
| completion_chars = 0 | ||
| request_tokens = estimate_prompt_tokens(body_to_send) | ||
| try: | ||
| async for chunk in r.aiter_bytes(): | ||
| completion_chars += len(chunk) | ||
| yield chunk | ||
| proxy_latency = (time.time() - proxy_start) * 1000.0 | ||
| stats["total_proxy_time_ms"] += proxy_latency | ||
| stats["avg_proxy_latency_ms"] = stats["total_proxy_time_ms"] / stats["total_requests"] | ||
| record_tool_usage(active_tool, request_tokens, completion_chars // 4, model_name, proxy_latency, route="litellm_fallback") |
There was a problem hiding this comment.
In the streaming path, completion_chars accumulates the raw bytes of the SSE stream chunks (including SSE protocol headers, JSON wrappers, and metadata) rather than the actual generated text characters. This leads to a massive overestimation of completion tokens when calling record_tool_usage. Consider parsing the SSE chunks to extract the actual content or using a more accurate estimation method.
| completion_chars = 0 | |
| request_tokens = estimate_prompt_tokens(body_to_send) | |
| try: | |
| async for chunk in r.aiter_bytes(): | |
| completion_chars += len(chunk) | |
| yield chunk | |
| proxy_latency = (time.time() - proxy_start) * 1000.0 | |
| stats["total_proxy_time_ms"] += proxy_latency | |
| stats["avg_proxy_latency_ms"] = stats["total_proxy_time_ms"] / stats["total_requests"] | |
| record_tool_usage(active_tool, request_tokens, completion_chars // 4, model_name, proxy_latency, route="litellm_fallback") | |
| completion_chars = 0 | |
| request_tokens = estimate_prompt_tokens(body_to_send) | |
| try: | |
| async for chunk in r.aiter_bytes(): | |
| completion_chars += len(chunk) | |
| yield chunk | |
| proxy_latency = (time.time() - proxy_start) * 1000.0 | |
| stats["total_proxy_time_ms"] += proxy_latency | |
| stats["avg_proxy_latency_ms"] = stats["total_proxy_time_ms"] / stats["total_requests"] | |
| # Note: completion_chars includes raw SSE metadata; consider parsing or adjusting this metric | |
| record_tool_usage(active_tool, request_tokens, completion_chars // 4, model_name, proxy_latency, route="litellm_fallback") |
| except Exception as e: | ||
| print(f"Request: model={model}, prompt='{prompt[:40]}...' failed/timed out as expected (API downstream might be simulated/unreachable): {e}") |
There was a problem hiding this comment.
The verification script catches all connection and timeout exceptions and prints them as 'expected', exiting with a success status code (0). This prevents the script from failing fast on actual downstream connection errors or misconfigurations in CI/CD environments. You should exit with a non-zero status code on unexpected connection failures.
| except Exception as e: | |
| print(f"Request: model={model}, prompt='{prompt[:40]}...' failed/timed out as expected (API downstream might be simulated/unreachable): {e}") | |
| except Exception as e: | |
| print(f"❌ ERROR: Request to model={model} failed or timed out: {e}", file=sys.stderr) | |
| sys.exit(1) |
| if diff == 0: | ||
| print("✅ SUCCESS: llm-routing-ollama was successfully skipped (cooled down) on the second request!") | ||
| else: | ||
| print(f"❌ FAILURE: llm-routing-ollama was NOT skipped (count increased by {diff})!") | ||
| else: | ||
| print("❌ FAILURE: First request did not even reach the triage router (check if all free models failed immediately without fallback).") |
There was a problem hiding this comment.
The verification script prints failure messages to stdout but does not exit with a non-zero status code when verification fails. This causes automated test runners and CI/CD pipelines to report the test as successful even when it fails. You should explicitly call sys.exit(1) on failures.
| if diff == 0: | |
| print("✅ SUCCESS: llm-routing-ollama was successfully skipped (cooled down) on the second request!") | |
| else: | |
| print(f"❌ FAILURE: llm-routing-ollama was NOT skipped (count increased by {diff})!") | |
| else: | |
| print("❌ FAILURE: First request did not even reach the triage router (check if all free models failed immediately without fallback).") | |
| if diff == 0: | |
| print("✅ SUCCESS: llm-routing-ollama was successfully skipped (cooled down) on the second request!") | |
| else: | |
| print(f"❌ FAILURE: llm-routing-ollama was NOT skipped (count increased by {diff})!") | |
| import sys | |
| sys.exit(1) | |
| else: | |
| print("❌ FAILURE: First request did not even reach the triage router (check if all free models failed immediately without fallback).") | |
| import sys | |
| sys.exit(1) |
| if diff == 0: | ||
| print("✅ SUCCESS: llm-routing-ollama was successfully cooled down and skipped on the second request!") | ||
| else: | ||
| print(f"❌ FAILURE: llm-routing-ollama was NOT cooled down (count increased by {diff})!") | ||
| else: | ||
| print("❌ FAILURE: First request did not even reach the triage router.") |
There was a problem hiding this comment.
The verification script prints failure messages to stdout but does not exit with a non-zero status code when verification fails. This causes automated test runners and CI/CD pipelines to report the test as successful even when it fails. You should explicitly call sys.exit(1) on failures.
| if diff == 0: | |
| print("✅ SUCCESS: llm-routing-ollama was successfully cooled down and skipped on the second request!") | |
| else: | |
| print(f"❌ FAILURE: llm-routing-ollama was NOT cooled down (count increased by {diff})!") | |
| else: | |
| print("❌ FAILURE: First request did not even reach the triage router.") | |
| if diff == 0: | |
| print("✅ SUCCESS: llm-routing-ollama was successfully cooled down and skipped on the second request!") | |
| else: | |
| print(f"❌ FAILURE: llm-routing-ollama was NOT cooled down (count increased by {diff})!") | |
| import sys | |
| sys.exit(1) | |
| else: | |
| print("❌ FAILURE: First request did not even reach the triage router.") | |
| import sys | |
| sys.exit(1) |
… breaker Valkey calls, generic error detail messages, SSE stream token counting, secure YAML rendering, and robust verification script exit codes
Comprehensive Fixes: Worktree Leak Cleanups, LiteLLM Authentication, and Code Review Feedback
This PR consolidates several sets of improvements addressing hardcoded workspace paths, LiteLLM API Gateway authentication, robust deployment procedures, and code safety enhancements.
Key Changes
1. Worktree Path & Configuration Leaks Resolved
os.path.expanduser, preventing hardcoded home directories from leaking into tracked scripts (verify_ollama_cooldown.py,verify_direct_ollama_cooldown.py,backup.sh,sync_gemini_token.py,test_antigravity.py, andtest_classifier_accuracy.py).start-stack.shwith environment variables ($WORKDIRand$HOME) mapped dynamically to the deployment host.2. LiteLLM Gateway Authentication Errors Fixed
pod.yamlwas hardcoded to a truncated placeholder string ("sk-lit...33bf"), causing dashboard spend log API queries to fail with 401 Unauthorized.$LITELLM_MASTER_KEYextracted from the deployment environment.3. Deployment Robustness & Path Escaping
\,|, and&) inside$WORKDIR,$HOME, and$LITELLM_MASTER_KEYinstart-stack.sh, protecting the deployment pipeline against syntax issues during template rendering.4. Explicit Routing Validation
verify_ollama_routing.pyverification utility to validate the actual returned routing model against the expected tier (expected_model), failing fast on model mismatch or downstream errors.5. Telemetry & Documentation Synchronization
README.mdto map precisely to the active exporters.6. Concurrent Processing & Triage Safety
httpx.AsyncClientinstantiation inside the main try/finally block ofexecute_proxyinrouter/main.py, guaranteeing connection closure on exceptions.annotations_lock).Summary by Sourcery
Introduce router-managed Ollama cooldown logic, tighten Ollama routing/gating behavior, and wire LiteLLM to use environment-driven authentication and updated fallback chains, while cleaning up scripts, paths, and docs.
New Features:
llm-routing-ollamaas a first-class LiteLLM model and integrate it into agent tier fallback chains.Bug Fixes:
Enhancements:
Build:
start-stack.shresolve the workspace dynamically, escape sed replacements safely, and template in the LiteLLM master key and host paths at deploy time.$HOMEinstead of hardcoded user directories.Documentation:
Tests:
Chores:
Summary by CodeRabbit
Release Notes
New Features
Bug Fixes
Documentation
Chores