diff --git a/codex-rs/app-server/src/attestation.rs b/codex-rs/app-server/src/attestation.rs index 206c38cee17e..c8203ec92b20 100644 --- a/codex-rs/app-server/src/attestation.rs +++ b/codex-rs/app-server/src/attestation.rs @@ -10,10 +10,12 @@ use codex_core::AttestationProvider; use codex_core::GenerateAttestationFuture; use serde::Serialize; use tokio::time::Duration; -use tokio::time::timeout; +use tokio::time::Instant; +use tokio::time::timeout_at; use tracing::warn; use crate::outgoing_message::OutgoingMessageSender; +use crate::outgoing_message::RequestSendError; use crate::thread_state::ThreadStateManager; const ATTESTATION_GENERATE_TIMEOUT: Duration = Duration::from_millis(100); @@ -70,16 +72,39 @@ async fn request_attestation_header_value_with_timeout( .first_attestation_capable_connection_for_thread(thread_id) .await?; - let connection_ids = [connection_id]; - let (request_id, rx) = outgoing - .send_request_to_connections( - Some(&connection_ids), + let deadline = Instant::now() + timeout_duration; + let (request_id, rx) = match outgoing + .send_request_to_connection_with_deadline( + connection_id, ServerRequestPayload::AttestationGenerate(AttestationGenerateParams {}), /*thread_id*/ None, + deadline, ) - .await; + .await + { + Ok(request) => request, + Err(RequestSendError::TimedOut) => { + warn!( + timeout_seconds = timeout_duration.as_secs_f64(), + "attestation generation request timed out before it was queued" + ); + return app_server_attestation_header_value( + AppServerAttestationStatus::Timeout, + /*token*/ None, + ); + } + Err(RequestSendError::Closed) => { + warn!( + "attestation generation request could not be queued because the client channel closed" + ); + return app_server_attestation_header_value( + AppServerAttestationStatus::RequestFailed, + /*token*/ None, + ); + } + }; - let result = match timeout(timeout_duration, rx).await { + let result = match timeout_at(deadline, rx).await { Ok(Ok(Ok(result))) => result, Ok(Ok(Err(err))) => { warn!( diff --git a/codex-rs/app-server/src/current_time.rs b/codex-rs/app-server/src/current_time.rs index 57be618af383..57d0821a9608 100644 --- a/codex-rs/app-server/src/current_time.rs +++ b/codex-rs/app-server/src/current_time.rs @@ -20,6 +20,7 @@ use tokio::time::timeout_at; use crate::outgoing_message::ConnectionId; use crate::outgoing_message::OutgoingMessageSender; +use crate::outgoing_message::RequestSendError; use crate::thread_state::ThreadStateManager; const CURRENT_TIME_REQUEST_TIMEOUT: Duration = Duration::from_secs(10); @@ -103,16 +104,28 @@ async fn request_current_time( .subscribed_connection_ids(thread_id) .await; let connection_id = require_single_current_time_connection(&connection_ids)?; - let connection_ids = [connection_id]; - let (request_id, rx) = outgoing - .send_request_to_connections( - Some(&connection_ids), + let (request_id, rx) = match outgoing + .send_request_to_connection_with_deadline( + connection_id, ServerRequestPayload::CurrentTimeRead(CurrentTimeReadParams { thread_id: thread_id.to_string(), }), /*thread_id*/ None, + deadline, ) - .await; + .await + { + Ok(request) => request, + Err(RequestSendError::TimedOut) => { + bail!( + "current-time request timed out after {}s", + CURRENT_TIME_REQUEST_TIMEOUT.as_secs() + ); + } + Err(RequestSendError::Closed) => { + bail!("current-time request could not be queued because the client channel closed"); + } + }; let result = match timeout_at(deadline, rx).await { Ok(Ok(Ok(result))) => result, diff --git a/codex-rs/app-server/src/outgoing_message.rs b/codex-rs/app-server/src/outgoing_message.rs index 16d5c2ff9d1d..029d56d8331f 100644 --- a/codex-rs/app-server/src/outgoing_message.rs +++ b/codex-rs/app-server/src/outgoing_message.rs @@ -21,6 +21,8 @@ use codex_protocol::request_permissions::RequestPermissionsResponse; use tokio::sync::Mutex; use tokio::sync::mpsc; use tokio::sync::oneshot; +use tokio::time::Instant; +use tokio::time::timeout_at; use tracing::Instrument; use tracing::Span; use tracing::warn; @@ -38,6 +40,12 @@ use codex_protocol::account::PlanType; pub(crate) type ClientRequestResult = std::result::Result; +#[derive(Debug, Eq, PartialEq)] +pub(crate) enum RequestSendError { + TimedOut, + Closed, +} + /// Stable identifier for a client request scoped to a transport connection. #[derive(Clone, Debug, Eq, Hash, PartialEq)] pub(crate) struct ConnectionRequestId { @@ -349,6 +357,53 @@ impl OutgoingMessageSender { (outgoing_message_id, rx_approve) } + pub(crate) async fn send_request_to_connection_with_deadline( + &self, + connection_id: ConnectionId, + request: ServerRequestPayload, + thread_id: Option, + deadline: Instant, + ) -> std::result::Result<(RequestId, oneshot::Receiver), RequestSendError> + { + let id = self.next_request_id(); + let request = request.request_with_id(id.clone()); + let outgoing_message = OutgoingMessage::Request(request.clone()); + let permit = match timeout_at(deadline, self.sender.reserve()).await { + Ok(Ok(permit)) => permit, + Ok(Err(_)) => return Err(RequestSendError::Closed), + Err(_) => return Err(RequestSendError::TimedOut), + }; + + let (callback_sender, callback_receiver) = oneshot::channel(); + let mut request_id_to_callback = + match timeout_at(deadline, self.request_id_to_callback.lock()).await { + Ok(request_id_to_callback) => request_id_to_callback, + Err(_) => return Err(RequestSendError::TimedOut), + }; + if Instant::now() >= deadline { + return Err(RequestSendError::TimedOut); + } + request_id_to_callback.insert( + id.clone(), + PendingCallbackEntry { + callback: callback_sender, + thread_id, + request: request.clone(), + }, + ); + drop(request_id_to_callback); + + permit.send(OutgoingEnvelope::ToConnection { + connection_id, + message: outgoing_message, + write_complete_tx: None, + }); + self.analytics_events_client + .track_server_request(connection_id.0, request); + + Ok((id, callback_receiver)) + } + pub(crate) async fn replay_requests_to_connection_for_thread( &self, connection_id: ConnectionId, @@ -756,6 +811,7 @@ mod tests { use pretty_assertions::assert_eq; use serde_json::json; use std::sync::Arc; + use tokio::time::Instant; use tokio::time::timeout; use uuid::Uuid; @@ -1086,6 +1142,43 @@ mod tests { } } + #[tokio::test] + async fn deadline_expires_while_waiting_for_outgoing_queue_capacity() { + let (tx, _rx) = mpsc::channel::(1); + tx.try_send(OutgoingEnvelope::Broadcast { + message: OutgoingMessage::AppServerNotification(ServerNotification::ConfigWarning( + ConfigWarningNotification { + summary: "queued".to_string(), + details: None, + path: None, + range: None, + }, + )), + }) + .expect("the queue should accept the filler message"); + let outgoing = + OutgoingMessageSender::new(tx, codex_analytics::AnalyticsEventsClient::disabled()); + + let error = outgoing + .send_request_to_connection_with_deadline( + ConnectionId(42), + ServerRequestPayload::ApplyPatchApproval(ApplyPatchApprovalParams { + conversation_id: ThreadId::new(), + call_id: "call-id".to_string(), + file_changes: HashMap::new(), + reason: None, + grant_root: None, + }), + /*thread_id*/ None, + Instant::now() + Duration::from_millis(10), + ) + .await + .expect_err("the full queue should make the request time out"); + + assert_eq!(error, RequestSendError::TimedOut); + assert!(outgoing.request_id_to_callback.lock().await.is_empty()); + } + #[tokio::test] async fn send_response_clears_registered_request_context() { let (tx, _rx) = mpsc::channel::(4);