| use std::future::Future; |
| use std::time::Duration; |
|
|
| use anyhow::Context; |
| use anyhow::Result; |
| use anyhow::bail; |
| use codex_exec_server::ExecServerClient; |
| use codex_exec_server::HttpHeader; |
| use codex_exec_server::HttpRedirectPolicy; |
| use codex_exec_server::HttpRequestBodyDeltaNotification; |
| use codex_exec_server::HttpRequestParams; |
| use codex_exec_server::HttpRequestResponse; |
| use codex_exec_server::InitializeParams; |
| use codex_exec_server::InitializeResponse; |
| use codex_exec_server::RemoteExecServerConnectArgs; |
| use codex_exec_server_protocol::JSONRPCMessage; |
| use codex_exec_server_protocol::JSONRPCNotification; |
| use codex_exec_server_protocol::JSONRPCRequest; |
| use codex_exec_server_protocol::JSONRPCResponse; |
| use codex_exec_server_protocol::MAX_HTTP_BODY_DELTA_BYTES; |
| use codex_exec_server_protocol::RequestId; |
| use codex_http_client::HttpClientFactory; |
| use codex_http_client::OutboundProxyPolicy; |
| use futures::SinkExt; |
| use futures::StreamExt; |
| use pretty_assertions::assert_eq; |
| use serde::Serialize; |
| use serde::de::DeserializeOwned; |
| use serde_json::from_slice; |
| use serde_json::from_str; |
| use serde_json::from_value; |
| use serde_json::to_string; |
| use serde_json::to_value; |
| use tokio::net::TcpListener; |
| use tokio::net::TcpStream; |
| use tokio::sync::oneshot; |
| use tokio::task::JoinHandle; |
| use tokio::time::timeout; |
| use tokio_tungstenite::WebSocketStream; |
| use tokio_tungstenite::accept_async; |
| use tokio_tungstenite::tungstenite::Message; |
|
|
| const CLIENT_NAME: &str = "test-exec-server-client"; |
| const HTTP_REQUEST_METHOD: &str = "http/request"; |
| const HTTP_REQUEST_BODY_DELTA_METHOD: &str = "http/request/bodyDelta"; |
| const INITIALIZE_METHOD: &str = "initialize"; |
| const INITIALIZED_METHOD: &str = "initialized"; |
| const TEST_TIMEOUT: Duration = Duration::from_secs(5); |
| const BYTE_BUDGET_TEST_TIMEOUT: Duration = Duration::from_secs(30); |
| const HTTP_BODY_DELTA_CHANNEL_CAPACITY: u64 = 256; |
| const HTTP_BODY_DELTA_BYTE_BUDGET: usize = 16 * 1024 * 1024; |
| const OVERFLOWING_BODY_DELTA_FRAMES: u64 = 1_024; |
|
|
| |
| |
| #[tokio::test] |
| async fn http_request_forces_buffered_request_params() -> Result<()> { |
| |
| |
| let server = spawn_scripted_exec_server(|mut peer| async move { |
| |
| |
| let (request_id, params) = peer.read_http_request().await?; |
| assert_eq!( |
| params, |
| HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/buffered".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "ignored-stream-id".to_string(), |
| stream_response: false, |
| } |
| ); |
|
|
| peer.write_response( |
| request_id, |
| HttpRequestResponse { |
| status: 200, |
| headers: Vec::new(), |
| body: b"buffered".to_vec().into(), |
| }, |
| ) |
| .await |
| }) |
| .await?; |
| let client = server.connect_client().await?; |
|
|
| |
| |
| let response = timeout( |
| TEST_TIMEOUT, |
| client.http_request(HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/buffered".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "ignored-stream-id".to_string(), |
| stream_response: true, |
| }), |
| ) |
| .await |
| .context("buffered http/request should complete")??; |
| assert_eq!( |
| response, |
| HttpRequestResponse { |
| status: 200, |
| headers: Vec::new(), |
| body: b"buffered".to_vec().into(), |
| } |
| ); |
|
|
| drop(client); |
| server.finish().await?; |
| Ok(()) |
| } |
|
|
| |
| |
| |
| #[tokio::test] |
| async fn http_response_body_stream_uses_generated_ids_and_receives_ordered_deltas() -> Result<()> { |
| |
| |
| let server = spawn_scripted_exec_server(|mut peer| async move { |
| let (request_id, params) = peer.read_http_request().await?; |
| assert_eq!( |
| params, |
| HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp".to_string(), |
| headers: vec![HttpHeader { |
| name: "accept".to_string(), |
| value: "text/event-stream".to_string(), |
| value_env_var: None, |
| }], |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "http-1".to_string(), |
| stream_response: true, |
| } |
| ); |
|
|
| |
| |
| peer.write_response( |
| request_id, |
| HttpRequestResponse { |
| status: 200, |
| headers: vec![HttpHeader { |
| name: "content-type".to_string(), |
| value: "text/event-stream".to_string(), |
| value_env_var: None, |
| }], |
| body: Vec::new().into(), |
| }, |
| ) |
| .await?; |
| for delta in [ |
| HttpRequestBodyDeltaNotification { |
| request_id: "http-1".to_string(), |
| seq: 1, |
| delta: b"hello ".to_vec().into(), |
| done: false, |
| error: None, |
| }, |
| HttpRequestBodyDeltaNotification { |
| request_id: "http-1".to_string(), |
| seq: 2, |
| delta: b"world".to_vec().into(), |
| done: false, |
| error: None, |
| }, |
| HttpRequestBodyDeltaNotification { |
| request_id: "http-1".to_string(), |
| seq: 3, |
| delta: b"!".to_vec().into(), |
| done: true, |
| error: None, |
| }, |
| ] { |
| peer.write_body_delta(delta).await?; |
| } |
|
|
| |
| let (request_id, params) = peer.read_http_request().await?; |
| assert_eq!( |
| params, |
| HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/reuse".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "http-2".to_string(), |
| stream_response: true, |
| } |
| ); |
| peer.write_response( |
| request_id, |
| HttpRequestResponse { |
| status: 204, |
| headers: Vec::new(), |
| body: Vec::new().into(), |
| }, |
| ) |
| .await |
| }) |
| .await?; |
| let client = server.connect_client().await?; |
|
|
| |
| let (response, mut body_stream) = timeout( |
| TEST_TIMEOUT, |
| client.http_request_stream(HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp".to_string(), |
| headers: vec![HttpHeader { |
| name: "accept".to_string(), |
| value: "text/event-stream".to_string(), |
| value_env_var: None, |
| }], |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "caller-stream-id".to_string(), |
| stream_response: false, |
| }), |
| ) |
| .await |
| .context("streamed http/request should return headers")??; |
| assert_eq!( |
| response, |
| HttpRequestResponse { |
| status: 200, |
| headers: vec![HttpHeader { |
| name: "content-type".to_string(), |
| value: "text/event-stream".to_string(), |
| value_env_var: None, |
| }], |
| body: Vec::new().into(), |
| } |
| ); |
|
|
| |
| let mut body = Vec::new(); |
| while let Some(chunk) = timeout(TEST_TIMEOUT, body_stream.recv()) |
| .await |
| .context("http response body delta should arrive")?? |
| { |
| body.extend_from_slice(&chunk); |
| } |
| assert_eq!(body, b"hello world!".to_vec()); |
|
|
| |
| |
| let (reuse_response, _reuse_body_stream) = timeout( |
| TEST_TIMEOUT, |
| client.http_request_stream(HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/reuse".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "caller-stream-id".to_string(), |
| stream_response: false, |
| }), |
| ) |
| .await |
| .context("second streamed http/request should return headers")??; |
| assert_eq!( |
| reuse_response, |
| HttpRequestResponse { |
| status: 204, |
| headers: Vec::new(), |
| body: Vec::new().into(), |
| } |
| ); |
|
|
| drop(client); |
| server.finish().await?; |
| Ok(()) |
| } |
|
|
| |
| |
| #[tokio::test] |
| async fn http_response_body_stream_drops_queued_terminal_before_next_generated_id() -> Result<()> { |
| |
| |
| let server = spawn_scripted_exec_server(|mut peer| async move { |
| let (request_id, params) = peer.read_http_request().await?; |
| assert_eq!( |
| params, |
| HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/queued-terminal".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "http-1".to_string(), |
| stream_response: true, |
| } |
| ); |
| peer.write_body_delta(HttpRequestBodyDeltaNotification { |
| request_id: "http-1".to_string(), |
| seq: 1, |
| delta: Vec::new().into(), |
| done: true, |
| error: None, |
| }) |
| .await?; |
| peer.write_response( |
| request_id, |
| HttpRequestResponse { |
| status: 200, |
| headers: Vec::new(), |
| body: Vec::new().into(), |
| }, |
| ) |
| .await?; |
|
|
| |
| |
| let (request_id, params) = peer.read_http_request().await?; |
| assert_eq!( |
| params, |
| HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/retry-queued-terminal".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "http-2".to_string(), |
| stream_response: true, |
| } |
| ); |
| peer.write_response( |
| request_id, |
| HttpRequestResponse { |
| status: 204, |
| headers: Vec::new(), |
| body: Vec::new().into(), |
| }, |
| ) |
| .await |
| }) |
| .await?; |
| let client = server.connect_client().await?; |
|
|
| |
| let (response, body_stream) = timeout( |
| TEST_TIMEOUT, |
| client.http_request_stream(HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/queued-terminal".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "caller-stream-id".to_string(), |
| stream_response: false, |
| }), |
| ) |
| .await |
| .context("streamed http/request should return headers")??; |
| assert_eq!( |
| response, |
| HttpRequestResponse { |
| status: 200, |
| headers: Vec::new(), |
| body: Vec::new().into(), |
| } |
| ); |
| drop(body_stream); |
|
|
| |
| |
| let params = HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/retry-queued-terminal".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "caller-stream-id".to_string(), |
| stream_response: false, |
| }; |
| let (reuse_response, _reuse_body_stream) = |
| timeout(TEST_TIMEOUT, client.http_request_stream(params)) |
| .await |
| .context("second streamed http/request should return headers")??; |
| assert_eq!( |
| reuse_response, |
| HttpRequestResponse { |
| status: 204, |
| headers: Vec::new(), |
| body: Vec::new().into(), |
| } |
| ); |
|
|
| drop(client); |
| server.finish().await?; |
| Ok(()) |
| } |
|
|
| |
| |
| #[tokio::test] |
| async fn http_response_body_stream_ignores_late_deltas_after_cancelled_request() -> Result<()> { |
| |
| |
| |
| let (request_seen_tx, request_seen_rx) = oneshot::channel(); |
| let server = spawn_scripted_exec_server(|mut peer| async move { |
| let (_request_id, params) = peer.read_http_request().await?; |
| assert_eq!( |
| params, |
| HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/cancel".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "http-1".to_string(), |
| stream_response: true, |
| } |
| ); |
| request_seen_tx |
| .send(()) |
| .expect("test should wait for the first request"); |
|
|
| |
| |
| let (request_id, params) = peer.read_http_request().await?; |
| assert_eq!( |
| params, |
| HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/retry-cancelled".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "http-2".to_string(), |
| stream_response: true, |
| } |
| ); |
| peer.write_body_delta(HttpRequestBodyDeltaNotification { |
| request_id: "http-1".to_string(), |
| seq: 1, |
| delta: b"stale".to_vec().into(), |
| done: false, |
| error: None, |
| }) |
| .await?; |
| peer.write_response( |
| request_id, |
| HttpRequestResponse { |
| status: 200, |
| headers: Vec::new(), |
| body: Vec::new().into(), |
| }, |
| ) |
| .await?; |
| peer.write_body_delta(HttpRequestBodyDeltaNotification { |
| request_id: "http-2".to_string(), |
| seq: 1, |
| delta: b"fresh".to_vec().into(), |
| done: true, |
| error: None, |
| }) |
| .await |
| }) |
| .await?; |
| let client = server.connect_client().await?; |
|
|
| |
| |
| let client_for_request = client.clone(); |
| let stream_task = tokio::spawn(async move { |
| let _ = client_for_request |
| .http_request_stream(HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/cancel".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "caller-stream-id".to_string(), |
| stream_response: false, |
| }) |
| .await; |
| }); |
| request_seen_rx |
| .await |
| .expect("server should observe the first http/request"); |
| stream_task.abort(); |
| let _ = stream_task.await; |
|
|
| |
| |
| let (response, mut body_stream) = timeout( |
| TEST_TIMEOUT, |
| client.http_request_stream(HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/retry-cancelled".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "caller-stream-id".to_string(), |
| stream_response: false, |
| }), |
| ) |
| .await |
| .context("second streamed http/request should return headers")??; |
| assert_eq!( |
| response, |
| HttpRequestResponse { |
| status: 200, |
| headers: Vec::new(), |
| body: Vec::new().into(), |
| } |
| ); |
| let mut body = Vec::new(); |
| while let Some(chunk) = timeout(TEST_TIMEOUT, body_stream.recv()) |
| .await |
| .context("fresh http response body delta should arrive")?? |
| { |
| body.extend_from_slice(&chunk); |
| } |
| assert_eq!(body, b"fresh".to_vec()); |
|
|
| drop(client); |
| server.finish().await?; |
| Ok(()) |
| } |
|
|
| |
| |
| #[tokio::test] |
| async fn http_response_body_stream_ignores_late_deltas_after_drop() -> Result<()> { |
| |
| |
| let (body_dropped_tx, body_dropped_rx) = oneshot::channel(); |
| let (stale_delta_sent_tx, stale_delta_sent_rx) = oneshot::channel(); |
| let server = spawn_scripted_exec_server(|mut peer| async move { |
| let (request_id, params) = peer.read_http_request().await?; |
| assert_eq!( |
| params, |
| HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/drop".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "http-1".to_string(), |
| stream_response: true, |
| } |
| ); |
| peer.write_response( |
| request_id, |
| HttpRequestResponse { |
| status: 200, |
| headers: Vec::new(), |
| body: Vec::new().into(), |
| }, |
| ) |
| .await?; |
| body_dropped_rx |
| .await |
| .expect("test should drop the first body stream"); |
| peer.write_body_delta(HttpRequestBodyDeltaNotification { |
| request_id: "http-1".to_string(), |
| seq: 1, |
| delta: b"stale".to_vec().into(), |
| done: false, |
| error: None, |
| }) |
| .await?; |
| stale_delta_sent_tx |
| .send(()) |
| .expect("test should wait for the stale delta"); |
|
|
| |
| |
| let (request_id, params) = peer.read_http_request().await?; |
| assert_eq!( |
| params, |
| HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/retry-dropped".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "http-2".to_string(), |
| stream_response: true, |
| } |
| ); |
| peer.write_response( |
| request_id, |
| HttpRequestResponse { |
| status: 200, |
| headers: Vec::new(), |
| body: Vec::new().into(), |
| }, |
| ) |
| .await?; |
| peer.write_body_delta(HttpRequestBodyDeltaNotification { |
| request_id: "http-2".to_string(), |
| seq: 1, |
| delta: b"fresh".to_vec().into(), |
| done: true, |
| error: None, |
| }) |
| .await |
| }) |
| .await?; |
| let client = server.connect_client().await?; |
|
|
| |
| |
| let (response, body_stream) = timeout( |
| TEST_TIMEOUT, |
| client.http_request_stream(HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/drop".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "caller-stream-id".to_string(), |
| stream_response: false, |
| }), |
| ) |
| .await |
| .context("streamed http/request should return headers")??; |
| assert_eq!( |
| response, |
| HttpRequestResponse { |
| status: 200, |
| headers: Vec::new(), |
| body: Vec::new().into(), |
| } |
| ); |
| drop(body_stream); |
| body_dropped_tx |
| .send(()) |
| .expect("server should wait for the body stream drop"); |
| stale_delta_sent_rx |
| .await |
| .expect("server should send one stale nonterminal delta"); |
|
|
| |
| |
| let (reuse_response, mut reuse_body_stream) = timeout( |
| TEST_TIMEOUT, |
| client.http_request_stream(HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/retry-dropped".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "caller-stream-id".to_string(), |
| stream_response: false, |
| }), |
| ) |
| .await |
| .context("second streamed http/request should return headers")??; |
| assert_eq!( |
| reuse_response, |
| HttpRequestResponse { |
| status: 200, |
| headers: Vec::new(), |
| body: Vec::new().into(), |
| } |
| ); |
| let mut body = Vec::new(); |
| while let Some(chunk) = timeout(TEST_TIMEOUT, reuse_body_stream.recv()) |
| .await |
| .context("fresh http response body delta should arrive")?? |
| { |
| body.extend_from_slice(&chunk); |
| } |
| assert_eq!(body, b"fresh".to_vec()); |
|
|
| drop(client); |
| server.finish().await?; |
| Ok(()) |
| } |
|
|
| |
| |
| #[tokio::test] |
| async fn http_response_body_stream_fails_when_transport_disconnects() -> Result<()> { |
| |
| |
| let server = spawn_scripted_exec_server(|mut peer| async move { |
| let (request_id, params) = peer.read_http_request().await?; |
| assert_eq!( |
| params, |
| HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/disconnect".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "http-1".to_string(), |
| stream_response: true, |
| } |
| ); |
| peer.write_response( |
| request_id, |
| HttpRequestResponse { |
| status: 200, |
| headers: Vec::new(), |
| body: Vec::new().into(), |
| }, |
| ) |
| .await |
| }) |
| .await?; |
| let client = server.connect_client().await?; |
|
|
| |
| let (_response, mut body_stream) = timeout( |
| TEST_TIMEOUT, |
| client.http_request_stream(HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/disconnect".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "caller-stream-id".to_string(), |
| stream_response: false, |
| }), |
| ) |
| .await |
| .context("streamed http/request should return headers")??; |
|
|
| |
| |
| let error = timeout(TEST_TIMEOUT, body_stream.recv()) |
| .await |
| .context("disconnect should wake http body stream")? |
| .expect_err("disconnect should fail the http body stream"); |
| let error_message = error.to_string(); |
| assert_eq!( |
| error_message.starts_with( |
| "exec-server protocol error: http response stream `http-1` failed: exec-server transport disconnected" |
| ), |
| true |
| ); |
|
|
| drop(client); |
| server.finish().await?; |
| Ok(()) |
| } |
|
|
| |
| |
| #[tokio::test] |
| async fn http_response_body_stream_rejects_oversized_delta() -> Result<()> { |
| let (finish_tx, finish_rx) = oneshot::channel(); |
| let server = spawn_scripted_exec_server(|mut peer| async move { |
| let (_request_id, params) = peer.read_http_request().await?; |
| assert_eq!( |
| params, |
| HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/oversized-delta".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "http-1".to_string(), |
| stream_response: true, |
| } |
| ); |
| peer.write_body_delta(HttpRequestBodyDeltaNotification { |
| request_id: "http-1".to_string(), |
| seq: 1, |
| delta: vec![0; MAX_HTTP_BODY_DELTA_BYTES + 1].into(), |
| done: false, |
| error: None, |
| }) |
| .await?; |
| finish_rx.await.expect("test should finish server task"); |
| Ok(()) |
| }) |
| .await?; |
| let client = server.connect_client().await?; |
|
|
| let request = HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/oversized-delta".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "caller-stream-id".to_string(), |
| stream_response: false, |
| }; |
| let result = timeout(TEST_TIMEOUT, client.http_request_stream(request)) |
| .await |
| .context("oversized body delta should close the executor transport")?; |
| let error = match result { |
| Ok(_) => bail!("oversized body delta should fail the request"), |
| Err(error) => error, |
| }; |
| let error = error.to_string(); |
| assert_eq!(error, "exec-server transport disconnected"); |
|
|
| finish_tx.send(()).expect("server task should stay active"); |
| drop(client); |
| server.finish().await?; |
| Ok(()) |
| } |
|
|
| |
| |
| #[tokio::test] |
| async fn http_response_body_stream_enforces_queued_byte_budget() -> Result<()> { |
| let (finish_tx, finish_rx) = oneshot::channel(); |
| let server = spawn_scripted_exec_server(|mut peer| async move { |
| let (request_id, params) = peer.read_http_request().await?; |
| assert_eq!( |
| params, |
| HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/byte-budget".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "http-1".to_string(), |
| stream_response: true, |
| } |
| ); |
| peer.write_response( |
| request_id, |
| HttpRequestResponse { |
| status: 200, |
| headers: Vec::new(), |
| body: Vec::new().into(), |
| }, |
| ) |
| .await?; |
|
|
| let frame_count = HTTP_BODY_DELTA_BYTE_BUDGET / MAX_HTTP_BODY_DELTA_BYTES + 1; |
| for seq in 1..=frame_count as u64 { |
| peer.write_body_delta(HttpRequestBodyDeltaNotification { |
| request_id: "http-1".to_string(), |
| seq, |
| delta: vec![0; MAX_HTTP_BODY_DELTA_BYTES].into(), |
| done: false, |
| error: None, |
| }) |
| .await?; |
| } |
|
|
| let (barrier_request_id, barrier_params) = peer.read_http_request().await?; |
| assert_eq!( |
| barrier_params, |
| HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/byte-budget-barrier".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "http-2".to_string(), |
| stream_response: true, |
| } |
| ); |
| peer.write_response( |
| barrier_request_id, |
| HttpRequestResponse { |
| status: 200, |
| headers: Vec::new(), |
| body: Vec::new().into(), |
| }, |
| ) |
| .await?; |
| peer.write_body_delta(HttpRequestBodyDeltaNotification { |
| request_id: "http-2".to_string(), |
| seq: 1, |
| delta: Vec::new().into(), |
| done: true, |
| error: None, |
| }) |
| .await?; |
| finish_rx.await.expect("test should finish server task"); |
| Ok(()) |
| }) |
| .await?; |
| let client = server.connect_client().await?; |
|
|
| let (_response, mut body_stream) = timeout( |
| TEST_TIMEOUT, |
| client.http_request_stream(HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/byte-budget".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "caller-stream-id".to_string(), |
| stream_response: false, |
| }), |
| ) |
| .await |
| .context("streamed http/request should return headers")??; |
|
|
| |
| |
| let (_response, mut barrier_stream) = timeout( |
| BYTE_BUDGET_TEST_TIMEOUT, |
| client.http_request_stream(HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/byte-budget-barrier".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "caller-barrier-id".to_string(), |
| stream_response: false, |
| }), |
| ) |
| .await |
| .context("barrier http/request should return headers")??; |
| assert_eq!( |
| timeout(TEST_TIMEOUT, barrier_stream.recv()) |
| .await |
| .context("barrier body stream should finish")??, |
| None |
| ); |
|
|
| let mut delivered_bytes = 0; |
| let error = loop { |
| match timeout(TEST_TIMEOUT, body_stream.recv()) |
| .await |
| .context("queued body stream should finish")? |
| { |
| Ok(Some(chunk)) => delivered_bytes += chunk.len(), |
| Ok(None) => bail!("byte-budget exhaustion should not look like clean EOF"), |
| Err(error) => break error, |
| } |
| }; |
| assert_eq!(delivered_bytes, HTTP_BODY_DELTA_BYTE_BUDGET); |
| assert!( |
| error |
| .to_string() |
| .contains("queued body deltas exceed 16777216 bytes") |
| ); |
|
|
| finish_tx.send(()).expect("server task should stay active"); |
| drop(client); |
| server.finish().await?; |
| Ok(()) |
| } |
|
|
| |
| |
| #[tokio::test] |
| async fn http_response_body_streams_share_queued_byte_budget() -> Result<()> { |
| let (finish_tx, finish_rx) = oneshot::channel(); |
| let server = spawn_scripted_exec_server(|mut peer| async move { |
| for (request_id, url) in [ |
| ("http-1", "https://example.test/mcp/shared-budget-one"), |
| ("http-2", "https://example.test/mcp/shared-budget-two"), |
| ] { |
| let (rpc_request_id, params) = peer.read_http_request().await?; |
| assert_eq!( |
| params, |
| HttpRequestParams { |
| method: "GET".to_string(), |
| url: url.to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: request_id.to_string(), |
| stream_response: true, |
| } |
| ); |
| peer.write_response( |
| rpc_request_id, |
| HttpRequestResponse { |
| status: 200, |
| headers: Vec::new(), |
| body: Vec::new().into(), |
| }, |
| ) |
| .await?; |
| } |
|
|
| let frames_per_stream = HTTP_BODY_DELTA_BYTE_BUDGET / MAX_HTTP_BODY_DELTA_BYTES / 2; |
| for request_id in ["http-1", "http-2"] { |
| for seq in 1..=frames_per_stream as u64 { |
| peer.write_body_delta(HttpRequestBodyDeltaNotification { |
| request_id: request_id.to_string(), |
| seq, |
| delta: vec![0; MAX_HTTP_BODY_DELTA_BYTES].into(), |
| done: false, |
| error: None, |
| }) |
| .await?; |
| } |
| } |
| peer.write_body_delta(HttpRequestBodyDeltaNotification { |
| request_id: "http-2".to_string(), |
| seq: frames_per_stream as u64 + 1, |
| delta: vec![0; MAX_HTTP_BODY_DELTA_BYTES].into(), |
| done: false, |
| error: None, |
| }) |
| .await?; |
|
|
| let (barrier_request_id, barrier_params) = peer.read_http_request().await?; |
| assert_eq!( |
| barrier_params, |
| HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/shared-budget-barrier".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "http-3".to_string(), |
| stream_response: true, |
| } |
| ); |
| peer.write_response( |
| barrier_request_id, |
| HttpRequestResponse { |
| status: 200, |
| headers: Vec::new(), |
| body: Vec::new().into(), |
| }, |
| ) |
| .await?; |
| for (request_id, seq) in [ |
| ("http-3", 1), |
| ("http-1", frames_per_stream as u64 + 1), |
| ("http-2", frames_per_stream as u64 + 2), |
| ] { |
| peer.write_body_delta(HttpRequestBodyDeltaNotification { |
| request_id: request_id.to_string(), |
| seq, |
| delta: Vec::new().into(), |
| done: true, |
| error: None, |
| }) |
| .await?; |
| } |
| finish_rx.await.expect("test should finish server task"); |
| Ok(()) |
| }) |
| .await?; |
| let client = server.connect_client().await?; |
|
|
| let (_response, mut first_stream) = timeout( |
| TEST_TIMEOUT, |
| client.http_request_stream(HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/shared-budget-one".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "caller-stream-one".to_string(), |
| stream_response: false, |
| }), |
| ) |
| .await |
| .context("first streamed http/request should return headers")??; |
| let (_response, mut second_stream) = timeout( |
| TEST_TIMEOUT, |
| client.http_request_stream(HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/shared-budget-two".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "caller-stream-two".to_string(), |
| stream_response: false, |
| }), |
| ) |
| .await |
| .context("second streamed http/request should return headers")??; |
|
|
| |
| |
| let (_response, mut barrier_stream) = timeout( |
| BYTE_BUDGET_TEST_TIMEOUT, |
| client.http_request_stream(HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/shared-budget-barrier".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "caller-barrier-id".to_string(), |
| stream_response: false, |
| }), |
| ) |
| .await |
| .context("barrier http/request should return headers")??; |
| assert_eq!( |
| timeout(TEST_TIMEOUT, barrier_stream.recv()) |
| .await |
| .context("barrier body stream should finish")??, |
| None |
| ); |
|
|
| let mut failed_stream_bytes = 0; |
| let error = loop { |
| match timeout(TEST_TIMEOUT, second_stream.recv()) |
| .await |
| .context("second body stream should finish")? |
| { |
| Ok(Some(chunk)) => failed_stream_bytes += chunk.len(), |
| Ok(None) => bail!("shared byte-budget exhaustion should not look like clean EOF"), |
| Err(error) => break error, |
| } |
| }; |
| assert_eq!( |
| (failed_stream_bytes, error.to_string()), |
| ( |
| HTTP_BODY_DELTA_BYTE_BUDGET / 2, |
| "exec-server protocol error: http response stream `http-2` failed: queued body deltas exceed 16777216 bytes".to_string(), |
| ) |
| ); |
|
|
| let mut surviving_stream_bytes = 0; |
| while let Some(chunk) = timeout(TEST_TIMEOUT, first_stream.recv()) |
| .await |
| .context("first body stream should finish")?? |
| { |
| surviving_stream_bytes += chunk.len(); |
| } |
| assert_eq!(surviving_stream_bytes, HTTP_BODY_DELTA_BYTE_BUDGET / 2); |
|
|
| finish_tx.send(()).expect("server task should stay active"); |
| drop(client); |
| server.finish().await?; |
| Ok(()) |
| } |
|
|
| |
| |
| #[tokio::test] |
| async fn http_response_body_stream_reports_disconnect_when_queue_is_full() -> Result<()> { |
| |
| |
| let server = spawn_scripted_exec_server(|mut peer| async move { |
| let (request_id, params) = peer.read_http_request().await?; |
| assert_eq!( |
| params, |
| HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/disconnect-full-queue".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "http-1".to_string(), |
| stream_response: true, |
| } |
| ); |
| for seq in 1..=HTTP_BODY_DELTA_CHANNEL_CAPACITY { |
| peer.write_body_delta(HttpRequestBodyDeltaNotification { |
| request_id: "http-1".to_string(), |
| seq, |
| delta: b"x".to_vec().into(), |
| done: false, |
| error: None, |
| }) |
| .await?; |
| } |
| peer.write_response( |
| request_id, |
| HttpRequestResponse { |
| status: 200, |
| headers: Vec::new(), |
| body: Vec::new().into(), |
| }, |
| ) |
| .await |
| }) |
| .await?; |
| let client = server.connect_client().await?; |
|
|
| |
| |
| let (_response, mut body_stream) = timeout( |
| TEST_TIMEOUT, |
| client.http_request_stream(HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/disconnect-full-queue".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "caller-stream-id".to_string(), |
| stream_response: false, |
| }), |
| ) |
| .await |
| .context("streamed http/request should return headers")??; |
|
|
| |
| |
| let mut chunks = 0; |
| let error = loop { |
| match timeout(TEST_TIMEOUT, body_stream.recv()) |
| .await |
| .context("disconnect should wake the full queued body stream")? |
| { |
| Ok(Some(_chunk)) => { |
| chunks += 1; |
| } |
| Ok(None) => bail!("disconnect with a full queue should not look like clean EOF"), |
| Err(error) => break error, |
| } |
| }; |
| assert_eq!( |
| ( |
| chunks, |
| error |
| .to_string() |
| .starts_with( |
| "exec-server protocol error: http response stream `http-1` failed: exec-server transport disconnected", |
| ), |
| ), |
| (HTTP_BODY_DELTA_CHANNEL_CAPACITY as usize, true) |
| ); |
|
|
| drop(client); |
| server.finish().await?; |
| Ok(()) |
| } |
|
|
| |
| |
| #[tokio::test] |
| async fn http_response_body_stream_reports_backpressure_truncation() -> Result<()> { |
| |
| |
| let (finish_tx, finish_rx) = oneshot::channel(); |
| let server = spawn_scripted_exec_server(|mut peer| async move { |
| let (request_id, params) = peer.read_http_request().await?; |
| assert_eq!( |
| params, |
| HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/backpressure".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "http-1".to_string(), |
| stream_response: true, |
| } |
| ); |
| for seq in 1..=OVERFLOWING_BODY_DELTA_FRAMES { |
| peer.write_body_delta(HttpRequestBodyDeltaNotification { |
| request_id: "http-1".to_string(), |
| seq, |
| delta: b"x".to_vec().into(), |
| done: false, |
| error: None, |
| }) |
| .await?; |
| } |
| peer.write_response( |
| request_id, |
| HttpRequestResponse { |
| status: 200, |
| headers: Vec::new(), |
| body: Vec::new().into(), |
| }, |
| ) |
| .await?; |
|
|
| |
| |
| finish_rx.await.expect("test should finish server task"); |
| Ok(()) |
| }) |
| .await?; |
| let client = server.connect_client().await?; |
|
|
| |
| |
| let (_response, mut body_stream) = timeout( |
| TEST_TIMEOUT, |
| client.http_request_stream(HttpRequestParams { |
| method: "GET".to_string(), |
| url: "https://example.test/mcp/backpressure".to_string(), |
| headers: Vec::new(), |
| body: None, |
| timeout_ms: None, |
| redirect_policy: HttpRedirectPolicy::Follow, |
| request_id: "caller-stream-id".to_string(), |
| stream_response: false, |
| }), |
| ) |
| .await |
| .context("streamed http/request should return headers")??; |
|
|
| |
| |
| let mut chunks = 0; |
| let error = loop { |
| match timeout(TEST_TIMEOUT, body_stream.recv()) |
| .await |
| .context("backpressure should close http body stream")? |
| { |
| Ok(Some(_chunk)) => { |
| chunks += 1; |
| } |
| Ok(None) => bail!("backpressure truncation should not look like clean EOF"), |
| Err(error) => break error, |
| } |
| }; |
| assert_eq!( |
| ( |
| chunks < OVERFLOWING_BODY_DELTA_FRAMES as usize, |
| error.to_string(), |
| ), |
| ( |
| true, |
| "exec-server protocol error: http response stream `http-1` failed: body delta channel filled before delivery".to_string(), |
| ) |
| ); |
|
|
| finish_tx |
| .send(()) |
| .expect("server task should wait for test completion"); |
| drop(client); |
| server.finish().await?; |
| Ok(()) |
| } |
|
|
| |
| |
| |
| |
| |
| struct ScriptedExecServer { |
| websocket_url: String, |
| task: JoinHandle<Result<()>>, |
| } |
|
|
| impl ScriptedExecServer { |
| |
| async fn connect_client(&self) -> Result<ExecServerClient> { |
| ExecServerClient::connect_websocket(RemoteExecServerConnectArgs::new( |
| self.websocket_url.clone(), |
| CLIENT_NAME.to_string(), |
| HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), |
| )) |
| .await |
| .context("client should connect to fake exec-server") |
| } |
|
|
| |
| async fn finish(self) -> Result<()> { |
| self.task |
| .await |
| .context("fake exec-server task should join")??; |
| Ok(()) |
| } |
| } |
|
|
| |
| async fn spawn_scripted_exec_server<F, Fut>(script: F) -> Result<ScriptedExecServer> |
| where |
| F: FnOnce(JsonRpcPeer) -> Fut + Send + 'static, |
| Fut: Future<Output = Result<()>> + Send + 'static, |
| { |
| let listener = TcpListener::bind("127.0.0.1:0") |
| .await |
| .context("fake exec-server should bind")?; |
| let websocket_url = format!("ws://{}", listener.local_addr()?); |
| let task = tokio::spawn(async move { |
| let (stream, _) = timeout(TEST_TIMEOUT, listener.accept()) |
| .await |
| .context("fake exec-server should accept a client")??; |
| let websocket = accept_async(stream) |
| .await |
| .context("fake exec-server websocket handshake should complete")?; |
| let mut peer = JsonRpcPeer { websocket }; |
| peer.complete_initialize().await?; |
| script(peer).await |
| }); |
| Ok(ScriptedExecServer { |
| websocket_url, |
| task, |
| }) |
| } |
|
|
| |
| struct JsonRpcPeer { |
| websocket: WebSocketStream<TcpStream>, |
| } |
|
|
| impl JsonRpcPeer { |
| |
| async fn complete_initialize(&mut self) -> Result<()> { |
| let request = self.read_request(INITIALIZE_METHOD).await?; |
| let params: InitializeParams = decode_request_params(&request)?; |
| assert_eq!( |
| params, |
| InitializeParams { |
| client_name: CLIENT_NAME.to_string(), |
| resume_session_id: None, |
| } |
| ); |
| self.write_response( |
| request.id, |
| InitializeResponse { |
| session_id: "session-1".to_string(), |
| environment_info: None, |
| }, |
| ) |
| .await?; |
| self.read_notification(INITIALIZED_METHOD).await?; |
| Ok(()) |
| } |
|
|
| |
| async fn read_http_request(&mut self) -> Result<(RequestId, HttpRequestParams)> { |
| let request = self.read_request(HTTP_REQUEST_METHOD).await?; |
| let params = decode_request_params(&request)?; |
| Ok((request.id, params)) |
| } |
|
|
| |
| async fn read_request(&mut self, expected_method: &str) -> Result<JSONRPCRequest> { |
| let message = self.read_message().await?; |
| let JSONRPCMessage::Request(request) = message else { |
| bail!("expected JSON-RPC request `{expected_method}`, got {message:?}"); |
| }; |
| if request.method != expected_method { |
| bail!( |
| "expected JSON-RPC request `{expected_method}`, got `{}`", |
| request.method |
| ); |
| } |
| Ok(request) |
| } |
|
|
| |
| async fn read_notification(&mut self, expected_method: &str) -> Result<JSONRPCNotification> { |
| let message = self.read_message().await?; |
| let JSONRPCMessage::Notification(notification) = message else { |
| bail!("expected JSON-RPC notification `{expected_method}`, got {message:?}"); |
| }; |
| if notification.method != expected_method { |
| bail!( |
| "expected JSON-RPC notification `{expected_method}`, got `{}`", |
| notification.method |
| ); |
| } |
| Ok(notification) |
| } |
|
|
| |
| async fn write_response<T>(&mut self, id: RequestId, result: T) -> Result<()> |
| where |
| T: Serialize, |
| { |
| self.write_message(JSONRPCMessage::Response(JSONRPCResponse { |
| id, |
| result: to_value(result)?, |
| })) |
| .await |
| } |
|
|
| |
| async fn write_body_delta(&mut self, delta: HttpRequestBodyDeltaNotification) -> Result<()> { |
| self.write_message(JSONRPCMessage::Notification(JSONRPCNotification { |
| method: HTTP_REQUEST_BODY_DELTA_METHOD.to_string(), |
| params: Some(to_value(delta)?), |
| })) |
| .await |
| } |
|
|
| |
| async fn read_message(&mut self) -> Result<JSONRPCMessage> { |
| let message = timeout(TEST_TIMEOUT, self.websocket.next()) |
| .await |
| .context("timed out waiting for JSON-RPC message")? |
| .context("client websocket closed before JSON-RPC message arrived")? |
| .context("failed to read websocket message")?; |
| match message { |
| Message::Text(text) => from_str(text.as_ref()).context("text JSON-RPC"), |
| Message::Binary(bytes) => from_slice(bytes.as_ref()).context("binary JSON-RPC"), |
| Message::Close(frame) => bail!("client websocket closed: {frame:?}"), |
| other => bail!("expected text or binary JSON-RPC message, got {other:?}"), |
| } |
| } |
|
|
| |
| async fn write_message(&mut self, message: JSONRPCMessage) -> Result<()> { |
| let encoded = to_string(&message)?; |
| timeout( |
| TEST_TIMEOUT, |
| self.websocket.send(Message::Text(encoded.into())), |
| ) |
| .await |
| .context("timed out writing JSON-RPC message")? |
| .context("failed to write JSON-RPC message") |
| } |
| } |
|
|
| |
| fn decode_request_params<T>(request: &JSONRPCRequest) -> Result<T> |
| where |
| T: DeserializeOwned, |
| { |
| let params = request |
| .params |
| .clone() |
| .context("JSON-RPC request should include params")?; |
| from_value(params).context("JSON-RPC request params should decode") |
| } |
|
|