Skip to content

feat(router): independent, default-on deployment affinity for the auto-router - #36146

Merged
tin-berri merged 1 commit into
litellm_internal_stagingfrom
litellm_autorouter_session_marker
Aug 8, 2026
Merged

feat(router): independent, default-on deployment affinity for the auto-router#36146
tin-berri merged 1 commit into
litellm_internal_stagingfrom
litellm_autorouter_session_marker

Conversation

@tin-berri

@tin-berri tin-berri commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

TLDR

Problem this solves:

  • An auto-router picks a model group, never a deployment, so a group fanned across several deployments spreads one conversation's turns across them and the provider prompt cache goes cold roughly every other turn
  • The only way to get deployment stickiness was 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 exclusive
  • Session deployment pins had no API key scoping, so two keys reusing the same client-supplied session id shared one pin

How it solves it:

  • New deployment_affinity on the auto-router config, on by default; session_affinity implies it, so a tier-pinned session always sticks to one deployment of one model, while deployment_affinity alone 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 before
  • The complexity router returns a TTL on the pre-routing hook response and the Router stamps it as a request-scoped internal metadata marker, written or cleared on every routing attempt so fallbacks to plain groups never carry a stale marker
  • DeploymentAffinityCheck treats 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 key

Important

Default-on behavior change. deployment_affinity defaults to true, so an auto-router whose callers send a session_id starts 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. Set deployment_affinity: false to keep the previous load-balanced behavior. Existing optional_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

  1. They send POST https://litellm-domain/v1/chat/completions with model smart-router and header x-litellm-session-id: chat-123
  2. The response header x-litellm-model-id reads deployment-a
  3. They send the next turn with the same session header
  4. x-litellm-model-id now reads deployment-b, and the provider bills the whole prompt at the uncached input rate again

After: 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

  1. They send the same POST with x-litellm-session-id: chat-123
  2. x-litellm-model-id reads deployment-a
  3. They send the next turn with the same session header
  4. x-litellm-model-id reads deployment-a again, and the provider response shows cached prompt tokens
  5. A caller on a different API key reusing chat-123 gets its own deployment choice and cannot steer or read the first caller's pin

