1 files changed, 32 insertions, 38 deletions

diff --git a/src/api/wait.rs b/src/api/wait.rs
index 0c47fbff..9403c4f2 100644
--- a/src/api/wait.rs
+++ b/src/api/wait.rs
@@ -193,7 +193,6 @@ pub(super) fn prompt_agent(
         )));
     };
 
-    let last_event_sequence = event_hub.current_sequence();
     let before_prompt = match agent_get(&request_id, &params.target, api_tx) {
         Ok(agent) => agent,
         Err(response) => {
@@ -202,7 +201,11 @@ pub(super) fn prompt_agent(
                 .map_err(std::io::Error::other);
         }
     };
+    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();
+    let last_event_sequence = event_hub.current_sequence();
     let prompt_response = dispatch_to_app_with_timeout(
         Request {
             id: request_id.clone(),
@@ -223,39 +226,38 @@ pub(super) fn prompt_agent(
         return agent_wait_not_running(request_id).map(Some);
     }
 
-    let wait_started = std::time::Instant::now();
+    let prompt_activity_observed = prompt_started_working
+        || matches!(
+            prompted.agent_status,
+            crate::api::schema::AgentStatus::Working | crate::api::schema::AgentStatus::Blocked
+        );
     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,
+
+    if !prompt_activity_observed {
+        let remaining_timeout_ms = remaining_timeout_ms(wait.timeout_ms, wait_started);
+        let (effect_timeout_ms, timeout_kind) = match remaining_timeout_ms {
+            Some(timeout_ms) if timeout_ms <= AGENT_PROMPT_EFFECT_TIMEOUT_MS => {
+                (timeout_ms, AgentWaitTimeoutKind::Status)
             }
+            _ => (
+                AGENT_PROMPT_EFFECT_TIMEOUT_MS,
+                AgentWaitTimeoutKind::PromptStalled {
+                    timeout_ms: AGENT_PROMPT_EFFECT_TIMEOUT_MS,
+                },
+            ),
         };
         let Some(outcome) = wait_for_resolved_agent(
             request_id.clone(),
             ResolvedAgentWait {
                 target: target.clone(),
-                until: all_agent_statuses(),
+                until: prompt_activity_statuses(),
                 timeout_ms: Some(effect_timeout_ms),
                 initial,
                 last_event_sequence,
-                after_state_change_seq,
-                accept_transient_status: false,
+                after_state_change_seq: Some(prompt_state_change_seq),
+                accept_transient_status: true,
                 timeout_kind,
             },
             stream,
@@ -270,10 +272,9 @@ pub(super) fn prompt_agent(
             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);
-        }
+    }
+    if agent_wait_matches(&initial, &until, None) {
+        return agent_prompt_success(request_id, initial).map(Some);
     }
 
     let Some(outcome) = wait_for_resolved_agent(
@@ -286,7 +287,7 @@ pub(super) fn prompt_agent(
             // 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,
+            after_state_change_seq: None,
             accept_transient_status: false,
             timeout_kind: AgentWaitTimeoutKind::Status,
         },
@@ -337,7 +338,7 @@ struct ResolvedAgentWait {
 #[derive(Clone, Copy)]
 enum AgentWaitTimeoutKind {
     Status,
-    PromptStalled { baseline: u64, timeout_ms: u64 },
+    PromptStalled { timeout_ms: u64 },
 }
 
 enum AgentWaitOutcome {
@@ -497,14 +498,10 @@ fn wait_for_resolved_agent(
     }
 }
 
-fn all_agent_statuses() -> Vec<crate::api::schema::AgentStatus> {
-    // Keep this exhaustive: every status is evidence that the sequence advanced.
+fn prompt_activity_statuses() -> Vec<crate::api::schema::AgentStatus> {
     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,
     ]
 }
 
@@ -617,15 +614,12 @@ fn agent_wait_timeout(
         AgentWaitTimeoutKind::Status => {
             ("timeout", "timed out waiting for agent status".to_string())
         }
-        AgentWaitTimeoutKind::PromptStalled {
-            baseline,
-            timeout_ms,
-        } => {
+        AgentWaitTimeoutKind::PromptStalled { 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}"
+                    "agent prompt produced no observed working or blocked state within {timeout_ms} ms; current status is {status}"
                 ),
             )
         }