Repository navigation
refactor(middleware): split middleware.rs into a folder + relocate token_bucket (Finding 2) - #1151
Merged
Merged
Conversation
…ken_bucket
Combines Finding 2 (relocate `worker/token_bucket.rs`) and a long-overdue
break-up of `middleware.rs` from `.claude/plans/2026-04-14-worker-module-followup-cleanup.md`
into a single PR. Pure structural refactor — zero behavior change, every
external import path stays valid via re-exports in `middleware/mod.rs`.
Why one PR
- `middleware.rs` was importing `crate::worker::token_bucket::TokenBucket`
via `pub use`, so the cleanest moment to relocate token_bucket is when
middleware itself becomes a folder. Doing the two changes in separate
PRs would create a transient `middleware.rs` → `middleware/mod.rs` →
`middleware.rs` flip that nobody needs.
Layout
Before:
model_gateway/src/middleware.rs (1110 lines, 8 concerns)
model_gateway/src/worker/token_bucket.rs
After:
model_gateway/src/middleware/mod.rs (re-exports for back-compat)
model_gateway/src/middleware/auth.rs (AuthConfig + auth_middleware)
model_gateway/src/middleware/concurrency.rs (TokenGuardBody, QueuedRequest,
QueueProcessor, ConcurrencyLimiter,
concurrency_limit_middleware)
model_gateway/src/middleware/logging.rs (RequestSpan, RequestLogger,
ResponseLogger, create_logging_layer)
model_gateway/src/middleware/metrics.rs (HttpMetricsLayer,
HttpMetricsMiddleware,
normalize_path_for_metrics +
is_dynamic_id + tests)
model_gateway/src/middleware/request_id.rs (RequestIdLayer,
RequestIdMiddleware,
generate_request_id, RequestId
re-export from smg_auth)
model_gateway/src/middleware/storage_context.rs (storage_context_middleware +
helpers + tests)
model_gateway/src/middleware/token_bucket.rs (relocated from worker/, git
rename so history follows)
model_gateway/src/middleware/wasm.rs (wasm_middleware)
What changed
- `model_gateway/src/middleware.rs`: deleted (1110 lines distributed across
the 8 files above, byte-for-byte identical except for per-file imports).
- `model_gateway/src/middleware/mod.rs`: new — registers all submodules and
re-exports the public surface (`AuthConfig`, `auth_middleware`,
`ConcurrencyLimiter`, `QueueProcessor`, `QueuedRequest`, `TokenGuardBody`,
`concurrency_limit_middleware`, `create_logging_layer`, `RequestLogger`,
`RequestSpan`, `ResponseLogger`, `HttpMetricsLayer`, `HttpMetricsMiddleware`,
`RequestId`, `RequestIdLayer`, `RequestIdMiddleware`,
`storage_context_middleware`, `TokenBucket`, `wasm_middleware`).
- `model_gateway/src/worker/token_bucket.rs` → `model_gateway/src/middleware/token_bucket.rs`:
git rename, file content unchanged. Tracks as a rename in the diff.
- `model_gateway/src/worker/mod.rs`: drops `pub mod token_bucket;`.
- Internal cross-module wiring inside the new folder:
* `concurrency.rs` imports `super::token_bucket::TokenBucket`.
* `logging.rs` imports `super::metrics::normalize_path_for_metrics` and
`super::request_id::RequestId`.
* `wasm.rs` imports `super::request_id::{generate_request_id, RequestId}`.
* `generate_request_id` is `pub(super)` so `wasm.rs` can use it without
leaking it past the module boundary.
* `normalize_path_for_metrics` is `pub(super)` so `logging.rs` can use it.
- Test modules followed their code: `build_storage_request_context_*`
tests live in `storage_context.rs`; the `normalize_path_*` and
`is_dynamic_id` tests live in `metrics.rs`.
Public surface preserved
- `crate::middleware::TokenBucket` still resolves (now via
`middleware/mod.rs:pub use token_bucket::TokenBucket`), so the existing
callers in `app_context.rs:19` and `service_discovery.rs:1240` are
untouched.
- Every `crate::middleware::{AuthConfig, QueuedRequest, …}` import in
`server.rs` still resolves via the re-exports.
- `RequestId` re-export from `smg_auth` is preserved (now in
`middleware/request_id.rs`).
- Net: no caller outside `worker/mod.rs` needs an import change.
What this unblocks
- `worker/` no longer hosts a primitive that has zero coupling to the
Worker abstraction — Finding 2 closed.
- `middleware.rs` is no longer a 1110-line bag of concerns; future
middleware additions can live in their own files instead of growing
the monolith.
- The folder layout sets a precedent for the next reviewer: one concern
per file, each file imports only what it uses.
What this does NOT do
- No new middleware trait. Tower's `Layer`/`Service` is already the
ecosystem-standard abstraction; introducing a project-local trait
would just be a wrapper.
- No behavior change. No test changes (tests moved alongside their
code, but their assertions are identical).
- No removal of public API. Every name that was public before is still
public.
Verification
- cargo check -p smg — clean.
- cargo clippy -p smg --all-targets -- -D warnings — clean.
- cargo clippy -p smg-golang --all-targets -- -D warnings — clean.
- cargo check -p smg-python — clean.
- cargo test -p smg --lib — 546 passed; 0 failed; 4 ignored.
- make fmt — clean.
- rustup run nightly cargo fmt -- --check — EXIT=0 (matches CI).
Refs: Finding 2 in .claude/plans/2026-04-14-worker-module-followup-cleanup.md
Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request refactors the monolithic middleware into a modular structure under model_gateway/src/middleware/, improving maintainability. The review feedback identifies potential token leaks in the concurrency limiter when requests are cancelled while queued or during execution. Additionally, an optimization was suggested for the WASM middleware to calculate timestamps once before the module execution loop to improve efficiency and consistency.
Comment on lines
+139
to
+164
| if self.token_bucket.try_acquire(1.0).is_ok() { | ||
| // Got token immediately | ||
| debug!("Queue: acquired token immediately for queued request"); | ||
| let _ = queued.permit_tx.send(Ok(())); | ||
| } else { | ||
| // Need to wait for token | ||
| let token_bucket = self.token_bucket.clone(); | ||
|
|
||
| // Spawn task only when we actually need to wait | ||
| #[expect( | ||
| clippy::disallowed_methods, | ||
| reason = "fire-and-forget permit acquisition: task is bounded by remaining_timeout and communicates via oneshot; dropping the JoinHandle detaches the task but it self-terminates" | ||
| )] | ||
| tokio::spawn(async move { | ||
| if token_bucket | ||
| .acquire_timeout(1.0, remaining_timeout) | ||
| .await | ||
| .is_ok() | ||
| { | ||
| debug!("Queue: acquired token after waiting"); | ||
| let _ = queued.permit_tx.send(Ok(())); | ||
| } else { | ||
| warn!("Queue: request timed out waiting for token"); | ||
| let _ = queued.permit_tx.send(Err(StatusCode::REQUEST_TIMEOUT)); | ||
| } | ||
| }); |
Contributor
There was a problem hiding this comment.
The current implementation of the queue processor can leak tokens if a request is cancelled while waiting in the queue or immediately after a token is acquired. If permit_tx.send() fails (because the receiver was dropped due to request cancellation), the acquired token is never returned to the bucket. Additionally, the spawned task should monitor for request cancellation while waiting for a token to avoid unnecessary acquisition. Note that data passed to the spawned task must have a 'static lifetime, so we use owned types or clones to ensure the data outlives the task.
// Try to acquire token for this request
if self.token_bucket.try_acquire(1.0).is_ok() {
// Got token immediately
debug!("Queue: acquired token immediately for queued request");
if queued.permit_tx.send(Ok(())).is_err() {
debug!("Queue: request cancelled, returning token");
self.token_bucket.return_tokens_sync(1.0);
}
} else {
// Need to wait for token
let token_bucket = self.token_bucket.clone();
let mut permit_tx = queued.permit_tx;
// Spawn task only when we actually need to wait
#[expect(
clippy::disallowed_methods,
reason = "fire-and-forget permit acquisition: task is bounded by remaining_timeout and communicates via oneshot; dropping the JoinHandle detaches the task but it self-terminates"
)]
tokio::spawn(async move {
tokio::select! {
_ = permit_tx.closed() => {
debug!("Queue: request cancelled while waiting for token");
}
res = token_bucket.acquire_timeout(1.0, remaining_timeout) => {
match res {
Ok(()) => {
debug!("Queue: acquired token after waiting");
if permit_tx.send(Ok(())).is_err() {
debug!("Queue: request cancelled after acquisition, returning token");
token_bucket.return_tokens_sync(1.0);
}
}
Err(_) => {
warn!("Queue: request timed out waiting for token");
let _ = permit_tx.send(Err(StatusCode::REQUEST_TIMEOUT));
}
}
}
}
});
}References
- Data passed to spawned background tasks must have a 'static lifetime. Use owned types or reference-counted pointers like Arc instead of passing references to ensure the data outlives the task.
Comment on lines
+238
to
+248
| if token_bucket.try_acquire(1.0).is_ok() { | ||
| debug!("Acquired token immediately"); | ||
| Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_ALLOWED); | ||
| let response = next.run(request).await; | ||
|
|
||
| // Wrap the response body with TokenGuardBody to return token when stream ends | ||
| // This ensures that for streaming responses, the token is only returned | ||
| // after the entire stream has been sent to the client. | ||
| let (parts, body) = response.into_parts(); | ||
| let guarded_body = TokenGuardBody::new(body, token_bucket, 1.0); | ||
| Response::from_parts(parts, Body::new(guarded_body)) |
Contributor
There was a problem hiding this comment.
There is a potential token leak in concurrency_limit_middleware. If the request is cancelled (e.g., the client disconnects) during the execution of next.run(request).await, the TokenGuardBody is never created, and the token acquired at line 238 is leaked. To fix this, you should use a RAII guard that returns the token on drop unless it is explicitly disarmed after the response is successfully wrapped in TokenGuardBody.
Comment on lines
+105
to
+118
| for module in modules_on_request { | ||
| let wasm_headers = build_wasm_headers_from_axum_headers(&headers); | ||
| let wasm_request = WasmRequest { | ||
| method: method_str.clone(), | ||
| path: path_str.clone(), | ||
| query: query_str.clone(), | ||
| headers: wasm_headers, | ||
| body: modified_body.clone(), | ||
| request_id: request_id.clone(), | ||
| now_epoch_ms: std::time::SystemTime::now() | ||
| .duration_since(std::time::UNIX_EPOCH) | ||
| .unwrap_or_else(|_| Duration::from_millis(0)) | ||
| .as_millis() as u64, | ||
| }; |
Contributor
There was a problem hiding this comment.
Calculating the current timestamp inside the loop for every WASM module is slightly inefficient and could lead to inconsistent timestamps across different modules for the same request. It is better to calculate the timestamp once before the loop.
Suggested change
| for module in modules_on_request { | |
| let wasm_headers = build_wasm_headers_from_axum_headers(&headers); | |
| let wasm_request = WasmRequest { | |
| method: method_str.clone(), | |
| path: path_str.clone(), | |
| query: query_str.clone(), | |
| headers: wasm_headers, | |
| body: modified_body.clone(), | |
| request_id: request_id.clone(), | |
| now_epoch_ms: std::time::SystemTime::now() | |
| .duration_since(std::time::UNIX_EPOCH) | |
| .unwrap_or_else(|_| Duration::from_millis(0)) | |
| .as_millis() as u64, | |
| }; | |
| let now_ms = std::time::SystemTime::now() | |
| .duration_since(std::time::UNIX_EPOCH) | |
| .unwrap_or_else(|_| Duration::from_millis(0)) | |
| .as_millis() as u64; | |
| for module in modules_on_request { | |
| let wasm_headers = build_wasm_headers_from_axum_headers(&headers); | |
| let wasm_request = WasmRequest { | |
| method: method_str.clone(), | |
| path: path_str.clone(), | |
| query: query_str.clone(), | |
| headers: wasm_headers, | |
| body: modified_body.clone(), | |
| request_id: request_id.clone(), | |
| now_epoch_ms: now_ms, | |
| }; |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
📝 WalkthroughWalkthroughThe pull request refactors a monolithic 1110-line middleware module into eight specialized submodules—auth, concurrency, logging, metrics, request_id, storage_context, token_bucket, and wasm—each handling distinct HTTP middleware concerns. The original middleware.rs is deleted and replaced with a modular structure via middleware/mod.rs, with approximately 1200 lines of new middleware logic distributed across the submodules. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning Review ran into problems🔥 ProblemsTimed out fetching pipeline failures after 30000ms Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/src/middleware/auth.rs`:
- Around line 39-51: The current strip_prefix("Bearer ") call is case-sensitive
and can miss valid tokens; update the token extraction around
request.headers().get(header::AUTHORIZATION) to perform a case-insensitive
"bearer " check (e.g., test h.len() > 7 and h[..7].eq_ignore_ascii_case("bearer
") then take &h[7..]) instead of strip_prefix, so the token variable, the
authorized computation (which uses Sha256::digest and expected_hash), and the
existing authorized logic remain unchanged but now accept any ASCII case
variation of the Bearer scheme.
In `@model_gateway/src/middleware/concurrency.rs`:
- Around line 76-84: The comment on poll_frame should use the repository
convention for non-unsafe assumptions: replace the "SAFETY:" marker with
"INVARIANT:" and keep the explanatory text that the inner body is never moved
and that Body is Unpin (type alias for UnsyncBoxBody which is Unpin); update the
comment above Pin::new(&mut this.inner).poll_frame(cx) to read "INVARIANT: We
never move the inner body, and Body is Unpin (it's a type alias for
UnsyncBoxBody which is Unpin)".
In `@model_gateway/src/middleware/logging.rs`:
- Around line 57-77: The RequestLogger::on_request currently attempts to read
RequestId from request.extensions but RequestIdLayer runs after TraceLayer, so
the ID will not be present here; either reorder middleware so RequestIdLayer
runs before TraceLayer if you need the ID recorded in the span at on_request
time, or (preferred per review) expand the existing comment near
RequestLogger/on_request to explicitly state that RequestId is intentionally
recorded later (because RequestIdLayer executes after TraceLayer) and document
where the request ID is actually attached (e.g., the layer or handler that
records it) so readers know why span.record("request_id", ...) may be empty.
In `@model_gateway/src/middleware/metrics.rs`:
- Around line 99-137: normalize_path_for_metrics can emit a trailing "{id}/" for
paths that end with a dynamic-id segment plus a slash; fix by tracking when you
replaced the current segment with "{id}" (e.g., a boolean like last_replaced set
inside the is_dynamic_id branch) and when handling the "if b == b'/'" branch
skip appending the trailing '/' if last_replaced is true and this slash is the
final byte of the path, otherwise append as before; ensure you clear
last_replaced when you append non-replacement segments so the flag only affects
the immediate trailing slash after an "{id}" replacement.
In `@model_gateway/src/middleware/wasm.rs`:
- Around line 230-240: Add a short clarifying comment in
model_gateway/src/middleware/wasm.rs inside the match for Action::Reject within
the OnResponse handling (the block that sets `status`, builds `final_response`,
assigns `*final_response.headers_mut() = headers`, and returns it) stating that
Reject during OnResponse intentionally preserves the current `headers` and
`body_bytes` (unlike OnRequest which returns an empty response) and explaining
the rationale (e.g., allow module to change status while keeping response
payload), so future readers know this semantic difference between OnRequest and
OnResponse.
- Around line 148-158: The unwrap_or_else on Request::builder() may swallow
build errors and drop the modified body; change the construction to match
Request::builder().method(method).uri(uri).body(Body::from(modified_body)) and
on Err(e) log the error (e.g. via tracing::error! or your logger) including e,
then fall back to creating a Request::new(Body::from(modified_body)) so the
modified_body is preserved; keep assigning *final_request.headers_mut() =
headers and *final_request.extensions_mut() = extensions before calling
next.run(final_request).await to ensure headers/extensions are preserved and the
error is visible in logs.
🪄 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: ASSERTIVE
Plan: Pro
Run ID: 923e6cd0-cff9-4871-9752-d8b924050451
📒 Files selected for processing (11)
model_gateway/src/middleware.rsmodel_gateway/src/middleware/auth.rsmodel_gateway/src/middleware/concurrency.rsmodel_gateway/src/middleware/logging.rsmodel_gateway/src/middleware/metrics.rsmodel_gateway/src/middleware/mod.rsmodel_gateway/src/middleware/request_id.rsmodel_gateway/src/middleware/storage_context.rsmodel_gateway/src/middleware/token_bucket.rsmodel_gateway/src/middleware/wasm.rsmodel_gateway/src/worker/mod.rs
💤 Files with no reviewable changes (2)
- model_gateway/src/worker/mod.rs
- model_gateway/src/middleware.rs
Comment on lines
+39
to
+51
| let token = request | ||
| .headers() | ||
| .get(header::AUTHORIZATION) | ||
| .and_then(|h| h.to_str().ok()) | ||
| .and_then(|h| h.strip_prefix("Bearer ")); | ||
|
|
||
| let authorized = token.is_some_and(|t| { | ||
| Sha256::digest(t.as_bytes()) | ||
| .as_slice() | ||
| .ct_eq(expected_hash) | ||
| .unwrap_u8() | ||
| == 1 | ||
| }); |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Consider case-insensitive "Bearer" prefix matching.
The strip_prefix("Bearer ") at line 43 is case-sensitive. Per RFC 6750, the "Bearer" scheme is case-insensitive. While most clients send "Bearer", some might send "bearer" or "BEARER".
♻️ Optional: case-insensitive Bearer prefix
let token = request
.headers()
.get(header::AUTHORIZATION)
.and_then(|h| h.to_str().ok())
- .and_then(|h| h.strip_prefix("Bearer "));
+ .and_then(|h| {
+ h.strip_prefix("Bearer ")
+ .or_else(|| h.strip_prefix("bearer "))
+ .or_else(|| h.strip_prefix("BEARER "))
+ });Or more robustly:
.and_then(|h| {
if h.len() > 7 && h[..7].eq_ignore_ascii_case("bearer ") {
Some(&h[7..])
} else {
None
}
})🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/src/middleware/auth.rs` around lines 39 - 51, The current
strip_prefix("Bearer ") call is case-sensitive and can miss valid tokens; update
the token extraction around request.headers().get(header::AUTHORIZATION) to
perform a case-insensitive "bearer " check (e.g., test h.len() > 7 and
h[..7].eq_ignore_ascii_case("bearer ") then take &h[7..]) instead of
strip_prefix, so the token variable, the authorized computation (which uses
Sha256::digest and expected_hash), and the existing authorized logic remain
unchanged but now accept any ASCII case variation of the Bearer scheme.
Comment on lines
+76
to
+84
| fn poll_frame( | ||
| self: Pin<&mut Self>, | ||
| cx: &mut Context<'_>, | ||
| ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> { | ||
| // SAFETY: We never move the inner body, and Body is Unpin | ||
| // (it's a type alias for UnsyncBoxBody which is Unpin) | ||
| let this = self.get_mut(); | ||
| Pin::new(&mut this.inner).poll_frame(cx) | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Clarify the pin projection safety comment.
The comment at lines 80-81 states "SAFETY" but this isn't an unsafe block. Per repository conventions from learnings, SAFETY: should be reserved for unsafe blocks, while INVARIANT: documents assumptions in safe code. Since Body being Unpin is an invariant the code relies on, consider updating the comment.
♻️ Use INVARIANT marker
fn poll_frame(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
- // SAFETY: We never move the inner body, and Body is Unpin
- // (it's a type alias for UnsyncBoxBody which is Unpin)
+ // INVARIANT: Body is Unpin (type alias for UnsyncBoxBody), so we can
+ // safely project through Pin without needing unsafe.
let this = self.get_mut();
Pin::new(&mut this.inner).poll_frame(cx)
}Based on learnings: "In Rust code across the repository, use the marker INVARIANT: to document assumptions in safe code. Reserve SAFETY: for explaining why unsafe blocks are sound."
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| fn poll_frame( | |
| self: Pin<&mut Self>, | |
| cx: &mut Context<'_>, | |
| ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> { | |
| // SAFETY: We never move the inner body, and Body is Unpin | |
| // (it's a type alias for UnsyncBoxBody which is Unpin) | |
| let this = self.get_mut(); | |
| Pin::new(&mut this.inner).poll_frame(cx) | |
| } | |
| fn poll_frame( | |
| self: Pin<&mut Self>, | |
| cx: &mut Context<'_>, | |
| ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> { | |
| // INVARIANT: Body is Unpin (type alias for UnsyncBoxBody), so we can | |
| // safely project through Pin without needing unsafe. | |
| let this = self.get_mut(); | |
| Pin::new(&mut this.inner).poll_frame(cx) | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/src/middleware/concurrency.rs` around lines 76 - 84, The
comment on poll_frame should use the repository convention for non-unsafe
assumptions: replace the "SAFETY:" marker with "INVARIANT:" and keep the
explanatory text that the inner body is never moved and that Body is Unpin (type
alias for UnsyncBoxBody which is Unpin); update the comment above Pin::new(&mut
this.inner).poll_frame(cx) to read "INVARIANT: We never move the inner body, and
Body is Unpin (it's a type alias for UnsyncBoxBody which is Unpin)".
Comment on lines
+57
to
+77
| impl<B> OnRequest<B> for RequestLogger { | ||
| fn on_request(&mut self, request: &Request<B>, span: &Span) { | ||
| let _enter = span.enter(); | ||
|
|
||
| // Try to get the request ID from extensions | ||
| // This will work if RequestIdLayer has already run | ||
| if let Some(request_id) = request.extensions().get::<RequestId>() { | ||
| span.record("request_id", request_id.0.as_str()); | ||
| } | ||
|
|
||
| let method = method_to_static_str(request.method().as_str()); | ||
| let path = normalize_path_for_metrics(request.uri().path()); | ||
| Metrics::record_http_request(method, &path); | ||
|
|
||
| // Log the request start | ||
| info!( | ||
| target: "smg::request", | ||
| "started processing request" | ||
| ); | ||
| } | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Middleware ordering dependency: RequestId may not be available in on_request.
Lines 33-34 note that RequestIdLayer runs after TraceLayer, so RequestId won't be in extensions when on_request fires. The code handles this gracefully (line 63-65 is conditional), but this means request IDs won't appear in the span's request_id field. If this is intentional, the comment at line 120-121 could be expanded to clarify that request IDs are recorded elsewhere.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/src/middleware/logging.rs` around lines 57 - 77, The
RequestLogger::on_request currently attempts to read RequestId from
request.extensions but RequestIdLayer runs after TraceLayer, so the ID will not
be present here; either reorder middleware so RequestIdLayer runs before
TraceLayer if you need the ID recorded in the span at on_request time, or
(preferred per review) expand the existing comment near RequestLogger/on_request
to explicitly state that RequestId is intentionally recorded later (because
RequestIdLayer executes after TraceLayer) and document where the request ID is
actually attached (e.g., the layer or handler that records it) so readers know
why span.record("request_id", ...) may be empty.
Comment on lines
+99
to
+137
| pub(super) fn normalize_path_for_metrics(path: &str) -> String { | ||
| let bytes = path.as_bytes(); | ||
| let mut segment_start = 0; | ||
| let mut segment_idx = 0; | ||
| let mut result: Option<String> = None; | ||
|
|
||
| for (pos, &b) in bytes.iter().enumerate() { | ||
| if b == b'/' || pos == bytes.len() - 1 { | ||
| // Determine segment end (include last char if not a slash) | ||
| let segment_end = if b == b'/' { pos } else { pos + 1 }; | ||
| let segment = &path[segment_start..segment_end]; | ||
|
|
||
| // Check segments after index 2 for dynamic IDs | ||
| if segment_idx > 2 && !segment.is_empty() && is_dynamic_id(segment) { | ||
| // Initialize result with everything before this segment | ||
| let result = result.get_or_insert_with(|| { | ||
| let mut s = String::with_capacity(path.len()); | ||
| s.push_str(&path[..segment_start]); | ||
| s | ||
| }); | ||
| result.push_str("{id}"); | ||
| } else if let Some(ref mut r) = result { | ||
| // Already normalizing, append this segment as-is | ||
| r.push_str(segment); | ||
| } | ||
|
|
||
| // Add slash after segment (except at end) | ||
| if b == b'/' { | ||
| if let Some(ref mut r) = result { | ||
| r.push('/'); | ||
| } | ||
| segment_start = pos + 1; | ||
| segment_idx += 1; | ||
| } | ||
| } | ||
| } | ||
|
|
||
| result.unwrap_or_else(|| path.to_owned()) | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Edge case: paths ending with slash may produce trailing {id}/.
The normalization logic handles most cases well, but when a path ends with a dynamic ID segment followed by a trailing slash (e.g., /v1/responses/resp_abc123def456/), the output might include {id}/ because the slash is appended after the segment replacement at line 128. This is likely acceptable for metrics purposes but worth noting.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/src/middleware/metrics.rs` around lines 99 - 137,
normalize_path_for_metrics can emit a trailing "{id}/" for paths that end with a
dynamic-id segment plus a slash; fix by tracking when you replaced the current
segment with "{id}" (e.g., a boolean like last_replaced set inside the
is_dynamic_id branch) and when handling the "if b == b'/'" branch skip appending
the trailing '/' if last_replaced is true and this slash is the final byte of
the path, otherwise append as before; ensure you clear last_replaced when you
append non-replacement segments so the flag only affects the immediate trailing
slash after an "{id}" replacement.
Comment on lines
+148
to
+158
| // Reconstruct request with modifications, preserving original extensions | ||
| let mut final_request = Request::builder() | ||
| .method(method) | ||
| .uri(uri) | ||
| .body(Body::from(modified_body)) | ||
| .unwrap_or_else(|_| Request::new(Body::empty())); | ||
| *final_request.headers_mut() = headers; | ||
| *final_request.extensions_mut() = extensions; | ||
|
|
||
| next.run(final_request).await | ||
| }; |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Request::builder() error fallback loses request data.
At line 153, unwrap_or_else(|_| Request::new(Body::empty())) silently replaces a failed request build with an empty request. While Request::builder() errors are rare (invalid URI/method), this could silently drop the modified body. Consider logging the error.
♻️ Log builder errors
let mut final_request = Request::builder()
.method(method)
.uri(uri)
.body(Body::from(modified_body))
- .unwrap_or_else(|_| Request::new(Body::empty()));
+ .unwrap_or_else(|e| {
+ error!("Failed to rebuild request after WASM processing: {}", e);
+ Request::new(Body::empty())
+ });📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| // Reconstruct request with modifications, preserving original extensions | |
| let mut final_request = Request::builder() | |
| .method(method) | |
| .uri(uri) | |
| .body(Body::from(modified_body)) | |
| .unwrap_or_else(|_| Request::new(Body::empty())); | |
| *final_request.headers_mut() = headers; | |
| *final_request.extensions_mut() = extensions; | |
| next.run(final_request).await | |
| }; | |
| // Reconstruct request with modifications, preserving original extensions | |
| let mut final_request = Request::builder() | |
| .method(method) | |
| .uri(uri) | |
| .body(Body::from(modified_body)) | |
| .unwrap_or_else(|e| { | |
| error!("Failed to rebuild request after WASM processing: {}", e); | |
| Request::new(Body::empty()) | |
| }); | |
| *final_request.headers_mut() = headers; | |
| *final_request.extensions_mut() = extensions; | |
| next.run(final_request).await | |
| }; |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/src/middleware/wasm.rs` around lines 148 - 158, The
unwrap_or_else on Request::builder() may swallow build errors and drop the
modified body; change the construction to match
Request::builder().method(method).uri(uri).body(Body::from(modified_body)) and
on Err(e) log the error (e.g. via tracing::error! or your logger) including e,
then fall back to creating a Request::new(Body::from(modified_body)) so the
modified_body is preserved; keep assigning *final_request.headers_mut() =
headers and *final_request.extensions_mut() = extensions before calling
next.run(final_request).await to ensure headers/extensions are preserved and the
error is visible in logs.
Comment on lines
+230
to
+240
| match action { | ||
| Action::Continue => {} | ||
| Action::Reject(status_code) => { | ||
| status = StatusCode::from_u16(status_code).unwrap_or(StatusCode::BAD_REQUEST); | ||
| let mut final_response = Response::builder() | ||
| .status(status) | ||
| .body(Body::from(body_bytes)) | ||
| .unwrap_or_else(|_| Response::new(Body::empty())); | ||
| *final_response.headers_mut() = headers; | ||
| return final_response; | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Clarify Action::Reject semantics in OnResponse phase.
When a WASM module returns Action::Reject(status_code) during OnResponse (lines 232-239), the response is returned with the new status but retains the current headers and body_bytes. This might be intentional (allow module to reject while preserving body), but differs from OnRequest where Reject returns an empty response. Consider adding a comment to clarify the intended behavior.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/src/middleware/wasm.rs` around lines 230 - 240, Add a short
clarifying comment in model_gateway/src/middleware/wasm.rs inside the match for
Action::Reject within the OnResponse handling (the block that sets `status`,
builds `final_response`, assigns `*final_response.headers_mut() = headers`, and
returns it) stating that Reject during OnResponse intentionally preserves the
current `headers` and `body_bytes` (unlike OnRequest which returns an empty
response) and explaining the rationale (e.g., allow module to change status
while keeping response payload), so future readers know this semantic difference
between OnRequest and OnResponse.
4 tasks
3 of 4 tasks
5 tasks done
5 tasks done
This was referenced Aug 13, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Combines Finding 2 (relocate
worker/token_bucket.rsout of the worker module) and the long-overdue break-up ofmiddleware.rsinto a single PR. Pure structural refactor — zero behavior change, and every external import path stays valid via re-exports inmiddleware/mod.rs.Why one PR
middleware.rswas importingcrate::worker::token_bucket::TokenBucketviapub use, so the cleanest moment to relocatetoken_bucketis when middleware itself becomes a folder. Doing the two changes in separate PRs would create a transientmiddleware.rs→middleware/mod.rs→middleware.rsflip that nobody needs.Layout
Before
After
What changed
model_gateway/src/middleware.rs: deleted (1110 lines distributed across the 8 files above, byte-for-byte identical except for per-file imports).model_gateway/src/middleware/mod.rs: new — registers all submodules and re-exports the public surface (AuthConfig,auth_middleware,ConcurrencyLimiter,QueueProcessor,QueuedRequest,TokenGuardBody,concurrency_limit_middleware,create_logging_layer,RequestLogger,RequestSpan,ResponseLogger,HttpMetricsLayer,HttpMetricsMiddleware,RequestId,RequestIdLayer,RequestIdMiddleware,storage_context_middleware,TokenBucket,wasm_middleware).model_gateway/src/worker/token_bucket.rs→model_gateway/src/middleware/token_bucket.rs: git rename, file content unchanged. Tracks as a 100% rename in the diff.model_gateway/src/worker/mod.rs: dropspub mod token_bucket;.concurrency.rsimportssuper::token_bucket::TokenBucketlogging.rsimportssuper::metrics::normalize_path_for_metricsandsuper::request_id::RequestIdwasm.rsimportssuper::request_id::{generate_request_id, RequestId}generate_request_idispub(super)sowasm.rscan use it without leaking it past the module boundarynormalize_path_for_metricsispub(super)sologging.rscan use itbuild_storage_request_context_*tests live instorage_context.rs; thenormalize_path_*andis_dynamic_idtests live inmetrics.rs.Public surface preserved
crate::middleware::TokenBucketstill resolves (now viamiddleware/mod.rs:pub use token_bucket::TokenBucket), so the existing callers inapp_context.rs:19andservice_discovery.rs:1240are untouched.crate::middleware::{AuthConfig, QueuedRequest, …}import inserver.rsstill resolves via the re-exports.RequestIdre-export fromsmg_authis preserved (now inmiddleware/request_id.rs).worker/mod.rsneeds an import change.What this unblocks
worker/no longer hosts a primitive that has zero coupling to the Worker abstraction — Finding 2 closed.middleware.rsis no longer a 1110-line bag of concerns; future middleware additions can live in their own files instead of growing the monolith.What this does NOT do
Layer/Serviceis already the ecosystem-standard abstraction; introducing a project-local trait would just be a wrapper around it.Test plan
cargo check -p smg— clean.cargo clippy -p smg --all-targets -- -D warnings— clean.cargo clippy -p smg-golang --all-targets -- -D warnings— clean.cargo check -p smg-python— clean.cargo test -p smg --lib— 546 passed; 0 failed; 4 ignored.make fmt— clean.rustup run nightly cargo fmt -- --check— EXIT=0 (matches CI; learned from refactor(worker): extract HashRing into its own module (Finding 5) #1129's format failure to always run nightly fmt before push).Refs: Finding 2 in
.claude/plans/2026-04-14-worker-module-followup-cleanup.mdSummary by CodeRabbit