diff --git a/rust/src/copilot_request_handler.rs b/rust/src/copilot_request_handler.rs index 961ae3876e..3b6e9369ab 100644 --- a/rust/src/copilot_request_handler.rs +++ b/rust/src/copilot_request_handler.rs @@ -40,6 +40,7 @@ use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; use tokio_util::sync::CancellationToken; use tracing::warn; +use self::http_response_reader::HttpResponseReader; use crate::generated::api_types::{ LlmInferenceHttpRequestChunkRequest, LlmInferenceHttpRequestStartRequest, LlmInferenceHttpRequestStartTransport, LlmInferenceHttpResponseChunkError, @@ -49,6 +50,8 @@ use crate::{ Client, ClientInner, JsonRpcRequest, JsonRpcResponse, RequestId, SessionId, error_codes, }; +mod http_response_reader; + const METHOD_HTTP_REQUEST_START: &str = "llmInference.httpRequestStart"; const METHOD_HTTP_REQUEST_CHUNK: &str = "llmInference.httpRequestChunk"; @@ -157,6 +160,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> + Send>>; @@ -916,32 +927,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 diff --git a/rust/src/copilot_request_handler/http_response_reader.rs b/rust/src/copilot_request_handler/http_response_reader.rs new file mode 100644 index 0000000000..96a5946a2a --- /dev/null +++ b/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, + pending: Bytes, + buffered: Vec, + error: Option, +} + +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, + ) -> Result { + 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 = (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, + visible: usize, + live: Arc, + } + + 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() + ); + } +} diff --git a/rust/tests/http_response_forwarding_test.rs b/rust/tests/http_response_forwarding_test.rs new file mode 100644 index 0000000000..1a0a2c77ce --- /dev/null +++ b/rust/tests/http_response_forwarding_test.rs @@ -0,0 +1,518 @@ +#![allow(clippy::unwrap_used)] + +use std::time::Duration; + +use async_trait::async_trait; +use base64::Engine; +use bytes::Bytes; +use futures_util::{StreamExt, stream}; +use github_copilot_sdk::{ + Client, ClientOptions, CopilotHttpRequest, CopilotHttpResponse, CopilotHttpResponseBody, + CopilotRequestContext, CopilotRequestError, CopilotRequestHandler, SDK_PROTOCOL_VERSION, + Transport, +}; +use http::HeaderMap; +use parking_lot::Mutex; +use serde_json::{Value, json}; +use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}; +use tokio::net::TcpListener; +use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf}; +use tokio::sync::mpsc; +use tokio::time::timeout; +use tokio_stream::wrappers::ReceiverStream; + +const DEADLINE: Duration = Duration::from_secs(5); +const CHUNK_SIZE: usize = 32 * 1024; + +async fn read_frame(reader: &mut BufReader) -> Value { + let mut length = None; + loop { + let mut line = String::new(); + assert_ne!(reader.read_line(&mut line).await.unwrap(), 0); + if line == "\r\n" { + break; + } + if let Some(value) = line.strip_prefix("Content-Length:") { + length = Some(value.trim().parse::().unwrap()); + } + } + let mut body = vec![0; length.unwrap()]; + reader.read_exact(&mut body).await.unwrap(); + serde_json::from_slice(&body).unwrap() +} + +async fn write_frame(writer: &mut OwnedWriteHalf, value: Value) { + let body = serde_json::to_vec(&value).unwrap(); + let mut frame = format!("Content-Length: {}\r\n\r\n", body.len()).into_bytes(); + frame.extend_from_slice(&body); + writer.write_all(&frame).await.unwrap(); +} + +/// A framed runtime peer exercising the public SDK client and real TCP transport. +struct Peer { + client: Client, + read: BufReader, + write: OwnedWriteHalf, + _work: tempfile::TempDir, +} + +impl Peer { + async fn start(handler: impl CopilotRequestHandler) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let work = tempfile::tempdir().unwrap(); + let options = ClientOptions::new() + .with_program(std::env::current_exe().unwrap()) + .with_cwd(work.path()) + .with_transport(Transport::External { + host: "127.0.0.1".into(), + port: listener.local_addr().unwrap().port(), + connection_token: None, + }) + .with_request_handler(handler); + let (client, (read, write)) = timeout(DEADLINE, async { + tokio::join!(Client::start(options), async { + let (socket, _) = listener.accept().await.unwrap(); + socket.set_nodelay(true).unwrap(); + let (read, mut write) = socket.into_split(); + let mut read = BufReader::new(read); + for method in ["connect", "llmInference.setProvider"] { + let message = read_frame(&mut read).await; + assert_eq!(message["method"], method); + let result = if method == "connect" { + json!({"ok": true, "version": "test", "protocolVersion": SDK_PROTOCOL_VERSION}) + } else { + json!({"success": true}) + }; + write_frame( + &mut write, + json!({"jsonrpc": "2.0", "id": message["id"], "result": result}), + ) + .await; + } + (read, write) + }) + }) + .await + .unwrap(); + Self { + client: client.unwrap(), + read, + write, + _work: work, + } + } + + async fn request(&mut self, url: &str) -> Value { + self.begin_request(url).await; + let head = self.next().await; + assert_eq!(head["method"], "llmInference.httpResponseStart"); + self.ack(&head).await; + head + } + + async fn begin_request(&mut self, url: &str) { + write_frame( + &mut self.write, + json!({ + "jsonrpc": "2.0", "id": 100000, + "method": "llmInference.httpRequestStart", + "params": {"requestId": "test", "method": "GET", "url": url, "headers": {}} + }), + ) + .await; + self.request_chunk(json!({"requestId": "test", "data": "", "end": true})) + .await; + } + + async fn request_chunk(&mut self, params: Value) { + write_frame( + &mut self.write, + json!({ + "jsonrpc": "2.0", "id": 100001, + "method": "llmInference.httpRequestChunk", "params": params + }), + ) + .await; + } + + async fn next(&mut self) -> Value { + timeout(DEADLINE, async { + loop { + let value = read_frame(&mut self.read).await; + if value.get("method").is_some() { + return value; + } + assert!(value.get("error").is_none(), "{value}"); + } + }) + .await + .expect("runtime response") + } + + async fn ack(&mut self, message: &Value) { + write_frame( + &mut self.write, + json!({"jsonrpc": "2.0", "id": message["id"], "result": {"accepted": true}}), + ) + .await; + } + + async fn stop(self) { + self.client.stop().await.unwrap(); + } +} + +struct BodyHandler(Mutex>); + +#[async_trait] +impl CopilotRequestHandler for BodyHandler { + async fn send_request( + &self, + _request: CopilotHttpRequest, + _context: &CopilotRequestContext, + ) -> Result { + Ok(CopilotHttpResponse::new( + 200, + None, + HeaderMap::new(), + self.0.lock().take().unwrap(), + )) + } +} + +async fn channel_peer() -> (Peer, mpsc::Sender>) { + let (tx, rx) = mpsc::channel(1); + let body = Box::pin(ReceiverStream::new(rx)); + let mut peer = Peer::start(BodyHandler(Mutex::new(Some(body)))).await; + peer.request("http://unused.test").await; + (peer, tx) +} + +fn data(message: &Value) -> Vec { + assert_eq!(message["method"], "llmInference.httpResponseChunk"); + assert_eq!(message["params"]["end"], false); + assert_eq!(message["params"]["binary"], true); + let decoded = base64::engine::general_purpose::STANDARD + .decode(message["params"]["data"].as_str().unwrap()) + .unwrap(); + assert!(!decoded.is_empty() && decoded.len() <= CHUNK_SIZE); + decoded +} + +#[tokio::test] +async fn reads_ahead_under_one_pinned_ack_with_bounded_backpressure() { + let (mut peer, tx) = channel_peer().await; + tx.send(Ok(Bytes::from_static(b"first"))).await.unwrap(); + let first = peer.next().await; + assert_eq!(data(&first), b"first"); + + timeout(DEADLINE, async { + for _ in 0..32 { + tx.send(Ok(Bytes::from(vec![b'x'; 1024]))).await.unwrap(); + } + // A permit proves the last queued fragment was consumed without an ACK. + drop(tx.reserve().await.unwrap()); + }) + .await + .expect("read-ahead while first ACK is outstanding"); + tx.send(Ok(Bytes::from_static(b"last"))).await.unwrap(); + assert!( + timeout(Duration::from_millis(20), tx.reserve()) + .await + .is_err(), + "full read-ahead buffer must stop polling the source" + ); + assert!( + timeout(Duration::from_millis(20), peer.next()) + .await + .is_err(), + "must not duplicate or overtake the outstanding write" + ); + peer.ack(&first).await; + let combined = peer.next().await; + assert_eq!(data(&combined), vec![b'x'; CHUNK_SIZE]); + assert_ne!(combined["id"], first["id"]); + drop(tx); + peer.ack(&combined).await; + let last = peer.next().await; + assert_eq!(data(&last), b"last"); + peer.ack(&last).await; + let end = peer.next().await; + assert_eq!(end["params"]["end"], true); + assert!(end["params"].get("error").is_none()); + peer.ack(&end).await; + peer.stop().await; +} + +#[tokio::test] +async fn sparse_bytes_flush_before_future_input_and_preserve_split_utf8() { + let (mut peer, tx) = channel_peer().await; + let mut received = Vec::new(); + // Both the UTF-8 code point and SSE delimiter cross source/write boundaries. + for byte in b"data: \xf0\x9f\x8c\x8d\n\n" { + tx.send(Ok(Bytes::copy_from_slice(&[*byte]))).await.unwrap(); + let message = peer.next().await; + received.extend(data(&message)); + peer.ack(&message).await; + } + assert_eq!(received, b"data: \xf0\x9f\x8c\x8d\n\n"); + drop(tx); + let end = peer.next().await; + assert_eq!(end["params"]["end"], true); + peer.ack(&end).await; + peer.stop().await; +} + +#[tokio::test] +async fn upstream_error_follows_buffered_partial_and_outstanding_ack() { + let (mut peer, tx) = channel_peer().await; + tx.send(Ok(Bytes::from_static(b"first"))).await.unwrap(); + let first = peer.next().await; + tx.send(Ok(Bytes::from_static(b"partial"))).await.unwrap(); + tx.send(Err(CopilotRequestError::message("upstream failed"))) + .await + .unwrap(); + timeout(DEADLINE, tx.closed()).await.unwrap(); + assert!( + timeout(Duration::from_millis(20), peer.next()) + .await + .is_err() + ); + peer.ack(&first).await; + let partial = peer.next().await; + assert_eq!(data(&partial), b"partial"); + assert!( + timeout(Duration::from_millis(20), peer.next()) + .await + .is_err() + ); + peer.ack(&partial).await; + let end = peer.next().await; + assert_eq!(end["params"]["end"], true); + assert_eq!(end["params"]["error"]["message"], "upstream failed"); + peer.ack(&end).await; + peer.stop().await; +} + +#[tokio::test] +async fn cancellation_drops_source_without_waiting_for_data_ack() { + let (mut peer, tx) = channel_peer().await; + tx.send(Ok(Bytes::from_static(b"first"))).await.unwrap(); + let first = peer.next().await; + assert_eq!(data(&first), b"first"); + peer.request_chunk(json!({"requestId": "test", "data": "", "cancel": true})) + .await; + let end = peer.next().await; + assert_eq!(end["params"]["end"], true); + assert_eq!(end["params"]["error"]["code"], "cancelled"); + timeout(DEADLINE, tx.closed()).await.unwrap(); + peer.ack(&end).await; + peer.stop().await; +} + +#[tokio::test] +async fn rejected_write_drops_source_and_reports_error() { + let (mut peer, tx) = channel_peer().await; + tx.send(Ok(Bytes::from_static(b"first"))).await.unwrap(); + let first = peer.next().await; + write_frame( + &mut peer.write, + json!({ + "jsonrpc": "2.0", "id": first["id"], + "error": {"code": -32603, "message": "write rejected"} + }), + ) + .await; + timeout(DEADLINE, tx.closed()).await.unwrap(); + let end = peer.next().await; + assert_eq!(end["params"]["end"], true); + assert!( + end["params"]["error"]["message"] + .as_str() + .unwrap() + .contains("write rejected") + ); + peer.ack(&end).await; + peer.stop().await; +} + +#[tokio::test] +async fn panicking_source_is_an_error_not_successful_eof() { + let body = stream::once(async { Ok(Bytes::from_static(b"first")) }).chain(stream::poll_fn( + |_| -> std::task::Poll>> { + panic!("failed source"); + }, + )); + let mut peer = Peer::start(BodyHandler(Mutex::new(Some(Box::pin(body))))).await; + peer.request("http://unused.test").await; + let first = peer.next().await; + assert_eq!(data(&first), b"first"); + peer.ack(&first).await; + let end = peer.next().await; + assert_eq!(end["params"]["end"], true); + assert_eq!( + end["params"]["error"]["message"], + "HTTP response body stream panicked" + ); + peer.ack(&end).await; + peer.stop().await; +} + +#[tokio::test] +async fn always_ready_source_does_not_starve_ack_or_cancellation() { + let body = stream::repeat_with(|| Ok(Bytes::from_static(b"x"))); + let mut peer = Peer::start(BodyHandler(Mutex::new(Some(Box::pin(body))))).await; + peer.request("http://unused.test").await; + for _ in 0..3 { + let chunk = peer.next().await; + assert!(data(&chunk).iter().all(|byte| *byte == b'x')); + peer.ack(&chunk).await; + } + let pending = peer.next().await; + assert!(!data(&pending).is_empty()); + peer.request_chunk(json!({"requestId": "test", "data": "", "cancel": true})) + .await; + let end = peer.next().await; + assert_eq!(end["params"]["error"]["code"], "cancelled"); + peer.ack(&end).await; + peer.stop().await; +} + +#[tokio::test] +async fn disconnected_consumer_drops_source_under_outstanding_ack() { + let (mut peer, tx) = channel_peer().await; + tx.send(Ok(Bytes::from_static(b"first"))).await.unwrap(); + assert_eq!(data(&peer.next().await), b"first"); + let Peer { + client, + read, + write, + _work, + } = peer; + drop(read); + drop(write); + timeout(DEADLINE, tx.closed()).await.unwrap(); + client.stop().await.unwrap(); +} + +struct ReturnAfterCancellation(BodyHandler); + +#[async_trait] +impl CopilotRequestHandler for ReturnAfterCancellation { + async fn send_request( + &self, + request: CopilotHttpRequest, + context: &CopilotRequestContext, + ) -> Result { + context.cancel.cancelled().await; + self.0.send_request(request, context).await + } +} + +#[tokio::test] +async fn cancellation_before_handler_returns_still_sends_head_before_terminal() { + for acknowledge_head in [false, true] { + let (tx, rx) = mpsc::channel(1); + let body = Box::pin(ReceiverStream::new(rx)); + let mut peer = + Peer::start(ReturnAfterCancellation(BodyHandler(Mutex::new(Some(body))))).await; + peer.begin_request("http://unused.test").await; + peer.request_chunk(json!({"requestId": "test", "data": "", "cancel": true})) + .await; + let head = peer.next().await; + assert_eq!(head["method"], "llmInference.httpResponseStart"); + assert_eq!(head["params"]["status"], 200); + if acknowledge_head { + peer.ack(&head).await; + } + let end = peer.next().await; + assert_eq!(end["method"], "llmInference.httpResponseChunk"); + assert_eq!(end["params"]["end"], true); + assert_eq!(end["params"]["error"]["code"], "cancelled"); + timeout(DEADLINE, tx.closed()).await.unwrap(); + peer.ack(&end).await; + assert!( + timeout(Duration::from_millis(20), peer.next()) + .await + .is_err() + ); + peer.stop().await; + } +} + +#[tokio::test] +async fn cancellation_while_head_ack_is_withheld_drops_source() { + let (tx, rx) = mpsc::channel(1); + let body = Box::pin(ReceiverStream::new(rx)); + let mut peer = Peer::start(BodyHandler(Mutex::new(Some(body)))).await; + peer.begin_request("http://unused.test").await; + let head = peer.next().await; + assert_eq!(head["method"], "llmInference.httpResponseStart"); + peer.request_chunk(json!({"requestId": "test", "data": "", "cancel": true})) + .await; + let end = peer.next().await; + assert_eq!(end["params"]["error"]["code"], "cancelled"); + timeout(DEADLINE, tx.closed()).await.unwrap(); + peer.ack(&end).await; + peer.stop().await; +} + +struct ForwardingHandler; + +#[async_trait] +impl CopilotRequestHandler for ForwardingHandler {} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn real_http_burst_preserves_headers_status_and_every_byte() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}/stream", listener.local_addr().unwrap()); + let expected: Vec = (0..CHUNK_SIZE * 8 + 7).map(|i| (i % 256) as u8).collect(); + let payload = expected.clone(); + let upstream = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let mut socket = BufReader::new(socket); + loop { + let mut line = String::new(); + assert_ne!(socket.read_line(&mut line).await.unwrap(), 0); + if line == "\r\n" { + break; + } + } + let mut response = b"HTTP/1.1 201 Created\r\nContent-Type: text/event-stream\r\nX-Test: one\r\nX-Test: two\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n".to_vec(); + for chunk in payload.chunks(1024) { + response.extend_from_slice(format!("{:x}\r\n", chunk.len()).as_bytes()); + response.extend_from_slice(chunk); + response.extend_from_slice(b"\r\n"); + } + response.extend_from_slice(b"0\r\n\r\n"); + socket.write_all(&response).await.unwrap(); + socket.shutdown().await.unwrap(); + }); + let mut peer = Peer::start(ForwardingHandler).await; + let head = peer.request(&url).await; + assert_eq!(head["params"]["status"], 201); + assert_eq!(head["params"]["statusText"], "Created"); + assert_eq!(head["params"]["headers"]["x-test"], json!(["one", "two"])); + let mut received = Vec::new(); + let mut writes = 0; + loop { + let message = peer.next().await; + if message["params"]["end"] == true { + assert!(message["params"].get("error").is_none(), "{message}"); + peer.ack(&message).await; + break; + } + received.extend(data(&message)); + writes += 1; + // Give the real HTTP transport turns to populate read-ahead under an ACK. + tokio::time::sleep(Duration::from_millis(2)).await; + peer.ack(&message).await; + } + assert_eq!(received, expected); + assert!( + writes < 32, + "bursty HTTP fragments were not combined: {writes}" + ); + timeout(DEADLINE, upstream).await.unwrap().unwrap(); + peer.stop().await; +}