fix: make Windows Codex prompt delay size-aware (#3552)

refs #3187

1 files changed, 53 insertions, 11 deletions

diff --git a/src/api/wait.rs b/src/api/wait.rs
index 9403c4f2..541779f8 100644
--- a/src/api/wait.rs
+++ b/src/api/wait.rs
@@ -9,8 +9,8 @@ use crate::api::schema::{
     SubscriptionEventEnvelope, SuccessResponse,
 };
 use crate::api::server::{
-    dispatch_to_app_with_timeout, should_stop_connection, APP_RESPONSE_TIMEOUT,
-    CONNECTION_POLL_INTERVAL,
+    dispatch_to_app_with_caller_timeout, dispatch_to_app_with_timeout, should_stop_connection,
+    APP_RESPONSE_TIMEOUT, CONNECTION_POLL_INTERVAL,
 };
 use crate::api::subscriptions::ActiveSubscription;
 use crate::api::subscriptions::{match_output, output_match_read_source};
@@ -176,7 +176,7 @@ pub(super) fn wait_for_agent(
 
 pub(super) fn prompt_agent(
     request_id: String,
-    params: crate::api::schema::AgentPromptParams,
+    mut params: crate::api::schema::AgentPromptParams,
     stream: &mut LocalStream,
     api_tx: &ApiRequestSender,
     event_hub: &EventHub,
@@ -193,7 +193,14 @@ pub(super) fn prompt_agent(
         )));
     };
 
-    let before_prompt = match agent_get(&request_id, &params.target, api_tx) {
+    let wait_started = std::time::Instant::now();
+    let before_prompt = match agent_get_for_prompt(
+        &request_id,
+        &params.target,
+        api_tx,
+        wait.timeout_ms,
+        wait_started,
+    ) {
         Ok(agent) => agent,
         Err(response) => {
             return serde_json::to_string(&response)
@@ -204,16 +211,24 @@ pub(super) fn prompt_agent(
     let prompt_started_working =
         before_prompt.agent_status == crate::api::schema::AgentStatus::Working;
     let target = params.target.clone();
-    let wait_started = std::time::Instant::now();
+    if let Some(prompt_wait) = params.wait.as_mut() {
+        prompt_wait.submission_deadline = wait
+            .timeout_ms
+            .map(|timeout_ms| wait_started + std::time::Duration::from_millis(timeout_ms));
+    }
     let last_event_sequence = event_hub.current_sequence();
-    let prompt_response = dispatch_to_app_with_timeout(
-        Request {
-            id: request_id.clone(),
-            method: Method::AgentPrompt(params),
-        },
+    let prompt_request = Request {
+        id: request_id.clone(),
+        method: Method::AgentPrompt(params),
+    };
+    #[cfg(windows)]
+    let prompt_response = dispatch_to_app_with_caller_timeout(
+        prompt_request,
         api_tx,
-        None,
+        remaining_timeout_ms(wait.timeout_ms, wait_started).map(std::time::Duration::from_millis),
     );
+    #[cfg(not(windows))]
+    let prompt_response = dispatch_to_app_with_timeout(prompt_request, api_tx, None);
     let Ok(prompted) = agent_from_response(&request_id, &prompt_response) else {
         return Ok(Some(prompt_response));
     };
@@ -561,6 +576,33 @@ fn agent_get(
     agent_from_response(request_id, &response)
 }
 
+fn agent_get_for_prompt(
+    request_id: &str,
+    target: &str,
+    api_tx: &ApiRequestSender,
+    total_timeout_ms: Option<u64>,
+    started: std::time::Instant,
+) -> Result<crate::api::schema::AgentInfo, ErrorResponse> {
+    let request = Request {
+        id: format!("{request_id}:agent"),
+        method: Method::AgentGet(crate::api::schema::AgentTarget {
+            target: target.to_string(),
+        }),
+    };
+    let remaining_ms = remaining_timeout_ms(total_timeout_ms, started);
+    let response = match remaining_ms {
+        Some(timeout_ms) if timeout_ms <= APP_RESPONSE_TIMEOUT.as_millis() as u64 => {
+            dispatch_to_app_with_caller_timeout(
+                request,
+                api_tx,
+                Some(std::time::Duration::from_millis(timeout_ms)),
+            )
+        }
+        _ => dispatch_to_app_with_timeout(request, api_tx, Some(APP_RESPONSE_TIMEOUT)),
+    };
+    agent_from_response(request_id, &response)
+}
+
 fn agent_from_response(
     request_id: &str,
     response: &str,