fix: bound stalled agent prompt waits

2 files changed, 151 insertions, 18 deletions

diff --git a/src/api/wait.rs b/src/api/wait.rs
index 2f8da2e3..4b9c853e 100644
--- a/src/api/wait.rs
+++ b/src/api/wait.rs
@@ -17,6 +17,8 @@ use crate::api::subscriptions::{match_output, output_match_read_source};
 use crate::api::{ApiRequestSender, EventHub};
 use crate::ipc::LocalStream;
 
+const AGENT_PROMPT_EFFECT_TIMEOUT_MS: u64 = 5_000;
+
 pub(super) fn wait_for_output(
     request_id: String,
     params: crate::api::schema::PaneWaitForOutputParams,
@@ -158,6 +160,7 @@ pub(super) fn wait_for_agent(
             last_event_sequence,
             after_state_change_seq: None,
             accept_transient_status: true,
+            timeout_kind: AgentWaitTimeoutKind::Status,
         },
         stream,
         api_tx,
@@ -219,19 +222,72 @@ pub(super) fn prompt_agent(
         return agent_wait_not_running(request_id).map(Some);
     }
 
-    // The prompt response is produced on the app thread after input is queued.
-    // Requiring a later sequence prevents pre-prompt completion from satisfying this wait.
+    let wait_started = std::time::Instant::now();
     let prompt_state_change_seq = prompted.state_change_seq;
+    let until = agent_wait_statuses(wait.until);
+    let mut initial = prompted;
+    let mut after_state_change_seq = Some(prompt_state_change_seq);
+
+    if initial.agent_status != crate::api::schema::AgentStatus::Working {
+        let effect_timeout_ms = wait
+            .timeout_ms
+            .map_or(AGENT_PROMPT_EFFECT_TIMEOUT_MS, |timeout_ms| {
+                timeout_ms.min(AGENT_PROMPT_EFFECT_TIMEOUT_MS)
+            });
+        let timeout_kind = if wait
+            .timeout_ms
+            .is_some_and(|timeout_ms| timeout_ms <= AGENT_PROMPT_EFFECT_TIMEOUT_MS)
+        {
+            AgentWaitTimeoutKind::Status
+        } else {
+            AgentWaitTimeoutKind::PromptStalled {
+                baseline: prompt_state_change_seq,
+                timeout_ms: effect_timeout_ms,
+            }
+        };
+        let Some(outcome) = wait_for_resolved_agent(
+            request_id.clone(),
+            ResolvedAgentWait {
+                target: target.clone(),
+                until: all_agent_statuses(),
+                timeout_ms: Some(effect_timeout_ms),
+                initial,
+                last_event_sequence,
+                after_state_change_seq,
+                accept_transient_status: false,
+                timeout_kind,
+            },
+            stream,
+            api_tx,
+            event_hub,
+            running,
+        )?
+        else {
+            return Ok(None);
+        };
+        initial = match outcome {
+            AgentWaitOutcome::Matched(agent) => *agent,
+            AgentWaitOutcome::Response(response) => return Ok(Some(response)),
+        };
+        after_state_change_seq = None;
+        if agent_wait_matches(&initial, &until, None) {
+            return agent_prompt_success(request_id, initial).map(Some);
+        }
+    }
+
     let Some(outcome) = wait_for_resolved_agent(
         request_id.clone(),
         ResolvedAgentWait {
             target,
-            until: agent_wait_statuses(wait.until),
-            timeout_ms: wait.timeout_ms,
-            initial: prompted,
+            until,
+            timeout_ms: remaining_timeout_ms(wait.timeout_ms, wait_started),
+            initial,
+            // Replay from before submission so terminal lifecycle events consumed by
+            // the activity gate still terminate this settled-state wait.
             last_event_sequence,
-            after_state_change_seq: Some(prompt_state_change_seq),
+            after_state_change_seq,
             accept_transient_status: false,
+            timeout_kind: AgentWaitTimeoutKind::Status,
         },
         stream,
         api_tx,
@@ -245,11 +301,24 @@ pub(super) fn prompt_agent(
         AgentWaitOutcome::Matched(agent) => *agent,
         AgentWaitOutcome::Response(response) => return Ok(Some(response)),
     };
+    agent_prompt_success(request_id, agent).map(Some)
+}
+
+fn remaining_timeout_ms(total_ms: Option<u64>, started: std::time::Instant) -> Option<u64> {
+    total_ms.map(|total_ms| {
+        let elapsed_ms = started.elapsed().as_millis().min(u128::from(u64::MAX)) as u64;
+        total_ms.saturating_sub(elapsed_ms)
+    })
+}
+
+fn agent_prompt_success(
+    request_id: String,
+    agent: crate::api::schema::AgentInfo,
+) -> std::io::Result<String> {
     serde_json::to_string(&SuccessResponse {
         id: request_id,
         result: ResponseResult::AgentPrompted { agent },
     })
-    .map(Some)
     .map_err(std::io::Error::other)
 }
 
@@ -261,6 +330,13 @@ struct ResolvedAgentWait {
     last_event_sequence: u64,
     after_state_change_seq: Option<u64>,
     accept_transient_status: bool,
+    timeout_kind: AgentWaitTimeoutKind,
+}
+
+#[derive(Clone, Copy)]
+enum AgentWaitTimeoutKind {
+    Status,
+    PromptStalled { baseline: u64, timeout_ms: u64 },
 }
 
 enum AgentWaitOutcome {
@@ -391,21 +467,46 @@ fn wait_for_resolved_agent(
         }
 
         if deadline.is_some_and(|deadline| std::time::Instant::now() >= deadline) {
-            return serde_json::to_string(&ErrorResponse {
-                id: request_id,
-                error: ErrorBody {
-                    code: "timeout".into(),
-                    message: "timed out waiting for agent status".into(),
-                },
-            })
-            .map(AgentWaitOutcome::Response)
-            .map(Some)
-            .map_err(std::io::Error::other);
+            let current = match agent_get(&request_id, &wait.target, api_tx) {
+                Ok(agent) => agent,
+                Err(response) => {
+                    return agent_wait_probe_error(response)
+                        .map(AgentWaitOutcome::Response)
+                        .map(Some);
+                }
+            };
+            if !agent_wait_identity_matches(
+                &current,
+                &expected_terminal_id,
+                expected_name.as_deref(),
+                expected_agent.as_deref(),
+            ) {
+                return agent_wait_not_running(request_id)
+                    .map(AgentWaitOutcome::Response)
+                    .map(Some);
+            }
+            if agent_wait_matches(&current, &wait.until, wait.after_state_change_seq) {
+                return Ok(Some(AgentWaitOutcome::Matched(Box::new(current))));
+            }
+            return agent_wait_timeout(request_id, wait.timeout_kind, &current)
+                .map(AgentWaitOutcome::Response)
+                .map(Some);
         }
         std::thread::sleep(CONNECTION_POLL_INTERVAL);
     }
 }
 
