app-server: bound client requests by their total deadline

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!(

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,

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<Result, JSONRPCErrorError>;
 
+#[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<ThreadId>,
+        deadline: Instant,
+    ) -> std::result::Result<(RequestId, oneshot::Receiver<ClientRequestResult>), 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::<OutgoingEnvelope>(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::<OutgoingEnvelope>(4);