From 66f49ea604315574d0b8a92d75dc196d5470df0a Mon Sep 17 00:00:00 2001 From: starr-openai Date: Wed, 18 Mar 2026 14:34:21 -0700 Subject: [PATCH] Remove outer handler mutex from exec-server RPC base Co-authored-by: Codex --- .../exec-server/src/client/local_backend.rs | 12 +++------- codex-rs/exec-server/src/server/handler.rs | 22 ++++++++++--------- codex-rs/exec-server/src/server/processor.rs | 5 ++--- codex-rs/exec-server/src/server/registry.rs | 12 +++------- 4 files changed, 20 insertions(+), 31 deletions(-) diff --git a/codex-rs/exec-server/src/client/local_backend.rs b/codex-rs/exec-server/src/client/local_backend.rs index 195847bf14..8f9a2481f8 100644 --- a/codex-rs/exec-server/src/client/local_backend.rs +++ b/codex-rs/exec-server/src/client/local_backend.rs @@ -1,7 +1,5 @@ use std::sync::Arc; -use tokio::sync::Mutex; - use crate::protocol::InitializeResponse; use crate::server::ExecServerHandler; @@ -9,24 +7,22 @@ use super::ExecServerError; #[derive(Clone)] pub(super) struct LocalBackend { - handler: Arc>, + handler: Arc, } impl LocalBackend { pub(super) fn new(handler: ExecServerHandler) -> Self { Self { - handler: Arc::new(Mutex::new(handler)), + handler: Arc::new(handler), } } pub(super) async fn shutdown(&self) { - self.handler.lock().await.shutdown().await; + self.handler.shutdown().await; } pub(super) async fn initialize(&self) -> Result { self.handler - .lock() - .await .initialize() .map_err(|error| ExecServerError::Server { code: error.code, @@ -36,8 +32,6 @@ impl LocalBackend { pub(super) async fn initialized(&self) -> Result<(), ExecServerError> { self.handler - .lock() - .await .initialized() .map_err(ExecServerError::Protocol) } diff --git a/codex-rs/exec-server/src/server/handler.rs b/codex-rs/exec-server/src/server/handler.rs index 6a8119b274..f10fa8ddd9 100644 --- a/codex-rs/exec-server/src/server/handler.rs +++ b/codex-rs/exec-server/src/server/handler.rs @@ -1,3 +1,6 @@ +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + use codex_app_server_protocol::JSONRPCErrorError; use crate::protocol::InitializeResponse; @@ -6,38 +9,37 @@ use crate::rpc::RpcNotificationSender; pub(crate) struct ExecServerHandler { _notifications: RpcNotificationSender, - initialize_requested: bool, - initialized: bool, + initialize_requested: AtomicBool, + initialized: AtomicBool, } impl ExecServerHandler { pub(crate) fn new(notifications: RpcNotificationSender) -> Self { Self { _notifications: notifications, - initialize_requested: false, - initialized: false, + initialize_requested: AtomicBool::new(false), + initialized: AtomicBool::new(false), } } pub(crate) async fn shutdown(&self) {} - pub(crate) fn initialize(&mut self) -> Result { - if self.initialize_requested { + pub(crate) fn initialize(&self) -> Result { + if self.initialize_requested.swap(true, Ordering::SeqCst) { return Err(crate::rpc::invalid_request( "initialize may only be sent once per connection".to_string(), )); } - self.initialize_requested = true; Ok(InitializeResponse { protocol_version: PROTOCOL_VERSION.to_string(), }) } - pub(crate) fn initialized(&mut self) -> Result<(), String> { - if !self.initialize_requested { + pub(crate) fn initialized(&self) -> Result<(), String> { + if !self.initialize_requested.load(Ordering::SeqCst) { return Err("received `initialized` notification before `initialize`".into()); } - self.initialized = true; + self.initialized.store(true, Ordering::SeqCst); Ok(()) } } diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index ddb7978be4..888bbf7f30 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -1,6 +1,5 @@ use std::sync::Arc; -use tokio::sync::Mutex; use tokio::sync::mpsc; use tracing::debug; use tracing::warn; @@ -21,7 +20,7 @@ pub(crate) async fn run_connection(connection: JsonRpcConnection) { let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(CHANNEL_CAPACITY); let notifications = RpcNotificationSender::new(outgoing_tx.clone()); - let handler = Arc::new(Mutex::new(ExecServerHandler::new(notifications))); + let handler = Arc::new(ExecServerHandler::new(notifications)); let outbound_task = tokio::spawn(async move { while let Some(message) = outgoing_rx.recv().await { @@ -101,7 +100,7 @@ pub(crate) async fn run_connection(connection: JsonRpcConnection) { } } - handler.lock().await.shutdown().await; + handler.shutdown().await; drop(outgoing_tx); let _ = outbound_task.await; } diff --git a/codex-rs/exec-server/src/server/registry.rs b/codex-rs/exec-server/src/server/registry.rs index 34aad179c6..100fe4a81a 100644 --- a/codex-rs/exec-server/src/server/registry.rs +++ b/codex-rs/exec-server/src/server/registry.rs @@ -1,26 +1,20 @@ use std::sync::Arc; -use tokio::sync::Mutex; - use crate::protocol::INITIALIZE_METHOD; use crate::protocol::INITIALIZED_METHOD; use crate::protocol::InitializeParams; use crate::rpc::RpcRouter; use crate::server::ExecServerHandler; -pub(crate) fn build_router() -> RpcRouter> { +pub(crate) fn build_router() -> RpcRouter { let mut router = RpcRouter::new(); router.request( INITIALIZE_METHOD, - |handler: Arc>, _params: InitializeParams| async move { - handler.lock().await.initialize() - }, + |handler: Arc, _params: InitializeParams| async move { handler.initialize() }, ); router.notification( INITIALIZED_METHOD, - |handler: Arc>, (): ()| async move { - handler.lock().await.initialized() - }, + |handler: Arc, (): ()| async move { handler.initialized() }, ); router }