From 8c0f7011674145da7e10627769fb221c74fbf7a0 Mon Sep 17 00:00:00 2001 From: Benjamin Carlsson Date: Sun, 11 Oct 2026 05:59:25 +0000 Subject: [PATCH] Add connection-bound reply handles for app-server requests (#53049) ## Why Server request IDs can be reused across connections. Retained replies need to target the connection that emitted the request, even if the UI selection later changes. ## What changed - Expose `AppServerReplyHandle` through `AppServerRequestHandle::server_request_reply_handle`, capturing the originating client and request ID for explicit `resolve` or `reject` calls. - Share single-use completion state across clones so only the first reply is submitted. Dropping a handle sends nothing; stale connections return I/O errors without retrying or switching connections. - Route existing client reply methods through request handles and use the new API in the managed worktree model test. ## Testing Add in-process and WebSocket tests for replies on separate connections with identical request IDs, rejection routing, duplicate completion, and stale handles after worker shutdown. Also check that debug output redacts handle details. GitOrigin-RevId: 940d0951ab67529b05919ca8ec7adb5813834b8c --- codex-rs/app-server-client/src/lib.rs | 123 ++++++++++++------ codex-rs/app-server-client/src/remote.rs | 96 ++++++++------ .../app-server-client/src/server_request.rs | 68 ++++++++++ .../src/server_request_tests.rs | 105 +++++++++++++++ .../src/tests/reply_tests.rs | 101 ++++++++++++++ .../tui/src/managed_worktree_model_tests.rs | 6 +- 6 files changed, 422 insertions(+), 77 deletions(-) create mode 100644 codex-rs/app-server-client/src/server_request.rs create mode 100644 codex-rs/app-server-client/src/server_request_tests.rs create mode 100644 codex-rs/app-server-client/src/tests/reply_tests.rs 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, _ => {},