fix(grpc): release request payloads at dispatch via RequestLease

model_gateway/src/routers/common/request_lease.rs

@@ -69,6 +69,9 @@ enum SerializedBody {
     None,
     Single(Bytes),
     Legs(Bytes, Bytes),
+    /// Wire size only: the serialization lives outside the lease (typed
+    /// proto dispatch).
+    Sized(usize),
 }
 
 impl SerializedBody {
@@ -79,6 +82,7 @@ impl SerializedBody {
             Self::None => 0,
             Self::Single(body) => body.len(),
             Self::Legs(prefill, decode) => prefill.len().max(decode.len()),
+            Self::Sized(len) => *len,
         }
     }
 }
@@ -144,10 +148,16 @@ impl<T> RequestLease<T> {
     pub(crate) fn body(&self) -> Option<Bytes> {
         match &self.lock().body {
             SerializedBody::Single(body) => Some(body.clone()),
-            SerializedBody::None | SerializedBody::Legs(..) => None,
+            SerializedBody::None | SerializedBody::Legs(..) | SerializedBody::Sized(_) => None,
         }
     }
 
+    /// Record the upstream wire size for the release metric when the wire
+    /// serialization lives outside the lease (typed proto dispatch).
+    pub(crate) fn note_upstream_len(&self, len: usize) {
+        self.lock().body = SerializedBody::Sized(len);
+    }
+
     /// Under `AfterDispatch`, free the parsed request and derivatives and
     /// count the serialized size as released early; under `AtRetryClose` a
     /// no-op — release happens when the lease drops.
@@ -185,6 +195,31 @@ impl<T> RequestLease<T> {
     }
 }
 
+/// Object-safe lease handle for pipelines that thread the release point
+/// through a heterogeneous request context (the gRPC stage pipeline).
+pub(crate) trait ErasedLease: Send + Sync {
+    /// True when the lease frees at dispatch (retries disabled).
+    fn releases_after_dispatch(&self) -> bool;
+    /// See [`RequestLease::note_upstream_len`].
+    fn note_upstream_len(&self, len: usize);
+    /// See [`RequestLease::release_dispatch`].
+    fn release_dispatch(&self);
+}
+
+impl<T: Send> ErasedLease for RequestLease<T> {
+    fn releases_after_dispatch(&self) -> bool {
+        self.release_point() == ReleasePoint::AfterDispatch
+    }
+
+    fn note_upstream_len(&self, len: usize) {
+        RequestLease::note_upstream_len(self, len);
+    }
+
+    fn release_dispatch(&self) {
+        RequestLease::release_dispatch(self);
+    }
+}
+
 /// Shared drop-probe idiom for release tests: a probed request type plus
 /// loopback stubs gated on the probe's weak count.
 #[cfg(test)]

model_gateway/src/routers/grpc/common/stages/dispatch_metadata.rs

@@ -41,6 +41,9 @@ impl PipelineStage for DispatchMetadataStage {
             RequestType::Embedding(req) => req.model.clone(),
             RequestType::Classify(req) => req.model.clone(),
             RequestType::Messages(req) => req.model.clone(),
+            // Runs before request execution, so the payload is never released
+            // here; the canonical model ID is the same value regardless.
+            RequestType::Released(_) => ctx.input.model_id.clone(),
         };
 
         let weight_version = ctx

model_gateway/src/routers/grpc/common/stages/helpers.rs

@@ -71,9 +71,10 @@ impl SamplingDefaultsMask {
                 min_p: true,
                 repetition_penalty: true,
             }),
-            RequestType::Responses(_) | RequestType::Embedding(_) | RequestType::Classify(_) => {
-                None
-            }
+            RequestType::Responses(_)
+            | RequestType::Embedding(_)
+            | RequestType::Classify(_)
+            | RequestType::Released(_) => None,
         }
     }
 

model_gateway/src/routers/grpc/common/stages/request_execution.rs

@@ -195,6 +195,19 @@ impl PipelineStage for RequestExecutionStage {
             sub_requests,
         ));
 
