From c7c51f80e13fee5948c094c1ca92f4202512b78a Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Tue, 10 Mar 2026 12:54:25 -0700 Subject: [PATCH] Address some conflict potential in tool registration Keep active-turn dynamic tool snapshots frozen to avoid cache churn, while tightening provider-owned tool request ownership so only the registering connection can answer them. Add a 5-second grace period for provider disconnects before failing any in-flight provider- owned dynamic tool calls, ignore late replies after those calls are already resolved, and extend websocket/stdio coverage for registration-after-start, hijack attempts, and provider disconnect during an in-flight call. --- codex-rs/app-server/README.md | 2 +- .../app-server/src/codex_message_processor.rs | 18 + codex-rs/app-server/src/dynamic_tools.rs | 4 + codex-rs/app-server/src/in_process.rs | 5 +- codex-rs/app-server/src/lib.rs | 4 +- codex-rs/app-server/src/message_processor.rs | 16 +- codex-rs/app-server/src/outgoing_message.rs | 373 +++++++++++++++++- .../app-server/src/server_request_error.rs | 35 ++ .../tests/suite/v2/dynamic_tools.rs | 22 +- .../tests/suite/v2/tool_provider_websocket.rs | 238 ++++++++++- codex-rs/core/src/codex.rs | 3 +- codex-rs/core/src/codex_tests.rs | 61 ++- codex-rs/core/src/state/turn.rs | 14 +- codex-rs/core/src/tools/js_repl/mod.rs | 2 +- codex-rs/core/src/tools/router.rs | 12 +- 15 files changed, 754 insertions(+), 55 deletions(-) diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 92b8c87417..268f854906 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -944,7 +944,7 @@ Clients can also register connection-owned dynamic tools with the experimental ` If `injectIntoContext` is `true`, the tool is exposed to the model as a function tool for future turns. If it is `false`, the tool remains callable only through `js_repl` via `codex.tool(...)` and is not injected into model context. -For provider-owned tools, the app-server routes `item/tool/call` only to the connection that registered the tool and applies that tool's `defaultTimeoutMs` while waiting for the client response. +For provider-owned tools, the app-server routes `item/tool/call` only to the connection that registered the tool, ignores replies from other connections for that pending request, applies that tool's `defaultTimeoutMs` while waiting for the client response, and fails any in-flight provider request about 5 seconds after that provider connection disconnects. When a dynamic tool is invoked during a turn, the server sends an `item/tool/call` JSON-RPC request to the client: diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 82f9d0723c..44b232629d 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -330,6 +330,7 @@ struct ThreadListFilters { // Duration before a ChatGPT login attempt is abandoned. const LOGIN_CHATGPT_TIMEOUT: Duration = Duration::from_secs(10 * 60); const APP_LIST_LOAD_TIMEOUT: Duration = Duration::from_secs(90); +const TOOL_PROVIDER_DISCONNECT_GRACE_PERIOD: Duration = Duration::from_secs(5); struct ActiveLogin { shutdown_handle: ShutdownHandle, login_id: Uuid, @@ -3213,6 +3214,23 @@ impl CodexMessageProcessor { .remove_connection(connection_id) .await { + let outgoing = Arc::clone(&self.outgoing); + tokio::spawn(async move { + tokio::time::sleep(TOOL_PROVIDER_DISCONNECT_GRACE_PERIOD).await; + let _ = outgoing + .cancel_pending_dynamic_tool_requests_for_connection( + connection_id, + JSONRPCErrorError { + code: INTERNAL_ERROR_CODE, + message: "dynamic tool provider is unavailable".to_string(), + data: Some(serde_json::json!({ + "reason": + crate::server_request_error::DYNAMIC_TOOL_PROVIDER_DISCONNECTED_PENDING_REQUEST_ERROR_REASON + })), + }, + ) + .await; + }); self.refresh_registered_dynamic_tools_for_loaded_threads() .await; } diff --git a/codex-rs/app-server/src/dynamic_tools.rs b/codex-rs/app-server/src/dynamic_tools.rs index e9c177c472..fc6037ff85 100644 --- a/codex-rs/app-server/src/dynamic_tools.rs +++ b/codex-rs/app-server/src/dynamic_tools.rs @@ -12,6 +12,7 @@ use tracing::error; use crate::outgoing_message::ClientRequestResult; use crate::outgoing_message::ThreadScopedOutgoingMessageSender; +use crate::server_request_error::is_dynamic_tool_provider_disconnected_server_request_error; use crate::server_request_error::is_turn_transition_server_request_error; pub(crate) async fn on_call_response( @@ -41,6 +42,9 @@ pub(crate) async fn on_call_response( let (response, _error) = match response { Ok(Ok(value)) => decode_response(value), Ok(Err(err)) if is_turn_transition_server_request_error(&err) => return, + Ok(Err(err)) if is_dynamic_tool_provider_disconnected_server_request_error(&err) => { + fallback_response("dynamic tool provider is unavailable") + } Ok(Err(err)) => { error!("request failed with client error: {err:?}"); fallback_response("dynamic tool request failed") diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index 4e3572ee5d..73c71e8570 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -544,12 +544,12 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { } Some(InProcessClientMessage::ServerRequestResponse { request_id, result }) => { outgoing_message_sender - .notify_client_response(request_id, result) + .notify_client_response(IN_PROCESS_CONNECTION_ID, request_id, result) .await; } Some(InProcessClientMessage::ServerRequestError { request_id, error }) => { outgoing_message_sender - .notify_client_error(request_id, error) + .notify_client_error(IN_PROCESS_CONNECTION_ID, request_id, error) .await; } Some(InProcessClientMessage::Shutdown { done_tx }) => { @@ -610,6 +610,7 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { }; outgoing_message_sender .notify_client_error( + IN_PROCESS_CONNECTION_ID, request_id, JSONRPCErrorError { code, diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 6d95c414e4..97d3792996 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -753,7 +753,7 @@ pub async fn run_main_with_transport( warn!("dropping response from unknown connection: {connection_id:?}"); continue; } - processor.process_response(response).await; + processor.process_response(connection_id, response).await; } JSONRPCMessage::Notification(notification) => { if !connections.contains_key(&connection_id) { @@ -767,7 +767,7 @@ pub async fn run_main_with_transport( warn!("dropping error from unknown connection: {connection_id:?}"); continue; } - processor.process_error(err).await; + processor.process_error(connection_id, err).await; } } } diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 82abd3466f..adbea4357d 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -399,16 +399,24 @@ impl MessageProcessor { } /// Handle a standalone JSON-RPC response originating from the peer. - pub(crate) async fn process_response(&mut self, response: JSONRPCResponse) { + pub(crate) async fn process_response( + &mut self, + connection_id: ConnectionId, + response: JSONRPCResponse, + ) { tracing::info!("<- response: {:?}", response); let JSONRPCResponse { id, result, .. } = response; - self.outgoing.notify_client_response(id, result).await + self.outgoing + .notify_client_response(connection_id, id, result) + .await } /// Handle an error object received from the peer. - pub(crate) async fn process_error(&mut self, err: JSONRPCError) { + pub(crate) async fn process_error(&mut self, connection_id: ConnectionId, err: JSONRPCError) { tracing::error!("<- error: {:?}", err); - self.outgoing.notify_client_error(err.id, err.error).await; + self.outgoing + .notify_client_error(connection_id, err.id, err.error) + .await; } async fn handle_client_request( diff --git a/codex-rs/app-server/src/outgoing_message.rs b/codex-rs/app-server/src/outgoing_message.rs index 86a3a4aefd..ceb47dd469 100644 --- a/codex-rs/app-server/src/outgoing_message.rs +++ b/codex-rs/app-server/src/outgoing_message.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::collections::HashSet; use std::sync::Arc; use std::sync::atomic::AtomicI64; use std::sync::atomic::Ordering; @@ -64,6 +65,7 @@ struct PendingCallbackEntry { callback: oneshot::Sender, thread_id: Option, request: ServerRequest, + allowed_responder_connections: Option>, } impl ThreadScopedOutgoingMessageSender { @@ -88,6 +90,7 @@ impl ThreadScopedOutgoingMessageSender { Some(self.connection_ids.as_slice()), payload, Some(self.thread_id), + false, ) .await } @@ -99,7 +102,7 @@ impl ThreadScopedOutgoingMessageSender { ) -> (RequestId, oneshot::Receiver) { let connection_ids = [connection_id]; self.outgoing - .send_request_to_connections(Some(&connection_ids), payload, Some(self.thread_id)) + .send_request_to_connections(Some(&connection_ids), payload, Some(self.thread_id), true) .await } @@ -164,7 +167,8 @@ impl OutgoingMessageSender { &self, request: ServerRequestPayload, ) -> (RequestId, oneshot::Receiver) { - self.send_request_to_connections(None, request, None).await + self.send_request_to_connections(None, request, None, false) + .await } fn next_request_id(&self) -> RequestId { @@ -176,6 +180,7 @@ impl OutgoingMessageSender { connection_ids: Option<&[ConnectionId]>, request: ServerRequestPayload, thread_id: Option, + restrict_responder_connections: bool, ) -> (RequestId, oneshot::Receiver) { let id = self.next_request_id(); let outgoing_message_id = id.clone(); @@ -190,6 +195,13 @@ impl OutgoingMessageSender { callback: tx_approve, thread_id, request: request.clone(), + allowed_responder_connections: restrict_responder_connections.then(|| { + connection_ids + .unwrap_or_default() + .iter() + .copied() + .collect::>() + }), }, ); } @@ -253,8 +265,13 @@ impl OutgoingMessageSender { } } - pub(crate) async fn notify_client_response(&self, id: RequestId, result: Result) { - let entry = self.take_request_callback(&id).await; + pub(crate) async fn notify_client_response( + &self, + connection_id: ConnectionId, + id: RequestId, + result: Result, + ) { + let entry = self.take_request_callback(&id, connection_id).await; match entry { Some((id, entry)) => { @@ -263,13 +280,18 @@ impl OutgoingMessageSender { } } None => { - warn!("could not find callback for {id:?}"); + warn!("could not find callback for {id:?} from connection {connection_id:?}"); } } } - pub(crate) async fn notify_client_error(&self, id: RequestId, error: JSONRPCErrorError) { - let entry = self.take_request_callback(&id).await; + pub(crate) async fn notify_client_error( + &self, + connection_id: ConnectionId, + id: RequestId, + error: JSONRPCErrorError, + ) { + let entry = self.take_request_callback(&id, connection_id).await; match entry { Some((id, entry)) => { @@ -279,13 +301,14 @@ impl OutgoingMessageSender { } } None => { - warn!("could not find callback for {id:?}"); + warn!("could not find callback for {id:?} from connection {connection_id:?}"); } } } pub(crate) async fn cancel_request(&self, id: &RequestId) -> bool { - self.take_request_callback(id).await.is_some() + let mut request_id_to_callback = self.request_id_to_callback.lock().await; + request_id_to_callback.remove(id).is_some() } pub(crate) async fn cancel_all_requests(&self, error: Option) { @@ -307,11 +330,60 @@ impl OutgoingMessageSender { } } + pub(crate) async fn cancel_pending_dynamic_tool_requests_for_connection( + &self, + connection_id: ConnectionId, + error: JSONRPCErrorError, + ) -> usize { + let entries = { + let mut request_id_to_callback = self.request_id_to_callback.lock().await; + let request_ids = request_id_to_callback + .iter() + .filter_map(|(request_id, entry)| { + (matches!(entry.request, ServerRequest::DynamicToolCall { .. }) + && entry + .allowed_responder_connections + .as_ref() + .is_some_and(|allowed| allowed.contains(&connection_id))) + .then_some(request_id.clone()) + }) + .collect::>(); + + let mut entries = Vec::with_capacity(request_ids.len()); + for request_id in request_ids { + if let Some(entry) = request_id_to_callback.remove(&request_id) { + entries.push(entry); + } + } + entries + }; + + let cancelled = entries.len(); + for entry in entries { + if let Err(err) = entry.callback.send(Err(error.clone())) { + let request_id = entry.request.id(); + warn!("could not notify callback for {request_id:?} due to: {err:?}"); + } + } + + cancelled + } + async fn take_request_callback( &self, id: &RequestId, + connection_id: ConnectionId, ) -> Option<(RequestId, PendingCallbackEntry)> { let mut request_id_to_callback = self.request_id_to_callback.lock().await; + let entry = request_id_to_callback.get(id)?; + if entry + .allowed_responder_connections + .as_ref() + .is_some_and(|allowed| !allowed.contains(&connection_id)) + { + warn!("rejecting response for {id:?} from unexpected connection {connection_id:?}"); + return None; + } request_id_to_callback.remove_entry(id) } @@ -814,7 +886,7 @@ mod tests { }; outgoing - .notify_client_error(request_id, error.clone()) + .notify_client_error(ConnectionId(0), request_id, error.clone()) .await; let result = timeout(Duration::from_secs(1), wait_for_result) @@ -940,4 +1012,285 @@ mod tests { .is_empty() ); } + + #[tokio::test] + async fn wrong_connection_cannot_resolve_targeted_request() { + let (tx, _rx) = mpsc::channel::(8); + let outgoing = Arc::new(OutgoingMessageSender::new(tx)); + let thread_id = ThreadId::new(); + let thread_outgoing = ThreadScopedOutgoingMessageSender::new( + outgoing.clone(), + vec![ConnectionId(7)], + thread_id, + ); + + let (request_id, wait_for_result) = thread_outgoing + .send_request_to_connection( + ConnectionId(7), + ServerRequestPayload::DynamicToolCall(DynamicToolCallParams { + thread_id: thread_id.to_string(), + turn_id: "turn-1".to_string(), + call_id: "call-0".to_string(), + tool: "tool".to_string(), + arguments: json!({}), + }), + ) + .await; + + outgoing + .notify_client_response(ConnectionId(9), request_id.clone(), json!({ "ok": true })) + .await; + let mut wait_for_result = Box::pin(wait_for_result); + assert!( + timeout(Duration::from_millis(100), &mut wait_for_result) + .await + .is_err(), + "wrong connection should not resolve targeted request", + ); + + let (request_id, wait_for_result) = thread_outgoing + .send_request_to_connection( + ConnectionId(7), + ServerRequestPayload::DynamicToolCall(DynamicToolCallParams { + thread_id: thread_id.to_string(), + turn_id: "turn-1".to_string(), + call_id: "call-1".to_string(), + tool: "tool".to_string(), + arguments: json!({}), + }), + ) + .await; + outgoing + .notify_client_response(ConnectionId(7), request_id, json!({ "ok": true })) + .await; + let result = timeout(Duration::from_secs(1), wait_for_result) + .await + .expect("authorized waiter should resolve") + .expect("authorized waiter should receive a callback"); + assert_eq!(result, Ok(json!({ "ok": true }))); + } + + #[tokio::test] + async fn wrong_connection_cannot_fail_targeted_request() { + let (tx, _rx) = mpsc::channel::(8); + let outgoing = Arc::new(OutgoingMessageSender::new(tx)); + let thread_id = ThreadId::new(); + let thread_outgoing = ThreadScopedOutgoingMessageSender::new( + outgoing.clone(), + vec![ConnectionId(11)], + thread_id, + ); + + let (request_id, wait_for_result) = thread_outgoing + .send_request_to_connection( + ConnectionId(11), + ServerRequestPayload::ToolRequestUserInput(ToolRequestUserInputParams { + thread_id: thread_id.to_string(), + turn_id: "turn-1".to_string(), + item_id: "call-0".to_string(), + questions: vec![], + }), + ) + .await; + let error = JSONRPCErrorError { + code: INTERNAL_ERROR_CODE, + message: "wrong connection".to_string(), + data: None, + }; + + outgoing + .notify_client_error(ConnectionId(12), request_id.clone(), error.clone()) + .await; + let mut wait_for_result = Box::pin(wait_for_result); + assert!( + timeout(Duration::from_millis(100), &mut wait_for_result) + .await + .is_err(), + "wrong connection should not fail targeted request", + ); + + let (request_id, wait_for_result) = thread_outgoing + .send_request_to_connection( + ConnectionId(11), + ServerRequestPayload::ToolRequestUserInput(ToolRequestUserInputParams { + thread_id: thread_id.to_string(), + turn_id: "turn-1".to_string(), + item_id: "call-1".to_string(), + questions: vec![], + }), + ) + .await; + outgoing + .notify_client_error(ConnectionId(11), request_id, error.clone()) + .await; + let result = timeout(Duration::from_secs(1), wait_for_result) + .await + .expect("authorized waiter should resolve") + .expect("authorized waiter should receive a callback"); + assert_eq!(result, Err(error)); + } + + #[tokio::test] + async fn disconnect_cancellation_only_resolves_dynamic_tool_requests() { + let (tx, _rx) = mpsc::channel::(8); + let outgoing = Arc::new(OutgoingMessageSender::new(tx)); + let thread_id = ThreadId::new(); + let thread_outgoing = ThreadScopedOutgoingMessageSender::new( + outgoing.clone(), + vec![ConnectionId(17)], + thread_id, + ); + + let (_dynamic_request_id, dynamic_tool_waiter) = thread_outgoing + .send_request_to_connection( + ConnectionId(17), + ServerRequestPayload::DynamicToolCall(DynamicToolCallParams { + thread_id: thread_id.to_string(), + turn_id: "turn-1".to_string(), + call_id: "call-0".to_string(), + tool: "tool".to_string(), + arguments: json!({}), + }), + ) + .await; + let (user_input_request_id, user_input_waiter) = thread_outgoing + .send_request_to_connection( + ConnectionId(17), + ServerRequestPayload::ToolRequestUserInput(ToolRequestUserInputParams { + thread_id: thread_id.to_string(), + turn_id: "turn-1".to_string(), + item_id: "call-1".to_string(), + questions: vec![], + }), + ) + .await; + let error = JSONRPCErrorError { + code: INTERNAL_ERROR_CODE, + message: "dynamic tool provider is unavailable".to_string(), + data: Some(serde_json::json!({ + "reason": + crate::server_request_error::DYNAMIC_TOOL_PROVIDER_DISCONNECTED_PENDING_REQUEST_ERROR_REASON + })), + }; + + assert_eq!( + outgoing + .cancel_pending_dynamic_tool_requests_for_connection( + ConnectionId(17), + error.clone(), + ) + .await, + 1 + ); + let dynamic_tool_result = timeout(Duration::from_secs(1), dynamic_tool_waiter) + .await + .expect("dynamic tool waiter should resolve") + .expect("dynamic tool waiter should receive a callback"); + assert_eq!(dynamic_tool_result, Err(error)); + + let mut user_input_waiter = Box::pin(user_input_waiter); + assert!( + timeout(Duration::from_millis(100), &mut user_input_waiter) + .await + .is_err(), + "non-dynamic targeted requests should remain pending", + ); + + outgoing + .notify_client_response( + ConnectionId(17), + user_input_request_id, + json!({ "ok": true }), + ) + .await; + let user_input_result = timeout(Duration::from_secs(1), user_input_waiter) + .await + .expect("user input waiter should resolve") + .expect("user input waiter should receive a callback"); + assert_eq!(user_input_result, Ok(json!({ "ok": true }))); + } + + #[tokio::test] + async fn late_response_after_disconnect_cancellation_is_ignored() { + let (tx, _rx) = mpsc::channel::(8); + let outgoing = Arc::new(OutgoingMessageSender::new(tx)); + let thread_id = ThreadId::new(); + let thread_outgoing = ThreadScopedOutgoingMessageSender::new( + outgoing.clone(), + vec![ConnectionId(23)], + thread_id, + ); + + let (request_id, wait_for_result) = thread_outgoing + .send_request_to_connection( + ConnectionId(23), + ServerRequestPayload::DynamicToolCall(DynamicToolCallParams { + thread_id: thread_id.to_string(), + turn_id: "turn-1".to_string(), + call_id: "call-0".to_string(), + tool: "tool".to_string(), + arguments: json!({}), + }), + ) + .await; + let error = JSONRPCErrorError { + code: INTERNAL_ERROR_CODE, + message: "dynamic tool provider is unavailable".to_string(), + data: Some(serde_json::json!({ + "reason": + crate::server_request_error::DYNAMIC_TOOL_PROVIDER_DISCONNECTED_PENDING_REQUEST_ERROR_REASON + })), + }; + + assert_eq!( + outgoing + .cancel_pending_dynamic_tool_requests_for_connection( + ConnectionId(23), + error.clone(), + ) + .await, + 1 + ); + let result = timeout(Duration::from_secs(1), wait_for_result) + .await + .expect("disconnect waiter should resolve") + .expect("disconnect waiter should receive a callback"); + assert_eq!(result, Err(error)); + + outgoing + .notify_client_response(ConnectionId(23), request_id, json!({ "ok": true })) + .await; + } + + #[tokio::test] + async fn non_targeted_thread_request_can_be_resolved_by_replayed_connection() { + let (tx, _rx) = mpsc::channel::(8); + let outgoing = Arc::new(OutgoingMessageSender::new(tx)); + let thread_id = ThreadId::new(); + let thread_outgoing = ThreadScopedOutgoingMessageSender::new( + outgoing.clone(), + vec![ConnectionId(1)], + thread_id, + ); + + let (request_id, wait_for_result) = thread_outgoing + .send_request(ServerRequestPayload::ToolRequestUserInput( + ToolRequestUserInputParams { + thread_id: thread_id.to_string(), + turn_id: "turn-1".to_string(), + item_id: "call-0".to_string(), + questions: vec![], + }, + )) + .await; + outgoing + .notify_client_response(ConnectionId(99), request_id, json!({ "ok": true })) + .await; + + let result = timeout(Duration::from_secs(1), wait_for_result) + .await + .expect("replayed waiter should resolve") + .expect("replayed waiter should receive a callback"); + assert_eq!(result, Ok(json!({ "ok": true }))); + } } diff --git a/codex-rs/app-server/src/server_request_error.rs b/codex-rs/app-server/src/server_request_error.rs index 524b50e44b..4649a4015f 100644 --- a/codex-rs/app-server/src/server_request_error.rs +++ b/codex-rs/app-server/src/server_request_error.rs @@ -1,6 +1,8 @@ use codex_app_server_protocol::JSONRPCErrorError; pub(crate) const TURN_TRANSITION_PENDING_REQUEST_ERROR_REASON: &str = "turnTransition"; +pub(crate) const DYNAMIC_TOOL_PROVIDER_DISCONNECTED_PENDING_REQUEST_ERROR_REASON: &str = + "dynamicToolProviderDisconnected"; pub(crate) fn is_turn_transition_server_request_error(error: &JSONRPCErrorError) -> bool { error @@ -11,8 +13,21 @@ pub(crate) fn is_turn_transition_server_request_error(error: &JSONRPCErrorError) == Some(TURN_TRANSITION_PENDING_REQUEST_ERROR_REASON) } +pub(crate) fn is_dynamic_tool_provider_disconnected_server_request_error( + error: &JSONRPCErrorError, +) -> bool { + error + .data + .as_ref() + .and_then(|data| data.get("reason")) + .and_then(serde_json::Value::as_str) + == Some(DYNAMIC_TOOL_PROVIDER_DISCONNECTED_PENDING_REQUEST_ERROR_REASON) +} + #[cfg(test)] mod tests { + use super::DYNAMIC_TOOL_PROVIDER_DISCONNECTED_PENDING_REQUEST_ERROR_REASON; + use super::is_dynamic_tool_provider_disconnected_server_request_error; use super::is_turn_transition_server_request_error; use codex_app_server_protocol::JSONRPCErrorError; use pretty_assertions::assert_eq; @@ -38,5 +53,25 @@ mod tests { }; assert_eq!(is_turn_transition_server_request_error(&error), false); + assert_eq!( + is_dynamic_tool_provider_disconnected_server_request_error(&error), + false + ); + } + + #[test] + fn dynamic_tool_provider_disconnect_error_is_detected() { + let error = JSONRPCErrorError { + code: -1, + message: "dynamic tool provider is unavailable".to_string(), + data: Some( + json!({ "reason": DYNAMIC_TOOL_PROVIDER_DISCONNECTED_PENDING_REQUEST_ERROR_REASON }), + ), + }; + + assert_eq!( + is_dynamic_tool_provider_disconnected_server_request_error(&error), + true + ); } } diff --git a/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs b/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs index cb04ecbc9a..ffe29abdb7 100644 --- a/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs +++ b/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs @@ -142,6 +142,17 @@ async fn tool_provider_registration_routes_stdio_dynamic_tool_calls() -> Result< let mut mcp = McpProcess::new(codex_home.path()).await?; timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + let thread_request_id = mcp + .send_thread_start_request(ThreadStartParams::default()) + .await?; + let thread_response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_request_id)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response::(thread_response)?; + let thread_id = thread.id.clone(); + let register_request_id = mcp .send_raw_request( "toolProvider/register", @@ -170,17 +181,6 @@ async fn tool_provider_registration_routes_stdio_dynamic_tool_calls() -> Result< .await??; let _: ToolProviderRegisterResponse = to_response(register_response)?; - let thread_request_id = mcp - .send_thread_start_request(ThreadStartParams::default()) - .await?; - let thread_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_request_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_response)?; - let thread_id = thread.id.clone(); - let turn_request_id = mcp .send_turn_start_request(TurnStartParams { thread_id: thread_id.clone(), diff --git a/codex-rs/app-server/tests/suite/v2/tool_provider_websocket.rs b/codex-rs/app-server/tests/suite/v2/tool_provider_websocket.rs index 1e00c642bd..11e604dee2 100644 --- a/codex-rs/app-server/tests/suite/v2/tool_provider_websocket.rs +++ b/codex-rs/app-server/tests/suite/v2/tool_provider_websocket.rs @@ -26,12 +26,18 @@ use codex_app_server_protocol::ToolProviderRegisterResponse; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; +use codex_protocol::models::FunctionCallOutputPayload; use core_test_support::responses; use futures::SinkExt; +use pretty_assertions::assert_eq; +use serde_json::Value; use serde_json::json; use tempfile::TempDir; +use tokio::time::Duration; +use tokio::time::sleep; use tokio::time::timeout; use tokio_tungstenite::tungstenite::Message as WebSocketMessage; +use wiremock::MockServer; #[tokio::test] async fn websocket_tool_provider_routes_dynamic_tool_calls_to_owner_connection() -> Result<()> { @@ -66,6 +72,16 @@ async fn websocket_tool_provider_routes_dynamic_tool_calls_to_owner_connection() let thread_init = read_response_for_id(&mut thread_client, 2).await?; assert_eq!(thread_init.id, RequestId::Integer(2)); + send_request( + &mut thread_client, + "thread/start", + 4, + Some(serde_json::to_value(ThreadStartParams::default())?), + ) + .await?; + let thread_start_response = read_response_for_id(&mut thread_client, 4).await?; + let ThreadStartResponse { thread, .. } = to_response(thread_start_response)?; + send_request( &mut provider, "toolProvider/register", @@ -91,16 +107,6 @@ async fn websocket_tool_provider_routes_dynamic_tool_calls_to_owner_connection() let register_response = read_response_for_id(&mut provider, 3).await?; let _: ToolProviderRegisterResponse = to_response(register_response)?; - send_request( - &mut thread_client, - "thread/start", - 4, - Some(serde_json::to_value(ThreadStartParams::default())?), - ) - .await?; - let thread_start_response = read_response_for_id(&mut thread_client, 4).await?; - let ThreadStartResponse { thread, .. } = to_response(thread_start_response)?; - send_request( &mut thread_client, "turn/start", @@ -143,6 +149,20 @@ async fn websocket_tool_provider_routes_dynamic_tool_calls_to_owner_connection() } ); + send_response( + &mut thread_client, + request_id.clone(), + serde_json::to_value(DynamicToolCallResponse { + content_items: vec![DynamicToolCallOutputContentItem::InputText { + text: "hijack".to_string(), + }], + success: true, + })?, + ) + .await?; + sleep(Duration::from_millis(250)).await; + assert_eq!(responses_bodies(&server).await?.len(), 1); + send_response( &mut provider, request_id, @@ -177,6 +197,155 @@ async fn websocket_tool_provider_routes_dynamic_tool_calls_to_owner_connection() } } + let bodies = responses_bodies(&server).await?; + let payload = bodies + .iter() + .find_map(|body| function_call_output_payload(body, call_id)) + .context("expected function_call_output in follow-up request")?; + assert_eq!( + payload, + FunctionCallOutputPayload::from_text("dynamic-ok".to_string()) + ); + + Ok(()) + } + .await; + + process + .kill() + .await + .context("failed to stop websocket app-server process")?; + test_result +} + +#[tokio::test] +async fn websocket_tool_provider_disconnect_fails_in_flight_request_after_grace_period() +-> Result<()> { + let call_id = "ws-provider-disconnect-call-1"; + let tool_name = "provider_tool"; + let tool_args = json!({ "city": "Paris" }); + let tool_call_arguments = serde_json::to_string(&tool_args)?; + let responses = vec![ + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_function_call(call_id, tool_name, &tool_call_arguments), + responses::ev_completed("resp-1"), + ]), + create_final_assistant_message_sse_response("Done")?, + ]; + let server = create_mock_responses_server_sequence_unchecked(responses).await; + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri(), "never")?; + + let bind_addr = reserve_local_addr()?; + let mut process = spawn_websocket_server(codex_home.path(), bind_addr).await?; + + let test_result = async { + let mut provider = connect_websocket(bind_addr).await?; + let mut thread_client = connect_websocket(bind_addr).await?; + + send_initialize_request(&mut provider, 1, "ws_provider").await?; + let provider_init = read_response_for_id(&mut provider, 1).await?; + assert_eq!(provider_init.id, RequestId::Integer(1)); + + send_initialize_request(&mut thread_client, 2, "ws_thread_client").await?; + let thread_init = read_response_for_id(&mut thread_client, 2).await?; + assert_eq!(thread_init.id, RequestId::Integer(2)); + + send_request( + &mut thread_client, + "thread/start", + 4, + Some(serde_json::to_value(ThreadStartParams::default())?), + ) + .await?; + let thread_start_response = read_response_for_id(&mut thread_client, 4).await?; + let ThreadStartResponse { thread, .. } = to_response(thread_start_response)?; + + send_request( + &mut provider, + "toolProvider/register", + 3, + Some(json!({ + "tool": { + "name": tool_name, + "description": "Provider tool", + "inputSchema": { + "type": "object", + "properties": { + "city": { "type": "string" } + }, + "required": ["city"], + "additionalProperties": false + }, + "defaultTimeoutMs": 30_000, + "injectIntoContext": true + } + })), + ) + .await?; + let register_response = read_response_for_id(&mut provider, 3).await?; + let _: ToolProviderRegisterResponse = to_response(register_response)?; + + send_request( + &mut thread_client, + "turn/start", + 5, + Some(serde_json::to_value(TurnStartParams { + thread_id: thread.id.clone(), + input: vec![V2UserInput::Text { + text: "Run the provider tool".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + })?), + ) + .await?; + let turn_start_response = read_response_for_id(&mut thread_client, 5).await?; + let _: TurnStartResponse = to_response::(turn_start_response)?; + + let request_id = loop { + match timeout(DEFAULT_READ_TIMEOUT, read_jsonrpc_message(&mut provider)).await?? { + JSONRPCMessage::Request(request) => match ServerRequest::try_from(request)? { + ServerRequest::DynamicToolCall { request_id, .. } => break request_id, + other => bail!("expected provider DynamicToolCall request, got {other:?}"), + }, + JSONRPCMessage::Notification(_) | JSONRPCMessage::Response(_) => {} + JSONRPCMessage::Error(error) => { + bail!("unexpected websocket error on provider client: {error:?}"); + } + } + }; + + provider + .close(None) + .await + .context("failed to close provider websocket")?; + drop(provider); + + let (bodies, payload) = wait_for_function_call_output(&server, call_id).await?; + assert_eq!( + payload, + FunctionCallOutputPayload::from_text( + "dynamic tool provider is unavailable".to_string(), + ) + ); + + let response_count = bodies.len(); + send_response( + &mut thread_client, + request_id, + serde_json::to_value(DynamicToolCallResponse { + content_items: vec![DynamicToolCallOutputContentItem::InputText { + text: "late-response".to_string(), + }], + success: true, + })?, + ) + .await?; + sleep(Duration::from_millis(250)).await; + assert_eq!(responses_bodies(&server).await?.len(), response_count); + Ok(()) } .await; @@ -200,3 +369,52 @@ async fn send_response( .context("failed to send websocket response frame")?; Ok(()) } + +async fn responses_bodies(server: &MockServer) -> Result> { + let requests = server + .received_requests() + .await + .context("failed to fetch received requests")?; + + requests + .into_iter() + .filter(|req| req.url.path().ends_with("/responses")) + .map(|req| { + req.body_json::() + .context("request body should be JSON") + }) + .collect() +} + +fn function_call_output_payload(body: &Value, call_id: &str) -> Option { + body.get("input") + .and_then(Value::as_array) + .and_then(|items| { + items.iter().find(|item| { + item.get("type").and_then(Value::as_str) == Some("function_call_output") + && item.get("call_id").and_then(Value::as_str) == Some(call_id) + }) + }) + .and_then(|item| item.get("output")) + .cloned() + .and_then(|output| serde_json::from_value(output).ok()) +} + +async fn wait_for_function_call_output( + server: &MockServer, + call_id: &str, +) -> Result<(Vec, FunctionCallOutputPayload)> { + timeout(Duration::from_secs(15), async { + loop { + let bodies = responses_bodies(server).await?; + if let Some(payload) = bodies + .iter() + .find_map(|body| function_call_output_payload(body, call_id)) + { + return Ok((bodies, payload)); + } + sleep(Duration::from_millis(100)).await; + } + }) + .await? +} diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 6df63c7069..23a37e4afc 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -6263,6 +6263,7 @@ async fn built_tools( skills_outcome: Option<&SkillLoadOutcome>, cancellation_token: &CancellationToken, ) -> CodexResult> { + let dynamic_tools = &turn_context.dynamic_tools; let mcp_connection_manager = sess.services.mcp_connection_manager.read().await; let has_mcp_servers = mcp_connection_manager.has_servers(); let mut mcp_tools = mcp_connection_manager @@ -6333,7 +6334,7 @@ async fn built_tools( .collect() }), app_tools, - turn_context.dynamic_tools.as_slice(), + dynamic_tools, ))) } diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 8e139dabc0..475e09d165 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -75,6 +75,7 @@ use tokio::time::sleep; use tracing_opentelemetry::OpenTelemetrySpanExt; use tracing_subscriber::prelude::*; +use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::mcp::CallToolResult as McpCallToolResult; use pretty_assertions::assert_eq; use rmcp::model::JsonObject; @@ -3618,6 +3619,64 @@ impl SessionTask for NeverEndingTask { } } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn update_registered_dynamic_tools_keeps_active_turn_state_frozen() { + let (session, turn_context, _rx) = make_session_and_context_with_rx().await; + let registered_tool = DynamicToolSpec { + name: "provider_tool".to_string(), + description: "Provider tool".to_string(), + input_schema: json!({ + "type": "object", + "properties": {}, + "additionalProperties": false + }), + inject_into_context: true, + provider_owned: true, + }; + let input = vec![UserInput::Text { + text: "hello".to_string(), + text_elements: Vec::new(), + }]; + + session + .spawn_task( + Arc::clone(&turn_context), + input, + NeverEndingTask { + kind: TaskKind::Regular, + listen_to_cancellation_token: false, + }, + ) + .await; + + session + .update_registered_dynamic_tools(vec![registered_tool.clone()]) + .await; + assert!(turn_context.dynamic_tools.is_empty()); + { + let state = session.state.lock().await; + assert_eq!( + state.session_configuration.effective_dynamic_tools(), + vec![registered_tool.clone()] + ); + } + + session.update_registered_dynamic_tools(Vec::new()).await; + assert!(turn_context.dynamic_tools.is_empty()); + + { + let state = session.state.lock().await; + assert!( + state + .session_configuration + .effective_dynamic_tools() + .is_empty() + ); + } + + session.abort_all_tasks(TurnAbortReason::Interrupted).await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[test_log::test] async fn abort_regular_task_emits_turn_aborted_only() { @@ -3976,7 +4035,7 @@ async fn fatal_tool_error_stops_turn_and_reports_error() { .collect(), ), app_tools, - turn_context.dynamic_tools.as_slice(), + &turn_context.dynamic_tools, ); let item = ResponseItem::CustomToolCall { id: None, diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index e2e141d387..4c6b14ccf3 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -79,13 +79,17 @@ pub(crate) struct TurnState { pending_request_permissions: HashMap>, pending_user_input: HashMap>, pending_elicitations: HashMap<(String, RequestId), oneshot::Sender>, - pending_dynamic_tools: HashMap>, + pending_dynamic_tools: HashMap, pending_input: Vec, granted_permissions: Option, pub(crate) tool_calls: u64, pub(crate) token_usage_at_turn_start: TokenUsage, } +pub(crate) struct PendingDynamicToolResponse { + response_tx: oneshot::Sender, +} + impl TurnState { pub(crate) fn insert_pending_approval( &mut self, @@ -165,14 +169,18 @@ impl TurnState { key: String, tx: oneshot::Sender, ) -> Option> { - self.pending_dynamic_tools.insert(key, tx) + self.pending_dynamic_tools + .insert(key, PendingDynamicToolResponse { response_tx: tx }) + .map(|pending| pending.response_tx) } pub(crate) fn remove_pending_dynamic_tool( &mut self, key: &str, ) -> Option> { - self.pending_dynamic_tools.remove(key) + self.pending_dynamic_tools + .remove(key) + .map(|pending| pending.response_tx) } pub(crate) fn push_pending_input(&mut self, input: ResponseInputItem) { diff --git a/codex-rs/core/src/tools/js_repl/mod.rs b/codex-rs/core/src/tools/js_repl/mod.rs index 38bcf53d04..f770de762f 100644 --- a/codex-rs/core/src/tools/js_repl/mod.rs +++ b/codex-rs/core/src/tools/js_repl/mod.rs @@ -1348,7 +1348,7 @@ impl JsReplManager { .collect(), ), None, - exec.turn.dynamic_tools.as_slice(), + &exec.turn.dynamic_tools, ); let payload = diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index 32f1531551..bd4d0e3388 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -168,7 +168,6 @@ impl ToolRouter { } = call; let payload_outputs_custom = matches!(payload, ToolPayload::Custom { .. }); let failure_call_id = call_id.clone(); - if source == ToolCallSource::Direct && turn.tools_config.js_repl_tools_only && !matches!(tool_name.as_str(), "js_repl" | "js_repl_reset") @@ -285,7 +284,7 @@ mod tests { .collect(), ), app_tools, - turn.dynamic_tools.as_slice(), + &turn.dynamic_tools, ); let call = ToolCall { @@ -338,7 +337,7 @@ mod tests { .collect(), ), app_tools, - turn.dynamic_tools.as_slice(), + &turn.dynamic_tools, ); let call = ToolCall { @@ -385,12 +384,7 @@ mod tests { let session = Arc::clone(&session); let turn = Arc::clone(&turn); - let router = ToolRouter::from_config( - &turn.tools_config, - None, - None, - turn.dynamic_tools.as_slice(), - ); + let router = ToolRouter::from_config(&turn.tools_config, None, None, &turn.dynamic_tools); let call = ToolCall { tool_name: "hidden_dynamic_tool".to_string(),