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

refs #3187

4 files changed, 102 insertions, 14 deletions

diff --git a/src/api/schema/agents.rs b/src/api/schema/agents.rs
index 9e1e7b15..ee496d5d 100644
--- a/src/api/schema/agents.rs
+++ b/src/api/schema/agents.rs
@@ -37,6 +37,9 @@ pub struct AgentPromptWaitOptions {
     pub until: Vec<AgentStatus>,
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub timeout_ms: Option<u64>,
+    #[serde(skip)]
+    #[schemars(skip)]
+    pub(crate) submission_deadline: Option<std::time::Instant>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
diff --git a/src/api/schema/tests.rs b/src/api/schema/tests.rs
index 06529f5e..693f0e85 100644
--- a/src/api/schema/tests.rs
+++ b/src/api/schema/tests.rs
@@ -133,6 +133,7 @@ fn agent_start_and_prompt_requests_round_trip() {
             wait: Some(AgentPromptWaitOptions {
                 until: vec![AgentStatus::Idle, AgentStatus::Done],
                 timeout_ms: Some(120_000),
+                submission_deadline: None,
             }),
         }),
     };
diff --git a/src/api/server.rs b/src/api/server.rs
index 64beec33..655d89c5 100644
--- a/src/api/server.rs
+++ b/src/api/server.rs
@@ -369,7 +369,7 @@ fn handle_request(
         );
     }
 
-    dispatch_to_app(request, api_tx, None, response_write_complete, None)
+    dispatch_to_app(request, api_tx, None, response_write_complete, None, None)
 }
 
 pub(crate) fn api_method_name(method: &Method) -> &'static str {
@@ -798,7 +798,22 @@ pub(super) fn dispatch_to_app_with_timeout(
     api_tx: &ApiRequestSender,
     timeout: Option<Duration>,
 ) -> String {
-    dispatch_to_app(request, api_tx, timeout, None, None)
+    dispatch_to_app(request, api_tx, timeout, None, None, None)
+}
+
+pub(super) fn dispatch_to_app_with_caller_timeout(
+    request: Request,
+    api_tx: &ApiRequestSender,
+    timeout: Option<Duration>,
+) -> String {
+    dispatch_to_app(
+        request,
+        api_tx,
+        timeout,
+        None,
+        None,
+        Some(("timeout", "timed out waiting for agent status")),
+    )
 }
 
 pub(super) fn dispatch_stream_open(
@@ -807,7 +822,7 @@ pub(super) fn dispatch_stream_open(
     timeout: Duration,
     active: Arc<AtomicBool>,
 ) -> String {
-    dispatch_to_app(request, api_tx, Some(timeout), None, Some(active))
+    dispatch_to_app(request, api_tx, Some(timeout), None, Some(active), None)
 }
 
 pub(super) fn dispatch_stream_frame(
@@ -821,6 +836,7 @@ pub(super) fn dispatch_stream_frame(
         Some(crate::app::pane_graphics::DIRECT_OUTER_TIMEOUT),
         None,
         Some(active),
+        None,
     )
 }
 
@@ -830,6 +846,7 @@ fn dispatch_to_app(
     timeout: Option<Duration>,
     response_write_complete: Option<std::sync::mpsc::Receiver<()>>,
     stream_active: Option<Arc<AtomicBool>>,
+    timeout_response: Option<(&str, &str)>,
 ) -> String {
     let request_id = request.id.clone();
     let request_active = stream_active.clone();
@@ -875,6 +892,11 @@ fn dispatch_to_app(
             if let Some(active) = request_active {
                 active.store(false, Ordering::Release);
             }
+            if err.kind() == std::io::ErrorKind::TimedOut {
+                if let Some((code, message)) = timeout_response {
+                    return error_response_json(request_id, code, message.into());
+                }
+            }
             error_response_json(
                 request_id,
                 "server_unavailable",
@@ -884,6 +906,26 @@ fn dispatch_to_app(
     }
 }
 
+#[cfg(test)]
+#[test]
+fn caller_timeout_dispatch_uses_timeout_error() {
+    let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
+    let response = dispatch_to_app_with_caller_timeout(
+        Request {
+            id: "prompt-timeout".into(),
+            method: Method::AgentPrompt(crate::api::schema::AgentPromptParams {
+                target: "reviewer".into(),
+                text: "review this".into(),
+                wait: None,
+            }),
+        },
+        &tx,
+        Some(Duration::ZERO),
+    );
+    let error: ErrorResponse = serde_json::from_str(&response).unwrap();
+    assert_eq!(error.error.code, "timeout");
+}
+
 fn error_response_json(id: String, code: &str, message: String) -> String {
     serde_json::to_string(&ErrorResponse {
         id,
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,