feat(router): coordinate session affinity across replicas [DYN-3249]

[PeaBrane]

Summary

• add immutable backend-neutral session claims with a process-local cache hot path for etcd and shared FileStore deployments
• resolve the authoritative worker/rank before scheduler bookkeeping and exact dispatch, with eventual delete/reset invalidation and terminal close
• document backend/lifetime semantics and add focused Rust plus two-frontend etcd/FileStore regression coverage

Each frontend previously kept an isolated affinity cache, so a later turn handled by another replica could select a different worker and lose prefix-cache reuse.

Fixes #11035

Validation

• cargo test -p dynamo-runtime --lib --no-default-features discovery::kv_store
• cargo test -p dynamo-runtime --lib --no-default-features pipeline::network::egress::push_router
• cargo test -p dynamo-llm --lib --no-default-features session_affinity
• cargo test -p dynamo-llm --lib --no-default-features kv_router::push_router
• .venv/bin/python -m pytest components/src/dynamo/common/tests/configuration/test_kv_router_args.py -q
• .venv/bin/python -m pytest tests/router/test_router_e2e_with_mockers.py::test_mocker_distributed_session_affinity -q
• cargo clippy --no-default-features -- -D warnings in lib/runtime and lib/llm
• fern check
• fern docs broken-links
• git diff --check

---

Summary by CodeRabbit

• New Features

  • Added clearer, more explicit session-affinity behavior for routers, including support for terminal requests that close a session binding.
  • Improved routing consistency so existing session bindings take priority over conflicting routing hints.
  • Expanded end-to-end coverage for distributed session affinity across supported storage backends.

• Documentation

  • Updated session-affinity docs to explain how session IDs, affinity TTL, and final requests behave.
  • Clarified router help text for session-affinity settings.

[github-actions]

🌿 Fern Docs Preview: https://nvidia-preview-019f172d-413a-742f-8da9-e58c8e96c17a.docs.buildwithfern.com/dynamo/dev

[devin-ai-integration]

Devin Review found 1 potential issue.

[coderabbitai]

Walkthrough

This PR replaces per-process in-memory session-affinity with a distributed claim layer. A new claim API (create_or_get_claim, close_claim, subscribe_claim_events) is added to the Discovery trait and implemented in KVStoreDiscovery (with a background watcher broadcasting Delete/Reset events) and stubbed as unsupported in KubeDiscoveryClient. AffinityCoordinator is refactored to derive a claim key, delegate to distributed claims for authoritative AffinityTarget resolution, maintain a local cache evicted on claim events, and expose ResolvedAffinity::into_stream(close_on_finish) with terminal-close semantics. Both push-router implementations are updated to use new_distributed coordinators, book_and_dispatch_exact, and close_on_finish. FileStore delete-event delivery is fixed for symlinked roots, the notify macOS backend is switched to kqueue, and an e2e test validates cross-replica affinity behavior.

Changes

Distributed Session Affinity

