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

@@ -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<RequestResult> {
         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<RequestResult> {
         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;

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<AppServerEvent> {
@@ -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<RequestResult> {
         self.request_json_rpc(jsonrpc_request_from_client_request(request))
             .await

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<AtomicBool>,
+}
+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;

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());
+}

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

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,
                 _ => {},