feat(router): make auto-router session affinity deployment-granular

[devin-ai-integration]

TLDR

Problem this solves:

• Auto-router session affinity pinned a model name, not a deployment
• Fanned model groups still spread a session across deployments
• That drops the provider prompt cache mid-session
• Session deployment pins were also shared across API keys

How it solves it:

• Auto-routed model groups now opt into the existing DeploymentAffinityCheck
• Pinning stays in one place, keyed by group plus session
• Session pins are now scoped by the caller's API key hash
• Off unless the auto-router already sets session_affinity

Relevant issues

Linear ticket

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 Anthropic API calls against a local proxy, no mocks. Model is anthropic/claude-haiku-4-5, the newest haiku-family model the Anthropic /v1/models endpoint lists. Each deployment in the fanned-out group has an explicit model_info.id so the x-litellm-model-id response header names the deployment that served the turn.

Config (/tmp/affinity_proof_config.yaml):

model_list:
  - model_name: target-group
    litellm_params:
      model: anthropic/claude-haiku-4-5
      api_key: os.environ/ANTHROPIC_API_KEY
    model_info:
      id: deployment-a
  - model_name: target-group
    litellm_params:
      model: anthropic/claude-haiku-4-5
      api_key: os.environ/ANTHROPIC_API_KEY
    model_info:
      id: deployment-b

  - model_name: smart-router
    litellm_params:
      model: auto_router/complexity_router
      complexity_router_config:
        session_affinity: true
        session_affinity_ttl_seconds: 3600
        tiers:
          SIMPLE: target-group
          MEDIUM: target-group
          COMPLEX: target-group
          REASONING: target-group
        default_model: target-group

general_settings:
  master_key: sk-affinity-proof

1. Start the proxy on this branch, at commit df1b93ea3dcb9486e59d3db414e87bc03c288083

sudo service postgresql start
(cd litellm/proxy && uv run --no-sync prisma db push --accept-data-loss --skip-generate)
uv run --no-sync litellm --config /tmp/affinity_proof_config.yaml --detailed_debug --port 4000 2>&1 | tee litellm.log

2. Six turns of one session, all carrying the same x-litellm-session-id

for i in 1 2 3 4 5 6; do
  echo -n "turn $i: "
  curl -s -D - -o /dev/null http://localhost:4000/v1/chat/completions \
    -H "Authorization: Bearer sk-affinity-proof" \
    -H "Content-Type: application/json" \
    -H "x-litellm-session-id: proof-session-A" \
    -d "{\"model\":\"smart-router\",\"messages\":[{\"role\":\"user\",\"content\":\"turn $i: say hi\"}],\"max_tokens\":16}" \
    | grep -i "x-litellm-model-id"
done

turn 1: x-litellm-model-id: deployment-a
turn 2: x-litellm-model-id: deployment-a
turn 3: x-litellm-model-id: deployment-a
turn 4: x-litellm-model-id: deployment-a
turn 5: x-litellm-model-id: deployment-a
turn 6: x-litellm-model-id: deployment-a

Every turn of the session lands on the same deployment, so the provider prompt cache stays warm

3. Same run, from the proxy log

DeploymentAffinityCheck: set session affinity mapping model_map_key=target-group deployment=deployment-a ttl=3600 session_id=proof-session-A   (x6)
DeploymentAffinityCheck: session-id affinity hit -> deployment=deployment-a session_id=proof-session-A                                        (x5)

Five hits against six sets: turn 1 has nothing cached so it picks freely, then every later turn reads the pin

4. Four fresh sessions, four turns each, same commit

for s in B C D E; do for i in 1 2 3 4; do
  echo -n "session $s turn $i: "
  curl -s -D - -o /dev/null http://localhost:4000/v1/chat/completions \
    -H "Authorization: Bearer sk-affinity-proof" \
    -H "Content-Type: application/json" \
    -H "x-litellm-session-id: proof-session-$s" \
    -d "{\"model\":\"smart-router\",\"messages\":[{\"role\":\"user\",\"content\":\"turn $i: say hi\"}],\"max_tokens\":16}" \
    | grep -i "x-litellm-model-id"
done; done

session B turn 1..4: deployment-b, deployment-b, deployment-b, deployment-b
session C turn 1..4: deployment-a, deployment-a, deployment-a, deployment-a
session D turn 1..4: deployment-b, deployment-b, deployment-b, deployment-b
session E turn 1..4: deployment-b, deployment-b, deployment-b, deployment-b