+        // Dispatch-phase release: response processing reads only the view
+        // captured at request building, so the context sheds its payload
+        // handle before the send. The lease then frees the parsed request
+        // now (retries disabled) or keeps it for replay until the retry
+        // window closes.
+        if let Some(lease) = ctx.input.request_lease.take() {
+            ctx.input.request_type.release_payload();
+            if lease.releases_after_dispatch() {
+                lease.note_upstream_len(execution_plan.wire_len());
+            }
+            lease.release_dispatch();
+        }
+
         // Extract dispatch metadata for tracing span
         let dispatch = ctx.state.dispatch.as_ref();
         let request_id = dispatch.map(|d| d.request_id.as_str()).unwrap_or("unknown");

model_gateway/src/routers/grpc/context.rs

@@ -33,6 +33,7 @@ use super::{
 };
 use crate::{
     middleware::TenantRequestMeta,
+    routers::common::request_lease::ErasedLease,
     worker::{RuntimeType, Worker, WorkerLoadGuard, WorkerRegistry},
 };
 
@@ -53,12 +54,19 @@ pub(crate) struct RequestInput {
     pub headers: Option<HeaderMap>,
     /// Canonical model ID used after aliases are resolved at request entry.
     pub model_id: String,
+    /// Captured at construction so it survives payload release.
+    pub streaming: bool,
     pub tenant_request_meta: Option<TenantRequestMeta>,
     /// Shared across every retry attempt of one logical request so
     /// `RateLimitReserveStage` reserves at most once. `None` for endpoints
     /// that haven't opted into tenant rate limiting yet (Responses,
     /// embeddings, classify).
     pub rate_limit_cell: Option<Arc<RateLimitCell>>,
+    /// Dispatch-phase owner of the parsed request (the router retry loop's
+    /// handle). When set, `RequestExecutionStage` drops the context's own
+    /// payload handle at dispatch and releases the lease; response stages
+    /// then read only the pre-extracted request view.
+    pub request_lease: Option<Arc<dyn ErasedLease>>,
 }
 
 /// Request type variants
@@ -71,9 +79,58 @@ pub(crate) enum RequestType {
     Embedding(Arc<EmbeddingRequest>),
     Classify(Arc<ClassifyRequest>),
     Messages(Arc<CreateMessageRequest>),
+    /// Payload dropped at dispatch; only the original kind survives for
+    /// post-dispatch stage dispatching.
+    Released(RequestKind),
+}
+
+/// Payload-free discriminant of [`RequestType`].
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) enum RequestKind {
+    Chat,
+    Generate,
+    Completion,
+    Responses,
+    Embedding,
+    Classify,
+    Messages,
+}
+
+impl std::fmt::Display for RequestKind {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let name = match self {
+            Self::Chat => "Chat",
+            Self::Generate => "Generate",
+            Self::Completion => "Completion",
+            Self::Responses => "Responses",
+            Self::Embedding => "Embedding",
+            Self::Classify => "Classify",
+            Self::Messages => "Messages",
+        };
+        write!(f, "{name}")
+    }
 }
 
 impl RequestType {
+    pub fn kind(&self) -> RequestKind {
+        match self {
+            Self::Chat(_) => RequestKind::Chat,
+            Self::Generate(_) => RequestKind::Generate,
+            Self::Completion(_) => RequestKind::Completion,
+            Self::Responses(_) => RequestKind::Responses,
+            Self::Embedding(_) => RequestKind::Embedding,
+            Self::Classify(_) => RequestKind::Classify,
+            Self::Messages(_) => RequestKind::Messages,
+            Self::Released(kind) => *kind,
+        }
+    }
+
+    /// Drop the context's payload handle at dispatch. The lease (or the
+    /// router loop) is the remaining owner; the typed accessors panic past
+    /// this point by construction.
+    pub fn release_payload(&mut self) {
+        *self = Self::Released(self.kind());
+    }
     /// Overwrite the request's own `model` field.
     ///
     /// Callers hold the request behind an `Arc` that the retry loop also
@@ -94,6 +151,7 @@ impl RequestType {
             Self::Embedding(request) => replace(&mut Arc::make_mut(request).model, model_id),
             Self::Classify(request) => replace(&mut Arc::make_mut(request).model, model_id),
             Self::Messages(request) => replace(&mut Arc::make_mut(request).model, model_id),
+            Self::Released(_) => {}
         }
     }
 
