feat(grpc): wire tenant token rate limiting into gRPC router

[XinyueZhang369]

Description

Problem

RateLimitManager::reserve() / settle_success() / close_reserved_only() (Phase 1, already merged on main) aren't called anywhere yet — nothing enforces tenant token-per-minute limits on the gRPC router's chat/generate/completion/messages endpoints.

Solution

Wire tenant rate limiting into the gRPC pipeline with a "guarded reserve-once" mechanism: a small RateLimitCell shared across all retry attempts of one logical request, checked by a new RateLimitReserveStage inserted right after preparation. The first attempt to reach the stage reserves (or is denied) and caches the outcome; every later retry attempt sees the cached outcome and skips straight through. RetryExecutor reruns the full pipeline per attempt unchanged — no prepare/dispatch split, no independently-constructed state that has to be kept in sync by hand.

Changes

• RateLimitCell / RateLimitOutcome / RateLimitReserveStage (new, routers/grpc/common/stages/rate_limit.rs), inserted into the Chat/Messages/Completion/Harmony pipelines right after preparation.
• router.rs: each route_*_impl constructs the cell, canonicalizes the model once before the retry loop (reused for both the reservation and every dispatch attempt, including the request body's own model field), stops retrying on a cached denial, and closes any reservation a non-2xx final response never got to settle.
• Non-streaming success settles inline with the response's real usage; streaming settles with real accumulated totals via a new reservation: Option<Arc<SharedReservationHandle>> parameter threaded through regular/streaming.rs and harmony/streaming.rs, with a ReservationAttachment layered onto AttachedBody as the Drop-based safety net for disconnect/preemption before that point.
• crates/mock_worker gained a [lib] target so its real (mock) gRPC server is spawnable in-process from model_gateway/tests/, enabling real end-to-end integration coverage (tests/tenant_rate_limiting_grpc_test.rs) rather than only unit tests against the pipeline stage in isolation.
• Several rounds of review-driven correctness fixes on top of the initial wiring — see commit messages for full detail:
  • Batched-completion undercount (only the first prompt's tokens were reserved), preemption/cancellation reservation leak, alias handover across retries, and an unusable Retry-After: 18446744073709551615 for impossible requests.
  • Alias-pinning left the request body itself uncanonicalized (response metadata/parser selection still saw the client's alias); a clean streaming EOF without an authoritative Complete frame was settling with 0 input tokens instead of keeping the reservation; n>1 requests were charging the shared prompt once per choice instead of once per request.
  • A clean EOF partway through an n>1 stream (some choices finished, others didn't) was accepted as fully authoritative instead of requiring every expected choice's Complete; Harmony PD mode's prefill-populated map could mask a decode phase that never actually finished.
  • cached_tokens had the same per-choice multiplication bug as prompt_tokens.
• Explicitly out of scope (matches the original Phase 2 cut): Responses endpoint, embeddings, classify (none use RetryExecutor today); OpenAI/Anthropic/Gemini routers (Phase 3).

Test Plan

• New integration tests (model_gateway/tests/tenant_rate_limiting_grpc_test.rs), driven against a real (mock) gRPC backend: a finite-wait 429 with Retry-After, an impossible-request 429 without Retry-After, non-streaming settle using the backend's real reported usage (not just the reserve-time estimate), and the feature-disabled no-op path.
• New unit tests throughout the touched modules: RateLimitCell::drop safety-net behavior (abandons on preemption, doesn't double-resolve after streaming handoff or denial), total_input_token_count batching, the u64::MAX Retry-After sentinel, and build_usage's per-choice prompt_tokens/cached_tokens max-not-sum behavior.
• Full lib suite (1345 tests) + this branch's integration suite (13 tests), all green.

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

[claude]

Thorough review of the tenant rate limiting wiring into the gRPC router. The "guarded reserve-once" mechanism via RateLimitCell / RateLimitReserveStage is clean: reserves exactly once per logical request across retry attempts, settles with real backend-reported usage (not just the reserve estimate), and has proper RAII safety nets for every lifecycle edge (preemption/Drop, streaming handoff, non-success close, disconnect via ReservationAttachment).

Key correctness points verified:

• n>1 handling: max() (not sum) for shared prompt/cached tokens, sum() for per-choice completion tokens, and expected-choices guard before settling streaming usage
• Harmony PD mode tracks decode_completed_indices separately from prefill-populated prompt_tokens to avoid masking a decode phase that never finished
• CAS-guarded resolution on SharedReservationHandle ensures first-resolver-wins across all paths (inline settle, streaming settle, ReservationAttachment abandon, cell Drop)
• Retry loop correctly short-circuits on Denied — never re-reserves or retries a rate-limit denial
• Integration tests (real mock gRPC backend) prove settle_success runs with real usage by choosing budget numbers that distinguish reserve-only from reserve+settle

0 issues found.