perf(rust): coalesce intercepted HTTP response chunks

rust/src/copilot_request_handler.rs

@@ -49,6 +49,10 @@ use crate::{
     Client, ClientInner, JsonRpcRequest, JsonRpcResponse, RequestId, SessionId, error_codes,
 };
 
+use self::http_response_reader::HttpResponseReader;
+
+mod http_response_reader;
+
 const METHOD_HTTP_REQUEST_START: &str = "llmInference.httpRequestStart";
 const METHOD_HTTP_REQUEST_CHUNK: &str = "llmInference.httpRequestChunk";
 
@@ -157,6 +161,14 @@ pub struct CopilotRequestContext {
 }
 
 /// Streaming response body: a sequence of byte chunks or a terminal error.
+///
+/// HTTP bytes are forwarded in order, but chunk boundaries are not preserved.
+/// The SDK reads ahead while awaiting runtime acknowledgements, using at most
+/// 64 KiB of raw-byte forwarding buffers plus one current source chunk. This
+/// excludes storage inside the source stream and RPC serialization. A source
+/// chunk (including its shared backing allocation) is not size-limited by this API.
+/// Bytes available after an acknowledgement are flushed without waiting for
+/// more input. WebSocket message boundaries are preserved separately.
 pub type CopilotHttpResponseBody =
     Pin<Box<dyn Stream<Item = Result<Bytes, CopilotRequestError>> + Send>>;
 
