fix(frontend): preserve response framing while tracking inflight requests

lib/llm/Cargo.toml

@@ -104,6 +104,7 @@ futures = { workspace = true }
 futures-util = { workspace = true }
 half = { version = "2", features = ["serde", "num-traits"] }
 hf-hub = { workspace = true, features = ["tokio", "rustls-tls", "ureq"] }
+http-body = "1"
 http-body-util = "0.1"
 ipnet = { workspace = true }
 libc = { workspace = true }

lib/llm/src/http/service/service_v2.rs

@@ -4,16 +4,19 @@
 use std::collections::{HashMap, HashSet};
 use std::env::var;
 use std::path::PathBuf;
+use std::pin::Pin;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::AtomicU8;
 use std::sync::atomic::AtomicU64;
 use std::sync::atomic::Ordering;
 use std::sync::{Arc, LazyLock};
+use std::task::{Context as TaskContext, Poll};
 use std::time::Duration;
 
-use axum::body::Body;
+use axum::body::{Body, Bytes, HttpBody};
 use axum::http::{Response, Version, header};
 use axum::response::IntoResponse;
+use http_body::{Frame, SizeHint};
 
 use super::Metrics;
 use super::RouteDoc;
@@ -171,12 +174,12 @@ impl InferenceReadinessState {
     }
 }
 
+/// Track admitted inference requests through response completion or cancellation.
 async fn track_inflight_inference(
     axum::extract::State(readiness): axum::extract::State<Arc<InferenceReadinessState>>,
     request: axum::extract::Request,
     next: axum::middleware::Next,
 ) -> axum::response::Response {
-    use futures::StreamExt;
     let state = &readiness.state;
 
     // Requests rejected during draining should not extend the drain window.
@@ -196,11 +199,45 @@ async fn track_inflight_inference(
     let (parts, body) = response.into_parts();
     // Keep the permit alive until the full response body, including streams,
     // finishes or is dropped.
-    let stream = body.into_data_stream().map(move |result| {
-        let _permit = &permit;
-        result
-    });
-    Response::from_parts(parts, Body::from_stream(stream))
+    Response::from_parts(
+        parts,
+        Body::new(InflightResponseBody {
+            body,
+            permit: Some(permit),
+        }),
+    )
+}
+
+// Preserve framing metadata and trailers while holding the response's admission.
+struct InflightResponseBody {
+    body: Body,
+    permit: Option<InflightPermit>,
+}
+
+impl HttpBody for InflightResponseBody {
+    type Data = Bytes;
+    type Error = axum::Error;
+
+    /// Release admission when the consumer observes EOF.
+    fn poll_frame(
+        self: Pin<&mut Self>,
+        cx: &mut TaskContext<'_>,
+    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
+        let this = self.get_mut();
+        let frame = Pin::new(&mut this.body).poll_frame(cx);
+        if matches!(frame, Poll::Ready(None)) {
+            this.permit.take();
+        }
+        frame
+    }
+
+    fn is_end_stream(&self) -> bool {
+        self.body.is_end_stream()
+    }
+
+    fn size_hint(&self) -> SizeHint {
+        self.body.size_hint()
+    }
 }
 
 /// HTTP service shared state
@@ -3069,6 +3106,52 @@ mod tests {
         assert_eq!(observer.inflight_count(), 0);
     }
 
+    #[tokio::test]
+    async fn inflight_body_preserves_trailers_and_holds_admission() {
+        use http_body_util::{BodyExt, StreamBody};
+
+        let observer = Arc::new(ServiceObserver::default());
+        let trailers = axum::http::HeaderMap::from_iter([(
+            header::HeaderName::from_static("x-checksum"),
+            header::HeaderValue::from_static("done"),
+        )]);
+        let frames = futures::stream::iter([
+            Ok::<_, std::convert::Infallible>(Frame::data(Bytes::from_static(b"data"))),
+            Ok(Frame::trailers(trailers.clone())),
+        ]);
+        let mut body = InflightResponseBody {
+            body: Body::new(StreamBody::new(frames)),
+            permit: Some(observer.acquire_inflight()),
+        };
+
+        assert_eq!(body.size_hint().exact(), None);
+        assert!(!body.is_end_stream());
+        assert_eq!(observer.inflight_count(), 1);
+        let frame = body.frame().await.unwrap().unwrap();
+        assert_eq!(frame.into_data().unwrap(), "data");
+        assert_eq!(observer.inflight_count(), 1);
+        let frame = body.frame().await.unwrap().unwrap();
+        assert_eq!(frame.into_trailers().unwrap(), trailers);
+        assert!(body.frame().await.is_none());
+        assert_eq!(observer.inflight_count(), 0);
+        drop(body);
+        assert_eq!(observer.inflight_count(), 0);
+    }
+
+    #[test]
+    fn inflight_body_preserves_length_and_releases_on_early_drop() {
+        let observer = Arc::new(ServiceObserver::default());
+        let body = InflightResponseBody {
+            body: Body::from("unread"),
+            permit: Some(observer.acquire_inflight()),
+        };
+        assert_eq!(body.size_hint().exact(), Some(6));
+        assert!(!body.is_end_stream());
+        assert_eq!(observer.inflight_count(), 1);
+        drop(body);
+        assert_eq!(observer.inflight_count(), 0);
+    }
+
     #[tokio::test]
     async fn test_admin_routes_only_on_admin_listener() {
         // Off by default, and bound to loopback when enabled.

lib/llm/tests/http-transcriptions.rs

@@ -334,7 +334,15 @@ async fn upload_too_large_for_the_request_plane_returns_payload_too_large() {
         .await
         .expect("transcription request failed");
     assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
-    let body: serde_json::Value = response.json().await.unwrap();
+    let content_length = response
+        .content_length()
+        .expect("the fixed-size 413 JSON body must retain its length");
+    let bytes = response
+        .bytes()
+        .await
+        .expect("incomplete 413 response body");
+    assert_eq!(content_length, bytes.len() as u64);
+    let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
     assert!(
         body["message"]
             .as_str()