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
This commit is contained in:
Adam Perry @ OpenAI
2026-08-01 05:57:45 +00:00
committed by copyberry
parent 6751b54cae
commit ee0247f95a
5 changed files with 413 additions and 208 deletions

View File

@@ -3,6 +3,7 @@ mod handler;
mod process_handler;
mod processor;
mod registry;
mod request_dispatcher;
mod session_registry;
mod transport;

View File

@@ -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<RpcServerOutboundMessage>) -> &'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());

View File

@@ -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<RpcRouter<ExecServerHandler>>,
handler: Arc<ExecServerHandler>,
outgoing_tx: mpsc::Sender<RpcServerOutboundMessage>,
disconnected_rx: watch::Receiver<bool>,
requests: RpcServerRequestSender,
telemetry: ExecServerTelemetry,
}
impl RequestDispatcher {
pub(super) fn new(
router: Arc<RpcRouter<ExecServerHandler>>,
handler: Arc<ExecServerHandler>,
outgoing_tx: mpsc::Sender<RpcServerOutboundMessage>,
disconnected_rx: watch::Receiver<bool>,
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<RpcServerOutboundMessage>) -> &'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;

View File

@@ -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);
}

View File

@@ -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::<ExecResponse>(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?;