feat(router): independent, default-on deployment affinity for the auto-router - #36146
Conversation
|
@greptileai can you reply here with a score and why that score? |
|
bugbot run |
|
Score: 4/5 This is a well-engineered feature with a clear problem statement, correct security hygiene, comprehensive tests, and a clean design. Here's the breakdown: Strengths:
Points off:
Both are minor relative to the overall quality. The architecture is sound, the security surface is handled correctly, and the real-Bedrock screenshots are compelling evidence the feature works end-to-end. |
Greptile SummaryThe PR adds deployment-level session affinity for auto-routed requests, scopes session pins by caller key, and makes concurrent pin claims first-writer-wins.
Confidence Score: 5/5The PR appears safe to merge because no blocking failure remains. No blocking failure remains.
|
| Filename | Overview |
|---|---|
| litellm/router.py | Registers request-scoped deployment affinity for eligible complexity routers and stamps or clears the affinity TTL marker during routing. |
| litellm/router_strategy/complexity_router/complexity_router.py | Exposes the configured session-affinity TTL in pre-routing responses when built-in session affinity is active. |
| litellm/router_utils/pre_call_checks/deployment_affinity_check.py | Adds caller-scoped session keys, marker-controlled affinity, and first-writer-wins deployment pin claims. |
| litellm/proxy/litellm_pre_call_utils.py | Prevents clients from injecting the internal deployment-affinity TTL marker. |
| tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py | Covers deployment stickiness, caller isolation, marker lifetime, TTL selection, and concurrent pin behavior. |
Reviews (2): Last reviewed commit: "feat(router): deployment-granular sessio..." | Re-trigger Greptile
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit 47e4d94. Configure here.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
47e4d94 to
42f3f0c
Compare
|
@greptileai can you reply here with a score and why that score? |
|
bugbot run |
|
Score: 4/5 What earns the 4: The design is solid throughout. The request-scoped marker ( What keeps it from 5: There's a narrow TOCTOU window in set_result: Final = await self.redis_cache.async_set_cache(key, value, nx=True, ttl=ttl)
existing: Final = None if set_result else await self.redis_cache.async_get_cache(key)
winner: Final = value if set_result or existing is None else existing
if winner == value and not set_result:
await self.redis_cache.async_set_cache(key, value, ttl=ttl) # plain SET, no NXIf the NX set fails (another pod won) but the key expires or is evicted between the NX attempt and the subsequent GET, Also minor: the in-memory-only These are corner cases in a best-effort feature; the PR is correct in all the scenarios that actually matter. |
42f3f0c to
7178f66
Compare
|
@greptileai can you reply here with a score and why that score? |
|
bugbot run |
|
Score: 4/5 The engineering here is solid—the Lua NX+GET script is properly atomic, the write-or-clear pattern on every attempt is correct, the API key scoping closes a real isolation gap, and the marker approach elegantly avoids the stale-pin problems described in the PR description. Tests are comprehensive. Three things hold it back from a 5: 1. run_script: Final = self.async_register_script(RedisCache.CLAIM_CACHE_SCRIPT)
result: Final = await run_script(keys=(key,), args=(...))
2. In-memory tier does not refresh TTL on re-claim, diverging from Redis The docstring says "Claiming with the value already stored refreshes its TTL on the redis tier (the in-memory tier keeps a live key's original expiry by design)." The Redis path does 3. pin_value: Final = DeploymentAffinityCacheValue(model_id=str(model_id))
...
if claimed_user_pin == pin_value:
|
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit 7178f66. Configure here.
|
|
||
| # 2) Session-id -> deployment affinity | ||
| if enable_session_id: | ||
| if session_affinity_active: |
There was a problem hiding this comment.
Medium: Session pin bypasses tag revocation
This marker now enables deployment filtering before get_deployments_for_tag() runs. If a session was pinned while the caller had a restricted tag, a later untagged request reduces the candidate list to that restricted deployment first; tag routing then sees no default deployment in the singleton list and returns the pinned deployment, allowing the caller to retain access after the tag is revoked. Apply affinity only to the deployments that remain after tag-based eligibility filtering, or validate the cached deployment against that complete eligible set before returning it.
PR overviewThis pull request adds independent, default-enabled deployment affinity to the auto-router, using session-based pins to keep requests on a selected deployment. It also integrates affinity configuration with the complexity router. The current implementation can preserve access to a previously pinned restricted deployment after the caller’s tag is revoked, creating an authorization bypass under that session state. Unbounded creation of session-affinity entries can also allow authenticated callers to consume shared Redis memory. Neither open issue has yet been addressed. Open issues (2)
Fixed/addressed: 0 · PR risk: 7/10 |
|
Script runner is now registered once at RedisCache init instead of per claim call. |
7178f66 to
10a15f0
Compare
|
In-memory TTL not sliding is long-standing allow_ttl_override semantics, shared by the pre-existing model pin; Redis stays authoritative multi-pod. |
|
The equality check only picks which debug line logs; claim correctness is enforced server-side by the Lua script. |
|
@greptileai can you reply here with a score and why that score? |
|
bugbot run |
|
Score: 4/5 This is a well-engineered feature with solid fundamentals. Here's why: What's strong:
What holds it to 4:
These are all minor against the overall quality of the design. The real-world Bedrock evidence in the PR description is compelling, and the test coverage across the cache, affinity check, and integration paths looks thorough. |
10a15f0 to
46a4170
Compare
|
Consolidated: the claim now lives in the affinity check via the existing script registration seam; cache layer is untouched. |
|
@greptileai can you reply here with a score and why that score? |
|
Score: 4/5 Here's the breakdown: What's well-done:
What holds it to 4 instead of 5:
None of these are blockers — the implementation is clean, the edge cases are handled, and the security properties are correct. But the default-on behavior change deserves a more prominent operator-facing signal, e.g. a logged warning at router init when |
e3ed042 to
e34e381
Compare
|
@greptileai can you reply here with a score and why that score? |
|
bugbot run |
|
Score: 4/5 This is a high-quality implementation with a clean design. Here's the breakdown: What's done well:
What causes the deduction:
Both are documented clearly and the impact is limited (deployment selection only, not tier/model selection), which is why this is a 4 and not a 3. The implementation itself is sound. |
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit e34e381. Configure here.
e34e381 to
f4d003c
Compare
|
@greptileai can you reply here with a score and why that score? |
|
bugbot run |
|
Score: 4/5 The implementation is technically solid and security-conscious. Here's the breakdown: What's well done:
What gives me pause (reason for -1):
Neither of these is a correctness bug — they're upgrade-experience concerns. The core routing path, Redis atomics, metadata lifecycle, and security properties are all correct. |
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit f4d003c. Configure here.
f4d003c to
7966a6d
Compare
|
@greptileai can you reply here with a score and why that score? |
|
bugbot run |
|
Score: 4/5 Here's why: What's well done:
What keeps this from a 5:
Overall this is a high-quality, well-reasoned PR that solves a real and subtle problem. The concerns are about upgrade smoothness and operational clarity rather than correctness. |
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit 7966a6d. Configure here.
|
@greptileai can you reply here with a score and why that score? |
|
bugbot run |
7966a6d to
8438a22
Compare
|
Score: 4/5 What's strong:
What to watch:
Overall: a well-engineered feature with good test coverage and correct security handling. The main real-world risk is the default-on behavior change catching existing auto-router operators off guard, which is why this isn't a 5/5. |
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit 8438a22. Configure here.
TLDR
Problem this solves:
session_affinity, which also freezes the tier, so a session that escalates to a stronger model can never come back down. Classification per turn and deployment stickiness were mutually exclusiveHow it solves it:
deployment_affinityon the auto-router config, on by default;session_affinityimplies it, so a tier-pinned session always sticks to one deployment of one model, whiledeployment_affinityalone pins per tier as sessions move between tiers: every turn is still classified on its own merits while a session returning to a model group lands on the deployment it used there beforeDeploymentAffinityChecktreats marker presence as session affinity for that request only, pins the deployment inside the routed group, and scopes session pin keys by the caller's hashed API keyImportant
Default-on behavior change.
deployment_affinitydefaults totrue, so an auto-router whose callers send asession_idstarts pinning deployments with no config change. Tiering, spend, and which model serves a turn are unaffected; only the choice among deployments of one model group changes. Setdeployment_affinity: falseto keep the previous load-balanced behavior. Existingoptional_pre_call_checks: ["session_affinity"]pins also miss once on upgrade, since the key gained an API key scope, and re-pin on the next request.User Flow
Before: a developer whose coding agent talks to an auto-router sees multi-turn sessions billed at full input price almost every turn, because turns bounce between deployments and the provider prompt cache never warms
smart-routerand headerx-litellm-session-id: chat-123x-litellm-model-idreadsdeployment-ax-litellm-model-idnow readsdeployment-b, and the provider bills the whole prompt at the uncached input rate againAfter: every turn of the session that lands on the same model group lands on the same deployment, so cached input is billed at the reduced rate
x-litellm-session-id: chat-123x-litellm-model-idreadsdeployment-ax-litellm-model-idreadsdeployment-aagain, and the provider response shows cached prompt tokenschat-123gets its own deployment choice and cannot steer or read the first caller's pinRelevant issues
model_group_aliastargets, forced affinity onto per-group configs that omitted it, and leaked pinning onto direct model group calls. A request-scoped marker makes all four impossible by construction_UNTRUSTED_METADATA_CONTROL_FIELDS) and excluded from provider-bound batch metadata (LITELLM_PROXY_INTERNAL_METADATA_KEYS)Linear ticket
Resolves LIT-5305
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Delays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
Screenshots / Proof of Fix
Real Bedrock calls against a local proxy backed by Postgres, no mocks.
cheap-tiercarries two deployments with explicitmodel_info.idsox-litellm-model-idnames the serving deployment, and the classifier runs on Nova Micro so tiers are decided by a real LLM call.Eight turns of one session, alternating a trivial ask with one that classifies up, so the run crosses tiers repeatedly and comes back:
Tiers alternate in both runs, so classification is untouched. Before, the four SIMPLE turns spread across B, A, A, B. After, all four land on cheap-B, and the config sets no affinity keys at all, which is the default-on path.
The third state, the one that was previously inexpressible, is
session_affinity: trueon the same rig: every turn including the COMPLEX ones is served by cheap-deployment-A, because freezing the tier at turn 1's SIMPLE sends complex prompts to the cheap model. That is the behavior this PR stops being the only way to get a stable deployment.First-turn deployment selection is random, so specific ids differ run to run. The invariant, identical within a session after and spread before, is what reproduces.
Type
🆕 New Feature
Changes
deployment_affinityonComplexityRouterConfig, default true, gated by its own predicate so it never switches the tier pin on.session_affinitykeeps its own predicate and its own default of falsePreRoutingHookResponsegainssession_affinity_ttl_seconds, andRouter.async_pre_routing_hookstamps or clears_session_deployment_affinity_ttlin the proxy-internal metadata bucket on every attempt, through the same write-or-clear owner asrouting_decisionDeploymentAffinityCheckreads the marker on both the filter and the pre-call persist paths, uses the marker TTL for the pin, and scopes session pin keys by hasheduser_api_key_hashwith anunscopedfallbackasync_register_scriptseam does get-or-set-or-refresh atomically on Redis, and the in-memory path is a synchronous check-and-set. Two overlapping first requests of one session can no longer flip a pinoptional_pre_call_checksentry is needed, and opting out skips the callback rather than registering a filter that can never fireThings a reviewer will ask about. The default is on because re-shuffling a conversation across deployments of the same model discards the provider cache for no benefit, and it is inert unless a session id is resolvable, so a caller that sends no session header gets no pin and no cache write. Set
deployment_affinity: falseto keep every turn load-balanced, which is what a deployment set with tight per-deployment rate limits wants.What does not change:
session_affinitystill defaults to false, no global flag is flipped, per-groupmodel_group_affinity_configsemantics are untouched, and direct calls to a model group are unaffected by an auto-router that targets it. Existing session pins foroptional_pre_call_checks: ["session_affinity"]users miss once on upgrade, since the key gained the API key scope, and re-pin gracefullyNote on CI: the osv-scan failure is inherited, this branch changes no lockfiles. The semantic-keyword test failures some local venvs show are missing-optional-dependency artifacts and reproduce on unmodified staging
QA runbook
--detailed_debugx-litellm-session-id: tiers alternate between cheap-tier and smart-tier while every cheap-tier turn names the same deploymentdeployment_affinity: falseto the router config and repeat: the cheap-tier turns spread againsession_affinity: trueinstead: every turn is served by one deployment of one tier, the pre-existing behaviorcheap-tierdirectly: deployments spread, no affinity hit linespytest tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py tests/test_litellm/router_strategy/test_complexity_router.py -qFinal Attestation
Note
Medium Risk
Default-on routing changes which deployment serves repeat sessions (tiering unchanged), and pin logic touches Redis/multi-pod concurrency; behavior is heavily tested but upgrades may reshuffle pins once due to API-key-scoped session keys.
Overview
Adds default-on deployment stickiness for auto-router sessions so multi-turn traffic that keeps landing on the same model group reuses the same deployment (prompt-cache friendly) without requiring
session_affinity, which still pins the tier.The complexity router gains
deployment_affinity(defaulttrue, independent ofsession_affinity, which implies deployment pinning). When deployment pinning applies, pre-routing responses carrysession_affinity_ttl_seconds, and the Router writes or clears internal metadata_session_deployment_affinity_ttlon every routing attempt (including fallbacks) via a shared_stamp_or_clear_metadata_keyhelper.DeploymentAffinityChecktreats that marker as session affinity for one request, uses the marker TTL for pins, scopes session pin cache keys by hashed API key, and replaces blind cache sets with first-writer-wins claims (Redis Lua get/set/refresh with in-memory fallback). Routers with deployment pinning auto-register the affinity callback through_ensure_deployment_affinity_callback.Proxy layers strip the new internal key from inbound metadata; OpenAPI/UI types document
deployment_affinity.Reviewed by Cursor Bugbot for commit 8438a22. Bugbot is set up for automated code reviews on this repo. Configure here.