Layer / File(s)	Summary
Discovery claim API types and trait methods lib/runtime/src/discovery/mod.rs	Adds ClaimPayload, ClaimPayloadFuture, ClaimOutcome, ClaimCloseOutcome, ClaimEvent types and default-unsupported create_or_get_claim, close_claim, subscribe_claim_events methods to the Discovery trait.
FileStore canonicalization and delete-event fixes lib/runtime/src/storage/kv/file.rs, lib/runtime/Cargo.toml, lib/runtime/src/storage/kv.rs	Canonicalizes directory paths in Directory::new, broadens delete-event matching to all Remove(_) variants, adds canonicalize_event_path for symlinked roots, switches notify to macos_kqueue, adds Manager::is_memory, and adds watcher/symlink tests.
KVStore claim watcher and create/get/close lib/runtime/src/discovery/kv_store.rs	Adds ClaimState with broadcast sender and OnceCell watcher gate; implements run_claim_watcher (reconnect loop emitting Delete/Reset), create_or_get_in_bucket (bounded competing inserts), close_claim (idempotent delete), and subscribe_claim_events; adds unit tests for races, disappearance, idempotency, and cancellation.
Kube discovery unsupported claim stubs lib/runtime/src/discovery/kube.rs	Adds one-time warn_claims_unsupported and returns Unsupported from create_or_get_claim and close_claim.
PushRouter book_and_dispatch_exact and peek_worker_for_request lib/runtime/src/pipeline/network/egress/push_router.rs	Adds book_and_dispatch_exact (optional round-robin advance, exact-instance dispatch with occupancy tracking) and peek_worker_for_request (non-booking worker peek).
AffinityCoordinator distributed claim integration lib/llm/src/session_affinity/coordinator.rs, lib/llm/src/session_affinity/mod.rs	Introduces ClaimCoordination; adds new_distributed; spawns claim-event listener evicting on Delete/Reset; changes acquisition to use claim_key, delegate to claims.resolve, commit to local cache; adds ResolvedAffinity::into_stream(close_on_finish) with CloseAction; adds session_final; reduces AffinityAcquire/AffinityInitialization/AffinityLease to pub(crate).
AffinityCoordinator unit tests lib/llm/src/session_affinity/tests.rs	Adds ClaimTestDiscovery in-memory double; updates all acquire/query/stream call sites; adds tests for winner caching, delete eviction, subscriber-lag, reset/disconnect, terminal close, initializer cancellation, and binding-override.
SessionAffinityPushRouter distributed wiring lib/llm/src/session_affinity/push_router.rs	Switches to new_distributed coordinator; replaces acquire helpers with resolve_affinity; refactors select_and_dispatch_prefill and generate to use book_and_dispatch_exact, advance_round_robin, close_on_finish, and resolved.into_stream.
KvPushRouter distributed wiring and selection refactor lib/llm/src/kv_router/push_router.rs, lib/llm/src/kv_router/push_router/selection.rs, lib/llm/src/kv_router/push_router/request_guard.rs	Wires new_distributed coordinator; changes select_with_affinity to return Option<ResolvedAffinity>; propagates close_on_finish; removes merge_affinity_pin in favor of direct affinity_pin precedence; updates unit tests.
E2E distributed session-affinity tests tests/router/common.py, tests/router/router_process.py, tests/router/test_router_e2e_with_mockers.py	Adds session_affinity_ttl_secs to FrontendRouterProcess; adds _test_distributed_session_affinity (two-router SSE test verifying claim override and cache eviction across etcd/file backends); adds test_mocker_distributed_session_affinity.
Documentation and CLI help updates docs/components/router/router-configuration.md, docs/agents/session-ids.md, docs/components/frontend/nvext.md, components/src/dynamo/common/configuration/groups/router_args.py	Rewrites Session Affinity docs for the distributed-claim contract; updates X-Dynamo-Session-ID and X-Dynamo-Session-Final contracts; revises --router-session-affinity-ttl-secs help text.

Estimated code review effort

🎯 5 (Critical) | ⏱️ ~120 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description has a summary and validation notes, but it omits the required Overview, Details, reviewer-start, and formatted Related Issues sections.	Rewrite the PR body to follow the repository template, adding Overview, Details, reviewer-start files, and a Related Issues block with Closes #11035.
Docstring Coverage	⚠️ Warning	Docstring coverage is 53.63% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: coordinating session affinity across replicas.
Linked Issues check	✅ Passed	The changes add distributed claims across replicas, backend support, invalidation, and terminal close handling, matching the bug's requirements.
Out of Scope Changes check	✅ Passed	I don't see unrelated feature work; the code, docs, and tests all support replica-coordinated session affinity and its backends.

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 3

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

lib/llm/src/session_affinity/push_router.rs (1)

146-158: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Preserve explicit DP rank in direct no-session routing.

Both direct fallback branches extract an AffinityTarget, but then only use worker_id. A request with prefill_dp_rank/dp_rank can be dispatched without the resolved rank, unlike the affinity path that passes/sets rank.