Each session is internally consistent and three of the four went to the deployment session A did not use, so the pin is per session rather than a global lock. First-turn selection is still random, which is why the split is 2 vs 3 rather than anything designed

5. The "before", same config against origin/litellm_internal_staging at commit ba917681461b1ad04d30f91da26e75b3521996f3, served on port 4001, one session id, eight turns

git worktree add /tmp/litellm-prefix origin/litellm_internal_staging --detach
cd /tmp/litellm-prefix && PYTHONPATH=/tmp/litellm-prefix \
  uv run --no-sync --project ~/repos/litellm litellm --config /tmp/affinity_proof_config.yaml --detailed_debug --port 4001

for i in 1 2 3 4 5 6 7 8; do
  echo -n "pre-fix turn $i: "
  curl -s -D - -o /dev/null http://localhost:4001/v1/chat/completions \
    -H "Authorization: Bearer sk-affinity-proof" \
    -H "Content-Type: application/json" \
    -H "x-litellm-session-id: pre-fix-session" \
    -d "{\"model\":\"smart-router\",\"messages\":[{\"role\":\"user\",\"content\":\"turn $i: say hi\"}],\"max_tokens\":16}" \
    | grep -i "x-litellm-model-id"
done

pre-fix turn 1: x-litellm-model-id: deployment-b
pre-fix turn 2: x-litellm-model-id: deployment-a
pre-fix turn 3: x-litellm-model-id: deployment-b
pre-fix turn 4: x-litellm-model-id: deployment-b
pre-fix turn 5: x-litellm-model-id: deployment-a
pre-fix turn 6: x-litellm-model-id: deployment-a
pre-fix turn 7: x-litellm-model-id: deployment-b
pre-fix turn 8: x-litellm-model-id: deployment-a

One session, four turns on each deployment, so the prompt cache goes cold roughly every other turn. The pre-fix log shows why: session_affinity_pin fires and routed_model is target-group on all eight turns, while set session affinity mapping appears zero times and no DeploymentAffinityCheck is registered. The tier's model name was pinned, the deployment was not

Deployment selection on a session's first turn is random, so the specific ids above will differ run to run; the invariants (identical within a session after the fix, spread before it) are what reproduce. Runs were single worker with the default in-memory cache, over /v1/chat/completions with the x-litellm-session-id header

Type

🆕 New Feature

Changes

For an agent workload with a large stable prefix, cache affinity is worth more than model choice: a provider-side prompt-cache hit costs roughly a tenth of the normal input rate, so a router that moves a session to a different model, or merely to a different deployment of the same model, can be a net loss even when the new model is cheaper

The complexity router's session_affinity only ever pinned the routed model name. When the tier entry is a model group fanned across several providers or deployments, deployment selection still spreads later turns of the same session across those deployments and the prefix cache goes cold, which is the exact failure this fixes

Approach

Two options were on the table: store (model, deployment_id) in the complexity router's own pin, or make auto-routing delegate deployment pinning to DeploymentAffinityCheck, which already pins a concrete deployment id per (model group, session id) and exists for precisely this implicit-prompt-caching reason

This PR takes the second one. The two components own disjoint questions: the complexity router decides which model group a session uses, and DeploymentAffinityCheck decides which deployment inside that group. Because the pre-routing hook rewrites the request's model before deployment resolution runs, the affinity check sees the routed group and pins within it with no coordination needed. There stays exactly one source of truth for the deployment id, and none of the affinity cache logic is duplicated. The first option would have meant re-implementing deployment filtering, cooldown awareness, and the post-call persist hook inside the complexity router, and would have left two mechanisms racing to answer the same question

Concretely, Router derives a model_group -> ttl mapping from its registered complexity routers, covering every tier value (scalar and pool form) plus default_model, for routers where session_affinity is on and no plugins are configured, which is the same gate async_pre_routing_hook uses for its own pin. That mapping is injected into DeploymentAffinityCheck as a callable rather than a snapshot, so deployments added at runtime are picked up, and it is OR'd into the effective flags so auto-routing can enable session pinning for a group without ever disabling something an operator configured. Duplicate groups take the minimum TTL, so a pin is never silently extended

def _get_effective_flags(self, model_group: str) -> tuple[bool, bool, bool]:
    auto_session_affinity = model_group in self.session_affinity_group_ttls()
    ...
    return (..., self.enable_session_id_affinity or auto_session_affinity)

The deployment pin honors the complexity router's session_affinity_ttl_seconds for those groups instead of the Router-wide deployment_affinity_ttl_seconds, so one auto-router config knob controls both halves of the stickiness

