diff options
Diffstat (limited to 'src/api/wait.rs')
| -rw-r--r-- | src/api/wait.rs | 430 |
1 files changed, 430 insertions, 0 deletions
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, ¶ms.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, ¶ms.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( + ¤t, + &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(¤t, &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"); + } +} |