fix(policies): honor the body rid over the header under manual and consistent_hashing

[ighutake-debug]

Description

Problem

With --routing-key-override, a request whose body carries a rid is documented to pin by the rid-derived key, with the routing-key header as the fallback, on any policy. Under manual and consistent_hashing the body rid was ignored and the header alone decided the worker.

PolicyRegistry::select_worker deliberately skips the sticky override for those two policies (routing_key_override_applies) because they key on the header themselves, but neither ManualPolicy::select_worker_impl nor ConsistentHashingPolicy::select_worker_impl consulted info.rid_key. The gRPC and HTTP pipelines already derive rid_key for every request under the override, so the key was available and simply not read.

Solution

Both policies now resolve their key as info.rid_key first, then the header (consistent_hashing keeps info.routing_key ahead of the raw header as before). PolicyRegistry::derive_rid_key returns None unless the override is enabled, and the key it returns is already capped and lineage-stripped, so:

• without --routing-key-override nothing changes;
• with it, manual / consistent_hashing pin on the same effective key the sticky override uses for every other policy, and the same key the routers use for keyed-load accounting (WorkerLoadGuard::with_key).

PD legs namespace the rid key exactly like header keys, so prefill and decode still stick independently.

Changes

• model_gateway/src/policies/manual.rs: consult info.rid_key before X-SMG-Routing-Key; module doc lists both key sources.
• model_gateway/src/policies/consistent_hashing.rs: rid_key → routing_key hint → raw header.
• model_gateway/src/policies/registry.rs: doc comments on select_worker / routing_key_override_applies describe the rid-first contract the key-native policies now uphold.
• Tests:
  • manual::tests::test_manual_rid_key_outranks_header_key — rotating header keys with one rid land on one worker; the pin is stored under the rid key.
  • manual::tests::test_manual_rid_key_namespaces_per_leg — same rid under prefill/decode gets independent entries.
  • consistent_hashing::tests::test_rid_key_outranks_routing_key_hint — rid wins over a validated header hint that hashes to a different worker.
  • registry::tests::rid_key_outranks_header_under_key_native_policies — end-to-end through PolicyRegistry::with_override for both PolicyConfig::Manual and PolicyConfig::ConsistentHashing, using derive_rid_key("conv_tN").

Test Plan

All four new tests were written first and failed against main for the expected reason (rid ignored, header decided), e.g.:

policies::manual::tests::test_manual_rid_key_outranks_header_key
  assertion `left == right` failed: rid-derived key must outrank the header key
  left: Some(0)  right: Some(1)

After the change:

cargo test -p smg --lib
test result: ok. 1754 passed; 0 failed; 5 ignored
cargo clippy -p smg --all-targets -- -D warnings   # clean
cargo +nightly fmt --all                            # clean

Manual repro from the issue (two workers, smg --policy manual --routing-key-override): pin key_a and key_b to different workers via x-smg-routing-key, then send -H 'x-smg-routing-key: key_a' with body "rid": "key_b" — now lands on key_b's worker.

Once #2462 lands, TestRoutingKeyPinning::test_body_rid_outranks_header_key can run under manual again.

Closes #2480

Checklist

• cargo +nightly fmt passes
• cargo clippy --all-targets --all-features -- -D warnings passes
• (Optional) Documentation updated
• (Optional) Please join us on Slack #sig-smg to discuss, review, and merge PRs

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Advanced

Run ID: 07d16029-c2eb-4f83-8fd2-8596d2ca95b8

📥 Commits

Reviewing files that changed from the base of the PR and between 07c8e82 and 6e6977f.

📒 Files selected for processing (3)

• model_gateway/src/policies/consistent_hashing.rs
• model_gateway/src/policies/manual.rs
• model_gateway/src/policies/registry.rs

Included review availability: Your plan provides up to 4 included reviews per hour; 2 remain after this review.

---

📝 Summary

Summary by CodeRabbit

• Bug Fixes

  • Improved worker selection consistency for conversations by prioritizing request-derived routing identifiers over rotating routing-key hints.
  • Prevented changing routing headers from unexpectedly breaking conversation pinning.
  • Applied consistent routing behavior across manual and consistent-hashing policies, including separate conversation-leg handling.

• Tests

  • Added coverage for routing-identifier precedence across conversation turns and policy types.

Walkthrough

Manual and consistent hashing policies now prefer rid_key over populated routing-key fields and X-SMG-Routing-Key headers. Tests cover precedence across turns, per-leg namespacing, and registry behavior with routing-key overrides.

Changes

RID-aware policy routing

Layer / File(s)	Summary
Policy key precedence model_gateway/src/policies/manual.rs, model_gateway/src/policies/consistent_hashing.rs	Manual and consistent hashing policies now prefer rid_key. Manual policy retains per-leg namespacing. Tests verify RID precedence and independent leg entries.
Registry policy validation model_gateway/src/policies/registry.rs	Documentation and tests verify RID precedence across both policies when routing-key overrides are enabled. Test imports include HashRing.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Bug fix · Severity of issue fixed: Medium

Suggested reviewers: catherinesue

Merge Risk: ⚪ Minimal · up to 6e697

RID-aware routing now consistently pins requests using the body-derived key when enabled, while retaining the documented fallback behavior. No current merge-blocking risk was identified.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: both policies now prioritize the body-derived rid over the header routing key.
Description check	✅ Passed	The description directly explains the problem, solution, affected policies, behavior with and without the override, tests, and validation results.
Linked Issues check	✅ Passed	Issue #2480 requires body-derived rid_key precedence under --routing-key-override for manual and consistent_hashing. The diff implements rid_key before the routing-key header in manual. It…
Out of Scope Changes check	✅ Passed	The changes stay within issue #2480. They update the two affected policies, related registry documentation, and automated tests for precedence and namespacing. No unrelated behavior or files are intro…
Docstring Coverage	✅ Passed	Docstring coverage is 87.50% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 16 functions across 3 files.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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