[smg] cleanup router RAII guards (#16560)

sgl-model-gateway/src/core/mod.rs

@@ -34,8 +34,8 @@ pub use job_queue::{Job, JobQueue, JobQueueConfig};
 pub use model_card::{ModelCard, ProviderType};
 pub use retry::{is_retryable_status, RetryExecutor};
 pub use worker::{
-    attach_guards_to_response, BasicWorker, ConnectionMode, HealthConfig, RuntimeType, Worker,
+    AttachedBody, BasicWorker, ConnectionMode, HealthConfig, RuntimeType, Worker, WorkerLoadGuard,
-    WorkerLoadGuard, WorkerType,
+    WorkerType,
 };
 pub use worker_builder::{BasicWorkerBuilder, DPAwareWorkerBuilder};
 pub use worker_manager::{LoadMonitor, WorkerManager};

sgl-model-gateway/src/core/worker.rs

@@ -1015,28 +1015,6 @@ impl WorkerLoadGuard {
         worker.increment_load();
         Self { worker }
     }
-
-    /// Attach this guard to a Response, tying the guard's lifetime to the response body.
-    ///
-    /// When the response body is fully consumed or dropped (e.g., client disconnects),
-    /// the guard is dropped and worker load is decremented automatically.
-    ///
-    /// This is the proper RAII pattern for SSE/streaming responses where the handler
-    /// returns immediately but the stream continues in a background task.
-    pub fn attach_to_response(
-        self,
-        response: axum::response::Response,
-    ) -> axum::response::Response {
-        let (parts, body) = response.into_parts();
-
-        // Wrap body with guard - guard drops when body drops
-        let guarded_body = GuardedBody {
-            inner: body,
-            _guard: self,
-        };
-
-        axum::response::Response::from_parts(parts, Body::new(guarded_body))
-    }
 }
 
 impl Drop for WorkerLoadGuard {
@@ -1045,65 +1023,45 @@ impl Drop for WorkerLoadGuard {
     }
 }
 
-/// Attach multiple guards to a Response (for dual prefill/decode workers)
+/// Body wrapper that holds an attached value.
-pub fn attach_guards_to_response(
-    guards: Vec<WorkerLoadGuard>,
-    response: axum::response::Response,
-) -> axum::response::Response {
-    let (parts, body) = response.into_parts();
-
-    let guarded_body = MultiGuardedBody {
-        inner: body,
-        _guards: guards,
-    };
-
-    axum::response::Response::from_parts(parts, Body::new(guarded_body))
-}
-
-/// Body wrapper that holds a WorkerLoadGuard
 ///
 /// When this body is dropped (stream ends or client disconnects),