Relevant issues

  • Auto-routed sessions get deployment-granular stickiness without freezing the tier, keyed per request instead of per model group
  • Pins are held per model group, so switching tiers leaves the previous group's pin undisturbed and returning to it reuses that deployment
  • Supersedes feat(router): make auto-router session affinity deployment-granular #36045: the group-name TTL map there silently failed for wildcard deployments and model_group_alias targets, 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
  • The marker key is inbound-stripped (_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

  • I have added meaningful tests
  • My PR passes all CI/CD checks (e.g., lint, format, unit tests)
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to 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-tier carries two deployments with explicit model_info.id so x-litellm-model-id names the serving deployment, and the classifier runs on Nova Micro so tiers are decided by a real LLM call.

model_list:
  - model_name: smart-router
    litellm_params:
      model: auto_router/complexity_router
      complexity_router_config:
        classifier_type: llm
        classifier_llm_config: {model: classifier-model, timeout_ms: 20000}
        tiers: {SIMPLE: cheap-tier, MEDIUM: cheap-tier, COMPLEX: smart-tier, REASONING: smart-tier}
        default_model: cheap-tier
  - model_name: cheap-tier
    litellm_params: {model: bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0, aws_region_name: us-east-1}
    model_info: {id: cheap-deployment-A}
  - model_name: cheap-tier
    litellm_params: {model: bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0, aws_region_name: us-east-1}
    model_info: {id: cheap-deployment-B}
  - model_name: smart-tier
    litellm_params: {model: bedrock/us.amazon.nova-lite-v1:0, aws_region_name: us-east-1}
    model_info: {id: smart-deployment-A}
  - model_name: classifier-model
    litellm_params: {model: bedrock/us.amazon.nova-micro-v1:0, aws_region_name: us-east-1}

Eight turns of one session, alternating a trivial ask with one that classifies up, so the run crosses tiers repeatedly and comes back:

SID="proof-session"
for i in 1 2 3 4 5 6 7 8; do
  if [ $((i % 2)) -eq 1 ]; then P="what is 2 plus $i"
  else P="Derive the amortized complexity of a Fibonacci heap decrease-key, probe $i, and justify each step rigorously"; fi
  curl -s -D - -o /dev/null http://127.0.0.1:4591/v1/chat/completions \
    -H "Authorization: Bearer sk-1234" -H "Content-Type: application/json" \
    -H "x-litellm-session-id: $SID" \
    -d "{\"model\":\"smart-router\",\"messages\":[{\"role\":\"user\",\"content\":\"$P\"}],\"max_tokens\":12}" \
    | grep -i "x-litellm-model-id"
done
turn classified before, on staging after, this branch
1 SIMPLE cheap-B cheap-B
2 COMPLEX smart-A smart-A
3 SIMPLE cheap-A cheap-B
4 COMPLEX smart-A smart-A
5 SIMPLE cheap-A cheap-B
6 COMPLEX smart-A smart-A
7 SIMPLE cheap-B cheap-B
8 COMPLEX smart-A smart-A

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: true on 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

  • New deployment_affinity on ComplexityRouterConfig, default true, gated by its own predicate so it never switches the tier pin on. session_affinity keeps its own predicate and its own default of false
  • PreRoutingHookResponse gains session_affinity_ttl_seconds, and Router.async_pre_routing_hook stamps or clears _session_deployment_affinity_ttl in the proxy-internal metadata bucket on every attempt, through the same write-or-clear owner as routing_decision
  • DeploymentAffinityCheck reads 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 hashed user_api_key_hash with an unscoped fallback
  • Pin writes are first-writer-wins claims: one Lua script registered through the existing async_register_script seam 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 pin
  • A complexity router with the deployment pin auto-registers the affinity callback, so no separate optional_pre_call_checks entry is needed, and opting out skips the callback rather than registering a filter that can never fire

Things 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: false to keep every turn load-balanced, which is what a deployment set with tight per-deployment rate limits wants.

What does not change: session_affinity still defaults to false, no global flag is flipped, per-group model_group_affinity_config semantics are untouched, and direct calls to a model group are unaffected by an auto-router that targets it. Existing session pins for optional_pre_call_checks: ["session_affinity"] users miss once on upgrade, since the key gained the API key scope, and re-pin gracefully

Note 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

  1. Start Postgres, push the schema, launch the proxy on this branch with the config above and --detailed_debug
  2. Run the eight-turn loop above with one x-litellm-session-id: tiers alternate between cheap-tier and smart-tier while every cheap-tier turn names the same deployment
  3. Add deployment_affinity: false to the router config and repeat: the cheap-tier turns spread again
  4. Add session_affinity: true instead: every turn is served by one deployment of one tier, the pre-existing behavior
  5. Send the same session id at cheap-tier directly: deployments spread, no affinity hit lines
  6. Generate a second virtual key and replay the same session id: its turns pin independently of the first key's deployment
  7. pytest tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py tests/test_litellm/router_strategy/test_complexity_router.py -q

Final Attestation

  • The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

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 (default true, independent of session_affinity, which implies deployment pinning). When deployment pinning applies, pre-routing responses carry session_affinity_ttl_seconds, and the Router writes or clears internal metadata _session_deployment_affinity_ttl on every routing attempt (including fallbacks) via a shared _stamp_or_clear_metadata_key helper.

DeploymentAffinityCheck treats 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.

@tin-berri

Copy link
Copy Markdown
Contributor Author

@greptileai can you reply here with a score and why that score?

@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@greptile-apps

greptile-apps Bot commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

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:

  • SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY is correctly added to both the inbound strip list (_UNTRUSTED_METADATA_CONTROL_FIELDS) and the proxy-internal exclusion list — an attacker cannot forge a deployment pin TTL from outside
  • _stamp_or_clear_metadata_key is the right abstraction: writes or clears on every attempt, so fallbacks can never resurrect a stale marker from an earlier attempt
  • _get_marker_session_affinity_ttl rejects strings, bools, zero, and negatives — the parametrized test covers all the malformed cases
  • _ensure_deployment_affinity_callback is idempotent and called from both the model_group_affinity_config path and init_complexity_router_deployment, so there's no double-registration
  • The init ordering fix (moving deployment_affinity_ttl_seconds and model_group_affinity_config before set_model_list) is necessary and correctly done
  • Test coverage is excellent: write path, read path, TTL, wildcard groups, key scoping, marker clearing on non-auto-routed fallbacks

Points off:

  1. One-time cache miss on upgrade for existing session_affinity users. The get_session_affinity_cache_key signature gained a required user_key parameter, and all existing cache entries under the old (unscoped) key format become dead. The PR description acknowledges "miss once and re-pin gracefully," but this is a silent behavioral change — no migration note, no log warning on first miss, no CHANGELOG entry. Users with long TTLs (hours/days) who rely on affinity for correctness (not just latency) will see their sessions spread exactly once after upgrading with no indication of why.

  2. _uses_session_affinity is computed twice on the hot path — once as a @property called from _with_session_deployment_affinity and once directly in async_pre_routing_hook as use_session_affinity: Final = self._uses_session_affinity. Not a functional issue, but since _uses_session_affinity re-evaluates bool(self.config.session_affinity and not self.config.plugins) on each access, the Final local should be the single call site.

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-apps

greptile-apps Bot commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

The PR adds deployment-level session affinity for auto-routed requests, scopes session pins by caller key, and makes concurrent pin claims first-writer-wins.

  • Propagates the complexity router’s session-affinity TTL through request-scoped internal metadata.
  • Registers deployment-affinity handling for eligible auto-router configurations.
  • Adds atomic Redis pin claims with a synchronized in-memory winner.
  • Extends regression coverage for TTL propagation, key scoping, marker cleanup, and concurrent claims.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains.

No blocking failure remains.

Important Files Changed

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

Comment thread litellm/router_utils/pre_call_checks/deployment_affinity_check.py Outdated
Comment thread litellm/router.py Outdated

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ 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

codecov Bot commented Aug 7, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 95.86777% with 5 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...utils/pre_call_checks/deployment_affinity_check.py 93.82% 5 Missing ⚠️

📢 Thoughts on this report? Let us know!

@tin-berri
tin-berri force-pushed the litellm_autorouter_session_marker branch from 47e4d94 to 42f3f0c Compare August 7, 2026 01:27
@tin-berri

Copy link
Copy Markdown
Contributor Author

@greptileai can you reply here with a score and why that score?

@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@greptile-apps

greptile-apps Bot commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

Score: 4/5

What earns the 4:

The design is solid throughout. The request-scoped marker (_session_deployment_affinity_ttl) stamped on every pre-routing attempt, with a clear on fallbacks, correctly solves the stale-pin-across-retries problem that plagued the per-group TTL map approach. Key scoping by user_api_key_hash is the right security boundary — two callers reusing the same client-supplied session ID cannot share a pin. The first-writer-wins async_claim_cache primitive properly eliminates the race where two concurrent first-turns of a session flip each other's pin. Stripping the marker from _UNTRUSTED_METADATA_CONTROL_FIELDS and the batch metadata sanitizer shows the right threat model. The proof screenshots with real Bedrock calls against an actual Postgres-backed proxy are unusually thorough.

What keeps it from 5:

There's a narrow TOCTOU window in async_claim_cache on the Redis path:

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 NX

If the NX set fails (another pod won) but the key expires or is evicted between the NX attempt and the subsequent GET, existing comes back None, winner becomes value, and the code then issues a plain SET (no NX) — potentially overwriting whichever pod actually holds the legitimate pin. For a best-effort routing optimization this is tolerable, but it means the "first-writer-wins" guarantee only holds under normal conditions, not under key expiry races. The fix would be to use a Lua script or GETSET/SET+GET pipeline instead of the separate NX then GET round-trips.

Also minor: the in-memory-only claimed if claimed is not None else value fallback silently swallows the case where TTL was so short the key expired between set and get, returning the caller's own value as though it won — harmless in practice but not strictly "first-writer-wins" semantics.

These are corner cases in a best-effort feature; the PR is correct in all the scenarios that actually matter.

Comment thread litellm/caching/dual_cache.py Outdated
@tin-berri
tin-berri force-pushed the litellm_autorouter_session_marker branch from 42f3f0c to 7178f66 Compare August 7, 2026 01:41
@tin-berri

Copy link
Copy Markdown
Contributor Author

@greptileai can you reply here with a score and why that score?

@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@greptile-apps

greptile-apps Bot commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

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. async_register_script called on every async_claim_cache invocation (perf)

run_script: Final = self.async_register_script(RedisCache.CLAIM_CACHE_SCRIPT)
result: Final = await run_script(keys=(key,), args=(...))

async_register_script runs on every claim, which is every routed request with session affinity. Script registration should be cached once (e.g., a class-level _claim_script slot initialized lazily), not re-registered per call.


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 EXPIRE when the existing value matches. The in-memory path takes the nx=True early-return, so TTL is never extended. In a multi-pod setup this is fine (Redis is authoritative), but in single-process (in-memory only) deployments the pin expires earlier than the caller expects.


3. claimed_pin == pin_value relies on a clean JSON round-trip through Redis

pin_value: Final = DeploymentAffinityCacheValue(model_id=str(model_id))
...
if claimed_user_pin == pin_value:

async_claim_cache returns json.loads(raw), which is a plain dict, not a DeploymentAffinityCacheValue. Pydantic's __eq__ won't match a dict. So claimed_user_pin == pin_value will almost always be False when Redis is in use—meaning the "pin already claimed" log branch is never reached and the "set affinity mapping" log is always printed even for losing claimers. The pin itself is correct (the NX semantics work), but the log interpretation and any future code branching on this equality will be wrong. You need either DeploymentAffinityCacheValue.model_validate(claimed_user_pin) == pin_value or compare the model_id field directly.

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ 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:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@veria-ai

veria-ai Bot commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

PR overview

This 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

@tin-berri

Copy link
Copy Markdown
Contributor Author

Script runner is now registered once at RedisCache init instead of per claim call.

@tin-berri
tin-berri force-pushed the litellm_autorouter_session_marker branch from 7178f66 to 10a15f0 Compare August 7, 2026 01:50
@tin-berri

Copy link
Copy Markdown
Contributor Author

In-memory TTL not sliding is long-standing allow_ttl_override semantics, shared by the pre-existing model pin; Redis stays authoritative multi-pod.

@tin-berri

Copy link
Copy Markdown
Contributor Author

The equality check only picks which debug line logs; claim correctness is enforced server-side by the Lua script.

@tin-berri

Copy link
Copy Markdown
Contributor Author

@greptileai can you reply here with a score and why that score?

@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@greptile-apps

greptile-apps Bot commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

Score: 4/5

This is a well-engineered feature with solid fundamentals. Here's why:

What's strong:

  • async_claim_cache is the right primitive. The Redis Lua script is correct — it atomically GET-or-SET, refreshes TTL on re-claim with the same value, and the in-memory path is race-safe in the asyncio single-thread model since set_cache(nx=True) runs synchronously.
  • Write-or-clear pattern via _stamp_or_clear_metadata_key is clean and prevents stale markers from leaking across fallback attempts.
  • API-key scoping on session pins is the right security fix — two callers with the same client-supplied session_id can no longer read each other's deployment pin.
  • Auto-registration of DeploymentAffinityCheck via _ensure_deployment_affinity_callback is a clean simplification; the refactored helper is invoked by both the model_group_affinity_config path and the new complexity-router path.
  • The marker is correctly added to both _UNTRUSTED_METADATA_CONTROL_FIELDS (strips inbound attacker-supplied values) and LITELLM_PROXY_INTERNAL_METADATA_KEYS (prevents it reaching providers).

What holds it to 4:

  1. MappingProxyType passed to model_copy(update=...) in _with_session_deployment_affinity. Pydantic v2's model_copy signature expects dict[str, Any] | None. At runtime this likely ducks through, but it's technically wrong and type checkers will flag it. A plain {"session_affinity_ttl_seconds": ...} dict is the correct arg.

  2. Float TTL truncation in async_claim_cacheint(ttl) if ttl else 0 truncates a float like 1.5 to 1, while the Lua EXPIRE comparison is > 0, so ttl=0.5 would store with no expiry. Fine for integer-second TTLs here, but fragile for future callers.

  3. The in-memory claimed is None fallback does a non-nx write after the nx write + get, meaning a key that expired between the two awaits can be overwritten without first-writer-wins semantics. For a best-effort optimization it's acceptable, but worth a comment.

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.

Comment thread litellm/router_utils/pre_call_checks/deployment_affinity_check.py
@tin-berri
tin-berri force-pushed the litellm_autorouter_session_marker branch from 10a15f0 to 46a4170 Compare August 7, 2026 02:03
@tin-berri

Copy link
Copy Markdown
Contributor Author

Consolidated: the claim now lives in the affinity check via the existing script registration seam; cache layer is untouched.

@tin-berri

Copy link
Copy Markdown
Contributor Author

@greptileai can you reply here with a score and why that score?

@greptile-apps

greptile-apps Bot commented Aug 8, 2026

Copy link
Copy Markdown
Contributor

Score: 4/5

Here's the breakdown:

What's well-done:

  • Correct design choice: The request-scoped marker approach (_session_deployment_affinity_ttl stamped per routing attempt, cleared on fallback) cleanly avoids the silent failures of the prior group-name TTL map (feat(router): make auto-router session affinity deployment-granular #36045) — wildcard deployments, alias targets, and leaked affinity onto direct group calls are all structurally impossible with this approach.
  • Security: The marker key is inbound-stripped (_UNTRUSTED_METADATA_CONTROL_FIELDS) and excluded from provider-bound batch metadata. Session pins are scoped by user_api_key_hash, so two callers reusing the same session_id can't steer each other's pin. Solid.
  • Atomicity: The Lua get-or-set-or-refresh script is the right tool for first-writer-wins on Redis. The in-memory fallback mirrors it exactly and degrades gracefully (pod-local stickiness through Redis blips rather than no stickiness at all).
  • _stamp_or_clear_metadata_key refactor: Good — the shared static helper makes it structurally impossible to write without clearing on fallback, which was the fragile part of the old _record_routing_decision inline logic.
  • Opt-out is first-class: deployment_affinity: false is clearly documented and suppresses the callback registration entirely rather than registering a filter that can never fire.

What holds it to 4 instead of 5:

  1. Default-on is a real behavior change for existing users. The PR acknowledges this under the [!IMPORTANT] block, but existing deployments that don't set the flag will silently start pinning deployments as soon as a caller sends a session_id. For users who have tight per-deployment rate limits, this could cause unexpected 429s before they diagnose why. The PR notes this in "Things a reviewer will ask about" but it's not surfaced prominently enough for operators who may not read every PR description on upgrade.

  2. One-time pin miss on upgrade for existing session_affinity users. Adding user_api_key_hash scope to the cache key is the right call for security, but existing pins under the old key are silently invalidated. The first turn after upgrade re-pins to a possibly different deployment, causing a prompt cache miss. Minor but real regression for current users.

  3. _uses_deployment_pin is suppressed when plugins are configured with no inline comment explaining why. A future plugin author could reasonably expect deployment affinity to still apply — worth a brief note that plugins may implement their own routing logic that conflicts with deployment pinning.

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 deployment_affinity is implicitly True and the router has multiple deployments per group.

Comment thread litellm/router_utils/pre_call_checks/deployment_affinity_check.py Outdated
@tin-berri
tin-berri force-pushed the litellm_autorouter_session_marker branch from e3ed042 to e34e381 Compare August 8, 2026 00:44
@tin-berri

Copy link
Copy Markdown
Contributor Author

@greptileai can you reply here with a score and why that score?

@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@greptile-apps

greptile-apps Bot commented Aug 8, 2026

Copy link
Copy Markdown
Contributor

Score: 4/5

This is a high-quality implementation with a clean design. Here's the breakdown:


What's done well:

  • Architecture: The marker-based approach (SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY stamped/cleared on every routing attempt) is the right abstraction. It makes fallback safety provable by construction rather than by convention.
  • Atomicity: The Redis Lua script for first-writer-wins is correct — GET-then-SET/EXPIRE in one atomic step, with a synchronous in-memory fallback that mirrors the same semantics.
  • Security: API key scoping of session pins and inbound stripping of the marker key (via _UNTRUSTED_METADATA_CONTROL_FIELDS and LITELLM_PROXY_INTERNAL_METADATA_KEYS) are exactly right.
  • Fault tolerance: Redis failures degrade to same-pod stickiness rather than unpinned sessions — the right tradeoff.
  • Test coverage: Thorough — concurrent pin races, legacy string pins, Redis outage fallback, wildcard group round-trips, write-or-clear per-attempt guarantees, malformed marker values. This is well above the average for this codebase.
  • Refactoring quality: _stamp_or_clear_metadata_key is a clean generalization of the routing_decision write-or-clear pattern. _ensure_deployment_affinity_callback removes the duplication between the model_group_affinity_config and complexity router init paths.

What causes the deduction:

  1. Backwards-incompatible session pin key format change (per rule b48b7341): Existing users of optional_pre_call_checks: ["session_affinity"] will get a one-time miss on upgrade because the cache key gained a user_api_key_hash scope component. The PR acknowledges this, but the miss is automatic with no flag to preserve the old format. The re-pin on the next request is graceful, but the behavioral change is not opt-in.

  2. Default-on behavior change (per rule b48b7341): deployment_affinity: true by default means every auto-router caller sending a session_id gets deployment pinning with no config change. The PR is transparent about this and the opt-out (deployment_affinity: false) is straightforward, but it still changes what happens for existing auto-router deployments on upgrade.

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.

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ 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.

@tin-berri
tin-berri enabled auto-merge (squash) August 8, 2026 01:02
@tin-berri
tin-berri disabled auto-merge August 8, 2026 01:02
@tin-berri
tin-berri force-pushed the litellm_autorouter_session_marker branch from e34e381 to f4d003c Compare August 8, 2026 01:14
@tin-berri

Copy link
Copy Markdown
Contributor Author

@greptileai can you reply here with a score and why that score?

@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@greptile-apps

greptile-apps Bot commented Aug 8, 2026

Copy link
Copy Markdown
Contributor

Score: 4/5

The implementation is technically solid and security-conscious. Here's the breakdown:

What's well done:

  • _stamp_or_clear_metadata_key correctly enforces write-or-clear on every attempt, so retries and fallbacks never carry a stale _session_deployment_affinity_ttl marker forward
  • The Lua first-writer-wins script is correct and atomic; the in-memory fallback mirrors it exactly (including the keepalive/TTL refresh on re-claim)
  • SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY is properly added to both LITELLM_PROXY_INTERNAL_METADATA_KEYS (inbound strip) and the callback_utils exclusion list (provider-bound batch metadata), closing both injection vectors
  • API key scoping for session pins is a meaningful security improvement over the previous unscoped keys
  • _ensure_deployment_affinity_callback has the guard that prevents double-registration when called from both the model_group_affinity_config path and init_complexity_router_deployment
  • Moving deployment_affinity_ttl_seconds and model_group_affinity_config assignment before set_model_list is the correct order — _ensure_deployment_affinity_callback can now be called during model list initialization without hitting unset attributes

What gives me pause (reason for -1):

  1. Default-on without a feature flag (rule b48b7341): deployment_affinity=True is a behaviour change for any auto-router caller that sends a session_id. The PR documents it and provides the escape hatch, and it's inert when no session ID is resolvable (most SDK-direct callers). But the pattern in this repo is to gate behaviour changes behind a flag rather than opt-out. The justification is strong — there is essentially no caller who benefits from re-shuffling deployments of the same model mid-conversation — but the upgrade path is still "callers that relied on load balancing across deployments mid-session must now explicitly set deployment_affinity: false."

  2. One-time pin miss for existing session_affinity users: Session pin keys now include the hashed API key, so all live pins evaporate on upgrade. For a production deployment serving an active multi-turn session, this causes a tier re-classification on the next turn. The PR acknowledges this but there's no migration path (can't forward-fill old-format keys).

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.

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ 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.

@tin-berri
tin-berri force-pushed the litellm_autorouter_session_marker branch from f4d003c to 7966a6d Compare August 8, 2026 18:08
@tin-berri

Copy link
Copy Markdown
Contributor Author

@greptileai can you reply here with a score and why that score?

@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@greptile-apps

greptile-apps Bot commented Aug 8, 2026

Copy link
Copy Markdown
Contributor

Score: 4/5

Here's why:

What's well done:

  1. Correct write-or-clear semantics_stamp_or_clear_metadata_key runs on every routing attempt (including the router_strategy is None early-return path), so fallbacks never carry a stale marker. This is the same discipline already applied to routing_decision.

  2. Atomic first-writer-wins — the Lua script (get → set-if-absent → refresh-on-match → return-winner) is correct. The in-memory fallback mirrors it synchronously, and both branches slide the TTL idle window rather than treating the TTL as an absolute session length.

  3. Security posture is solidSESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY is added to both _UNTRUSTED_METADATA_CONTROL_FIELDS (stripped from inbound requests) and LITELLM_PROXY_INTERNAL_METADATA_KEYS (excluded from provider-bound logging). API key scoping on session pin keys prevents cross-caller collisions on shared session IDs.

  4. Clean separation_uses_tier_pin / _uses_deployment_pin properties, _with_session_deployment_affinity wrapper, and _ensure_deployment_affinity_callback() deduplication are all small, testable, and clearly named.

  5. Inert without a session_id — the default-on flag does nothing for callers that don't send a session header, which is the right shape for a behavioral default change.

What keeps this from a 5:

  1. Upgrade regression on existing session_affinity pins — the session cache key gained an API key scope component, so every existing pin silently misses on the first post-upgrade request and re-pins. Documented in the PR description, but it's a customer-visible one-time regression with no migration path and no log line that says "your pin expired due to upgrade."

  2. _ensure_deployment_affinity_callback registers a callback with all global flags = False — the instance looks disabled from its constructor args and only activates via marker presence. A future developer auditing the callback chain will see a DeploymentAffinityCheck(enable_user_key_affinity=False, enable_responses_api_affinity=False, enable_session_id_affinity=False) and may reasonably conclude it's a no-op. The _get_effective_flags and _get_marker_session_affinity_ttl interaction makes it non-obvious that this instance can still route.

  3. Default-on behavioral change — any auto-router caller already sending session_id headers gets deployment pinning with no config change required (or warned about). The PR is transparent about this, but it touches the critical path for a feature (auto-router) that's already in production use.

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.

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ 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.

@tin-berri
tin-berri enabled auto-merge (squash) August 8, 2026 18:33
@tin-berri

Copy link
Copy Markdown
Contributor Author

@greptileai can you reply here with a score and why that score?

@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@tin-berri
tin-berri force-pushed the litellm_autorouter_session_marker branch from 7966a6d to 8438a22 Compare August 8, 2026 18:35
@greptile-apps

greptile-apps Bot commented Aug 8, 2026

Copy link
Copy Markdown
Contributor

Score: 4/5


What's strong:

  • Design is correct. The request-scoped marker approach (_session_deployment_affinity_ttl) cleanly solves the "stale marker survives fallback" problem that the superseded feat(router): make auto-router session affinity deployment-granular #36045 had. Every routing attempt writes or clears, so retries/fallbacks can never carry a stale value forward.

  • First-writer-wins claim is solid. The Lua script is well-reasoned: get-or-set-or-keepalive in one atomic op, with a clean in-memory fallback that degrades gracefully on Redis outage rather than blowing up or leaving the session unpinned. The TTL-slide-for-winner-only logic is correct.

  • Security posture is correct. The marker key is in LITELLM_PROXY_INTERNAL_METADATA_KEYS (excluded from provider-bound metadata) and in _UNTRUSTED_METADATA_CONTROL_FIELDS (inbound-stripped), so it can't be injected or leaked. Session pins are scoped by hashed API key, so two callers reusing the same client-supplied session_id cannot read or steer each other's pin.

  • Test coverage is thorough. The new tests cover: scoped pin write/read, cross-scope invisibility, marker TTL honored over callback default, malformed markers rejected, concurrent first-request pin claims, TTL keepalive for winner only, late-attached Redis, Redis outage fallback, wildcard group round-trip, and the stamp/clear behavior. That's the right set.

  • __init__ assignment order fix is necessary and correct. Moving self.deployment_affinity_ttl_seconds and self.model_group_affinity_config before set_model_list() was required once init_complexity_router_deployment started calling _ensure_deployment_affinity_callback, which reads both fields.


What to watch:

  • Default-on is a real behavior change for existing auto-router users. Any caller sending x-litellm-session-id on an auto-router immediately starts getting deployment pins after upgrade — no config change needed. The PR documents the opt-out (deployment_affinity: false) and notes the change in the PR description, but teams relying on strict load balancing across deployments for rate-limiting purposes will be surprised unless they read the release notes. The existing session_affinity pin key also gains an API key scope, causing a one-time miss for existing users. Both are documented but both are the kind of thing that shows up as a support ticket.

  • _uses_deployment_pin property is computed fresh every call. It's pure (reads two config fields and one bool), so this is harmless, but it's called on the hot pre-routing path. A @cached_property would be cleaner since config is set at construction and never mutated.

  • _with_session_deployment_affinity is called in two branches (session-affinity early-return path and the general return response at the bottom), but the normal path's call is unconditional. This is correct since _uses_deployment_pin gates both, but the duplication is a subtle maintenance surface: if a third return path is added later it's easy to miss wrapping it.

  • _CLAIM_PIN_SCRIPT decoded result is trusted as the authoritative pin value and written to in-memory. If Redis returns a value written by a different pod using the legacy bare-string format, _pinned_model_id handles it via the isinstance(stored, str) branch — that's covered — but the in-memory tier ends up holding a bare string, which is the "legacy" shape. Harmless in practice since _pinned_model_id reads both, but worth noting.

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.

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ 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.

@tin-berri
tin-berri merged commit e35ee4e into litellm_internal_staging Aug 8, 2026
83 checks passed
@tin-berri
tin-berri deleted the litellm_autorouter_session_marker branch August 8, 2026 20:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants