feat: make agent automation a first-class cli primitive

9 files changed, 565 insertions, 55 deletions

diff --git a/src/api/mod.rs b/src/api/mod.rs
index 618fa4f9..badf1707 100644
--- a/src/api/mod.rs
+++ b/src/api/mod.rs
@@ -45,6 +45,7 @@ pub(crate) fn request_changes_ui(request: &Request) -> bool {
             | Method::AgentFocus(_)
             | Method::AgentStart(_)
             | Method::AgentPrompt(_)
+            | Method::AgentSendKeys(_)
             | Method::PaneSplit(_)
             | Method::PaneSwap(_)
             | Method::PaneMove(_)
diff --git a/src/api/schema.rs b/src/api/schema.rs
index 41f372e3..efaee54c 100644
--- a/src/api/schema.rs
+++ b/src/api/schema.rs
@@ -109,8 +109,8 @@ pub enum Method {
     AgentRead(AgentReadParams),
     #[serde(rename = "agent.explain")]
     AgentExplain(AgentTarget),
-    #[serde(rename = "agent.send")]
-    AgentSend(AgentSendParams),
+    #[serde(rename = "agent.send_keys")]
+    AgentSendKeys(AgentSendKeysParams),
     #[serde(rename = "agent.rename")]
     AgentRename(AgentRenameParams),
     #[serde(rename = "agent.focus")]
@@ -119,6 +119,8 @@ pub enum Method {
     AgentStart(AgentStartParams),
     #[serde(rename = "agent.prompt")]
     AgentPrompt(AgentPromptParams),
+    #[serde(rename = "agent.wait")]
+    AgentWait(AgentWaitParams),
     #[serde(rename = "pane.split")]
     PaneSplit(PaneSplitParams),
     #[serde(rename = "pane.swap")]
diff --git a/src/api/schema/agents.rs b/src/api/schema/agents.rs
index 8592e56b..59949293 100644
--- a/src/api/schema/agents.rs
+++ b/src/api/schema/agents.rs
@@ -17,9 +17,26 @@ pub struct AgentReadParams {
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
-pub struct AgentSendParams {
+pub struct AgentSendKeysParams {
     pub target: String,
-    pub text: String,
+    pub keys: Vec<String>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
+pub struct AgentWaitParams {
+    pub target: String,
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub until: Vec<AgentStatus>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub timeout_ms: Option<u64>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
+pub struct AgentPromptWaitOptions {
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub until: Vec<AgentStatus>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub timeout_ms: Option<u64>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
@@ -45,6 +62,8 @@ pub struct AgentStartParams {
 pub struct AgentPromptParams {
     pub target: String,
     pub text: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub wait: Option<AgentPromptWaitOptions>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
diff --git a/src/api/schema/events.rs b/src/api/schema/events.rs
index 25ccc8ec..50b5d522 100644
--- a/src/api/schema/events.rs
+++ b/src/api/schema/events.rs
@@ -520,6 +520,10 @@ pub enum EventData {
         workspace_id: String,
         #[serde(default, skip_serializing_if = "Option::is_none")]
         agent: Option<String>,
+        #[serde(default, skip_serializing_if = "std::ops::Not::not")]
+        released: bool,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        final_status: Option<AgentStatus>,
     },
     PaneAgentStatusChanged {
         pane_id: String,
diff --git a/src/api/schema/integrations.rs b/src/api/schema/integrations.rs
index f20b9095..5a267ff1 100644
--- a/src/api/schema/integrations.rs
+++ b/src/api/schema/integrations.rs
@@ -29,6 +29,25 @@ pub enum IntegrationTarget {
     Mastracode,
 }
 
+impl IntegrationTarget {
+    pub(crate) const ALL: [Self; 14] = [
+        Self::Pi,
+        Self::Omp,
+        Self::Claude,
+        Self::Codex,
+        Self::Copilot,
+        Self::Devin,
+        Self::Droid,
+        Self::Kimi,
+        Self::Opencode,
+        Self::Kilo,
+        Self::Hermes,
+        Self::Qodercli,
+        Self::Cursor,
+        Self::Mastracode,
+    ];
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
 pub struct IntegrationInstallResult {
     pub messages: Vec<String>,
diff --git a/src/api/schema/response.rs b/src/api/schema/response.rs
index ad058191..e1f47efd 100644
--- a/src/api/schema/response.rs
+++ b/src/api/schema/response.rs
@@ -103,7 +103,6 @@ pub enum ResponseResult {
     },
     AgentPrompted {
         agent: AgentInfo,
-        baseline_state_change_seq: u64,
     },
     AgentList {
         agents: Vec<AgentInfo>,
diff --git a/src/api/schema/tests.rs b/src/api/schema/tests.rs
index 232c2b73..236d4f12 100644
--- a/src/api/schema/tests.rs
+++ b/src/api/schema/tests.rs
@@ -86,6 +86,7 @@ fn agent_start_and_prompt_requests_round_trip() {
         method: Method::AgentPrompt(AgentPromptParams {
             target: "reviewer".into(),
             text: "review this".into(),
+            wait: None,
         }),
     };
     let prompt_json = serde_json::to_value(&prompt).unwrap();
@@ -94,6 +95,31 @@ fn agent_start_and_prompt_requests_round_trip() {
         serde_json::from_value::<Request>(prompt_json).unwrap(),
         prompt
     );
+
+    let prompt_and_wait = Request {
+        id: "prompt-and-wait".into(),
+        method: Method::AgentPrompt(AgentPromptParams {
+            target: "reviewer".into(),
+            text: "review this".into(),
+            wait: Some(AgentPromptWaitOptions {
+                until: vec![AgentStatus::Idle, AgentStatus::Done],
+                timeout_ms: Some(120_000),
+            }),
+        }),
+    };
+    let prompt_and_wait_json = serde_json::to_value(&prompt_and_wait).unwrap();
+    assert_eq!(
+        prompt_and_wait_json["params"]["wait"]["until"],
+        serde_json::json!(["idle", "done"])
+    );
+    assert_eq!(
+        prompt_and_wait_json["params"]["wait"]["timeout_ms"],
+        120_000
+    );
+    assert_eq!(
+        serde_json::from_value::<Request>(prompt_and_wait_json).unwrap(),
+        prompt_and_wait
+    );
 }
 
 #[test]
diff --git a/src/api/server.rs b/src/api/server.rs
index ce538853..0a36e2fa 100644
--- a/src/api/server.rs
+++ b/src/api/server.rs
@@ -14,7 +14,7 @@ use crate::api::schema::{
     ErrorBody, ErrorResponse, Method, Request, ResponseResult, ServerCapabilities, SuccessResponse,
 };
 use crate::api::subscriptions::ActiveSubscription;
-use crate::api::wait::{wait_for_event, wait_for_output};
+use crate::api::wait::{prompt_agent, wait_for_agent, wait_for_event, wait_for_output};
 use crate::api::{request_changes_ui, socket_path, ApiRequestMessage, ApiRequestSender, EventHub};
 use crate::ipc::{
     bind_local_listener, is_connection_closed_error, local_stream_peer_closed,
@@ -218,62 +218,42 @@ fn handle_connection(
             result
         }
         Method::EventsWait(params) => {
-            let Some(response) = wait_for_event(
+            let response = wait_for_event(
                 request_id.clone(),
                 params,
                 &mut stream,
                 api_tx,
                 event_hub,
                 running,
-            )?
-            else {
-                crate::logging::api_request_completed(
-                    &request_id,
-                    method,
-                    "client_disconnected",
-                    changes_ui,
-                );
-                return Ok(());
-            };
-            let result = write_text_line_allow_disconnect(&mut stream, &response);
-            match &result {
-                Ok(()) => crate::logging::api_request_completed(
-                    &request_id,
-                    method,
-                    api_response_outcome(&response),
-                    changes_ui,
-                ),
-                Err(err) => {
-                    crate::logging::api_request_failed(&request_id, method, &err.to_string())
-                }
-            }
-            result
+            )?;
+            finish_wait_response(&mut stream, response, &request_id, method, changes_ui)
+        }
+        Method::AgentPrompt(params) => {
+            let response = prompt_agent(
+                request_id.clone(),
+                params,
+                &mut stream,
+                api_tx,
+                event_hub,
+                running,
+            )?;
+            finish_wait_response(&mut stream, response, &request_id, method, changes_ui)
+        }
+        Method::AgentWait(params) => {
+            let response = wait_for_agent(
+                request_id.clone(),
+                params,
+                &mut stream,
+                api_tx,
+                event_hub,
+                running,
+            )?;
+            finish_wait_response(&mut stream, response, &request_id, method, changes_ui)
         }
         Method::PaneWaitForOutput(params) => {
-            let Some(response) =
-                wait_for_output(request_id.clone(), params, &mut stream, api_tx, running)?
-            else {
-                crate::logging::api_request_completed(
-                    &request_id,
-                    method,
-                    "client_disconnected",
-                    changes_ui,
-                );
-                return Ok(());
-            };
-            let result = write_text_line_allow_disconnect(&mut stream, &response);
-            match &result {
-                Ok(()) => crate::logging::api_request_completed(
-                    &request_id,
-                    method,
-                    api_response_outcome(&response),
-                    changes_ui,
-                ),
-                Err(err) => {
-                    crate::logging::api_request_failed(&request_id, method, &err.to_string())
-                }
-            }
-            result
+            let response =
+                wait_for_output(request_id.clone(), params, &mut stream, api_tx, running)?;
+            finish_wait_response(&mut stream, response, &request_id, method, changes_ui)
         }
         method_body => {
             let (response_write_tx, response_write_rx) = std::sync::mpsc::channel();
@@ -304,6 +284,35 @@ fn handle_connection(
     }
 }
 
+fn finish_wait_response(
+    stream: &mut LocalStream,
+    response: Option<String>,
+    request_id: &str,
+    method: &'static str,
+    changes_ui: bool,
+) -> std::io::Result<()> {
+    let Some(response) = response else {
+        crate::logging::api_request_completed(
+            request_id,
+            method,
+            "client_disconnected",
+            changes_ui,
+        );
+        return Ok(());
+    };
+    let result = write_text_line_allow_disconnect(stream, &response);
+    match &result {
+        Ok(()) => crate::logging::api_request_completed(
+            request_id,
+            method,
+            api_response_outcome(&response),
+            changes_ui,
+        ),
+        Err(err) => crate::logging::api_request_failed(request_id, method, &err.to_string()),
+    }
+    result
+}
+
 fn handle_request(
     request: Request,
     api_tx: &ApiRequestSender,
@@ -367,11 +376,12 @@ fn api_method_name(method: &Method) -> &'static str {
         Method::AgentGet(_) => "agent.get",
         Method::AgentRead(_) => "agent.read",
         Method::AgentExplain(_) => "agent.explain",
-        Method::AgentSend(_) => "agent.send",
+        Method::AgentSendKeys(_) => "agent.send_keys",
         Method::AgentRename(_) => "agent.rename",
         Method::AgentFocus(_) => "agent.focus",
         Method::AgentStart(_) => "agent.start",
         Method::AgentPrompt(_) => "agent.prompt",
+        Method::AgentWait(_) => "agent.wait",
         Method::PaneSplit(_) => "pane.split",
         Method::PaneSwap(_) => "pane.swap",
         Method::PaneMove(_) => "pane.move",
diff --git a/src/api/wait.rs b/src/api/wait.rs
index d556fa96..2f8da2e3 100644
--- a/src/api/wait.rs
+++ b/src/api/wait.rs
@@ -126,6 +126,404 @@ pub(super) fn wait_for_output(
     }
 }
 
+pub(super) fn wait_for_agent(
+    request_id: String,
+    params: crate::api::schema::AgentWaitParams,
+    stream: &mut LocalStream,
+    api_tx: &ApiRequestSender,
+    event_hub: &EventHub,
+    running: &Arc<AtomicBool>,
+) -> std::io::Result<Option<String>> {
+    let last_event_sequence = event_hub.current_sequence();
+    let initial = match agent_get(&request_id, &params.target, api_tx) {
+        Ok(agent) => agent,
+        Err(response) => {
+            return serde_json::to_string(&response)
+                .map(Some)
+                .map_err(std::io::Error::other);
+        }
+    };
+    let until = agent_wait_statuses(params.until);
+    if agent_wait_matches(&initial, &until, None) {
+        return agent_wait_success(request_id, initial).map(Some);
+    }
+
+    match wait_for_resolved_agent(
+        request_id.clone(),
+        ResolvedAgentWait {
+            target: params.target,
+            until,
+            timeout_ms: params.timeout_ms,
+            initial,
+            last_event_sequence,
+            after_state_change_seq: None,
+            accept_transient_status: true,
+        },
+        stream,
+        api_tx,
+        event_hub,
+        running,
+    )? {
+        Some(AgentWaitOutcome::Matched(agent)) => agent_wait_success(request_id, *agent).map(Some),
+        Some(AgentWaitOutcome::Response(response)) => Ok(Some(response)),
+        None => Ok(None),
+    }
+}
+
+pub(super) fn prompt_agent(
+    request_id: String,
+    params: crate::api::schema::AgentPromptParams,
+    stream: &mut LocalStream,
+    api_tx: &ApiRequestSender,
+    event_hub: &EventHub,
+    running: &Arc<AtomicBool>,
+) -> std::io::Result<Option<String>> {
+    let Some(wait) = params.wait.clone() else {
+        return Ok(Some(dispatch_to_app_with_timeout(
+            Request {
+                id: request_id,
+                method: Method::AgentPrompt(params),
+            },
+            api_tx,
+            None,
+        )));
+    };
+
+    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) => {
+            return serde_json::to_string(&response)
+                .map(Some)
+                .map_err(std::io::Error::other);
+        }
+    };
+    let target = params.target.clone();
+    let prompt_response = dispatch_to_app_with_timeout(
+        Request {
+            id: request_id.clone(),
+            method: Method::AgentPrompt(params),
+        },
+        api_tx,
+        None,
+    );
+    let Ok(prompted) = agent_from_response(&request_id, &prompt_response) else {
+        return Ok(Some(prompt_response));
+    };
+    if !agent_wait_identity_matches(
+        &prompted,
+        &before_prompt.terminal_id,
+        before_prompt.name.as_deref().filter(|name| *name == target),
+        before_prompt.agent.as_deref(),
+    ) {
+        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 prompt_state_change_seq = prompted.state_change_seq;
+    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,
+            last_event_sequence,
+            after_state_change_seq: Some(prompt_state_change_seq),
+            accept_transient_status: false,
+        },
+        stream,
+        api_tx,
+        event_hub,
+        running,
+    )?
+    else {
+        return Ok(None);
+    };
+    let agent = match outcome {
+        AgentWaitOutcome::Matched(agent) => *agent,
+        AgentWaitOutcome::Response(response) => return Ok(Some(response)),
+    };
+    serde_json::to_string(&SuccessResponse {
+        id: request_id,
+        result: ResponseResult::AgentPrompted { agent },
+    })
+    .map(Some)
+    .map_err(std::io::Error::other)
+}
+
+struct ResolvedAgentWait {
+    target: String,
+    until: Vec<crate::api::schema::AgentStatus>,
+    timeout_ms: Option<u64>,
+    initial: crate::api::schema::AgentInfo,
+    last_event_sequence: u64,
+    after_state_change_seq: Option<u64>,
+    accept_transient_status: bool,
+}
+
+enum AgentWaitOutcome {
+    Matched(Box<crate::api::schema::AgentInfo>),
+    Response(String),
+}
+
+fn wait_for_resolved_agent(
+    request_id: String,
+    wait: ResolvedAgentWait,
+    stream: &mut LocalStream,
+    api_tx: &ApiRequestSender,
+    event_hub: &EventHub,
+    running: &Arc<AtomicBool>,
+) -> std::io::Result<Option<AgentWaitOutcome>> {
+    let deadline = wait
+        .timeout_ms
+        .map(|ms| std::time::Instant::now() + std::time::Duration::from_millis(ms));
+    let expected_terminal_id = wait.initial.terminal_id.clone();
+    let expected_name = wait
+        .initial
+        .name
+        .as_ref()
+        .filter(|name| name.as_str() == wait.target)
+        .cloned();
+    let expected_agent = wait.initial.agent.clone();
+    let pane_id = wait.initial.pane_id.clone();
+    let mut last_event_sequence = wait.last_event_sequence;
+
+    loop {
+        if should_stop_connection(stream, running)? {
+            return Ok(None);
+        }
+
+        let mut should_probe = false;
+        let mut matched_event_status = None;
+        for (sequence, event) in event_hub.events_after(last_event_sequence) {
+            last_event_sequence = sequence;
+            match event.data {
+                EventData::PaneAgentDetected {
+                    pane_id: event_pane,
+                    agent,
+                    released,
+                    final_status,
+                    ..
+                } if event_pane == pane_id => {
+                    if released {
+                        if let Some(status) = final_status
+                            .filter(|status| wait.until.contains(status))
+                            .or(matched_event_status)
+                        {
+                            let mut matched = wait.initial.clone();
+                            matched.agent_status = status;
+                            return Ok(Some(AgentWaitOutcome::Matched(Box::new(matched))));
+                        }
+                        return agent_wait_not_running(request_id)
+                            .map(AgentWaitOutcome::Response)
+                            .map(Some);
+                    }
+                    if agent.is_some() && expected_agent.is_some() && agent != expected_agent {
+                        return agent_wait_not_running(request_id)
+                            .map(AgentWaitOutcome::Response)
+                            .map(Some);
+                    }
+                    should_probe = true;
+                }
+                EventData::PaneAgentStatusChanged {
+                    pane_id: event_pane,
+                    agent_status,
+                    ..
+                } if event_pane == pane_id => {
+                    if wait.accept_transient_status && wait.until.contains(&agent_status) {
+                        matched_event_status = Some(agent_status);
+                    }
+                    should_probe = true;
+                }
+                EventData::PaneUpdated { pane } if pane.pane_id == pane_id => should_probe = true,
+                EventData::PaneMoved {
+                    previous_pane_id, ..
+                } if previous_pane_id == pane_id => {
+                    return agent_wait_not_running(request_id)
+                        .map(AgentWaitOutcome::Response)
+                        .map(Some);
+                }
+                EventData::PaneClosed {
+                    pane_id: event_pane,
+                    ..
+                }
+                | EventData::PaneExited {
+                    pane_id: event_pane,
+                    ..
+                } if event_pane == pane_id => {
+                    return agent_wait_not_running(request_id)
+                        .map(AgentWaitOutcome::Response)
+                        .map(Some);
+                }
+                _ => {}
+            }
+        }
+
+        if should_probe {
+            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 let Some(status) = matched_event_status {
+                let mut matched = current;
+                matched.agent_status = status;
+                return Ok(Some(AgentWaitOutcome::Matched(Box::new(matched))));
+            }
+            if agent_wait_matches(&current, &wait.until, wait.after_state_change_seq) {
+                return Ok(Some(AgentWaitOutcome::Matched(Box::new(current))));
+            }
+        }
+
+        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);
+        }
+        std::thread::sleep(CONNECTION_POLL_INTERVAL);
+    }
+}
+
+fn agent_wait_statuses(
+    until: Vec<crate::api::schema::AgentStatus>,
+) -> Vec<crate::api::schema::AgentStatus> {
+    if until.is_empty() {
+        vec![
+            crate::api::schema::AgentStatus::Idle,
+            crate::api::schema::AgentStatus::Done,
+            crate::api::schema::AgentStatus::Blocked,
+        ]
+    } else {
+        until
+    }
+}
+
+fn agent_wait_identity_matches(
+    agent: &crate::api::schema::AgentInfo,
+    expected_terminal_id: &str,
+    expected_name: Option<&str>,
+    expected_agent: Option<&str>,
+) -> bool {
+    agent.terminal_id == expected_terminal_id
+        && expected_name.is_none_or(|name| agent.name.as_deref() == Some(name))
+        && match (expected_agent, agent.agent.as_deref()) {
+            (Some(expected), Some(current)) => expected == current,
+            (Some(_), None) => agent.name.is_some(),
+            (None, _) => true,
+        }
+}
+
+fn agent_wait_matches(
+    agent: &crate::api::schema::AgentInfo,
+    until: &[crate::api::schema::AgentStatus],
+    after_state_change_seq: Option<u64>,
+) -> bool {
+    until.contains(&agent.agent_status)
+        && after_state_change_seq.is_none_or(|baseline| agent.state_change_seq > baseline)
+}
+
+fn agent_get(
+    request_id: &str,
+    target: &str,
+    api_tx: &ApiRequestSender,
+) -> Result<crate::api::schema::AgentInfo, ErrorResponse> {
+    let response = dispatch_to_app_with_timeout(
+        Request {
+            id: format!("{request_id}:agent"),
+            method: Method::AgentGet(crate::api::schema::AgentTarget {
+                target: target.to_string(),
+            }),
+        },
+        api_tx,
+        Some(APP_RESPONSE_TIMEOUT),
+    );
+    agent_from_response(request_id, &response)
+}
+
+fn agent_from_response(
+    request_id: &str,
+    response: &str,
+) -> Result<crate::api::schema::AgentInfo, ErrorResponse> {
+    let value: serde_json::Value = serde_json::from_str(response).map_err(|_| ErrorResponse {
+        id: request_id.into(),
+        error: ErrorBody {
+            code: "internal_error".into(),
+            message: "failed to decode agent response".into(),
+        },
+    })?;
+    if value.get("error").is_some() {
+        let error = serde_json::from_value(value["error"].clone()).map_err(|_| ErrorResponse {
+            id: request_id.into(),
+            error: ErrorBody {
+                code: "internal_error".into(),
+                message: "failed to decode agent error".into(),
+            },
+        })?;
+        return Err(ErrorResponse {
+            id: request_id.into(),
+            error,
+        });
+    }
+    serde_json::from_value(value["result"]["agent"].clone()).map_err(|_| ErrorResponse {
+        id: request_id.into(),
+        error: ErrorBody {
+            code: "internal_error".into(),
+            message: "failed to decode agent result".into(),
+        },
+    })
+}
+
+fn agent_wait_success(
+    request_id: String,
+    agent: crate::api::schema::AgentInfo,
+) -> std::io::Result<String> {
+    serde_json::to_string(&SuccessResponse {
+        id: request_id,
+        result: ResponseResult::AgentInfo { agent },
+    })
+    .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,
+        error: ErrorBody {
+            code: "agent_not_running".into(),
+            message: "agent is no longer running in the target pane".into(),
+        },
+    })
+    .map_err(std::io::Error::other)
+}
+
+fn agent_wait_probe_error(response: ErrorResponse) -> std::io::Result<String> {
+    if response.error.code == "agent_not_found" {
+        return agent_wait_not_running(response.id);
+    }
+    serde_json::to_string(&response).map_err(std::io::Error::other)
+}
+
 pub(super) fn wait_for_event(
     request_id: String,
     params: EventsWaitParams,
@@ -246,3 +644,35 @@ fn wait_matched_response(request_id: &str, event: serde_json::Value) -> String {
     })
     .unwrap()
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn agent_wait_probe_only_translates_agent_disappearance() {
+        let disappeared = agent_wait_probe_error(ErrorResponse {
+            id: "wait".into(),
+            error: ErrorBody {
+                code: "agent_not_found".into(),
+                message: "missing".into(),
+            },
+        })
+        .unwrap();
+        let disappeared: ErrorResponse = serde_json::from_str(&disappeared).unwrap();
+        assert_eq!(disappeared.id, "wait");
+        assert_eq!(disappeared.error.code, "agent_not_running");
+
+        let unavailable = agent_wait_probe_error(ErrorResponse {
+            id: "wait".into(),
+            error: ErrorBody {
+                code: "server_unavailable".into(),
+                message: "timed out waiting for app response".into(),
+            },
+        })
+        .unwrap();
+        let unavailable: ErrorResponse = serde_json::from_str(&unavailable).unwrap();
+        assert_eq!(unavailable.id, "wait");
+        assert_eq!(unavailable.error.code, "server_unavailable");
+    }
+}