From ee0247f95a6fe2b094ba2253d82cae2a2b4c2dff Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Sat, 1 Aug 2026 05:57:45 +0000 Subject: [PATCH] Extract exec-server request dispatching (#36440) ## What changed - Move JSON-RPC request, notification, response, error, and malformed-message handling into a dedicated `RequestDispatcher`. - Keep the connection loop responsible for receiving events and closing the connection when dispatch reports a terminal condition. ## Testing - Add an integration test confirming that ordinary requests are processed serially by default, including when a blocking `process/read` queues later requests. GitOrigin-RevId: 29d1358d4524edd492ff3855b29f23c42c8b3390 --- codex-rs/exec-server/src/server.rs | 1 + codex-rs/exec-server/src/server/processor.rs | 232 ++---------------- .../src/server/request_dispatcher.rs | 219 +++++++++++++++++ .../src/server/request_dispatcher_tests.rs | 64 +++++ codex-rs/exec-server/tests/process.rs | 105 ++++++++ 5 files changed, 413 insertions(+), 208 deletions(-) create mode 100644 codex-rs/exec-server/src/server/request_dispatcher.rs create mode 100644 codex-rs/exec-server/src/server/request_dispatcher_tests.rs diff --git a/codex-rs/exec-server/src/server.rs b/codex-rs/exec-server/src/server.rs index ba58b36dc9..3db6b58a4a 100644 --- a/codex-rs/exec-server/src/server.rs +++ b/codex-rs/exec-server/src/server.rs @@ -3,6 +3,7 @@ mod handler; mod process_handler; mod processor; mod registry; +mod request_dispatcher; mod session_registry; mod transport; diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index 016147e756..72ec510bf3 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -1,8 +1,7 @@ use std::sync::Arc; -use std::time::Instant; +use codex_exec_server_protocol::JSONRPCMessage; use tokio::sync::mpsc; -use tracing::Instrument; use tracing::debug; use tracing::warn; @@ -10,14 +9,13 @@ use crate::ExecServerRuntimePaths; use crate::connection::CHANNEL_CAPACITY; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; -use crate::rpc::RpcCallError; use crate::rpc::RpcNotificationSender; use crate::rpc::RpcServerOutboundMessage; use crate::rpc::encode_server_message; -use crate::rpc::invalid_request; -use crate::rpc::method_not_found; use crate::server::ExecServerHandler; use crate::server::registry::build_router; +use crate::server::request_dispatcher::RequestDispatcher; +use crate::server::request_dispatcher::RequestTaskResult; use crate::server::session_registry::SessionRegistry; use crate::telemetry::ConnectionTransport; use crate::telemetry::ExecServerTelemetry; @@ -90,7 +88,7 @@ async fn run_connection( let JsonRpcConnection { outgoing_tx: json_outgoing_tx, mut incoming_rx, - mut disconnected_rx, + disconnected_rx, task_handles: connection_tasks, transport: _transport, } = connection; @@ -120,127 +118,32 @@ async fn run_connection( } }); + let mut dispatcher = RequestDispatcher::new( + router, + Arc::clone(&handler), + outgoing_tx.clone(), + disconnected_rx, + requests.clone(), + telemetry, + ); + // 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; } - match event { + let result = match event { JsonRpcConnectionEvent::MalformedMessage { reason } => { - warn!("ignoring malformed exec-server message: {reason}"); - if outgoing_tx - .send(RpcServerOutboundMessage::Error { - request_id: codex_exec_server_protocol::RequestId::Integer(-1), - error: invalid_request(reason), - }) - .await - .is_err() - { - break; - } + dispatcher.handle_malformed_message(reason).await } JsonRpcConnectionEvent::Message(message) => match message { - codex_exec_server_protocol::JSONRPCMessage::Request(request) => { - let request_started_at = Instant::now(); - if let Some((method, route)) = router.request_route(request.method.as_str()) { - let request_span = request_span(method, &request); - let message = tokio::select! { - message = route(Arc::clone(&handler), request).instrument(request_span.clone()) => message, - _ = disconnected_rx.changed() => { - request_span.record("result", "disconnected"); - telemetry.request_completed( - method, - "disconnected", - request_started_at.elapsed(), - ); - 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"); - telemetry.request_completed( - method, - "disconnected", - request_started_at.elapsed(), - ); - break; - } - request_span.record("result", result); - telemetry.request_completed(method, result, request_started_at.elapsed()); - drop(request_span); - } else { - let method = "unknown"; - let request_span = request_span(method, &request); - if outgoing_tx - .send(RpcServerOutboundMessage::Error { - request_id: request.id, - error: method_not_found(format!( - "exec-server stub does not implement `{}` yet", - request.method - )), - }) - .await - .is_err() - { - request_span.record("result", "disconnected"); - telemetry.request_completed( - method, - "disconnected", - request_started_at.elapsed(), - ); - break; - } - request_span.record("result", "error"); - telemetry.request_completed(method, "error", request_started_at.elapsed()); - } - } - codex_exec_server_protocol::JSONRPCMessage::Notification(notification) => { - let Some(route) = router.notification_route(notification.method.as_str()) - else { - warn!( - "closing exec-server connection after unexpected notification: {}", - notification.method - ); - break; - }; - let result = tokio::select! { - result = route(Arc::clone(&handler), notification) => result, - _ = disconnected_rx.changed() => { - debug!( - "exec-server transport disconnected while handling notification" - ); - break; - } - }; - if let Err(err) = result { - warn!("closing exec-server connection after protocol error: {err}"); - break; - } - } - codex_exec_server_protocol::JSONRPCMessage::Response(response) => { - if !requests.complete(response.id.clone(), Ok(response.result)) { - warn!( - "closing exec-server connection after unexpected client response: {:?}", - response.id - ); - break; - } - } - codex_exec_server_protocol::JSONRPCMessage::Error(error) => { - if !requests.complete(error.id.clone(), Err(RpcCallError::Server(error.error))) - { - warn!( - "closing exec-server connection after unexpected client error: {:?}", - error.id - ); - break; - } + JSONRPCMessage::Request(request) => dispatcher.dispatch_request(request).await, + JSONRPCMessage::Notification(notification) => { + dispatcher.handle_notification(notification).await } + JSONRPCMessage::Response(response) => dispatcher.handle_response(response), + JSONRPCMessage::Error(error) => dispatcher.handle_error(error), }, JsonRpcConnectionEvent::Disconnected { reason } => { if let Some(reason) = reason { @@ -248,10 +151,14 @@ async fn run_connection( } break; } + }; + if result == RequestTaskResult::ConnectionClosed { + break; } } requests.close(); + drop(dispatcher); handler.shutdown().await; drop(handler); drop(requests); @@ -263,38 +170,6 @@ async fn run_connection( let _ = outbound_task.await; } -fn request_span( - span_name: &str, - request: &codex_exec_server_protocol::JSONRPCRequest, -) -> tracing::Span { - let method = request.method.as_str(); - let span = tracing::info_span!( - "codex.exec_server.request", - otel.kind = "server", - otel.name = span_name, - method, - result = tracing::field::Empty, - ); - if let Some(trace) = &request.trace - && !codex_otel::set_parent_from_w3c_trace_context(&span, trace) - { - warn!(method, "ignoring invalid inbound exec-server trace carrier"); - } - span -} - -fn request_result(message: &Option) -> &'static str { - match message { - Some(RpcServerOutboundMessage::Error { .. }) => "error", - Some( - RpcServerOutboundMessage::Request(_) - | RpcServerOutboundMessage::Response { .. } - | RpcServerOutboundMessage::Notification(_), - ) - | None => "success", - } -} - #[cfg(test)] mod tests { use std::collections::HashMap; @@ -307,11 +182,6 @@ mod tests { use codex_exec_server_protocol::JSONRPCResponse; use codex_exec_server_protocol::RequestId; use codex_utils_path_uri::PathUri; - use opentelemetry::trace::SpanId; - use opentelemetry::trace::TraceId; - use opentelemetry::trace::TracerProvider as _; - use opentelemetry_sdk::trace::InMemorySpanExporter; - use opentelemetry_sdk::trace::SdkTracerProvider; use pretty_assertions::assert_eq; use serde::Serialize; use serde::de::DeserializeOwned; @@ -323,10 +193,7 @@ mod tests { use tokio::io::duplex; use tokio::task::JoinHandle; use tokio::time::timeout; - use tracing_subscriber::filter::filter_fn; - use tracing_subscriber::prelude::*; - use super::request_span; use super::run_connection; use crate::ExecServerRuntimePaths; use crate::ProcessId; @@ -350,57 +217,6 @@ mod tests { use crate::protocol::TerminateResponse; use crate::server::session_registry::SessionRegistry; - #[test] - fn request_span_uses_bounded_name_wire_method_and_inbound_trace_parent() { - let span_exporter = InMemorySpanExporter::default(); - let tracer_provider = SdkTracerProvider::builder() - .with_simple_exporter(span_exporter.clone()) - .build(); - let tracer = tracer_provider.tracer("exec-server-test"); - let subscriber = tracing_subscriber::registry().with( - tracing_opentelemetry::layer() - .with_tracer(tracer) - .with_filter(filter_fn(codex_otel::OtelProvider::trace_export_filter)), - ); - let trace_id = TraceId::from_hex("00000000000000000000000000000001").expect("trace id"); - let parent_span_id = SpanId::from_hex("0000000000000002").expect("span id"); - let trace = codex_protocol::protocol::W3cTraceContext { - traceparent: Some(format!("00-{trace_id}-{parent_span_id}-01")), - tracestate: None, - }; - - let method = "custom/method"; - tracing::subscriber::with_default(subscriber, || { - tracing::callsite::rebuild_interest_cache(); - let request = JSONRPCRequest { - id: RequestId::Integer(1), - method: method.to_string(), - params: None, - trace: Some(trace), - }; - let request_span = request_span("unknown", &request); - request_span.in_scope(|| {}); - drop(request_span); - }); - - tracer_provider.force_flush().expect("flush traces"); - let spans = span_exporter.get_finished_spans().expect("span export"); - let request_span = spans - .iter() - .find(|span| span.name.as_ref() == "unknown") - .expect("unknown method span"); - assert_eq!( - request_span - .attributes - .iter() - .find(|attribute| attribute.key.as_str() == "method") - .map(|attribute| attribute.value.clone()), - Some(opentelemetry::Value::String(method.into())) - ); - assert_eq!(request_span.span_context.trace_id(), trace_id); - assert_eq!(request_span.parent_span_id, parent_span_id); - } - #[tokio::test] async fn connection_accepts_pipelined_scalar_requests() { let registry = SessionRegistry::new(crate::ExecServerTelemetry::default()); diff --git a/codex-rs/exec-server/src/server/request_dispatcher.rs b/codex-rs/exec-server/src/server/request_dispatcher.rs new file mode 100644 index 0000000000..1d4bcbae37 --- /dev/null +++ b/codex-rs/exec-server/src/server/request_dispatcher.rs @@ -0,0 +1,219 @@ +use std::sync::Arc; +use std::time::Instant; + +use codex_exec_server_protocol::JSONRPCError; +use codex_exec_server_protocol::JSONRPCNotification; +use codex_exec_server_protocol::JSONRPCRequest; +use codex_exec_server_protocol::JSONRPCResponse; +use codex_exec_server_protocol::RequestId; +use tokio::sync::mpsc; +use tokio::sync::watch; +use tracing::Instrument; +use tracing::debug; +use tracing::warn; + +use crate::rpc::RpcCallError; +use crate::rpc::RpcRouter; +use crate::rpc::RpcServerOutboundMessage; +use crate::rpc::invalid_request; +use crate::rpc::method_not_found; +use crate::rpc_server_requests::RpcServerRequestSender; +use crate::server::ExecServerHandler; +use crate::telemetry::ExecServerTelemetry; + +pub(super) struct RequestDispatcher { + router: Arc>, + handler: Arc, + outgoing_tx: mpsc::Sender, + disconnected_rx: watch::Receiver, + requests: RpcServerRequestSender, + telemetry: ExecServerTelemetry, +} + +impl RequestDispatcher { + pub(super) fn new( + router: Arc>, + handler: Arc, + outgoing_tx: mpsc::Sender, + disconnected_rx: watch::Receiver, + requests: RpcServerRequestSender, + telemetry: ExecServerTelemetry, + ) -> Self { + Self { + router, + handler, + outgoing_tx, + disconnected_rx, + requests, + telemetry, + } + } + + pub(super) async fn handle_malformed_message(&self, reason: String) -> RequestTaskResult { + warn!("ignoring malformed exec-server message: {reason}"); + if self + .outgoing_tx + .send(RpcServerOutboundMessage::Error { + request_id: RequestId::Integer(-1), + error: invalid_request(reason), + }) + .await + .is_err() + { + return RequestTaskResult::ConnectionClosed; + } + + RequestTaskResult::Completed + } + + pub(super) async fn handle_notification( + &mut self, + notification: JSONRPCNotification, + ) -> RequestTaskResult { + let Some(route) = self.router.notification_route(notification.method.as_str()) else { + warn!( + "closing exec-server connection after unexpected notification: {}", + notification.method + ); + return RequestTaskResult::ConnectionClosed; + }; + let result = tokio::select! { + result = route(Arc::clone(&self.handler), notification) => result, + _ = self.disconnected_rx.changed() => { + debug!("exec-server transport disconnected while handling notification"); + return RequestTaskResult::ConnectionClosed; + } + }; + if let Err(error) = result { + warn!("closing exec-server connection after protocol error: {error}"); + return RequestTaskResult::ConnectionClosed; + } + + RequestTaskResult::Completed + } + + pub(super) fn handle_response(&self, response: JSONRPCResponse) -> RequestTaskResult { + if !self + .requests + .complete(response.id.clone(), Ok(response.result)) + { + warn!( + "closing exec-server connection after unexpected client response: {:?}", + response.id + ); + return RequestTaskResult::ConnectionClosed; + } + + RequestTaskResult::Completed + } + + pub(super) fn handle_error(&self, error: JSONRPCError) -> RequestTaskResult { + if !self + .requests + .complete(error.id.clone(), Err(RpcCallError::Server(error.error))) + { + warn!( + "closing exec-server connection after unexpected client error: {:?}", + error.id + ); + return RequestTaskResult::ConnectionClosed; + } + + RequestTaskResult::Completed + } + + pub(super) async fn dispatch_request(&mut self, request: JSONRPCRequest) -> RequestTaskResult { + let started_at = Instant::now(); + let Some((method, route)) = self.router.request_route(request.method.as_str()) else { + let method = "unknown"; + let span = request_span(method, &request); + if self + .outgoing_tx + .send(RpcServerOutboundMessage::Error { + request_id: request.id, + error: method_not_found(format!( + "exec-server stub does not implement `{}` yet", + request.method + )), + }) + .await + .is_err() + { + span.record("result", "disconnected"); + self.telemetry + .request_completed(method, "disconnected", started_at.elapsed()); + return RequestTaskResult::ConnectionClosed; + } + span.record("result", "error"); + self.telemetry + .request_completed(method, "error", started_at.elapsed()); + return RequestTaskResult::Completed; + }; + + let span = request_span(method, &request); + let message = tokio::select! { + message = route(Arc::clone(&self.handler), request).instrument(span.clone()) => message, + _ = self.disconnected_rx.changed() => { + span.record("result", "disconnected"); + self.telemetry + .request_completed(method, "disconnected", started_at.elapsed()); + debug!("exec-server transport disconnected while handling request"); + return RequestTaskResult::ConnectionClosed; + } + }; + let result = request_result(&message); + if let Some(message) = message + && self.outgoing_tx.send(message).await.is_err() + { + span.record("result", "disconnected"); + self.telemetry + .request_completed(method, "disconnected", started_at.elapsed()); + return RequestTaskResult::ConnectionClosed; + } + span.record("result", result); + self.telemetry + .request_completed(method, result, started_at.elapsed()); + drop(span); + + RequestTaskResult::Completed + } +} + +#[derive(Eq, PartialEq)] +pub(super) enum RequestTaskResult { + Completed, + ConnectionClosed, +} + +fn request_span(span_name: &str, request: &JSONRPCRequest) -> tracing::Span { + let method = request.method.as_str(); + let span = tracing::info_span!( + "codex.exec_server.request", + otel.kind = "server", + otel.name = span_name, + method, + result = tracing::field::Empty, + ); + if let Some(trace) = &request.trace + && !codex_otel::set_parent_from_w3c_trace_context(&span, trace) + { + warn!(method, "ignoring invalid inbound exec-server trace carrier"); + } + span +} + +fn request_result(message: &Option) -> &'static str { + match message { + Some(RpcServerOutboundMessage::Error { .. }) => "error", + Some( + RpcServerOutboundMessage::Request(_) + | RpcServerOutboundMessage::Response { .. } + | RpcServerOutboundMessage::Notification(_), + ) + | None => "success", + } +} + +#[cfg(test)] +#[path = "request_dispatcher_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/server/request_dispatcher_tests.rs b/codex-rs/exec-server/src/server/request_dispatcher_tests.rs new file mode 100644 index 0000000000..619c6eda9f --- /dev/null +++ b/codex-rs/exec-server/src/server/request_dispatcher_tests.rs @@ -0,0 +1,64 @@ +use codex_exec_server_protocol::JSONRPCRequest; +use codex_exec_server_protocol::RequestId; +use opentelemetry::trace::SpanId; +use opentelemetry::trace::TraceId; +use opentelemetry::trace::TracerProvider as _; +use opentelemetry_sdk::trace::InMemorySpanExporter; +use opentelemetry_sdk::trace::SdkTracerProvider; +use pretty_assertions::assert_eq; +use tracing_subscriber::filter::filter_fn; +use tracing_subscriber::prelude::*; + +use super::request_span; + +/// Request spans retain the wire method and inbound trace while bounding their exported names. +#[test] +fn request_span_uses_bounded_name_wire_method_and_inbound_trace_parent() { + let span_exporter = InMemorySpanExporter::default(); + let tracer_provider = SdkTracerProvider::builder() + .with_simple_exporter(span_exporter.clone()) + .build(); + let tracer = tracer_provider.tracer("exec-server-test"); + let subscriber = tracing_subscriber::registry().with( + tracing_opentelemetry::layer() + .with_tracer(tracer) + .with_filter(filter_fn(codex_otel::OtelProvider::trace_export_filter)), + ); + let trace_id = TraceId::from_hex("00000000000000000000000000000001").expect("trace id"); + let parent_span_id = SpanId::from_hex("0000000000000002").expect("span id"); + let trace = codex_protocol::protocol::W3cTraceContext { + traceparent: Some(format!("00-{trace_id}-{parent_span_id}-01")), + tracestate: None, + }; + + let method = "custom/method"; + tracing::subscriber::with_default(subscriber, || { + tracing::callsite::rebuild_interest_cache(); + let request = JSONRPCRequest { + id: RequestId::Integer(1), + method: method.to_string(), + params: None, + trace: Some(trace), + }; + let request_span = request_span("unknown", &request); + request_span.in_scope(|| {}); + drop(request_span); + }); + + tracer_provider.force_flush().expect("flush traces"); + let spans = span_exporter.get_finished_spans().expect("span export"); + let request_span = spans + .iter() + .find(|span| span.name.as_ref() == "unknown") + .expect("unknown method span"); + assert_eq!( + request_span + .attributes + .iter() + .find(|attribute| attribute.key.as_str() == "method") + .map(|attribute| attribute.value.clone()), + Some(opentelemetry::Value::String(method.into())) + ); + assert_eq!(request_span.span_context.trace_id(), trace_id); + assert_eq!(request_span.parent_span_id, parent_span_id); +} diff --git a/codex-rs/exec-server/tests/process.rs b/codex-rs/exec-server/tests/process.rs index 68409c3018..273c022eaa 100644 --- a/codex-rs/exec-server/tests/process.rs +++ b/codex-rs/exec-server/tests/process.rs @@ -2,6 +2,7 @@ mod common; +use codex_exec_server::EnvironmentInfo; use codex_exec_server::ExecResponse; use codex_exec_server::InitializeParams; use codex_exec_server::InitializeResponse; @@ -81,6 +82,110 @@ async fn exec_server_starts_process_over_websocket() -> anyhow::Result<()> { Ok(()) } +/// Ordinary requests run one at a time when concurrent processing is not enabled. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_server_runs_ordinary_requests_serially_by_default() -> anyhow::Result<()> { + let mut server = exec_server().await?; + let initialize_id = server + .send_request( + "initialize", + serde_json::to_value(InitializeParams { + client_name: "exec-server-test".to_string(), + resume_session_id: None, + })?, + ) + .await?; + let _ = server + .wait_for_event(|event| { + matches!( + event, + JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &initialize_id + ) + }) + .await?; + server + .send_notification("initialized", serde_json::json!({})) + .await?; + + let process_start_id = server + .send_request( + "process/start", + serde_json::json!({ + "processId": "proc-serial-read", + "argv": ["/bin/sh", "-c", "parent=$PPID; while kill -0 \"$parent\" 2>/dev/null; do sleep 1; done"], + "cwd": PathUri::from_host_native_path(std::env::current_dir()?)?, + "env": {}, + "tty": false, + "pipeStdin": false, + "arg0": null + }), + ) + .await?; + let _ = server + .wait_for_event(|event| { + matches!( + event, + JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &process_start_id + ) + }) + .await?; + + let read_id = server + .send_request( + "process/read", + serde_json::json!({ + "processId": "proc-serial-read", + "afterSeq": null, + "maxBytes": null, + "waitMs": 250 + }), + ) + .await?; + let queued_environment_info_id = server + .send_request("environment/info", serde_json::json!({})) + .await?; + let queued_start_id = server + .send_request( + "process/start", + serde_json::json!({ + "processId": "proc-serial-queued", + "argv": ["/bin/sh", "-c", "parent=$PPID; while kill -0 \"$parent\" 2>/dev/null; do sleep 1; done"], + "cwd": PathUri::from_host_native_path(std::env::current_dir()?)?, + "env": {}, + "tty": false, + "pipeStdin": false, + "arg0": null + }), + ) + .await?; + + let JSONRPCMessage::Response(JSONRPCResponse { id, .. }) = server.next_event().await? else { + panic!("expected the blocked process/read to finish before the queued process/start"); + }; + assert_eq!(id, read_id); + let JSONRPCMessage::Response(JSONRPCResponse { id, result }) = server.next_event().await? + else { + panic!("expected the queued environment/info response after process/read"); + }; + assert_eq!(id, queued_environment_info_id); + let _: EnvironmentInfo = serde_json::from_value(result)?; + let JSONRPCMessage::Response(JSONRPCResponse { id, result }) = server.next_event().await? + else { + panic!("expected the queued process/start response after process/read"); + }; + assert_eq!(id, queued_start_id); + assert_eq!( + serde_json::from_value::(result)?, + ExecResponse { + process_id: ProcessId::from("proc-serial-queued"), + sandbox_type: Some(ProcessSandboxType::None), + } + ); + + server.shutdown().await?; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn exec_server_defaults_omitted_pipe_stdin_to_closed_stdin() -> anyhow::Result<()> { let mut server = exec_server().await?;