-/// the guard is dropped, decrementing worker load.
+/// the attached value is dropped automatically. This is useful for RAII guards
-struct GuardedBody {
+/// like WorkerLoadGuard that need to be tied to a response body's lifetime.
+pub struct AttachedBody<T> {
     inner: Body,
-    _guard: WorkerLoadGuard,
+    _attached: T,
 }
 
-/// Body wrapper that holds multiple WorkerLoadGuards (for dual prefill/decode)
+impl<T> AttachedBody<T> {
-struct MultiGuardedBody {
+    pub fn new(inner: Body, attached: T) -> Self {
-    inner: Body,
+        Self {
-    _guards: Vec<WorkerLoadGuard>,
+            inner,
+            _attached: attached,
+        }
+    }
 }
 
-impl http_body::Body for GuardedBody {
+impl<T: Send + Unpin + 'static> AttachedBody<T> {
+    pub fn wrap_response(
+        response: axum::response::Response,
+        attached: T,
+    ) -> axum::response::Response {
+        let (parts, body) = response.into_parts();
+        axum::response::Response::from_parts(parts, Body::new(Self::new(body, attached)))
+    }
+}
+
+impl<T: Send + Unpin + 'static> http_body::Body for AttachedBody<T> {
     type Data = bytes::Bytes;
     type Error = axum::Error;
 
     fn poll_frame(
-        mut self: std::pin::Pin<&mut Self>,
+        self: std::pin::Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
     ) -> std::task::Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
-        std::pin::Pin::new(&mut self.inner).poll_frame(cx)
+        let this = self.get_mut();
-    }
+        std::pin::Pin::new(&mut this.inner).poll_frame(cx)
-
-    fn is_end_stream(&self) -> bool {
-        self.inner.is_end_stream()
-    }
-
-    fn size_hint(&self) -> http_body::SizeHint {
-        self.inner.size_hint()
-    }
-}
-
-impl http_body::Body for MultiGuardedBody {
-    type Data = bytes::Bytes;
-    type Error = axum::Error;
-
-    fn poll_frame(
-        mut self: std::pin::Pin<&mut Self>,
-        cx: &mut std::task::Context<'_>,
-    ) -> std::task::Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
-        std::pin::Pin::new(&mut self.inner).poll_frame(cx)
     }
 
     fn is_end_stream(&self) -> bool {

sgl-model-gateway/src/routers/grpc/context.rs

@@ -13,7 +13,7 @@ use super::{
     proto_wrapper::{ProtoEmbedComplete, ProtoRequest, ProtoStream},
 };
 use crate::{
-    core::{attach_guards_to_response, Worker, WorkerLoadGuard},
+    core::{Worker, WorkerLoadGuard},
     protocols::{
         chat::{ChatCompletionRequest, ChatCompletionResponse},
         classify::{ClassifyRequest, ClassifyResponse},
@@ -158,47 +158,29 @@ pub(crate) struct DispatchMetadata {
 /// Load guards for worker load tracking
 /// Automatically decrements load when dropped
 pub(crate) enum LoadGuards {
-    Single(WorkerLoadGuard),
+    Single {
+        _guard: WorkerLoadGuard,
+    },
     Dual {
-        prefill: WorkerLoadGuard,
+        _prefill: WorkerLoadGuard,
-        decode: WorkerLoadGuard,
+        _decode: WorkerLoadGuard,
     },
 }
 
 impl From<&WorkerSelection> for LoadGuards {
     fn from(selection: &WorkerSelection) -> Self {
         match selection {
-            WorkerSelection::Single { worker } => {
+            WorkerSelection::Single { worker } => LoadGuards::Single {
-                LoadGuards::Single(WorkerLoadGuard::new(worker.clone()))
+                _guard: WorkerLoadGuard::new(worker.clone()),
-            }
+            },
             WorkerSelection::Dual { prefill, decode } => LoadGuards::Dual {
-                prefill: WorkerLoadGuard::new(prefill.clone()),
+                _prefill: WorkerLoadGuard::new(prefill.clone()),
-                decode: WorkerLoadGuard::new(decode.clone()),
+                _decode: WorkerLoadGuard::new(decode.clone()),
             },
         }
     }
 }
 
-impl LoadGuards {
-    /// Attach these load guards to a Response, tying their lifetime to the response body.
-    ///
-    /// When the response body is fully consumed or dropped (e.g., client disconnects),
-    /// the guards are dropped and worker load is decremented automatically.
-    ///
-    /// This is the proper RAII pattern for SSE/streaming responses.
-    pub fn attach_to_response(
-        self,
-        response: axum::response::Response,
-    ) -> axum::response::Response {
-        let guards = match self {
-            LoadGuards::Single(guard) => vec![guard],
-            LoadGuards::Dual { prefill, decode } => vec![prefill, decode],
-        };
-
-        attach_guards_to_response(guards, response)
-    }
-}
-
 /// Response processing state (Step 6)
 #[derive(Default)]
 pub(crate) struct ResponseState {

sgl-model-gateway/src/routers/grpc/harmony/stages/response_processing.rs

@@ -7,12 +7,15 @@ use axum::response::Response;
 use tracing::error;
 
 use super::super::{HarmonyResponseProcessor, HarmonyStreamingProcessor};
-use crate::routers::{
+use crate::{
+    core::AttachedBody,
+    routers::{
         error,
         grpc::{
             common::stages::PipelineStage,
             context::{FinalResponse, RequestContext, RequestType},
         },
+    },
 };
 
 /// Harmony Response Processing stage: Parse and format Harmony responses
@@ -81,7 +84,7 @@ impl PipelineStage for HarmonyResponseProcessingStage {
 
                     // Attach load guards to response body for proper RAII lifecycle
                     let response = match ctx.state.load_guards.take() {
-                        Some(guards) => guards.attach_to_response(response),
+                        Some(guards) => AttachedBody::wrap_response(response, guards),
                         None => response,
                     };
 

sgl-model-gateway/src/routers/grpc/regular/stages/chat/response_processing.rs

@@ -9,13 +9,16 @@ use async_trait::async_trait;
 use axum::response::Response;
 use tracing::error;
 
-use crate::routers::{
+use crate::{
+    core::AttachedBody,
+    routers::{
         error,
         grpc::{
             common::stages::PipelineStage,
             context::{FinalResponse, RequestContext},
             regular::{processor, streaming},
         },
+    },
 };
 
 /// Chat response processing stage
@@ -100,7 +103,7 @@ impl ChatResponseProcessingStage {
 
             // Attach load guards to response body for proper RAII lifecycle
             let response = match ctx.state.load_guards.take() {
-                Some(guards) => guards.attach_to_response(response),
+                Some(guards) => AttachedBody::wrap_response(response, guards),
                 None => response,
             };
 

sgl-model-gateway/src/routers/grpc/regular/stages/generate/response_processing.rs

@@ -6,13 +6,16 @@ use async_trait::async_trait;
 use axum::response::Response;
 use tracing::error;
 
-use crate::routers::{
+use crate::{
+    core::AttachedBody,
+    routers::{
         error,
         grpc::{
             common::stages::PipelineStage,
             context::{FinalResponse, RequestContext},
             regular::{processor, streaming},
         },
+    },
 };
 
 /// Generate response processing stage
@@ -100,7 +103,7 @@ impl GenerateResponseProcessingStage {
 
             // Attach load guards to response body for proper RAII lifecycle
             let response = match ctx.state.load_guards.take() {
-                Some(guards) => guards.attach_to_response(response),
+                Some(guards) => AttachedBody::wrap_response(response, guards),
                 None => response,
             };
 

sgl-model-gateway/src/routers/http/pd_router.rs

@@ -835,7 +835,7 @@ impl PDRouter {
         prefill: Arc<dyn Worker>,
         decode: Arc<dyn Worker>,
     ) -> Response {
-        use crate::core::attach_guards_to_response;
+        use crate::core::AttachedBody;
 
         let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
 
@@ -885,7 +885,7 @@ impl PDRouter {
         // Attach load guards to response body for proper RAII lifecycle
         // Guards are dropped when response body is consumed or client disconnects
         let guards = vec![WorkerLoadGuard::new(prefill), WorkerLoadGuard::new(decode)];
-        attach_guards_to_response(guards, response)
+        AttachedBody::wrap_response(response, guards)
     }
 
     // Helper to process non-streaming decode response with logprob merging

sgl-model-gateway/src/routers/http/router.rs

@@ -16,7 +16,7 @@ use crate::{
     app_context::AppContext,
     config::types::RetryConfig,
     core::{
-        is_retryable_status, ConnectionMode, RetryExecutor, Worker, WorkerLoadGuard,
+        is_retryable_status, AttachedBody, ConnectionMode, RetryExecutor, Worker, WorkerLoadGuard,
         WorkerRegistry, WorkerType, UNKNOWN_MODEL_ID,
     },
     observability::{
@@ -629,7 +629,7 @@ impl Router {
             // Attach load guard to response body for proper RAII lifecycle
             // Guard is dropped when response body is consumed or client disconnects
             if let Some(guard) = load_guard {
-                response = guard.attach_to_response(response);
+                response = AttachedBody::wrap_response(response, guard);
             }
             response
         }

sgl-model-gateway/tests/load_guard_raii_test.rs

@@ -11,7 +11,7 @@ use axum::{body::Body, response::Response};
 use bytes::Bytes;
 use futures_util::StreamExt;
 use http_body_util::BodyExt;
-use smg::core::{attach_guards_to_response, BasicWorkerBuilder, Worker, WorkerLoadGuard};
+use smg::core::{AttachedBody, BasicWorkerBuilder, Worker, WorkerLoadGuard};
 use tokio::sync::mpsc;
 use tokio_stream::wrappers::UnboundedReceiverStream;
 
@@ -40,7 +40,7 @@ async fn test_guard_dropped_when_response_body_consumed() {
     let guard = WorkerLoadGuard::new(worker.clone());
     assert_eq!(worker.load(), 1);
 
-    let guarded_response = guard.attach_to_response(response);
+    let guarded_response = AttachedBody::wrap_response(response, guard);
 
     // Load should still be 1 (guard is in the body)
     assert_eq!(worker.load(), 1);
@@ -65,7 +65,7 @@ async fn test_guard_dropped_when_response_dropped_without_consumption() {
         let guard = WorkerLoadGuard::new(worker.clone());
         assert_eq!(worker.load(), 1);
 
-        let _guarded_response = guard.attach_to_response(response);
+        let _guarded_response = AttachedBody::wrap_response(response, guard);
 
         // Load is still 1
         assert_eq!(worker.load(), 1);
@@ -89,7 +89,7 @@ async fn test_streaming_guard_dropped_when_stream_ends() {
     let guard = WorkerLoadGuard::new(worker.clone());
     assert_eq!(worker.load(), 1);
 
-    let guarded_response = guard.attach_to_response(response);
+    let guarded_response = AttachedBody::wrap_response(response, guard);
 
     // Spawn a task to consume the response
     let worker_clone = worker.clone();
@@ -136,7 +136,7 @@ async fn test_streaming_guard_dropped_on_client_disconnect() {
     let guard = WorkerLoadGuard::new(worker.clone());
     assert_eq!(worker.load(), 1);
 
-    let guarded_response = guard.attach_to_response(response);
+    let guarded_response = AttachedBody::wrap_response(response, guard);
 
     // Start consuming but drop early (simulate client disconnect)
     {
@@ -176,8 +176,7 @@ async fn test_multiple_guards_all_dropped() {
         assert_eq!(worker1.load(), 1);
         assert_eq!(worker2.load(), 1);
 
-        // Attach both guards using attach_guards_to_response
+        let _response = AttachedBody::wrap_response(response, vec![guard1, guard2]);
-        let _response = attach_guards_to_response(vec![guard1, guard2], response);
 
         // Both loads are 1
         assert_eq!(worker1.load(), 1);
@@ -201,7 +200,7 @@ async fn test_guard_with_empty_body() {
         let guard = WorkerLoadGuard::new(worker.clone());
         assert_eq!(worker.load(), 1);
 
-        let guarded_response = guard.attach_to_response(response);
+        let guarded_response = AttachedBody::wrap_response(response, guard);
 
         // Consume empty body
         let body = guarded_response.into_body();