@@ -916,32 +928,57 @@ async fn stream_http_response(
     exchange: &CopilotRequestExchange,
     cancel: &CancellationToken,
 ) -> Result<(), CopilotRequestError> {
-    exchange
-        .start_response(response.status, response.status_text, response.headers)
-        .await?;
+    tokio::select! {
+        biased;
+        // The RPC enqueues its complete frame before its first suspension.
+        // Poll it first even if already cancelled: the writer actor then commits
+        // the head before the terminal error, without waiting for the head ACK.
+        result = exchange.start_response(response.status, response.status_text, response.headers) => {
+            result?;
+        }
+        _ = cancel.cancelled() => {
+            drop(response.body);
+            return exchange
+                .error_response("Request cancelled by runtime", Some("cancelled".to_string()))
+                .await;
+        }
+    }
 
-    let mut body = response.body;
-    loop {
-        tokio::select! {
-            _ = cancel.cancelled() => {
-                return exchange
-                    .error_response("Request cancelled by runtime", Some("cancelled".to_string()))
-                    .await;
+    let forward = async {
+        let mut reader = HttpResponseReader::new(response.body);
+        let mut chunk = Vec::new();
+        loop {
+            match reader.next_chunk(&mut chunk).await {
+                Ok(true) => {}
+                Ok(false) => return exchange.end_response().await,
+                Err(error) => return exchange.error_response(error.to_string(), None).await,
             }
-            next = body.next() => match next {
-                Some(Ok(chunk)) => {
-                    for piece in chunk.chunks(32 * 1024) {
-                        exchange.write_binary(piece).await?;
+
+            // Keep the same acknowledged write alive while polling the upstream.
+            // Prefer its completion so a ready source cannot delay the next write.
+            let write = exchange.write_binary(&chunk);
+            tokio::pin!(write);
+            loop {
+                tokio::select! {
+                    biased;
+                    result = &mut write => {
+                        result?;
+                        break;
                     }
+                    () = reader.read_more(), if reader.can_read() => {}
                 }
-                Some(Err(e)) => {
-                    return exchange.error_response(e.to_string(), None).await;
-                }
-                None => break,
             }
         }
+    };
+    tokio::select! {
+        biased;
+        _ = cancel.cancelled() => {
+            exchange
+                .error_response("Request cancelled by runtime", Some("cancelled".to_string()))
+                .await
+        }
+        result = forward => result,
     }
-    exchange.end_response().await
 }
 
 /// Forward runtime→upstream WebSocket messages until the runtime closes its side

rust/src/copilot_request_handler/http_response_reader.rs

@@ -0,0 +1,271 @@
+use std::panic::AssertUnwindSafe;
+
+use bytes::{Buf, Bytes};
+use futures_util::{FutureExt, StreamExt};
+
+use super::{CopilotHttpResponseBody, CopilotRequestError};
+
+const CHUNK_SIZE: usize = 32 * 1024;
+
+/// One write of read-ahead, not a queue of `Bytes` slices: even a tiny slice can
+/// retain an arbitrarily large allocation. Only the current source frame may
+/// retain such an allocation, and it is released as soon as its bytes are copied.
+pub(super) struct HttpResponseReader {
+    body: Option<CopilotHttpResponseBody>,
+    pending: Bytes,
+    buffered: Vec<u8>,
+    error: Option<CopilotRequestError>,
+}
+
+impl HttpResponseReader {
+    pub(super) fn new(body: CopilotHttpResponseBody) -> Self {
+        Self {
+            body: Some(body),
+            pending: Bytes::new(),
+            buffered: Vec::new(),
+            error: None,
+        }
+    }
+
+    pub(super) fn can_read(&self) -> bool {
+        self.buffered.len() < CHUNK_SIZE && (!self.pending.is_empty() || self.body.is_some())
+    }
+
+    /// Cancel-safe: no suspension occurs between acquiring bytes and saving them.
+    pub(super) async fn read_more(&mut self) {
+        // Custom streams can yield arbitrarily many ready (even empty) frames.
+        tokio::task::consume_budget().await;
+        if self.pending.is_empty() {
+            let Some(body) = &mut self.body else {
+                return;
+            };
+            match AssertUnwindSafe(body.next()).catch_unwind().await {
+                Ok(Some(Ok(bytes))) => self.pending = bytes,
+                Ok(Some(Err(error))) => {
+                    self.error = Some(error);
+                    self.body = None;
+                }
+                Ok(None) => self.body = None,
+                Err(_) => {
+                    self.error = Some(CopilotRequestError::message(
+                        "HTTP response body stream panicked",
+                    ));
+                    self.body = None;
+                }
+            }
+        }
+        if self.pending.is_empty() {
+            // An empty Bytes can still retain its source allocation.
+            self.pending = Bytes::new();
+            return;
+        }
+        if self.buffered.capacity() == 0 {
+            self.buffered.reserve_exact(CHUNK_SIZE);
+        }
+        let count = self.pending.len().min(CHUNK_SIZE - self.buffered.len());
+        self.buffered.extend_from_slice(&self.pending[..count]);
+        self.pending.advance(count);
+        if self.pending.is_empty() {
+            self.pending = Bytes::new();
+        }
+    }
+
+    pub(super) async fn next_chunk(
+        &mut self,
+        output: &mut Vec<u8>,
+    ) -> Result<bool, CopilotRequestError> {
+        while self.buffered.is_empty() && self.can_read() {
+            self.read_more().await;
+        }
+        // Flush partial bytes before an upstream failure or EOF, without waiting
+        // to fill the buffer. Reuse the two bounded allocations on later writes.
+        if !self.buffered.is_empty() {
+            output.clear();
+            std::mem::swap(output, &mut self.buffered);
+            return Ok(true);
+        }
+        if let Some(error) = self.error.take() {
+            return Err(error);
+        }
+        Ok(false)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::time::Duration;
+
+    use futures_util::stream;
+
+    use super::*;
+
+    #[tokio::test]
+    async fn preserves_empty_large_fragmented_and_partial_bodies() {
+        for size in [0, 1, CHUNK_SIZE - 1, CHUNK_SIZE, CHUNK_SIZE * 3 + 7] {
+            let expected: Vec<u8> = (0..size).map(|i| (i % 256) as u8).collect();
+            for fragment_size in [1, 1024, CHUNK_SIZE, CHUNK_SIZE * 4] {
+                let fragments: Vec<_> = expected
+                    .chunks(fragment_size)
+                    .map(Bytes::copy_from_slice)
+                    .collect();
+                let body = stream::iter(
+                    [Bytes::new()]
+                        .into_iter()
+                        .chain(fragments)
+                        .chain([Bytes::new()])
+                        .map(Ok),
+                );
+                let mut reader = HttpResponseReader::new(Box::pin(body));
+                let mut output = Vec::new();
+                let mut actual = Vec::new();
+                while reader.next_chunk(&mut output).await.unwrap() {
+                    assert!(output.len() <= CHUNK_SIZE);
+                    actual.extend_from_slice(&output);
+                    while reader.can_read() {
+                        reader.read_more().await;
+                    }
+                    assert!(reader.buffered.capacity() <= CHUNK_SIZE);
+                    assert!(output.capacity() <= CHUNK_SIZE);
+                }
+                assert_eq!(actual, expected);
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn flushes_partial_bytes_without_polling_pending_input() {
+        let body = stream::once(async { Ok(Bytes::from_static(b"data: first\n\n")) })
+            .chain(stream::pending());
+        let mut reader = HttpResponseReader::new(Box::pin(body));
+        let mut output = Vec::new();
+        assert!(
+            reader
+                .next_chunk(&mut output)
+                .now_or_never()
+                .unwrap()
+                .unwrap()
+        );
+        assert_eq!(output, b"data: first\n\n");
+        assert!(reader.read_more().now_or_never().is_none());
+    }
+
+    #[tokio::test]
+    async fn bounds_read_ahead_even_for_single_byte_fragments() {
+        let polls = Arc::new(AtomicUsize::new(0));
+        let counter = polls.clone();
+        let body = stream::repeat_with(move || {
+            counter.fetch_add(1, Ordering::SeqCst);
+            Ok(Bytes::from_static(b"x"))
+        });
+        let mut reader = HttpResponseReader::new(Box::pin(body));
+        let mut output = Vec::new();
+        assert!(reader.next_chunk(&mut output).await.unwrap());
+        while reader.can_read() {
+            reader.read_more().await;
+        }
+        assert_eq!(polls.load(Ordering::SeqCst), CHUNK_SIZE + 1);
+        assert_eq!(reader.buffered.len(), CHUNK_SIZE);
+        assert_eq!(reader.buffered.capacity(), CHUNK_SIZE);
+        assert_eq!(output.capacity(), CHUNK_SIZE);
+        assert!(reader.pending.is_empty());
+    }
+
+    struct BackingAllocation {
+        data: Vec<u8>,
+        visible: usize,
+        live: Arc<AtomicUsize>,
+    }
+
+    impl AsRef<[u8]> for BackingAllocation {
+        fn as_ref(&self) -> &[u8] {
+            &self.data[..self.visible]
+        }
+    }
+
+    impl Drop for BackingAllocation {
+        fn drop(&mut self) {
+            self.live.fetch_sub(1, Ordering::SeqCst);
+        }
+    }
+
+    #[tokio::test]
+    async fn releases_large_backing_allocations_including_empty_frames() {
+        let live = Arc::new(AtomicUsize::new(0));
+        let counter = live.clone();
+        let body = stream::iter([0, 1, 1, CHUNK_SIZE * 2]).map(move |visible| {
+            counter.fetch_add(1, Ordering::SeqCst);
+            Ok(Bytes::from_owner(BackingAllocation {
+                data: vec![42; CHUNK_SIZE * 64],
+                visible,
+                live: counter.clone(),
+            }))
+        });
+        let mut reader = HttpResponseReader::new(Box::pin(body));
+        reader.read_more().await;
+        assert_eq!(live.load(Ordering::SeqCst), 0);
+        reader.read_more().await;
+        reader.read_more().await;
+        assert_eq!(live.load(Ordering::SeqCst), 0);
+        reader.read_more().await;
+        assert_eq!(live.load(Ordering::SeqCst), 1);
+        assert_eq!(reader.pending.len(), CHUNK_SIZE + 2);
+        let mut output = Vec::new();
+        assert!(reader.next_chunk(&mut output).await.unwrap());
+        reader.read_more().await;
+        assert_eq!(live.load(Ordering::SeqCst), 1);
+        assert!(reader.next_chunk(&mut output).await.unwrap());
+        reader.read_more().await;
+        assert_eq!(live.load(Ordering::SeqCst), 0);
+    }
+
+    #[tokio::test]
+    async fn flushes_last_partial_before_upstream_error_or_panic() {
+        for panic in [false, true] {
+            let body = stream::iter([Ok(Bytes::from_static(b"partial"))]).chain(stream::once(
+                async move {
+                    assert!(!panic, "failing user stream");
+                    Err(CopilotRequestError::message("upstream failed"))
+                },
+            ));
+            let mut reader = HttpResponseReader::new(Box::pin(body));
+            reader.read_more().await;
+            reader.read_more().await;
+            let mut output = Vec::new();
+            assert!(reader.next_chunk(&mut output).await.unwrap());
+            assert_eq!(output, b"partial");
+            let error = reader.next_chunk(&mut output).await.unwrap_err();
+            assert_eq!(
+                error.to_string(),
+                if panic {
+                    "HTTP response body stream panicked"
+                } else {
+                    "upstream failed"
+                }
+            );
+        }
+    }
+
+    #[tokio::test]
+    async fn dropping_reader_drops_pending_source() {
+        let (tx, rx) = tokio::sync::mpsc::channel(1);
+        let mut reader =
+            HttpResponseReader::new(Box::pin(tokio_stream::wrappers::ReceiverStream::new(rx)));
+        assert!(reader.read_more().now_or_never().is_none());
+        drop(reader);
+        assert!(tx.is_closed());
+    }
+
+    #[tokio::test]
+    async fn always_ready_empty_frames_yield_to_cancellation() {
+        let mut reader =
+            HttpResponseReader::new(Box::pin(stream::repeat_with(|| Ok(Bytes::new()))));
+        let mut output = Vec::new();
+        assert!(
+            tokio::time::timeout(Duration::from_millis(10), reader.next_chunk(&mut output))
+                .await
+                .is_err()
+        );
+    }
+}