Fallback behavior

DeploymentAffinityCheck.async_filter_deployments runs after health, cooldown, and blocked-deployment filtering, and already returns the full healthy list when the pinned deployment is not in it. So an unhealthy, cooling-down, rate-limited, or removed-from-config deployment falls back to normal selection instead of failing the request, and this PR adds no new failure mode. Note that TPM/RPM enforcement runs after the callback filter, so a pinned deployment can still be selected and then rejected by a later rate-limit check; that is pre-existing behavior for every affinity user and is untouched here

One known limitation, pre-existing but now reachable by auto-router users: when _get_stable_model_map_key_from_deployments cannot derive a stable key for the group (for example Azure deployments without base_model set), session pinning silently degrades to normal selection. There is a test asserting it degrades rather than erroring

Behavior changes worth calling out

session_affinity still defaults to False and no default is flipped. For operators who already set it, the pin now also sticks to a deployment, which is the point of the change

DeploymentAffinityCheck's session cache key had no API-key scoping, so two virtual keys reusing the same client-supplied session_id shared a pin. It is now scoped by the hashed user_api_key_hash, falling back to unscoped for SDK usage with no authenticated caller, matching what the complexity router already does for its own pin. Existing proxy pins miss once after upgrade and re-pin, which is graceful

While gating the session-only write path, a latent issue surfaced: async_pre_call_deployment_hook wrote an API-key affinity entry whenever a user key was resolvable, even for groups with only session affinity enabled. Those writes were never read. The write is now gated on enable_user_key

The provider handed to DeploymentAffinityCheck is built over a weakref.ref to the Router rather than a bound method, and returns an empty mapping once the referent is gone. The callback is registered into the module-level litellm.callbacks list, which is never pruned, so a bound method would have retained the Router and everything it owns for the process lifetime. That matters for the per-request routers built from a caller-supplied user_config and for proxy config reloads, which is the same reason _live_routers is a WeakSet. There are GC regression tests covering every registration path

Sync note for whoever documents this: there is no synchronous complexity pre-routing hook at all. Router.get_available_deployment never invokes async_pre_routing_hook, so auto-routing (and therefore this pin) is async-only. Nothing was added to the sync path

For the docs page covering the affinity combo: enabling session_affinity on an auto-router no longer requires separately adding session_affinity to optional_pre_call_checks to get deployment stickiness, and the TTL that governs the deployment pin for auto-routed groups is the auto-router's session_affinity_ttl_seconds

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

Link to Devin session: https://app.devin.ai/sessions/86f1017360074499a6f46d06fa1bd8d9

[devin-ai-integration]

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

• Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
• Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

• Disable automatic comment, CI, and merge conflict monitoring

[greptile-apps]

Greptile Summary

The PR makes auto-router session affinity deployment-granular and scopes session pins by API key hash. The latest changes use a weak reference for the dynamic group-TTL provider so globally registered callbacks do not retain discarded Router instances

• Derives per-model-group affinity TTLs from eligible complexity-router configurations
• Enables deployment affinity automatically for auto-routed groups with session affinity
• Uses API-key-scoped session cache keys and applies the configured session TTL
• Adds lifecycle, routing, cooldown, expiration, and API-key isolation regression coverage

Confidence Score: 5/5

The PR appears safe to merge

No blocking failure remains

Important Files Changed

Filename	Overview
litellm/router.py	Derives complexity-router affinity groups, registers the deployment-affinity callback, and exposes the TTL mapping through a weakref-backed provider
litellm/router_utils/pre_call_checks/deployment_affinity_check.py	Enables per-group auto-router affinity, scopes session keys by caller key hash, and applies group-specific TTLs
tests/test_litellm/router_strategy/test_complexity_router.py	Covers complexity-router group discovery and callback registration
tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py	Adds deployment pinning, isolation, fallback, expiration, configuration, and Router garbage-collection regression tests

_{Reviews (2): Last reviewed commit: "fix(router): weaken optional affinity pr..." | Re-trigger Greptile}

[codecov]

Codecov Report

❌ Patch coverage is 98.30508% with 1 line in your changes missing coverage. Please review.

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

📢 Thoughts on this report? Let us know!

[codspeed-hq]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing litellm_auto_router_deployment_affinity (f8f82c5) with litellm_internal_staging (ba91768)}

[devin-ai-integration]

@greptileai two commits landed after your review: the group-TTL provider handed to DeploymentAffinityCheck is now weakref-backed so the callback stops retaining the Router, plus GC regression tests