fix(concurrency): return permits on request cancellation

[Moersity]

Description

Problem

The concurrency middleware only transferred an acquired token into TokenGuardBody after next.run(request).await completed. If the client disconnected or the request future was cancelled before a response body was created, the acquired token was never returned. Queued requests could also leak a token when the oneshot receiver disappeared after acquisition. With a zero refill rate, each leak permanently reduced concurrency capacity.

Solution

Represent each successful acquisition as an RAII TokenPermit. The middleware owns the permit while the request future runs, then moves it into the response body guard. Dropping the request future, response body, or an undeliverable queued permit now returns the token synchronously.

Changes

• Add TokenPermit to couple token acquisition and release.
• Hold the permit across next.run so cancellation returns the token.
• Send permits, rather than unit acknowledgements, through the queue oneshot.
• Add regression tests for request cancellation, dropped response bodies, and cancelled queued requests.

Test Plan

• cargo +nightly fmt --all
• cargo clippy --all-targets --all-features -- -D warnings
• env -u RUST_LOG cargo test
• cargo test -p smg middleware::concurrency::tests --lib
• Verify aborting a task after it acquires the only zero-refill token restores bucket availability from 0 to 1.
• Verify dropping a guarded response body restores the acquired token.
• Verify dropping a queued request receiver before permit delivery returns the acquired token.

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: Pro Plus

Run ID: df29533a-8383-4a47-8e78-cbc6358ea15f

📥 Commits

Reviewing files that changed from the base of the PR and between 937bd70 and 8d9d237.

📒 Files selected for processing (1)

• model_gateway/src/middleware/concurrency.rs

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

• model_gateway/src/middleware/concurrency.rs

---

📝 Walkthrough

Summary by CodeRabbit

• Bug Fixes
  • Improved request concurrency handling so capacity is reliably restored when responses complete or requests are cancelled.
  • Fixed queued requests cancelled before processing, preventing them from reducing available capacity.
  • Added safeguards to keep request limits accurate across interrupted, cancelled, and completed requests.
  • Improved handling of response-body cleanup to ensure capacity is returned even when responses are dropped early.

Walkthrough

The concurrency middleware now uses owned RAII token permits. Queue and middleware paths transfer permits through responses and wrap response bodies with shared logic. Tests cover token release after body drops, task cancellation, and queued request cancellation.

Changes

Concurrency permit lifecycle

Layer / File(s)	Summary
Permit ownership and response wrapping model_gateway/src/middleware/concurrency.rs	TokenPermit owns acquired tokens and returns them when dropped. TokenGuardBody stores the permit, and run_with_permit centralizes response wrapping.
Queued permit transfer model_gateway/src/middleware/concurrency.rs	Immediate and delayed queue acquisition send TokenPermit values through the permit channel.
Middleware execution and cancellation tests model_gateway/src/middleware/concurrency.rs	Immediate and queued middleware paths use shared permit-wrapped execution. Tests cover body drops, task cancellation, and cancelled queued requests.

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

Sequence Diagram(s)

sequenceDiagram
  participant RequestMiddleware
  participant PermitQueue
  participant TokenPermit
  participant TokenGuardBody
  RequestMiddleware->>PermitQueue: request permit
  PermitQueue->>TokenPermit: acquire tokens
  PermitQueue-->>RequestMiddleware: send owned permit
  RequestMiddleware->>TokenGuardBody: wrap response with permit
  TokenGuardBody-->>TokenPermit: drop on completion or cancellation
  TokenPermit->>PermitQueue: return tokens

Loading

Suggested reviewers: catherinesue, key4ng

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the concurrency fix for returning permits when requests are cancelled.
Description check	✅ Passed	The description directly explains the token leak problem, RAII solution, implementation changes, and regression tests.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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 unit tests (beta)

• Create PR with unit tests

---

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: 2

🤖 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/middleware/concurrency.rs`:
- Around line 75-77: Update TokenGuardBody::new so it cannot create a
TokenPermit without first acquiring tokens: remove the constructor or make it
call TokenBucket::try_acquire and return Result<Self, ()>. Keep TokenPermit-only
wrapping private within the module, and run the type-design-analyzer for the
Rust type invariant and encapsulation review.
- Around line 320-329: Update the concurrency test around the spawned task to
synchronize on successful permit acquisition instead of relying on
tokio::task::yield_now(). Signal from inside the task immediately after
TokenPermit::try_acquire succeeds, await that signal before asserting
bucket.available_tokens() is zero, and preserve the existing abort and
token-restoration assertions.

🪄 Autofix

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: 19977f10-0132-42a1-a7cf-5dee0b1edfd2

📥 Commits

Reviewing files that changed from the base of the PR and between 3c4e508 and 937bd70.

📒 Files selected for processing (1)

• model_gateway/src/middleware/concurrency.rs