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.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 5d071550-b01c-44e4-83da-9ff002f312e1

📥 Commits

Reviewing files that changed from the base of the PR and between 43779b4 and f47310d.

📒 Files selected for processing (23)

• crates/mock_worker/Cargo.toml
• crates/mock_worker/src/lib.rs
• crates/mock_worker/src/main.rs
• model_gateway/Cargo.toml
• model_gateway/src/app_context.rs
• model_gateway/src/rate_limit/rejection.rs
• model_gateway/src/routers/grpc/common/response_formatting.rs
• model_gateway/src/routers/grpc/common/stages/helpers.rs
• model_gateway/src/routers/grpc/common/stages/mod.rs
• model_gateway/src/routers/grpc/common/stages/rate_limit.rs
• model_gateway/src/routers/grpc/context.rs
• model_gateway/src/routers/grpc/harmony/stages/response_processing.rs
• model_gateway/src/routers/grpc/harmony/streaming.rs
• model_gateway/src/routers/grpc/pipeline.rs
• model_gateway/src/routers/grpc/regular/responses/streaming.rs
• model_gateway/src/routers/grpc/regular/stages/chat/response_processing.rs
• model_gateway/src/routers/grpc/regular/stages/completion/response_processing.rs
• model_gateway/src/routers/grpc/regular/stages/generate/response_processing.rs
• model_gateway/src/routers/grpc/regular/stages/messages/response_processing.rs
• model_gateway/src/routers/grpc/regular/streaming.rs
• model_gateway/src/routers/grpc/router.rs
• model_gateway/tests/common/mod.rs
• model_gateway/tests/tenant_rate_limiting_grpc_test.rs

🚧 Files skipped from review as they are similar to previous changes (22)

• model_gateway/Cargo.toml
• crates/mock_worker/src/main.rs
• crates/mock_worker/Cargo.toml
• model_gateway/src/routers/grpc/regular/stages/generate/response_processing.rs
• model_gateway/src/routers/grpc/common/stages/mod.rs
• model_gateway/src/routers/grpc/regular/stages/chat/response_processing.rs
• model_gateway/src/routers/grpc/common/response_formatting.rs
• model_gateway/src/routers/grpc/common/stages/helpers.rs
• model_gateway/src/routers/grpc/regular/stages/messages/response_processing.rs
• model_gateway/src/app_context.rs
• model_gateway/src/rate_limit/rejection.rs
• model_gateway/src/routers/grpc/harmony/stages/response_processing.rs
• crates/mock_worker/src/lib.rs
• model_gateway/tests/common/mod.rs
• model_gateway/src/routers/grpc/regular/stages/completion/response_processing.rs
• model_gateway/src/routers/grpc/router.rs
• model_gateway/src/routers/grpc/harmony/streaming.rs
• model_gateway/src/routers/grpc/regular/responses/streaming.rs
• model_gateway/src/routers/grpc/common/stages/rate_limit.rs
• model_gateway/src/routers/grpc/pipeline.rs
• model_gateway/src/routers/grpc/regular/streaming.rs
• model_gateway/src/routers/grpc/context.rs

---

📝 Walkthrough

Summary by CodeRabbit

• New Features

  • Added tenant-aware rate limiting for chat, completion, message, and generation requests.
  • Added consistent rate-limit handling for streaming and non-streaming responses.
  • Requests exceeding finite limits now provide retry guidance when applicable.

• Bug Fixes

  • Prevented double-counting shared prompt and cached token usage.
  • Prevented rate-limit denials from being retried.
  • Improved reservation cleanup for disconnected or incomplete streams.
  • Omitted retry guidance when a request cannot succeed within the limit.

• Tests

  • Added end-to-end coverage for enforcement, settlement, streaming, and disabled limits.

Walkthrough

The change exposes mock_worker as a library and adds tenant rate-limit reservation handling to gRPC pipelines. Reservations persist across retries, settle from authoritative usage, transfer through streaming responses, and receive end-to-end coverage.