@@ -107,22 +165,14 @@ impl RequestType {
             Self::Embedding(r) => r.rid.as_deref(),
             Self::Classify(r) => r.rid.as_deref(),
             Self::Messages(r) => r.rid.as_deref(),
-            Self::Responses(_) => None,
+            Self::Responses(_) | Self::Released(_) => None,
         }
     }
 }
 
 impl std::fmt::Display for RequestType {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self {
-            Self::Chat(_) => write!(f, "Chat"),
-            Self::Generate(_) => write!(f, "Generate"),
-            Self::Completion(_) => write!(f, "Completion"),
-            Self::Responses(_) => write!(f, "Responses"),
-            Self::Embedding(_) => write!(f, "Embedding"),
-            Self::Classify(_) => write!(f, "Classify"),
-            Self::Messages(_) => write!(f, "Messages"),
-        }
+        self.kind().fmt(f)
     }
 }
 
@@ -281,6 +331,19 @@ impl ExecutionPlan {
             },
         }
     }
+
+    /// Serialized wire size of the built request(s), for the release metric.
+    pub(crate) fn wire_len(&self) -> usize {
+        match self {
+            Self::Single(request) => request.wire_len(),
+            Self::PrefillDecode(request) | Self::EncodePrefillDecode { request } => {
+                request.wire_len()
+            }
+            Self::Batch { requests, .. } => {
+                requests.iter().map(ProtoGenerateRequest::wire_len).sum()
+            }
+        }
+    }
 }
 
 /// Output from preparation stage (Step 1)
@@ -491,7 +554,7 @@ pub(crate) struct ResponseState {
     pub skip_special_tokens: Option<bool>,
 
     /// Response-phase view of the request, set by request building so
-    /// response processing never reads the payload after dispatch.
+    /// response processing never reads the (possibly released) payload.
     pub request_view: Option<super::regular::views::RequestView>,
 
     /// Execution result (streams from workers)
@@ -528,13 +591,25 @@ impl RequestContext {
             model_id.push_str(&canonical_model_id);
             request_type.set_model(&model_id);
         }
+        let streaming = match &request_type {
+            RequestType::Chat(req) => req.stream,
+            RequestType::Generate(req) => req.stream,
+            RequestType::Completion(req) => req.stream,
+            RequestType::Responses(req) => req.stream.unwrap_or(false),
+            RequestType::Messages(req) => req.stream.unwrap_or(false),
+            // Embeddings and classification never stream.
+            RequestType::Embedding(_) | RequestType::Classify(_) => false,
+            RequestType::Released(_) => false,
+        };
         Self {
             input: RequestInput {
                 request_type,
                 headers,
                 model_id,
+                streaming,
                 tenant_request_meta: None,
                 rate_limit_cell: None,
+                request_lease: None,
             },
             components,
             state: ProcessingState::default(),
@@ -701,17 +776,10 @@ impl RequestContext {
         }
     }
 
-    /// Check if request is streaming
+    /// Check if request is streaming (captured at construction, so valid
+    /// after payload release).
     pub fn is_streaming(&self) -> bool {
-        match &self.input.request_type {
-            RequestType::Chat(req) => req.stream,
-            RequestType::Generate(req) => req.stream,
-            RequestType::Completion(req) => req.stream,
-            RequestType::Responses(req) => req.stream.unwrap_or(false),
-            RequestType::Messages(req) => req.stream.unwrap_or(false),
-            RequestType::Embedding(_) => false, // Embeddings are never streaming
-            RequestType::Classify(_) => false,  // Classification is never streaming
-        }
+        self.input.streaming
     }
 
     /// Get the cached tokenizer, cloning the Arc (cheap 8-byte clone)

model_gateway/src/routers/grpc/harmony/stages/request_building.rs

@@ -100,7 +100,8 @@ impl PipelineStage for HarmonyRequestBuildingStage {
             | RequestType::Completion(_)
             | RequestType::Embedding(_)
             | RequestType::Classify(_)
-            | RequestType::Messages(_)) => {
+            | RequestType::Messages(_)
+            | RequestType::Released(_)) => {
                 error!(
                     function = "HarmonyRequestBuildingStage::execute",
                     request_type = %request_type,

model_gateway/src/routers/grpc/harmony/stages/response_processing.rs

@@ -170,7 +170,8 @@ impl PipelineStage for HarmonyResponseProcessingStage {
             | RequestType::Completion(_)
             | RequestType::Embedding(_)
             | RequestType::Classify(_)
-            | RequestType::Messages(_)) => {
+            | RequestType::Messages(_)
+            | RequestType::Released(_)) => {
                 error!(
                     function = "HarmonyResponseProcessingStage::execute",
                     request_type = %request_type,