Download codex-rs/exec-server/tests/http_client.rs from SaylorTwift/codex: direct link, hf CLI and curl.
- Browser
- Download file 54.9 kB
-
https://huggingface.co/SaylorTwift/codex/resolve/main/codex-rs/exec-server/tests/http_client.rs
- Command line
-
hf download hf://SaylorTwift/codex/codex-rs/exec-server/tests/http_client.rs
-
curl -L -o http_client.rs https://huggingface.co/SaylorTwift/codex/resolve/main/codex-rs/exec-server/tests/http_client.rs
54.9 kB
| 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; | |
| /// What this tests: the buffered HTTP helper always sends a buffered | |
| /// `http/request`, even when a caller accidentally provides streaming flags. | |
| async fn http_request_forces_buffered_request_params() -> Result<()> { | |
| // Phase 1: start a fake WebSocket exec-server so the test covers the | |
| // public client connection path without depending on the HTTP runner. | |
| let server = spawn_scripted_exec_server(|mut peer| async move { | |
| // Phase 2: verify the buffered helper forces buffered mode before it | |
| // sends the JSON-RPC call. | |
| 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?; | |
| // Phase 3: call the buffered helper with streaming-only fields populated | |
| // and assert callers still receive the buffered response body. | |
| 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(()) | |
| } | |
| /// What this tests: streamed executor HTTP response frames are routed by the | |
| /// client's generated request id, delivered in sequence, and concatenated by | |
| /// the caller. | |
| async fn http_response_body_stream_uses_generated_ids_and_receives_ordered_deltas() -> Result<()> { | |
| // Phase 1: script two requests. The caller supplies reusable ids, but the | |
| // client replaces them with connection-local ids on the wire. | |
| 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, | |
| } | |
| ); | |
| // Phase 2: return headers first, then body notifications in the order | |
| // the public body stream should expose them. | |
| 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?; | |
| } | |
| // Phase 3: accept the next generated request id after EOF. | |
| 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?; | |
| // Phase 4: start a streaming HTTP request through the public client API. | |
| 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(), | |
| } | |
| ); | |
| // Phase 5: drain the body stream and verify the caller-visible byte order. | |
| 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()); | |
| // Phase 6: start another stream through the public API to validate cleanup | |
| // after EOF without reaching into the client routing table. | |
| 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(()) | |
| } | |
| /// What this tests: dropping a body stream with a queued terminal frame removes | |
| /// the old route while the next stream gets a fresh generated id. | |
| async fn http_response_body_stream_drops_queued_terminal_before_next_generated_id() -> Result<()> { | |
| // Phase 1: send terminal EOF before the header response so the public body | |
| // stream starts with EOF already queued but unread. | |
| 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?; | |
| // Phase 2: accept another stream after the client drops the unread | |
| // body. The second request receives a distinct generated id. | |
| 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?; | |
| // Phase 3: drop the body stream without reading the queued EOF frame. | |
| 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); | |
| // Phase 4: start another stream through the public API. The caller-provided | |
| // id is ignored, so the request uses the next generated route id. | |
| 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(()) | |
| } | |
| /// What this tests: cancelling a streaming HTTP request while it is waiting for | |
| /// headers drops its route, and a later stream gets a fresh generated id. | |
| async fn http_response_body_stream_ignores_late_deltas_after_cancelled_request() -> Result<()> { | |
| // Phase 1: coordinate cancellation after the fake server observes the | |
| // first request but before it returns headers. The server later sends a | |
| // stale delta for the cancelled id before serving the fresh stream. | |
| 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"); | |
| // Phase 2: the next stream uses a new generated id. A late body delta | |
| // for the cancelled id is ignored by the client-side router. | |
| 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?; | |
| // Phase 3: start a streaming request and abort the caller future while it | |
| // is blocked waiting for response headers. | |
| 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; | |
| // Phase 4: start a new stream immediately. It receives only the fresh body | |
| // bytes for its generated id. | |
| 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(()) | |
| } | |
| /// What this tests: dropping a returned body stream before EOF removes its | |
| /// route and prevents stale body deltas from reaching the next stream. | |
| async fn http_response_body_stream_ignores_late_deltas_after_drop() -> Result<()> { | |
| // Phase 1: script two requests. The first returns only headers; after the | |
| // client drops its body receiver, the server sends a stale body delta. | |
| 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"); | |
| // Phase 2: accept the next request with a new generated id. The new | |
| // stream must receive only fresh body bytes. | |
| 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?; | |
| // Phase 3: receive headers for the first stream, then drop the body stream | |
| // without reading any body frames. | |
| 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"); | |
| // Phase 4: start the next stream immediately. The caller-provided id is | |
| // ignored, and the fresh generated id isolates it from stale bytes. | |
| 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(()) | |
| } | |
| /// What this tests: an in-flight streamed HTTP body is failed when the shared | |
| /// JSON-RPC transport disconnects before a terminal body frame. | |
| async fn http_response_body_stream_fails_when_transport_disconnects() -> Result<()> { | |
| // Phase 1: return response headers for a streaming request, then drop the | |
| // fake server transport without sending EOF. | |
| 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?; | |
| // Phase 2: start a streaming HTTP request and receive headers. | |
| 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")??; | |
| // Phase 3: assert transport disconnect wakes the body stream with a | |
| // terminal error instead of hanging. | |
| 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(()) | |
| } | |
| /// What this tests: an executor cannot make the orchestrator decode and retain | |
| /// a body frame larger than the response-stream wire contract allows. | |
| 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(()) | |
| } | |
| /// What this tests: frame-count backpressure cannot hide an unbounded amount | |
| /// of executor-controlled body bytes across the orchestrator's stream queues. | |
| 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")??; | |
| // Receiving this terminal notification proves the earlier byte-budget | |
| // notifications have all passed through the ordered notification handler. | |
| 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(()) | |
| } | |
| /// What this tests: every response stream on one executor connection shares | |
| /// the same queued-body byte budget. | |
| 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")??; | |
| // This terminal notification is ordered after both streams contend for the | |
| // budget, so neither stream is drained before the overflow is observed. | |
| 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(()) | |
| } | |
| /// What this tests: transport disconnect still records a terminal stream | |
| /// failure even when the client-side body-delta queue is already full. | |
| async fn http_response_body_stream_reports_disconnect_when_queue_is_full() -> Result<()> { | |
| // Phase 1: fill the queued body-delta route exactly to capacity before the | |
| // response headers arrive, then drop the transport without sending EOF. | |
| 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?; | |
| // Phase 2: start the streaming request and receive headers while the | |
| // queue is already full. | |
| 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")??; | |
| // Phase 3: drain the queued chunks and assert the transport disconnect is | |
| // still reported as an error rather than a clean EOF. | |
| 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(()) | |
| } | |
| /// What this tests: body-delta backpressure closes the public body stream as | |
| /// an error rather than letting callers accept a truncated body as clean EOF. | |
| async fn http_response_body_stream_reports_backpressure_truncation() -> Result<()> { | |
| // Phase 1: send enough body frames before headers to overflow the bounded | |
| // client-side route while the public request future is still pending. | |
| 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?; | |
| // Phase 2: keep the transport connected so the body stream reports the | |
| // backpressure failure rather than a disconnect. | |
| finish_rx.await.expect("test should finish server task"); | |
| Ok(()) | |
| }) | |
| .await?; | |
| let client = server.connect_client().await?; | |
| // Phase 3: start the streaming request; the server overfills the route | |
| // before returning the body stream to this consumer. | |
| 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")??; | |
| // Phase 4: drain queued chunks and assert the truncated stream ends in an | |
| // explicit error, not a clean EOF. | |
| 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(()) | |
| } | |
| /// Fake WebSocket exec-server used by the integration tests. | |
| /// | |
| /// The helper exercises `ExecServerClient::connect_websocket`, including the | |
| /// initialize handshake, while each test controls the exact JSON-RPC traffic | |
| /// that follows. | |
| struct ScriptedExecServer { | |
| websocket_url: String, | |
| task: JoinHandle<Result<()>>, | |
| } | |
| impl ScriptedExecServer { | |
| /// Connects the public exec-server client to this fake WebSocket endpoint. | |
| 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") | |
| } | |
| /// Waits for the scripted fake server to finish. | |
| async fn finish(self) -> Result<()> { | |
| self.task | |
| .await | |
| .context("fake exec-server task should join")??; | |
| Ok(()) | |
| } | |
| } | |
| /// Starts a fake exec-server that accepts one WebSocket client. | |
| 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, | |
| }) | |
| } | |
| /// JSON-RPC peer for the fake exec-server WebSocket. | |
| struct JsonRpcPeer { | |
| websocket: WebSocketStream<TcpStream>, | |
| } | |
| impl JsonRpcPeer { | |
| /// Completes and validates the client initialize handshake. | |
| 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(()) | |
| } | |
| /// Reads one typed `http/request` call from the client. | |
| 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)) | |
| } | |
| /// Reads a JSON-RPC request and validates its method. | |
| 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) | |
| } | |
| /// Reads a JSON-RPC notification and validates its method. | |
| 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) | |
| } | |
| /// Sends a successful JSON-RPC response. | |
| 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 | |
| } | |
| /// Sends one streamed HTTP body notification. | |
| 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 | |
| } | |
| /// Reads one WebSocket JSON-RPC message. | |
| 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:?}"), | |
| } | |
| } | |
| /// Writes one WebSocket JSON-RPC message. | |
| 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") | |
| } | |
| } | |
| /// Decodes a request params object into its typed protocol payload. | |
| 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") | |
| } | |