From bb13cf171d1401b02b0df2d53a5ce03d2760aa86 Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Wed, 24 Jun 2026 22:32:31 +0000 Subject: [PATCH] exec-server: handle post-init requests concurrently --- codex-rs/exec-server/README.md | 4 + codex-rs/exec-server/src/local_process.rs | 29 + codex-rs/exec-server/src/server/handler.rs | 30 +- .../exec-server/src/server/process_handler.rs | 4 + codex-rs/exec-server/src/server/processor.rs | 693 +++++++++++++++--- 5 files changed, 670 insertions(+), 90 deletions(-) diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index b1f9279186..0d7cd56e1c 100644 --- a/codex-rs/exec-server/README.md +++ b/codex-rs/exec-server/README.md @@ -115,6 +115,10 @@ Each connection follows this sequence: 3. Send `initialized`. 4. Call process or filesystem RPCs. +After initialization, a client may issue requests concurrently. When it does, responses may arrive +out of order, so it must match them by request id and await an earlier response before issuing any +operation that depends on it. + If the server receives any notification other than `initialized`, it replies with an error using request id `-1`. diff --git a/codex-rs/exec-server/src/local_process.rs b/codex-rs/exec-server/src/local_process.rs index 973019808a..0091b5ed1d 100644 --- a/codex-rs/exec-server/src/local_process.rs +++ b/codex-rs/exec-server/src/local_process.rs @@ -205,6 +205,14 @@ impl LocalProcess { } } + pub(crate) async fn clear_incomplete_starts(&self) { + self.inner + .processes + .lock() + .await + .retain(|_, process| !matches!(process, ProcessEntry::Starting(_))); + } + pub(crate) fn set_notification_sender(&self, notifications: Option) { let mut notification_sender = self .inner @@ -993,6 +1001,27 @@ mod tests { assert_eq!(error, expected); } + #[tokio::test] + async fn clear_incomplete_starts_removes_starting_processes() { + let backend = LocalProcess::default(); + let process_id = ProcessId::from("incomplete-start"); + backend.inner.processes.lock().await.insert( + process_id.clone(), + ProcessEntry::Starting(Arc::new(ProcessStart)), + ); + + backend.clear_incomplete_starts().await; + + assert!( + !backend + .inner + .processes + .lock() + .await + .contains_key(&process_id) + ); + } + #[test] fn child_env_defaults_to_exact_env() { let params = test_exec_params(HashMap::from([("ONLY_THIS".to_string(), "1".to_string())])); diff --git a/codex-rs/exec-server/src/server/handler.rs b/codex-rs/exec-server/src/server/handler.rs index 561cf0ff63..7c96e7c11a 100644 --- a/codex-rs/exec-server/src/server/handler.rs +++ b/codex-rs/exec-server/src/server/handler.rs @@ -42,6 +42,7 @@ use crate::protocol::FsWalkResponse; use crate::protocol::FsWriteFileParams; use crate::protocol::FsWriteFileResponse; use crate::protocol::HttpRequestParams; +use crate::protocol::INITIALIZE_METHOD; use crate::protocol::InitializeParams; use crate::protocol::InitializeResponse; use crate::protocol::ReadParams; @@ -98,7 +99,12 @@ impl ExecServerHandler { self.background_tasks.close(); self.background_tasks.wait().await; self.file_system.shutdown().await; - if let Some(session) = self.session() { + if let Some(session) = self.session() + && session.is_session_attached() + { + // The connection processor cancels and joins request tasks before shutdown, so any + // remaining `Starting` entries are stale reservations left by cancelled start futures. + session.process().clear_incomplete_starts().await; session.detach().await; } } @@ -108,6 +114,28 @@ impl ExecServerHandler { .is_none_or(|session| session.is_session_attached()) } + /// Returns an ordering error for a request admitted before initialization is complete. + /// + /// Call this from the connection's receive loop before spawning the request handler. Latching + /// the error there prevents a later `initialized` notification from making an earlier request + /// appear valid when its task eventually starts polling. The `initialize` request is exempt. + pub(crate) fn request_initialization_error(&self, method: &str) -> Option { + if method == INITIALIZE_METHOD { + return None; + } + if !self.initialize_requested.load(Ordering::SeqCst) { + return Some(invalid_request(format!( + "client must call initialize before invoking `{method}`" + ))); + } + if !self.initialized.load(Ordering::SeqCst) { + return Some(invalid_request(format!( + "client must send initialized before invoking `{method}`" + ))); + } + None + } + pub(crate) async fn initialize( &self, params: InitializeParams, diff --git a/codex-rs/exec-server/src/server/process_handler.rs b/codex-rs/exec-server/src/server/process_handler.rs index 9ff6690411..a5c9ca21ec 100644 --- a/codex-rs/exec-server/src/server/process_handler.rs +++ b/codex-rs/exec-server/src/server/process_handler.rs @@ -33,6 +33,10 @@ impl ProcessHandler { self.process.shutdown().await; } + pub(crate) async fn clear_incomplete_starts(&self) { + self.process.clear_incomplete_starts().await; + } + pub(crate) fn set_notification_sender(&self, notifications: Option) { self.process.set_notification_sender(notifications); } diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index 3eedfb2f20..7081ea4387 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -1,6 +1,9 @@ +use std::future::Future; use std::sync::Arc; use tokio::sync::mpsc; +use tokio::sync::watch; +use tokio::task::JoinSet; use tracing::Instrument; use tracing::debug; use tracing::warn; @@ -9,6 +12,7 @@ use crate::ExecServerRuntimePaths; use crate::connection::CHANNEL_CAPACITY; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; +use crate::protocol::INITIALIZE_METHOD; use crate::rpc::RpcNotificationSender; use crate::rpc::RpcServerOutboundMessage; use crate::rpc::encode_server_message; @@ -18,6 +22,16 @@ use crate::server::ExecServerHandler; use crate::server::registry::build_router; use crate::server::session_registry::SessionRegistry; +type RequestTaskResult = Result<(), mpsc::error::SendError>; +type RequestTaskJoinResult = Result; + +enum ConnectionActivity { + Incoming(Option), + RequestTask(RequestTaskJoinResult), + RequestTasksDrained, + Disconnected, +} + #[derive(Clone)] pub(crate) struct ConnectionProcessor { session_registry: Arc, @@ -79,14 +93,39 @@ async fn run_connection( } }); - // Process inbound events sequentially to preserve initialize/initialized ordering. - while let Some(event) = incoming_rx.recv().await { - if !handler.is_session_attached() { - debug!("exec-server connection evicted after session resume"); - break; - } + // Run requests independently so one slow request does not block the connection, up to the + // transport channel capacity per connection. + let mut request_tasks = JoinSet::::new(); + 'connection: loop { + let event = match wait_for_connection_activity( + &mut incoming_rx, + &mut disconnected_rx, + &mut request_tasks, + ) + .await + { + ConnectionActivity::Incoming(event) => event, + ConnectionActivity::RequestTask(Ok(Ok(()))) + | ConnectionActivity::RequestTasksDrained => continue 'connection, + ConnectionActivity::RequestTask(Ok(Err(_))) => { + debug!("closing exec-server connection after response channel closed"); + break 'connection; + } + ConnectionActivity::RequestTask(Err(err)) => { + warn!(error = %err, "exec-server request task failed"); + break 'connection; + } + ConnectionActivity::Disconnected => { + debug!("exec-server transport disconnected"); + break 'connection; + } + }; match event { - JsonRpcConnectionEvent::MalformedMessage { reason } => { + Some(_) if !handler.is_session_attached() => { + warn!("exec-server connection evicted after session resume"); + break 'connection; + } + Some(JsonRpcConnectionEvent::MalformedMessage { reason }) => { warn!("ignoring malformed exec-server message: {reason}"); if outgoing_tx .send(RpcServerOutboundMessage::Error { @@ -96,69 +135,102 @@ async fn run_connection( .await .is_err() { - break; + break 'connection; } } - JsonRpcConnectionEvent::Message(message) => match message { + Some(JsonRpcConnectionEvent::Message(message)) => match message { codex_exec_server_protocol::JSONRPCMessage::Request(request) => { - if let Some(route) = router.request_route(request.method.as_str()) { - let request_span = request_span(request.method.as_str(), &request); - let message = tokio::select! { - message = route(Arc::clone(&handler), request).instrument(request_span.clone()) => message, - _ = disconnected_rx.changed() => { - request_span.record("result", "disconnected"); - debug!("exec-server transport disconnected while handling request"); - break; - } - }; - let result = request_result(&message); - if let Some(message) = message - && outgoing_tx.send(message).await.is_err() - { - request_span.record("result", "disconnected"); - break; - } - request_span.record("result", result); + // Capture protocol-ordering violations before spawning. Otherwise a later + // `initialized` notification could make an early request appear valid before + // this task is first polled. + let is_initialize = request.method == INITIALIZE_METHOD; + let route = router.request_route(request.method.as_str()); + let initialization_error = route.and_then(|_| { + handler.request_initialization_error(request.method.as_str()) + }); + let span_name = if route.is_some() { + request.method.as_str() } else { - let request_span = request_span("unknown", &request); - if outgoing_tx - .send(RpcServerOutboundMessage::Error { - request_id: request.id, + "unknown" + }; + let request_span = request_span(span_name, &request); + let request_id = request.id.clone(); + let request_method = request.method.clone(); + let routed_response = if initialization_error.is_none() { + route.map(|route| route(Arc::clone(&handler), request)) + } else { + None + }; + + let outgoing_tx = outgoing_tx.clone(); + let task_span = request_span.clone(); + let request_task = async move { + let message = if let Some(error) = initialization_error { + Some(RpcServerOutboundMessage::Error { request_id, error }) + } else if let Some(response) = routed_response { + response.await + } else { + Some(RpcServerOutboundMessage::Error { + request_id, error: method_not_found(format!( - "exec-server stub does not implement `{}` yet", - request.method + "exec-server stub does not implement `{request_method}` yet" )), }) - .await - .is_err() - { - request_span.record("result", "disconnected"); - break; + }; + let result = request_result(&message); + if let Some(message) = message { + // The sole receiver belongs to the outbound encoder task. A send error + // means the connection cannot deliver any more responses. + outgoing_tx.send(message).await?; } - request_span.record("result", "error"); + request_span.record("result", result); + Ok(()) + } + .instrument(task_span); + + if is_initialize { + // `initialize` claims connection state inside its route. Await it before + // admitting another frame so duplicate handshakes are ordered by the wire, + // rather than by which spawned task the scheduler polls first. + let Some(result) = + await_or_disconnect(request_task, &mut disconnected_rx).await + else { + debug!("exec-server transport disconnected while handling initialize"); + break 'connection; + }; + if result.is_err() { + debug!("closing exec-server connection after response channel closed"); + break 'connection; + } + } else { + request_tasks.spawn(request_task); } } codex_exec_server_protocol::JSONRPCMessage::Notification(notification) => { + // Notifications stay inline because `initialized` is an ordering barrier: + // later requests must latch state after it runs, and an invalid notification + // must close the connection before another frame is admitted. It is currently + // the only notification route and performs no asynchronous work. let Some(route) = router.notification_route(notification.method.as_str()) else { warn!( "closing exec-server connection after unexpected notification: {}", notification.method ); - break; + break 'connection; }; - let result = tokio::select! { - result = route(Arc::clone(&handler), notification) => result, - _ = disconnected_rx.changed() => { - debug!( - "exec-server transport disconnected while handling notification" - ); - break; - } + let Some(result) = await_or_disconnect( + route(Arc::clone(&handler), notification), + &mut disconnected_rx, + ) + .await + else { + debug!("exec-server transport disconnected while handling notification"); + break 'connection; }; if let Err(err) = result { warn!("closing exec-server connection after protocol error: {err}"); - break; + break 'connection; } } codex_exec_server_protocol::JSONRPCMessage::Response(response) => { @@ -166,25 +238,33 @@ async fn run_connection( "closing exec-server connection after unexpected client response: {:?}", response.id ); - break; + break 'connection; } codex_exec_server_protocol::JSONRPCMessage::Error(error) => { warn!( "closing exec-server connection after unexpected client error: {:?}", error.id ); - break; + break 'connection; } }, - JsonRpcConnectionEvent::Disconnected { reason } => { + Some(JsonRpcConnectionEvent::Disconnected { reason }) => { if let Some(reason) = reason { debug!("exec-server connection disconnected: {reason}"); } - break; + break 'connection; + } + None => { + debug!("exec-server incoming event channel closed"); + break 'connection; } } } + // Abort and await requests before handler shutdown clears any incomplete process starts and + // detaches the session. Long polls therefore cannot delay resume, and no request can mutate the + // session after its `Starting` entries have been swept. + request_tasks.shutdown().await; handler.shutdown().await; drop(handler); drop(outgoing_tx); @@ -205,7 +285,8 @@ fn request_span( otel.kind = "server", otel.name = span_name, method, - result = tracing::field::Empty, + // An aborted request drops the span with this fallback. Completed requests overwrite it. + result = "disconnected", ); if let Some(trace) = &request.trace && !codex_otel::set_parent_from_w3c_trace_context(&span, trace) @@ -225,12 +306,57 @@ fn request_result(message: &Option) -> &'static str { } } +async fn wait_for_connection_activity( + incoming_rx: &mut mpsc::Receiver, + disconnected_rx: &mut watch::Receiver, + request_tasks: &mut JoinSet, +) -> ConnectionActivity { + let has_request_tasks = !request_tasks.is_empty(); + let can_receive_event = request_tasks.len() < CHANNEL_CAPACITY; + // All three futures are cancellation safe, so the branches that lose this race retain their + // next event or task completion for the following iteration. At the request limit, stop + // draining the bounded incoming channel until a task completes, propagating backpressure to + // the transport. Ready task completions cannot starve input because no new request tasks are + // added while this wait reaps them. + tokio::select! { + biased; + _ = disconnected_rx.changed() => ConnectionActivity::Disconnected, + result = request_tasks.join_next(), if has_request_tasks => { + match result { + Some(result) => ConnectionActivity::RequestTask(result), + None => ConnectionActivity::RequestTasksDrained, + } + } + // Keep incoming events last so a disconnect or terminal task failure stops the connection + // before it admits another request. + event = incoming_rx.recv(), if can_receive_event => ConnectionActivity::Incoming(event), + } +} + +async fn await_or_disconnect( + future: F, + disconnected_rx: &mut watch::Receiver, +) -> Option +where + F: Future, +{ + // `watch::Receiver::changed` is cancellation safe. The notification future is intentionally + // dropped on disconnect because its result can no longer be delivered and connection teardown + // starts immediately. + tokio::select! { + output = future => Some(output), + _ = disconnected_rx.changed() => None, + } +} + #[cfg(test)] mod tests { use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; + use codex_exec_server_protocol::JSONRPCError; + use codex_exec_server_protocol::JSONRPCErrorError; use codex_exec_server_protocol::JSONRPCMessage; use codex_exec_server_protocol::JSONRPCNotification; use codex_exec_server_protocol::JSONRPCRequest; @@ -251,16 +377,26 @@ mod tests { use tokio::io::DuplexStream; use tokio::io::Lines; use tokio::io::duplex; + use tokio::sync::mpsc; + use tokio::sync::oneshot; + use tokio::sync::watch; use tokio::task::JoinHandle; + use tokio::task::JoinSet; use tokio::time::timeout; use tracing_subscriber::filter::filter_fn; use tracing_subscriber::prelude::*; + use super::ConnectionActivity; + use super::RequestTaskResult; use super::request_span; use super::run_connection; + use super::wait_for_connection_activity; use crate::ExecServerRuntimePaths; use crate::ProcessId; + use crate::connection::CHANNEL_CAPACITY; use crate::connection::JsonRpcConnection; + use crate::connection::JsonRpcConnectionEvent; + use crate::connection::JsonRpcTransport; use crate::protocol::ENVIRONMENT_INFO_METHOD; use crate::protocol::EXEC_METHOD; use crate::protocol::EXEC_READ_METHOD; @@ -349,7 +485,158 @@ mod tests { send_request(&mut writer, /*id*/ 2, ENVIRONMENT_INFO_METHOD, &()).await; send_request(&mut writer, /*id*/ 3, ENVIRONMENT_INFO_METHOD, &()).await; - let _: EnvironmentInfo = read_response(&mut lines, /*expected_id*/ 2).await; + let (first_id, _first_response) = + read_response_with_id::(&mut lines).await; + let (second_id, _second_response) = + read_response_with_id::(&mut lines).await; + let mut response_ids = [first_id, second_id]; + response_ids.sort(); + assert_eq!(response_ids, [RequestId::Integer(2), RequestId::Integer(3)]); + + drop(writer); + drop(lines); + timeout(Duration::from_secs(1), task) + .await + .expect("processor should exit") + .expect("processor should join"); + } + + #[tokio::test] + async fn initialized_before_initialize_closes_connection() { + let registry = SessionRegistry::new(); + let (mut writer, mut lines, task) = + spawn_test_connection(registry, "initialized-before-initialize"); + + send_notification(&mut writer, INITIALIZED_METHOD, &()).await; + + timeout(Duration::from_secs(1), task) + .await + .expect("processor should reject initialized before initialize") + .expect("processor should join"); + assert_eq!(lines.next_line().await.expect("read connection EOF"), None); + drop(writer); + } + + #[tokio::test] + async fn duplicate_initialize_requests_use_wire_order() { + let registry = SessionRegistry::new(); + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY); + let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY); + let (_disconnected_tx, disconnected_rx) = watch::channel(false); + for (id, resume_session_id) in [(1, None), (2, Some("must-not-be-used"))] { + incoming_tx + .send(JsonRpcConnectionEvent::Message(JSONRPCMessage::Request( + JSONRPCRequest { + id: RequestId::Integer(id), + method: INITIALIZE_METHOD.to_string(), + params: Some( + serde_json::to_value(InitializeParams { + client_name: format!("client-{id}"), + resume_session_id: resume_session_id.map(str::to_string), + }) + .expect("serialize initialize params"), + ), + trace: None, + }, + ))) + .await + .expect("incoming channel should remain open"); + } + let connection = JsonRpcConnection { + outgoing_tx, + incoming_rx, + disconnected_rx, + task_handles: Vec::new(), + transport: JsonRpcTransport::Plain, + }; + let task = tokio::spawn(run_connection(connection, registry, test_runtime_paths())); + + let mut first_response = None; + let mut second_error = None; + for _ in 0..2 { + let message = timeout(Duration::from_secs(1), outgoing_rx.recv()) + .await + .expect("initialize request should complete") + .expect("outgoing channel should remain open"); + match message { + JSONRPCMessage::Response(JSONRPCResponse { + id: RequestId::Integer(1), + result, + }) => { + first_response = Some( + serde_json::from_value::(result) + .expect("decode initialize response"), + ); + } + JSONRPCMessage::Error( + error @ JSONRPCError { + id: RequestId::Integer(2), + .. + }, + ) => { + second_error = Some(error); + } + other => panic!("unexpected initialize result: {other:?}"), + } + } + + assert!(first_response.is_some()); + assert_eq!( + second_error, + Some(JSONRPCError { + id: RequestId::Integer(2), + error: JSONRPCErrorError { + code: -32600, + message: "initialize may only be sent once per connection".to_string(), + data: None, + }, + }) + ); + + drop(incoming_tx); + drop(outgoing_rx); + timeout(Duration::from_secs(1), task) + .await + .expect("processor should exit") + .expect("processor should join"); + } + + #[tokio::test] + async fn request_before_initialized_returns_error_without_closing_connection() { + let registry = SessionRegistry::new(); + let (mut writer, mut lines, task) = + spawn_test_connection(registry, "request-before-initialized"); + + send_request( + &mut writer, + /*id*/ 1, + INITIALIZE_METHOD, + &InitializeParams { + client_name: "exec-server-test".to_string(), + resume_session_id: None, + }, + ) + .await; + let _: InitializeResponse = read_response(&mut lines, /*expected_id*/ 1).await; + + send_request(&mut writer, /*id*/ 2, ENVIRONMENT_INFO_METHOD, &()).await; + // Pipeline `initialized` behind the invalid request. The request must still fail based on + // its position in the incoming stream, even if its spawned task runs after the notification. + send_notification(&mut writer, INITIALIZED_METHOD, &()).await; + assert_eq!( + read_error(&mut lines).await, + JSONRPCError { + id: RequestId::Integer(2), + error: JSONRPCErrorError { + code: -32600, + data: None, + message: "client must send initialized before invoking `environment/info`" + .to_string(), + }, + } + ); + + send_request(&mut writer, /*id*/ 3, ENVIRONMENT_INFO_METHOD, &()).await; let _: EnvironmentInfo = read_response(&mut lines, /*expected_id*/ 3).await; drop(writer); @@ -361,7 +648,178 @@ mod tests { } #[tokio::test] - async fn transport_disconnect_detaches_session_during_in_flight_read() { + async fn in_flight_read_does_not_block_independent_request() { + let registry = SessionRegistry::new(); + let (mut writer, mut lines, task) = + spawn_test_connection(Arc::clone(®istry), "concurrent-read"); + + send_request( + &mut writer, + /*id*/ 1, + INITIALIZE_METHOD, + &InitializeParams { + client_name: "exec-server-test".to_string(), + resume_session_id: None, + }, + ) + .await; + let _: InitializeResponse = read_response(&mut lines, /*expected_id*/ 1).await; + send_notification(&mut writer, INITIALIZED_METHOD, &()).await; + + let process_id = ProcessId::from("proc-concurrent-read"); + send_request( + &mut writer, + /*id*/ 2, + EXEC_METHOD, + &exec_params_with_argv(process_id.clone(), long_running_process_argv()), + ) + .await; + let _: ExecResponse = read_response(&mut lines, /*expected_id*/ 2).await; + + send_request( + &mut writer, + /*id*/ 3, + EXEC_READ_METHOD, + &ReadParams { + process_id: process_id.clone(), + after_seq: None, + max_bytes: None, + wait_ms: Some(600_000), + }, + ) + .await; + send_request(&mut writer, /*id*/ 4, ENVIRONMENT_INFO_METHOD, &()).await; + + // The ordered transport admits request 3 first, but its process stays alive without output + // and its read deadline is ten minutes away. Receiving response 4 next therefore proves + // that the read did not block request handling without relying on a short wall-clock + // assertion. + let _: EnvironmentInfo = read_response(&mut lines, /*expected_id*/ 4).await; + + timeout( + Duration::from_secs(2), + terminate_process(&mut writer, &mut lines, /*request_id*/ 5, process_id), + ) + .await + .expect("process should terminate"); + + drop(writer); + drop(lines); + timeout(Duration::from_secs(1), task) + .await + .expect("processor should exit") + .expect("processor should join"); + } + + #[tokio::test] + async fn response_sink_failure_closes_connection() { + let registry = SessionRegistry::new(); + let (mut writer, mut lines, task) = spawn_test_connection(registry, "closed-response-sink"); + + send_request( + &mut writer, + /*id*/ 1, + INITIALIZE_METHOD, + &InitializeParams { + client_name: "exec-server-test".to_string(), + resume_session_id: None, + }, + ) + .await; + let _: InitializeResponse = read_response(&mut lines, /*expected_id*/ 1).await; + send_notification(&mut writer, INITIALIZED_METHOD, &()).await; + + drop(lines); + send_request(&mut writer, /*id*/ 2, ENVIRONMENT_INFO_METHOD, &()).await; + timeout(Duration::from_secs(1), task) + .await + .expect("processor should exit after the response sink fails") + .expect("processor should join"); + drop(writer); + } + + #[tokio::test] + async fn request_task_failure_wakes_idle_connection_wait() { + let (_incoming_tx, mut incoming_rx) = mpsc::channel(1); + let (_disconnected_tx, mut disconnected_rx) = watch::channel(false); + let mut request_tasks = JoinSet::::new(); + request_tasks.spawn(async { panic!("intentional request task panic") }); + + let activity = timeout( + Duration::from_secs(1), + wait_for_connection_activity( + &mut incoming_rx, + &mut disconnected_rx, + &mut request_tasks, + ), + ) + .await + .expect("request task failure should wake the idle connection wait"); + let ConnectionActivity::RequestTask(Err(err)) = activity else { + panic!("expected failed request task activity"); + }; + assert!(err.is_panic()); + } + + #[tokio::test] + async fn incoming_events_wait_for_request_capacity() { + let (incoming_tx, mut incoming_rx) = mpsc::channel(1); + let (_disconnected_tx, mut disconnected_rx) = watch::channel(false); + let mut request_tasks = JoinSet::::new(); + incoming_tx + .send(JsonRpcConnectionEvent::MalformedMessage { + reason: "queued".to_string(), + }) + .await + .expect("incoming channel should remain open"); + + for _ in 1..CHANNEL_CAPACITY { + request_tasks.spawn(std::future::pending::()); + } + let (release_tx, release_rx) = oneshot::channel(); + request_tasks.spawn(async move { + release_rx.await.expect("request task should be released"); + Ok(()) + }); + assert_eq!(request_tasks.len(), CHANNEL_CAPACITY); + + let mut activity = Box::pin(wait_for_connection_activity( + &mut incoming_rx, + &mut disconnected_rx, + &mut request_tasks, + )); + tokio::select! { + biased; + _ = &mut activity => panic!("incoming event bypassed request backpressure"), + _ = tokio::task::yield_now() => {} + } + + release_tx + .send(()) + .expect("request task should remain live"); + let ConnectionActivity::RequestTask(Ok(Ok(()))) = activity.await else { + panic!("expected the released request task to complete"); + }; + + let activity = wait_for_connection_activity( + &mut incoming_rx, + &mut disconnected_rx, + &mut request_tasks, + ) + .await; + let ConnectionActivity::Incoming(Some(JsonRpcConnectionEvent::MalformedMessage { reason })) = + activity + else { + panic!("expected the queued incoming event after capacity opened"); + }; + assert_eq!(reason, "queued"); + + request_tasks.shutdown().await; + } + + #[tokio::test(start_paused = true)] + async fn transport_disconnect_with_in_flight_read_allows_session_resume() { + // Keep the test-only detached-session TTL from racing session resume on a loaded host. let registry = SessionRegistry::new(); let (mut first_writer, mut first_lines, first_task) = spawn_test_connection(Arc::clone(®istry), "first"); @@ -380,12 +838,12 @@ mod tests { read_response(&mut first_lines, /*expected_id*/ 1).await; send_notification(&mut first_writer, INITIALIZED_METHOD, &()).await; - let process_id = ProcessId::from("proc-long-poll"); + let process_id = ProcessId::from("proc-disconnect-read"); send_request( &mut first_writer, /*id*/ 2, EXEC_METHOD, - &exec_params(process_id.clone()), + &exec_params_with_argv(process_id.clone(), long_running_process_argv()), ) .await; let _: ExecResponse = read_response(&mut first_lines, /*expected_id*/ 2).await; @@ -398,12 +856,21 @@ mod tests { process_id: process_id.clone(), after_seq: None, max_bytes: None, - wait_ms: Some(5_000), + wait_ms: Some(600_000), }, ) .await; + // The malformed frame is processed inline after request 3. Its response is an ordered + // barrier proving the long read was admitted to the task set before disconnect. + send_malformed_message(&mut first_writer).await; + let _ = read_error(&mut first_lines).await; + drop(first_writer); - tokio::time::sleep(Duration::from_millis(25)).await; + drop(first_lines); + timeout(Duration::from_secs(1), first_task) + .await + .expect("first processor should exit") + .expect("first processor should join"); let (mut second_writer, mut second_lines, second_task) = spawn_test_connection(Arc::clone(®istry), "second"); @@ -417,30 +884,20 @@ mod tests { }, ) .await; - let second_initialize_response = timeout( - Duration::from_secs(1), - read_response::(&mut second_lines, /*expected_id*/ 1), - ) - .await - .expect("resume initialize should not wait for the old read to finish"); + let second_initialize_response: InitializeResponse = + read_response(&mut second_lines, /*expected_id*/ 1).await; assert_eq!( second_initialize_response.session_id, initialize_response.session_id ); - timeout(Duration::from_secs(1), first_task) - .await - .expect("first processor should exit") - .expect("first processor should join"); send_notification(&mut second_writer, INITIALIZED_METHOD, &()).await; - - send_request( + terminate_process( &mut second_writer, - /*id*/ 2, - EXEC_TERMINATE_METHOD, - &TerminateParams { process_id }, + &mut second_lines, + /*request_id*/ 2, + process_id, ) .await; - let _: TerminateResponse = read_response(&mut second_lines, /*expected_id*/ 2).await; drop(second_writer); drop(second_lines); @@ -499,6 +956,13 @@ mod tests { .await; } + async fn send_malformed_message(writer: &mut DuplexStream) { + writer + .write_all(b"not-json\n") + .await + .expect("write malformed message"); + } + async fn write_message(writer: &mut DuplexStream, message: &JSONRPCMessage) { let encoded = serde_json::to_vec(message).expect("serialize JSON-RPC message"); writer.write_all(&encoded).await.expect("write request"); @@ -509,29 +973,80 @@ mod tests { lines: &mut Lines>, expected_id: i64, ) -> T { - let line = lines - .next_line() - .await - .expect("read response") - .expect("response line"); - match serde_json::from_str::(&line).expect("decode JSON-RPC response") { + let (id, response) = read_response_with_id(lines).await; + assert_eq!(id, RequestId::Integer(expected_id)); + response + } + + async fn read_response_with_id( + lines: &mut Lines>, + ) -> (RequestId, T) { + match read_message(lines).await { JSONRPCMessage::Response(JSONRPCResponse { id, result }) => { - assert_eq!(id, RequestId::Integer(expected_id)); - serde_json::from_value(result).expect("decode response result") + let response = serde_json::from_value(result).expect("decode response result"); + (id, response) } JSONRPCMessage::Error(error) => panic!("unexpected JSON-RPC error: {error:?}"), other => panic!("expected JSON-RPC response, got {other:?}"), } } - fn exec_params(process_id: ProcessId) -> ExecParams { + async fn read_error(lines: &mut Lines>) -> JSONRPCError { + match read_message(lines).await { + JSONRPCMessage::Error(error) => error, + other => panic!("expected JSON-RPC error, got {other:?}"), + } + } + + async fn terminate_process( + writer: &mut DuplexStream, + lines: &mut Lines>, + request_id: i64, + process_id: ProcessId, + ) { + send_request( + writer, + request_id, + EXEC_TERMINATE_METHOD, + &TerminateParams { process_id }, + ) + .await; + + loop { + match read_message(lines).await { + JSONRPCMessage::Response(JSONRPCResponse { id, result }) + if id == RequestId::Integer(request_id) => + { + let _: TerminateResponse = + serde_json::from_value(result).expect("decode terminate response"); + return; + } + JSONRPCMessage::Response(_) | JSONRPCMessage::Notification(_) => {} + JSONRPCMessage::Error(error) => { + panic!("unexpected JSON-RPC error: {error:?}") + } + other => panic!("expected JSON-RPC response or notification, got {other:?}"), + } + } + } + + async fn read_message(lines: &mut Lines>) -> JSONRPCMessage { + let line = lines + .next_line() + .await + .expect("read response") + .expect("response line"); + serde_json::from_str(&line).expect("decode JSON-RPC message") + } + + fn exec_params_with_argv(process_id: ProcessId, argv: Vec) -> ExecParams { let mut env = HashMap::new(); if let Some(path) = std::env::var_os("PATH") { env.insert("PATH".to_string(), path.to_string_lossy().into_owned()); } ExecParams { process_id, - argv: sleep_then_print_argv(), + argv, cwd: PathUri::from_host_native_path(std::env::current_dir().expect("cwd")) .expect("cwd URI"), env_policy: None, @@ -545,18 +1060,18 @@ mod tests { } } - fn sleep_then_print_argv() -> Vec { + fn long_running_process_argv() -> Vec { if cfg!(windows) { vec![ std::env::var("COMSPEC").unwrap_or_else(|_| "cmd.exe".to_string()), "/C".to_string(), - "ping -n 3 127.0.0.1 >NUL && echo late".to_string(), + "ping -n 3601 127.0.0.1 >NUL".to_string(), ] } else { vec![ "/bin/sh".to_string(), "-c".to_string(), - "sleep 1; printf late".to_string(), + "sleep 3600".to_string(), ] } }