diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index ca41cb0037..92fe7b4662 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/codex-rs/app-server-client/src/lib.rs @@ -22,6 +22,8 @@ mod path; mod remote; +mod server_request; +pub use server_request::AppServerReplyHandle; use std::error::Error; use std::fmt; @@ -549,26 +551,9 @@ impl InProcessAppServerClient { request_id: RequestId, result: JsonRpcResult, ) -> IoResult<()> { - let (response_tx, response_rx) = oneshot::channel(); - self.command_tx - .send(ClientCommand::ResolveServerRequest { - request_id, - result, - response_tx, - }) + self.request_handle() + .resolve_server_request(request_id, result) .await - .map_err(|_| { - IoError::new( - ErrorKind::BrokenPipe, - "in-process app-server worker channel is closed", - ) - })?; - response_rx.await.map_err(|_| { - IoError::new( - ErrorKind::BrokenPipe, - "in-process app-server resolve channel is closed", - ) - })? } /// Rejects a pending server request with JSON-RPC error payload. @@ -577,26 +562,9 @@ impl InProcessAppServerClient { request_id: RequestId, error: JSONRPCErrorError, ) -> IoResult<()> { - let (response_tx, response_rx) = oneshot::channel(); - self.command_tx - .send(ClientCommand::RejectServerRequest { - request_id, - error, - response_tx, - }) + self.request_handle() + .reject_server_request(request_id, error) .await - .map_err(|_| { - IoError::new( - ErrorKind::BrokenPipe, - "in-process app-server worker channel is closed", - ) - })?; - response_rx.await.map_err(|_| { - IoError::new( - ErrorKind::BrokenPipe, - "in-process app-server reject channel is closed", - ) - })? } /// Returns the next in-process event, or `None` when worker exits. @@ -643,6 +611,61 @@ impl InProcessAppServerClient { } impl InProcessAppServerRequestHandle { + pub(crate) async fn resolve_server_request( + &self, + request_id: RequestId, + result: JsonRpcResult, + ) -> IoResult<()> { + let (response_tx, response_rx) = oneshot::channel(); + self.command_tx + .send(ClientCommand::ResolveServerRequest { + request_id, + result, + response_tx, + }) + .await + .map_err(|_| { + IoError::new( + ErrorKind::BrokenPipe, + "in-process app-server worker channel is closed", + ) + })?; + response_rx.await.map_err(|_| { + IoError::new( + ErrorKind::BrokenPipe, + "in-process app-server resolve channel is closed", + ) + })? + } + + /// Rejects a pending server request with JSON-RPC error payload. + pub(crate) async fn reject_server_request( + &self, + request_id: RequestId, + error: JSONRPCErrorError, + ) -> IoResult<()> { + let (response_tx, response_rx) = oneshot::channel(); + self.command_tx + .send(ClientCommand::RejectServerRequest { + request_id, + error, + response_tx, + }) + .await + .map_err(|_| { + IoError::new( + ErrorKind::BrokenPipe, + "in-process app-server worker channel is closed", + ) + })?; + response_rx.await.map_err(|_| { + IoError::new( + ErrorKind::BrokenPipe, + "in-process app-server reject channel is closed", + ) + })? + } + pub async fn request(&self, request: ClientRequest) -> IoResult { let (response_tx, response_rx) = oneshot::channel(); self.command_tx @@ -676,6 +699,28 @@ impl InProcessAppServerRequestHandle { } impl AppServerRequestHandle { + pub(crate) async fn resolve_server_request( + &self, + request_id: RequestId, + result: JsonRpcResult, + ) -> IoResult<()> { + match self { + Self::InProcess(handle) => handle.resolve_server_request(request_id, result).await, + Self::Remote(handle) => handle.resolve_server_request(request_id, result).await, + } + } + + pub(crate) async fn reject_server_request( + &self, + request_id: RequestId, + error: JSONRPCErrorError, + ) -> IoResult<()> { + match self { + Self::InProcess(handle) => handle.reject_server_request(request_id, error).await, + Self::Remote(handle) => handle.reject_server_request(request_id, error).await, + } + } + pub async fn request(&self, request: ClientRequest) -> IoResult { match self { Self::InProcess(handle) => handle.request(request).await, @@ -789,6 +834,8 @@ impl AppServerClient { #[cfg(test)] mod tests { + #[path = "reply_tests.rs"] + mod replies; use super::*; use codex_app_server_protocol::AccountUpdatedNotification; use codex_app_server_protocol::ConfigRequirementsReadResponse; diff --git a/codex-rs/app-server-client/src/remote.rs b/codex-rs/app-server-client/src/remote.rs index 6b5571bf03..f95bee50c5 100644 --- a/codex-rs/app-server-client/src/remote.rs +++ b/codex-rs/app-server-client/src/remote.rs @@ -573,26 +573,9 @@ impl RemoteAppServerClient { request_id: RequestId, result: JsonRpcResult, ) -> IoResult<()> { - let (response_tx, response_rx) = oneshot::channel(); - self.command_tx - .send(RemoteClientCommand::ResolveServerRequest { - request_id, - result, - response_tx, - }) + self.request_handle() + .resolve_server_request(request_id, result) .await - .map_err(|_| { - IoError::new( - ErrorKind::BrokenPipe, - "remote app-server worker channel is closed", - ) - })?; - response_rx.await.map_err(|_| { - IoError::new( - ErrorKind::BrokenPipe, - "remote app-server resolve channel is closed", - ) - })? } pub async fn reject_server_request( @@ -600,26 +583,9 @@ impl RemoteAppServerClient { request_id: RequestId, error: JSONRPCErrorError, ) -> IoResult<()> { - let (response_tx, response_rx) = oneshot::channel(); - self.command_tx - .send(RemoteClientCommand::RejectServerRequest { - request_id, - error, - response_tx, - }) + self.request_handle() + .reject_server_request(request_id, error) .await - .map_err(|_| { - IoError::new( - ErrorKind::BrokenPipe, - "remote app-server worker channel is closed", - ) - })?; - response_rx.await.map_err(|_| { - IoError::new( - ErrorKind::BrokenPipe, - "remote app-server reject channel is closed", - ) - })? } pub async fn next_event(&mut self) -> Option { @@ -658,6 +624,60 @@ impl RemoteAppServerClient { } impl RemoteAppServerRequestHandle { + pub(crate) async fn resolve_server_request( + &self, + request_id: RequestId, + result: JsonRpcResult, + ) -> IoResult<()> { + let (response_tx, response_rx) = oneshot::channel(); + self.command_tx + .send(RemoteClientCommand::ResolveServerRequest { + request_id, + result, + response_tx, + }) + .await + .map_err(|_| { + IoError::new( + ErrorKind::BrokenPipe, + "remote app-server worker channel is closed", + ) + })?; + response_rx.await.map_err(|_| { + IoError::new( + ErrorKind::BrokenPipe, + "remote app-server resolve channel is closed", + ) + })? + } + + pub(crate) async fn reject_server_request( + &self, + request_id: RequestId, + error: JSONRPCErrorError, + ) -> IoResult<()> { + let (response_tx, response_rx) = oneshot::channel(); + self.command_tx + .send(RemoteClientCommand::RejectServerRequest { + request_id, + error, + response_tx, + }) + .await + .map_err(|_| { + IoError::new( + ErrorKind::BrokenPipe, + "remote app-server worker channel is closed", + ) + })?; + response_rx.await.map_err(|_| { + IoError::new( + ErrorKind::BrokenPipe, + "remote app-server reject channel is closed", + ) + })? + } + pub async fn request(&self, request: ClientRequest) -> IoResult { self.request_json_rpc(jsonrpc_request_from_client_request(request)) .await diff --git a/codex-rs/app-server-client/src/server_request.rs b/codex-rs/app-server-client/src/server_request.rs new file mode 100644 index 0000000000..464b26b7f2 --- /dev/null +++ b/codex-rs/app-server-client/src/server_request.rs @@ -0,0 +1,68 @@ +//! Retained reply capability: a server request ID and the worker that emitted it. +use crate::AppServerRequestHandle; +use codex_app_server_protocol::JSONRPCErrorError; +use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::Result as JsonRpcResult; +use std::fmt; +use std::io; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + +/// A reply always targets its originating connection, regardless of later UI selection. +/// Clones share a single-use completion state, so only the first resolve or reject is +/// submitted. Dropping a handle sends nothing. It does not keep a shutdown worker alive +/// or reconnect it; a stale connection returns its normal I/O error. No retries occur. +#[derive(Clone)] +pub struct AppServerReplyHandle { + client: AppServerRequestHandle, + request_id: RequestId, + claimed: Arc, +} +impl fmt::Debug for AppServerReplyHandle { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("AppServerReplyHandle([redacted])") + } +} +impl AppServerRequestHandle { + /// Capture at event ingress, using an ID from this connection's event stream. + /// Request IDs may be reused by other connections and must not select a client. + pub fn server_request_reply_handle(&self, request_id: RequestId) -> AppServerReplyHandle { + AppServerReplyHandle { + client: self.clone(), + request_id, + claimed: Arc::new(AtomicBool::new(false)), + } + } +} +impl AppServerReplyHandle { + fn claim(&self) -> io::Result<()> { + self.claimed + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .map(|_| ()) + .map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "app-server reply handle already used", + ) + }) + } + /// Resolve at the caller's explicit request. Failure can have an unknown outcome. + pub async fn resolve(&self, result: JsonRpcResult) -> io::Result<()> { + self.claim()?; + self.client + .resolve_server_request(self.request_id.clone(), result) + .await + } + /// Reject on the originating connection; never fall back to another connection. + pub async fn reject(&self, error: JSONRPCErrorError) -> io::Result<()> { + self.claim()?; + self.client + .reject_server_request(self.request_id.clone(), error) + .await + } +} + +#[cfg(test)] +#[path = "server_request_tests.rs"] +mod tests; diff --git a/codex-rs/app-server-client/src/server_request_tests.rs b/codex-rs/app-server-client/src/server_request_tests.rs new file mode 100644 index 0000000000..24e9dd1e48 --- /dev/null +++ b/codex-rs/app-server-client/src/server_request_tests.rs @@ -0,0 +1,105 @@ +//! Worker-channel ownership checks, including stale handles and identical request IDs. +use super::*; +use crate::ClientCommand; +use crate::InProcessAppServerRequestHandle; +use pretty_assertions::assert_eq; +use tokio::sync::mpsc; + +#[tokio::test] +async fn in_process_replies_keep_their_worker_and_request_id() { + let (first_tx, mut first_rx) = mpsc::channel(/*buffer*/ 2); + let (second_tx, mut second_rx) = mpsc::channel(/*buffer*/ 2); + let first = AppServerRequestHandle::InProcess(InProcessAppServerRequestHandle { + command_tx: first_tx, + }); + let second = AppServerRequestHandle::InProcess(InProcessAppServerRequestHandle { + command_tx: second_tx, + }); + let id = RequestId::String("private-request".into()); + let reply = first.server_request_reply_handle(id.clone()); + let stale = first.server_request_reply_handle(RequestId::String("stale-request".into())); + let other = second.server_request_reply_handle(id.clone()); + drop(first); + assert_eq!(format!("{reply:?}"), "AppServerReplyHandle([redacted])"); + assert_eq!(reply.request_id, id); + let (result, ()) = tokio::join!(reply.resolve(serde_json::json!({"first": true})), async { + let ClientCommand::ResolveServerRequest { + request_id, + result, + response_tx, + } = first_rx.recv().await.unwrap() + else { + panic!("wrong command") + }; + assert_eq!( + (request_id, result), + (id.clone(), serde_json::json!({"first": true})) + ); + response_tx.send(Ok(())).unwrap(); + }); + result.unwrap(); + assert!(second_rx.try_recv().is_err()); + let error = JSONRPCErrorError { + code: -32000, + message: "declined".into(), + data: None, + }; + let (result, ()) = tokio::join!(other.reject(error.clone()), async { + let ClientCommand::RejectServerRequest { + request_id, + error: received, + response_tx, + } = second_rx.recv().await.unwrap() + else { + panic!("wrong command") + }; + assert_eq!((request_id, received), (id, error)); + response_tx.send(Ok(())).unwrap(); + }); + result.unwrap(); + assert!(first_rx.try_recv().is_err()); + drop(first_rx); + assert_eq!( + stale + .resolve(serde_json::json!({})) + .await + .unwrap_err() + .kind(), + io::ErrorKind::BrokenPipe + ); + drop(other); + assert!(second_rx.try_recv().is_err()); +} + +#[tokio::test] +async fn cloned_reply_handle_only_submits_once() { + let (command_tx, mut command_rx) = mpsc::channel(/*buffer*/ 2); + let client = AppServerRequestHandle::InProcess(InProcessAppServerRequestHandle { command_tx }); + let reply = client.server_request_reply_handle(RequestId::Integer(1)); + let duplicate = reply.clone(); + let error = JSONRPCErrorError { + code: -32000, + message: "declined".into(), + data: None, + }; + + let (resolved, rejected, ()) = tokio::join!( + reply.resolve(serde_json::json!({"accepted": true})), + duplicate.reject(error), + async { + match command_rx.recv().await.unwrap() { + ClientCommand::ResolveServerRequest { response_tx, .. } + | ClientCommand::RejectServerRequest { response_tx, .. } => { + response_tx.send(Ok(())).unwrap(); + } + _ => panic!("wrong command"), + } + } + ); + let outcomes = [resolved, rejected].map(|result| result.map_err(|error| error.kind())); + assert!( + outcomes == [Ok(()), Err(io::ErrorKind::InvalidInput)] + || outcomes == [Err(io::ErrorKind::InvalidInput), Ok(())] + ); + assert!(command_rx.try_recv().is_err()); +} diff --git a/codex-rs/app-server-client/src/tests/reply_tests.rs b/codex-rs/app-server-client/src/tests/reply_tests.rs new file mode 100644 index 0000000000..0fe5a0f7ac --- /dev/null +++ b/codex-rs/app-server-client/src/tests/reply_tests.rs @@ -0,0 +1,101 @@ +//! Real WebSocket workers retain independent replies even when request IDs collide. +use super::*; +use pretty_assertions::assert_eq; + +#[tokio::test] +async fn remote_reply_handles_keep_original_connection() { + let mut clients = Vec::new(); + let mut replies = Vec::new(); + let mut stale_replies = Vec::new(); + let mut observed = Vec::new(); + for label in ["first", "second"] { + let (done_tx, done_rx) = oneshot::channel(); + let websocket_url = start_test_remote_server(move |mut websocket| async move { + expect_remote_initialize(&mut websocket).await; + write_websocket_message(&mut websocket, JSONRPCMessage::Request(JSONRPCRequest { + id: RequestId::Integer(1), method: "item/tool/requestUserInput".into(), + params: Some(serde_json::json!({"threadId":"same-thread", "turnId":"same-turn", "itemId":"same-item", "questions":[], "isBlocking":true})), + trace: None, + })).await; + let response = read_websocket_message(&mut websocket).await; + done_tx.send((label, response)).unwrap(); + while websocket.next().await.is_some() {} + }).await; + let mut client = AppServerClient::Remote( + RemoteAppServerClient::connect(test_remote_connect_args(websocket_url)) + .await + .unwrap(), + ); + let AppServerEvent::ServerRequest(request) = client.next_event().await.unwrap() else { + panic!("expected request") + }; + replies.push( + client + .request_handle() + .server_request_reply_handle(request.id().clone()), + ); + stale_replies.push( + client + .request_handle() + .server_request_reply_handle(RequestId::Integer(2)), + ); + clients.push(client); + observed.push(done_rx); + } + replies[1] + .reject(JSONRPCErrorError { + code: -32000, + message: "second".into(), + data: None, + }) + .await + .unwrap(); + replies[0] + .resolve(serde_json::json!({"first": true})) + .await + .unwrap(); + for receiver in observed { + let (label, message) = timeout(Duration::from_secs(/*secs*/ 5), receiver) + .await + .unwrap() + .unwrap(); + match message { + JSONRPCMessage::Response(response) => assert_eq!( + (label, response.id, response.result), + ( + "first", + RequestId::Integer(1), + serde_json::json!({"first": true}) + ) + ), + JSONRPCMessage::Error(error) => assert_eq!( + (label, error.id, error.error.message), + ("second", RequestId::Integer(1), "second".into()) + ), + _ => panic!("expected response or rejection"), + } + } + for client in clients { + client.shutdown().await.unwrap(); + } + for reply in replies { + assert_eq!( + reply + .resolve(serde_json::json!({})) + .await + .unwrap_err() + .kind(), + ErrorKind::InvalidInput + ); + } + for reply in stale_replies { + assert_eq!( + reply + .resolve(serde_json::json!({})) + .await + .unwrap_err() + .kind(), + ErrorKind::BrokenPipe + ); + } +} diff --git a/codex-rs/tui/src/managed_worktree_model_tests.rs b/codex-rs/tui/src/managed_worktree_model_tests.rs index ae36b611a5..2f66b7cf68 100644 --- a/codex-rs/tui/src/managed_worktree_model_tests.rs +++ b/codex-rs/tui/src/managed_worktree_model_tests.rs @@ -61,7 +61,11 @@ async fn model_call( let ServerRequest::McpServerElicitationRequest { request_id, params } = *request else { panic!("unexpected server request: {request:?}"); }; assert_eq!(params.thread_id, thread); approvals += 1; - server.resolve_server_request(request_id, json!({"action":"accept","content":null,"_meta":null})).await?; + server + .request_handle() + .server_request_reply_handle(request_id) + .resolve(json!({"action":"accept","content":null,"_meta":null})) + .await?; } AppServerEvent::ServerNotification(notification) if matches!(notification.as_ref(), ServerNotification::TurnCompleted(event) if event.thread_id == thread) => break, _ => {},