Changes

Tenant rate limiting

Layer / File(s)	Summary
Mock worker library surface crates/mock_worker/*, model_gateway/Cargo.toml	The mock worker is available as a library dependency for in-process tests.
Reservation contracts and admission stage model_gateway/src/app_context.rs, model_gateway/src/rate_limit/*, model_gateway/src/routers/grpc/common/*	Requests store shared reservation state. The new stage reserves estimated input tokens, caches outcomes across retries, and handles reservation cleanup. Usage aggregation counts shared prompt and cache tokens once.
Pipeline and retry orchestration model_gateway/src/routers/grpc/pipeline.rs, model_gateway/src/routers/grpc/router.rs	All supported gRPC pipelines receive rate-limit dependencies. Retries share reservation state, rate-limit denials are not retried, and successful responses settle actual usage.
Streaming handoff and settlement model_gateway/src/routers/grpc/harmony/*, model_gateway/src/routers/grpc/regular/*	Streaming processors transfer reservations to response bodies and settle them only after authoritative completion messages are received for all expected choices.
End-to-end validation model_gateway/tests/common/mod.rs, model_gateway/tests/tenant_rate_limiting_grpc_test.rs	Tests use a real router and in-process mock worker to verify denials, retry headers, backend-reported usage, and disabled rate limiting.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Router
  participant RequestPipeline
  participant RateLimitManager
  participant StreamingProcessor
  participant ResponseBody
  Router->>RequestPipeline: execute request with shared reservation state
  RequestPipeline->>RateLimitManager: reserve tenant input-token budget
  RateLimitManager-->>RequestPipeline: admitted reservation or denial
  RequestPipeline->>StreamingProcessor: process response with reservation
  StreamingProcessor->>ResponseBody: attach reservation lifecycle
  ResponseBody-->>RateLimitManager: settle complete usage or abandon unresolved reservation

Loading

Possibly related PRs

• smg-project/smg#337: Overlaps in shared token tracking and usage accounting.
• smg-project/smg#519: Overlaps in prompt and cache token aggregation.

Suggested reviewers: catherinesue, slin1237

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes wiring tenant token rate limiting into the gRPC router.
Description check	✅ Passed	The description directly explains the tenant rate-limiting implementation, affected pipelines, cleanup behavior, tests, and scope.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch xz/grpc-rate-limit-reserve-guard

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (7)

model_gateway/src/routers/grpc/regular/stages/chat/response_processing.rs (1)

125-134: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

🟡 Nit: Five response-processing stages repeat the same four-arm attachment match. Each stage combines optional load guards with an optional ReservationAttachment using the same match (guards, reservation) block. No shared helper exists for that combination, so the pattern was copied once per stage. A future third attachment kind would require eight arms in five places.

Add one helper, for example AttachedBody::wrap_with_optional(response, guards, reservation), and call it from each site.

• model_gateway/src/routers/grpc/regular/stages/chat/response_processing.rs#L125-L134: replace the four-arm match with the shared helper call.
• model_gateway/src/routers/grpc/regular/stages/completion/response_processing.rs#L106-L115: replace the four-arm match with the shared helper call.
• model_gateway/src/routers/grpc/regular/stages/generate/response_processing.rs#L117-L126: replace the four-arm match with the shared helper call.
• model_gateway/src/routers/grpc/regular/stages/messages/response_processing.rs#L112-L121: replace the four-arm match with the shared helper call.
• model_gateway/src/routers/grpc/harmony/stages/response_processing.rs#L100-L110: replace the four-arm match with the shared helper call.

🤖 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 `@model_gateway/src/routers/grpc/regular/stages/chat/response_processing.rs`
around lines 125 - 134, Factor the repeated optional-attachment logic into an
AttachedBody helper such as wrap_with_optional(response, guards, reservation),
preserving the existing behavior for every Some/None combination. Replace the
four-arm match at
model_gateway/src/routers/grpc/regular/stages/chat/response_processing.rs:125-134,
completion/response_processing.rs:106-115,
generate/response_processing.rs:117-126,
messages/response_processing.rs:112-121, and
model_gateway/src/routers/grpc/harmony/stages/response_processing.rs:100-110
with calls to the shared helper.

model_gateway/src/routers/grpc/pipeline.rs (1)

1613-1755: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

🟡 Nit: Reservation settlement has no unit coverage.

The tests cover the reserve stage well: single admission across retries, cached denial, disabled manager, and a missing cell. They do not cover settle_reservation. The zero-usage case flagged at Lines 508-514 would be caught by a test that asserts the settled amount for a response without usage.

Do you want me to add those settlement tests?

🤖 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 `@model_gateway/src/routers/grpc/pipeline.rs` around lines 1613 - 1755, Add
unit coverage for settle_reservation, including a response without usage that
settles a zero token amount. Extend the existing rate_limit_reserve_tests module
with setup using the visible RateLimitManager and reservation context, then
assert the manager records the expected settled amount for zero usage while
preserving the existing reserve-stage tests.

Source: Coding guidelines

model_gateway/src/routers/grpc/regular/streaming.rs (1)

2992-3006: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

🟡 Nit: The completion settle gate reads a metrics map, and the field doc no longer matches the field.

Two small issues in the same computation:

1. saw_complete is derived from reasoning_tokens.len(). That map exists to aggregate reasoning-token metrics. It happens to receive exactly one entry per completed index today, so the count is correct. It is not a completion counter, and nothing enforces that it stays one. A later change that skips the insert for zero reasoning tokens, or moves it into the Chunk arm, would silently break the reservation gate. Track completed indices in a dedicated HashSet<u32>, as process_generate_streaming already does at Line 907.

2. The field doc at Lines 70-73 states that saw_complete records whether a Complete message "was ever seen". Line 2998 sets it only when every expected choice completed. Align the doc with the stricter meaning, or rename the field to all_choices_completed.

♻️ Proposed change: dedicated counter

     let mut reasoning_tokens: HashMap<u32, u32> = HashMap::new();
+    let mut completed_indices: HashSet<u32> = HashSet::new();

                     let index = index_offset + complete.index();
+                    completed_indices.insert(index);
                     total_prompt = total_prompt.max(complete.prompt_tokens());

-        let expected_choices = completion_request.n.unwrap_or(1).max(1);
-        let saw_complete = reasoning_tokens.len() as u32 >= expected_choices;
+        let expected_choices = completion_request.n.unwrap_or(1).max(1);
+        let saw_complete = completed_indices.len() as u32 >= expected_choices;

🤖 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 `@model_gateway/src/routers/grpc/regular/streaming.rs` around lines 2992 -
3006, Track completed choice indices in a dedicated HashSet<u32> within the
streaming completion flow, following the existing process_generate_streaming
pattern, and derive saw_complete from that set reaching expected_choices instead
of reasoning_tokens.len(). Update the saw_complete field documentation to
describe that all expected choices completed, or rename the field consistently
if that better matches the established API.

model_gateway/src/routers/grpc/router.rs (1)

550-614: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

🟡 Nit: Four dispatch methods now repeat the same reservation scaffolding.

route_chat_impl, route_generate_impl, route_messages_impl, and route_completion_impl each repeat six steps: canonicalize the model ID, rewrite the body model, clone the retry inputs, create the RateLimitCell, suppress retries on Denied, and close the reservation after the retry loop. Any future change to the reservation lifecycle must land in four places.

Two smaller cleanups are available without a large refactor:

1. Extract the denial predicate, for example fn reservation_denied(cell: &RateLimitCell) -> bool, and call it from all four should_retry closures.
2. resolve_retry_config resolves the alias itself, and it now receives an already-canonical ID at Line 567. The second resolve is redundant. Consider a retry_config_for_canonical variant that skips it.

A full generic extraction over the four request and response types is likely not worth the type gymnastics here.

🤖 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 `@model_gateway/src/routers/grpc/router.rs` around lines 550 - 614, Reduce
duplicated reservation logic across route_chat_impl, route_generate_impl,
route_messages_impl, and route_completion_impl by extracting a shared
reservation_denied(&RateLimitCell) predicate and using it in each should_retry
closure. Add a retry configuration helper for already-canonical model IDs, then
use it after resolve_canonical_model_id so retry configuration does not resolve
the alias again; leave the broader dispatch flow unchanged.

model_gateway/tests/tenant_rate_limiting_grpc_test.rs (1)

94-174: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy lift

🟡 Nit: Confirm equivalent end-to-end coverage exists for other pipelines and PD dispatch.

build_router only wires regular_mode(vec![]) with a Grpc/TokenSpeed worker, and every test in this file drives route_generate. The PR wires tenant rate limiting into Chat, Messages, Completion, Generate, and Harmony pipelines, and PD disaggregation introduces a distinct dual-dispatch path (separate prefill/decode reservation and settlement) alongside regular routing.

Confirm whether equivalent end-to-end reservation/settlement coverage exists elsewhere for the other pipelines and for PD mode, since Harmony's multi-choice completion gating and PD's split dispatch could plausibly diverge from the single-dispatch Generate behavior asserted here.

As per coding guidelines, "Account for dual-dispatch complexity introduced by PD disaggregation in addition to regular routing" and "Run the pr-test-analyzer agent to verify that tests adequately cover new or changed functionality."

🤖 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 `@model_gateway/tests/tenant_rate_limiting_grpc_test.rs` around lines 94 - 174,
Extend the tenant rate-limiting integration coverage beyond build_router’s
regular TokenSpeed route_generate path: add equivalent reservation/settlement
tests for Chat, Messages, Completion, and Harmony, including Harmony
multi-choice behavior, and add PD disaggregation tests covering separate
prefill/decode dispatch reservations and settlement. Reuse the existing router
setup and assertions where possible, and verify the resulting test suite with
the pr-test-analyzer agent.

Source: Coding guidelines

model_gateway/tests/common/mod.rs (1)

328-473: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

🟡 Nit: Extract shared AppContext construction logic instead of duplicating it a third time.

create_test_context_with_tokenizer_registry (Lines 328-473) repeats nearly all of the ~140-line tail already present in the unchanged create_test_context_with_parsers (Lines 481-620): registries, storage backends, worker monitor, job queue, workflow engines, OpenAI-mode worker registration, and MCP orchestrator setup. The only real differences are the tokenizer registry source, the parser factories, and now the rate-limit manager.

Consider factoring the shared tail into a private helper that accepts these three pieces as parameters, so future AppContext field additions (like this PR's rate_limit_manager) only need to be wired in one place instead of drifting across multiple near-identical test helpers.

♻️ Sketch of a shared helper

+fn build_test_app_context(
+    config: RouterConfig,
+    client: reqwest::Client,
+    tokenizer_registry: Arc<TokenizerRegistry>,
+    reasoning_parser_factory: Option<ReasoningParserFactory>,
+    tool_parser_factory: Option<ToolParserFactory>,
+    rate_limiter: Option<Arc<TokenBucket>>,
+    rate_limit_manager: Option<Arc<smg::rate_limit::RateLimitManager>>,
+) -> Pin<Box<dyn Future<Output = Arc<AppContext>> + Send>> {
+    // ... shared registries/storage/monitor/job-queue/workflow/MCP init,
+    // moved out of create_test_context_with_tokenizer_registry and
+    // create_test_context_with_parsers.
+}

🤖 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 `@model_gateway/tests/common/mod.rs` around lines 328 - 473, Extract the shared
AppContext setup from create_test_context_with_tokenizer_registry and
create_test_context_with_parsers into one private helper. Parameterize the
helper with the tokenizer registry, reasoning parser factory, and tool parser
factory, and centralize rate_limit_manager plus all shared registries, storage,
monitor, queue, workflow, worker, and MCP initialization there. Keep both public
helpers as thin wrappers that supply their respective inputs and preserve their
existing behavior.

model_gateway/src/routers/grpc/common/stages/rate_limit.rs (1)

111-114: 🗄️ Data Integrity & Integration | 🔵 Trivial

🟡 Nit — Track HTTP/gRPC rate-limit parity as a follow-up.

The doc comment states Responses, embeddings, and classify have not opted into rate_limit_cell, and the PR scope also excludes non-gRPC routers. This is disclosed, intentional, phased scope, so it is not a defect in this change. Track closing this gap in a follow-up so the HTTP and gRPC routers eventually enforce the same tenant rate-limit contract.

As per coding guidelines, "Ensure HTTP and gRPC routers implement the same API contract across both code paths."

🤖 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 `@model_gateway/src/routers/grpc/common/stages/rate_limit.rs` around lines 111
- 114, Track a follow-up to align rate-limit enforcement between HTTP and gRPC
routers, including Responses, embeddings, and classify endpoints that currently
lack rate_limit_cell opt-in. Preserve the current scoped behavior in the
rate-limit reservation stage and document the follow-up using the project’s
established issue or TODO tracking mechanism.

Source: Coding guidelines

🤖 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 `@model_gateway/src/routers/grpc/common/stages/rate_limit.rs`:
- Around line 144-149: In the missing-tenant-metadata branch of the rate-limit
`execute()` flow, add a `tracing::warn!` describing that rate limiting is
failing open because `ctx.input.tenant_request_meta` is absent, then preserve
the existing `Ok(None)` return.

---

Nitpick comments:
In `@model_gateway/src/routers/grpc/common/stages/rate_limit.rs`:
- Around line 111-114: Track a follow-up to align rate-limit enforcement between
HTTP and gRPC routers, including Responses, embeddings, and classify endpoints
that currently lack rate_limit_cell opt-in. Preserve the current scoped behavior
in the rate-limit reservation stage and document the follow-up using the
project’s established issue or TODO tracking mechanism.

In `@model_gateway/src/routers/grpc/pipeline.rs`:
- Around line 1613-1755: Add unit coverage for settle_reservation, including a
response without usage that settles a zero token amount. Extend the existing
rate_limit_reserve_tests module with setup using the visible RateLimitManager
and reservation context, then assert the manager records the expected settled
amount for zero usage while preserving the existing reserve-stage tests.

In `@model_gateway/src/routers/grpc/regular/stages/chat/response_processing.rs`:
- Around line 125-134: Factor the repeated optional-attachment logic into an
AttachedBody helper such as wrap_with_optional(response, guards, reservation),
preserving the existing behavior for every Some/None combination. Replace the
four-arm match at
model_gateway/src/routers/grpc/regular/stages/chat/response_processing.rs:125-134,
completion/response_processing.rs:106-115,
generate/response_processing.rs:117-126,
messages/response_processing.rs:112-121, and
model_gateway/src/routers/grpc/harmony/stages/response_processing.rs:100-110
with calls to the shared helper.

In `@model_gateway/src/routers/grpc/regular/streaming.rs`:
- Around line 2992-3006: Track completed choice indices in a dedicated
HashSet<u32> within the streaming completion flow, following the existing
process_generate_streaming pattern, and derive saw_complete from that set
reaching expected_choices instead of reasoning_tokens.len(). Update the
saw_complete field documentation to describe that all expected choices
completed, or rename the field consistently if that better matches the
established API.

In `@model_gateway/src/routers/grpc/router.rs`:
- Around line 550-614: Reduce duplicated reservation logic across
route_chat_impl, route_generate_impl, route_messages_impl, and
route_completion_impl by extracting a shared reservation_denied(&RateLimitCell)
predicate and using it in each should_retry closure. Add a retry configuration
helper for already-canonical model IDs, then use it after
resolve_canonical_model_id so retry configuration does not resolve the alias
again; leave the broader dispatch flow unchanged.

In `@model_gateway/tests/common/mod.rs`:
- Around line 328-473: Extract the shared AppContext setup from
create_test_context_with_tokenizer_registry and create_test_context_with_parsers
into one private helper. Parameterize the helper with the tokenizer registry,
reasoning parser factory, and tool parser factory, and centralize
rate_limit_manager plus all shared registries, storage, monitor, queue,
workflow, worker, and MCP initialization there. Keep both public helpers as thin
wrappers that supply their respective inputs and preserve their existing
behavior.

In `@model_gateway/tests/tenant_rate_limiting_grpc_test.rs`:
- Around line 94-174: Extend the tenant rate-limiting integration coverage
beyond build_router’s regular TokenSpeed route_generate path: add equivalent
reservation/settlement tests for Chat, Messages, Completion, and Harmony,
including Harmony multi-choice behavior, and add PD disaggregation tests
covering separate prefill/decode dispatch reservations and settlement. Reuse the
existing router setup and assertions where possible, and verify the resulting
test suite with the pr-test-analyzer agent.

🪄 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: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 5a8f2db8-4007-4989-93fc-0cd14799b003

📥 Commits

Reviewing files that changed from the base of the PR and between 386be7a and 9f78090.

📒 Files selected for processing (22)

• crates/mock_worker/Cargo.toml
• crates/mock_worker/src/lib.rs
• crates/mock_worker/src/main.rs
• model_gateway/Cargo.toml
• model_gateway/src/app_context.rs
• model_gateway/src/rate_limit/rejection.rs
• model_gateway/src/routers/grpc/common/response_formatting.rs
• model_gateway/src/routers/grpc/common/stages/mod.rs
• model_gateway/src/routers/grpc/common/stages/rate_limit.rs
• model_gateway/src/routers/grpc/context.rs
• model_gateway/src/routers/grpc/harmony/stages/response_processing.rs
• model_gateway/src/routers/grpc/harmony/streaming.rs
• model_gateway/src/routers/grpc/pipeline.rs
• model_gateway/src/routers/grpc/regular/responses/streaming.rs
• model_gateway/src/routers/grpc/regular/stages/chat/response_processing.rs
• model_gateway/src/routers/grpc/regular/stages/completion/response_processing.rs
• model_gateway/src/routers/grpc/regular/stages/generate/response_processing.rs
• model_gateway/src/routers/grpc/regular/stages/messages/response_processing.rs
• model_gateway/src/routers/grpc/regular/streaming.rs
• model_gateway/src/routers/grpc/router.rs
• model_gateway/tests/common/mod.rs
• model_gateway/tests/tenant_rate_limiting_grpc_test.rs

[coderabbitai]

🧹 Nitpick comments (2)

model_gateway/tests/common/mod.rs (1)

506-518: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

🟡 Nit: create_test_context_with_parsers hardcodes None for the rate-limit manager.

create_test_context_with_tokenizer_registry derives the manager from config, but this helper ignores tenant rate-limit configuration. A test that enables tenant rate limits through this helper gets no manager, so the request path silently skips reservation and the test passes for the wrong reason. Derive the manager from config here as well, or document that this helper intentionally disables tenant rate limiting.

♻️ Proposed change

 pub fn create_test_context_with_parsers(
     config: RouterConfig,
 ) -> Pin<Box<dyn Future<Output = Arc<AppContext>> + Send>> {
+    // Tenant rate limiting stays disabled here on purpose: callers of this
+    // helper do not exercise the reservation path.
     Box::pin(build_test_app_context(
         config,
         Arc::new(TokenizerRegistry::new()),
         Some(ReasoningParserFactory::new()),
         Some(ToolParserFactory::new()),
         None,
         empty_mcp_config(),
     ))
 }

🤖 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 `@model_gateway/tests/common/mod.rs` around lines 506 - 518, Update
create_test_context_with_parsers to derive and pass the rate-limit manager from
config, matching create_test_context_with_tokenizer_registry, instead of
hardcoding None; preserve the helper’s existing parser factory initialization.

model_gateway/src/routers/grpc/common/stages/helpers.rs (1)

84-105: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win

🟡 Nit — Add focused tests for combined response guards.

attach_response_guards() exercises (Some, Some), (Some, None), (None, Some), and (None, None) via the streaming response paths, but the existing tests exercise load guards and reservation handles through separate RAII/unit tests rather than through AttachedBody::wrap_response with the shared (guards, ReservationAttachment) combo. Add a small focused test that drops the combined wrapped body and asserts both the load guard and unresolved reservation are cleaned up.

🤖 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 `@model_gateway/src/routers/grpc/common/stages/helpers.rs` around lines 84 -
105, Add a focused test for attach_response_guards covering the
(Some(LoadGuards), Some(ReservationAttachment)) case, wrapping a response
through AttachedBody::wrap_response, then dropping the body and asserting both
the load guard and unresolved reservation are released. Reuse the existing RAII
test fixtures and assertions, keeping the test scoped to combined guard cleanup.

Source: Coding guidelines

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

Nitpick comments:
In `@model_gateway/src/routers/grpc/common/stages/helpers.rs`:
- Around line 84-105: Add a focused test for attach_response_guards covering the
(Some(LoadGuards), Some(ReservationAttachment)) case, wrapping a response
through AttachedBody::wrap_response, then dropping the body and asserting both
the load guard and unresolved reservation are released. Reuse the existing RAII
test fixtures and assertions, keeping the test scoped to combined guard cleanup.

In `@model_gateway/tests/common/mod.rs`:
- Around line 506-518: Update create_test_context_with_parsers to derive and
pass the rate-limit manager from config, matching
create_test_context_with_tokenizer_registry, instead of hardcoding None;
preserve the helper’s existing parser factory initialization.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 3f524414-635f-477d-ad22-b6ad059d6762

📥 Commits

Reviewing files that changed from the base of the PR and between 9f78090 and 43779b4.

📒 Files selected for processing (12)

• model_gateway/src/routers/grpc/common/stages/helpers.rs
• model_gateway/src/routers/grpc/common/stages/rate_limit.rs
• model_gateway/src/routers/grpc/harmony/stages/response_processing.rs
• model_gateway/src/routers/grpc/pipeline.rs
• model_gateway/src/routers/grpc/regular/stages/chat/response_processing.rs
• model_gateway/src/routers/grpc/regular/stages/completion/response_processing.rs
• model_gateway/src/routers/grpc/regular/stages/generate/response_processing.rs
• model_gateway/src/routers/grpc/regular/stages/messages/response_processing.rs
• model_gateway/src/routers/grpc/regular/streaming.rs
• model_gateway/src/routers/grpc/router.rs
• model_gateway/tests/common/mod.rs
• model_gateway/tests/tenant_rate_limiting_grpc_test.rs

🚧 Files skipped from review as they are similar to previous changes (10)

• model_gateway/src/routers/grpc/harmony/stages/response_processing.rs
• model_gateway/src/routers/grpc/router.rs
• model_gateway/src/routers/grpc/regular/stages/chat/response_processing.rs
• model_gateway/src/routers/grpc/regular/stages/completion/response_processing.rs
• model_gateway/src/routers/grpc/common/stages/rate_limit.rs
• model_gateway/src/routers/grpc/regular/stages/generate/response_processing.rs
• model_gateway/src/routers/grpc/regular/stages/messages/response_processing.rs
• model_gateway/tests/tenant_rate_limiting_grpc_test.rs
• model_gateway/src/routers/grpc/pipeline.rs
• model_gateway/src/routers/grpc/regular/streaming.rs