Suggested fix

         let Some(session_id) = session_id else {
             let explicit = explicit_target(&request, RequestPhase::Prefill)?;
-            let Some(pinned_worker) = explicit else {
+            let Some(target) = explicit else {
                 return Err(invalid_argument(
                     "worker ID required for prefill request in Direct routing mode",
                 ));
             };
+            let rank = target.dp_rank;
             return self
                 .inner
                 .select_and_dispatch_exact(
                     request,
-                    Some(pinned_worker.worker_id),
-                    move |request, worker_id| prepare(request, worker_id, None),
+                    Some(target.worker_id),
+                    move |request, worker_id| prepare(request, worker_id, rank),
                 )
                 .await;
         };

For the generate branch, similarly normalize the request routing rank before inner.direct(...) when target.dp_rank is set.

Also applies to: 241-248

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/llm/src/session_affinity/push_router.rs` around lines 146 - 158, The
direct no-session routing path is dropping the explicit affinity rank and only
forwarding worker_id. Update the branches that use explicit_target in
push_router’s no-session handling so the resolved dp_rank/prefill_dp_rank is
preserved and applied before dispatch, matching the affinity-path behavior. In
the prefill branch, pass the rank through the select_and_dispatch_exact/prepare
flow, and in the generate branch normalize the request routing rank before
calling inner.direct when target.dp_rank is present.

🧹 Nitpick comments (1)

lib/runtime/src/discovery/kv_store.rs (1)

43-44: 🩺 Stability & Availability | 🔵 Trivial

Track the hidden-watch-failure gap.

This TODO describes a stale-cache risk in the new claim invalidation path. Please link it to a GitHub issue or add a follow-up contract once Bucket::watch can surface backend errors.

Do you want me to draft the follow-up issue text?

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/runtime/src/discovery/kv_store.rs` around lines 43 - 44, The TODO in
Bucket::watch should be tracked with a concrete follow-up, since it currently
leaves hidden backend failures in the claim invalidation path unaddressed.
Update the note near Bucket::watch in kv_store.rs to reference a GitHub issue or
add a follow-up contract describing how etcd reconnect/compaction and FileStore
overflow errors will be surfaced and converted into Reset once that support
exists.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@lib/llm/src/session_affinity/coordinator.rs`:
- Around line 720-760: `CloseAction::run` is closing the shared session claim
unconditionally after `finish()` drops only the current lease, which can evict a
binding that still has active requests. Update the close path in
`CloseAction`/`AffinityTrackedStream` to guard the eviction and `claims.close`
call behind a check that no other active leases remain for the same claim key,
using the existing `coordinator`/`claims` state and the `claim_key` to verify
exclusivity before terminal close.

In `@lib/llm/src/session_affinity/tests.rs`:
- Around line 435-443: The reacquire in session_affinity/tests.rs does not
verify that `ClaimEvent::Reset` rehydrates the authoritative claim from
`ClaimTestDiscovery::claims`; update the test around the
`coordinator.acquire(...).resolve(...)` flow to assert the reopened claim still
resolves to `target(8, Some(0))` and rejects the new `target(99, Some(0))`
proposal. Use the existing `coordinator`, `session_id`, and `ClaimTestDiscovery`
setup to prove the post-Reset claim comes from discovery rather than the local
cache.

In `@lib/runtime/src/storage/kv/file.rs`:
- Around line 643-648: The regression test
external_delete_is_observed_under_noncanonical_root currently uses t.path() for
both FileStore instances, so it never exercises a non-canonical root. Update the
test to create a symlinked root from the tempdir and construct one of the
FileStore values with that symlink path while keeping the other on the canonical
path, so the watcher/creator interaction in FileStore::new and the external
delete observation path are covered.

---

Outside diff comments:
In `@lib/llm/src/session_affinity/push_router.rs`:
- Around line 146-158: The direct no-session routing path is dropping the
explicit affinity rank and only forwarding worker_id. Update the branches that
use explicit_target in push_router’s no-session handling so the resolved
dp_rank/prefill_dp_rank is preserved and applied before dispatch, matching the
affinity-path behavior. In the prefill branch, pass the rank through the
select_and_dispatch_exact/prepare flow, and in the generate branch normalize the
request routing rank before calling inner.direct when target.dp_rank is present.

---

Nitpick comments:
In `@lib/runtime/src/discovery/kv_store.rs`:
- Around line 43-44: The TODO in Bucket::watch should be tracked with a concrete
follow-up, since it currently leaves hidden backend failures in the claim
invalidation path unaddressed. Update the note near Bucket::watch in kv_store.rs
to reference a GitHub issue or add a follow-up contract describing how etcd
reconnect/compaction and FileStore overflow errors will be surfaced and
converted into Reset once that support exists.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 1211536f-d8df-45a5-a3c1-3987f8d67269

📥 Commits

Reviewing files that changed from the base of the PR and between 8f2431e and 065abee.

⛔ Files ignored due to path filters (2)

• Cargo.lock is excluded by !**/*.lock
• lib/bindings/python/Cargo.lock is excluded by !**/*.lock

📒 Files selected for processing (21)

• components/src/dynamo/common/configuration/groups/router_args.py
• docs/agents/session-ids.md
• docs/components/frontend/nvext.md
• docs/components/router/router-configuration.md
• lib/llm/src/kv_router/push_router.rs
• lib/llm/src/kv_router/push_router/request_guard.rs
• lib/llm/src/kv_router/push_router/selection.rs
• lib/llm/src/session_affinity/coordinator.rs
• lib/llm/src/session_affinity/mod.rs
• lib/llm/src/session_affinity/push_router.rs
• lib/llm/src/session_affinity/tests.rs
• lib/runtime/Cargo.toml
• lib/runtime/src/discovery/kube.rs
• lib/runtime/src/discovery/kv_store.rs
• lib/runtime/src/discovery/mod.rs
• lib/runtime/src/pipeline/network/egress/push_router.rs
• lib/runtime/src/storage/kv.rs
• lib/runtime/src/storage/kv/file.rs
• tests/router/common.py
• tests/router/router_process.py
• tests/router/test_router_e2e_with_mockers.py