+fn all_agent_statuses() -> Vec<crate::api::schema::AgentStatus> {
+    // Keep this exhaustive: every status is evidence that the sequence advanced.
+    vec![
+        crate::api::schema::AgentStatus::Idle,
+        crate::api::schema::AgentStatus::Working,
+        crate::api::schema::AgentStatus::Blocked,
+        crate::api::schema::AgentStatus::Done,
+        crate::api::schema::AgentStatus::Unknown,
+    ]
+}
+
 fn agent_wait_statuses(
     until: Vec<crate::api::schema::AgentStatus>,
 ) -> Vec<crate::api::schema::AgentStatus> {
@@ -506,6 +607,38 @@ fn agent_wait_success(
     .map_err(std::io::Error::other)
 }
 
+fn agent_wait_timeout(
+    request_id: String,
+    kind: AgentWaitTimeoutKind,
+    current: &crate::api::schema::AgentInfo,
+) -> std::io::Result<String> {
+    let (code, message) = match kind {
+        AgentWaitTimeoutKind::Status => {
+            ("timeout", "timed out waiting for agent status".to_string())
+        }
+        AgentWaitTimeoutKind::PromptStalled {
+            baseline,
+            timeout_ms,
+        } => {
+            let status = format!("{:?}", current.agent_status).to_ascii_lowercase();
+            (
+                "agent_prompt_stalled",
+                format!(
+                    "agent prompt produced no observed state change within {timeout_ms} ms; status is {status} and state_change_seq remained {baseline}"
+                ),
+            )
+        }
+    };
+    serde_json::to_string(&ErrorResponse {
+        id: request_id,
+        error: ErrorBody {
+            code: code.into(),
+            message,
+        },
+    })
+    .map_err(std::io::Error::other)
+}
+
 fn agent_wait_not_running(request_id: String) -> std::io::Result<String> {
     serde_json::to_string(&ErrorResponse {
         id: request_id,
diff --git a/src/cli/spec.rs b/src/cli/spec.rs
index 7b2977ac..7fa08550 100644
--- a/src/cli/spec.rs
+++ b/src/cli/spec.rs
@@ -355,7 +355,7 @@ fn agent_command() -> Command {
                         .help("Fail after this many milliseconds"),
                 )
                 .after_help(
-                    "Without --until, --wait matches the first idle, done, or blocked state observed after submission. It does not track turns: if the agent is already working, that active turn's completion may match. Use --until unknown explicitly when needed. Without --timeout, it waits indefinitely.",
+                    "When submission starts from a non-working state, --wait first requires an observed state change within 5000ms; otherwise it returns agent_prompt_stalled. A shorter --timeout returns timeout instead. It then matches idle, done, or blocked by default, or any exact --until state. It does not track turns: if the agent is already working, that active turn's completion may match. Without --timeout, the settled-state wait is indefinite.",
                 ),
         )
         .subcommand(