use std::fs; use std::process::Stdio; use std::time::Duration; use anyhow::{Context, Result}; use chrono::Utc; use serde_json::Value; use tokio::io::{AsyncBufReadExt, AsyncReadExt, BufReader}; use tokio::process::Command; use tokio::sync::{mpsc, oneshot}; use tokio::time::sleep; use crate::adapters::{OutputMode, build_command}; use crate::types::{ AgentRequest, RunResult, RunnerEvent, StreamSource, UsageAlert, UsageAlertKind, }; #[derive(Debug)] enum InternalEvent { ToolConversationId(String), Chunk(StreamSource, String), StreamClosed, } pub struct RunningHandle { pub receiver: mpsc::UnboundedReceiver, cancel: Option>, } impl RunningHandle { pub fn cancel(&mut self) { if let Some(cancel) = self.cancel.take() { let _ = cancel.send(()); } } } pub fn spawn(request: AgentRequest) -> RunningHandle { let (event_tx, event_rx) = mpsc::unbounded_channel(); let (cancel_tx, cancel_rx) = oneshot::channel(); tokio::spawn(async move { if let Err(err) = run(request, event_tx.clone(), cancel_rx).await { let _ = event_tx.send(RunnerEvent::Failed(err.to_string())); } }); RunningHandle { receiver: event_rx, cancel: Some(cancel_tx), } } async fn run( request: AgentRequest, event_tx: mpsc::UnboundedSender, mut cancel_rx: oneshot::Receiver<()>, ) -> Result<()> { let started_at = Utc::now(); let _ = event_tx.send(RunnerEvent::Started(started_at)); let command_spec = build_command(&request)?; if let Some(tool_conversation_id) = &command_spec.tool_conversation_id_hint { let _ = event_tx.send(RunnerEvent::ToolConversationId( tool_conversation_id.clone(), )); } let mut child = Command::new(&command_spec.program) .args(&command_spec.args) .current_dir(&request.cwd) .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::piped()) .spawn() .with_context(|| format!("failed to spawn {}", request.tool.command_name()))?; let stdout = child .stdout .take() .context("failed to capture child stdout")?; let stderr = child .stderr .take() .context("failed to capture child stderr")?; let (internal_tx, mut internal_rx) = mpsc::unbounded_channel(); match command_spec.output_mode { OutputMode::PlainText => { tokio::spawn(read_plain_stream( stdout, StreamSource::Stdout, internal_tx.clone(), )); } OutputMode::AgyJson | OutputMode::CodexJson => { tokio::spawn(read_structured_stream( stdout, command_spec.output_mode, internal_tx.clone(), )); } } tokio::spawn(read_plain_stream( stderr, StreamSource::Stderr, internal_tx.clone(), )); let mut stdout_buf = String::new(); let mut stderr_buf = String::new(); let mut closed_streams = 0usize; let mut exit_code = None; let mut cancelled = false; let mut usage_alert = None; let mut tool_conversation_id = command_spec.tool_conversation_id_hint.clone(); loop { tokio::select! { maybe_event = internal_rx.recv() => { match maybe_event { Some(InternalEvent::ToolConversationId(id)) => { if tool_conversation_id.as_deref() != Some(id.as_str()) { tool_conversation_id = Some(id.clone()); let _ = event_tx.send(RunnerEvent::ToolConversationId(id)); } } Some(InternalEvent::Chunk(source, text)) => { if let Some(alert) = detect_usage_issue(&text) && usage_alert.is_none() { usage_alert = Some(alert.clone()); let _ = event_tx.send(RunnerEvent::Alert(alert)); } match source { StreamSource::Stdout => stdout_buf.push_str(&text), StreamSource::Stderr => stderr_buf.push_str(&text), } let _ = event_tx.send(RunnerEvent::Chunk { source, text }); } Some(InternalEvent::StreamClosed) => { closed_streams += 1; } None => break, } } _ = &mut cancel_rx, if !cancelled => { cancelled = true; let _ = child.kill().await; } _ = sleep(Duration::from_millis(50)) => { if exit_code.is_none() { if let Some(status) = child.try_wait().context("failed waiting on child process")? { exit_code = status.code(); } } } } if exit_code.is_some() && closed_streams >= 2 { break; } } if exit_code.is_none() { let status = child .wait() .await .context("failed waiting on child process")?; exit_code = status.code(); } let ended_at = Utc::now(); let assistant_response = if let Some(path) = &command_spec.response_capture_file { let response = fs::read_to_string(path).unwrap_or_else(|_| stdout_buf.clone()); let _ = fs::remove_file(path); response } else { stdout_buf.clone() }; let result = RunResult { stdout: stdout_buf, stderr: stderr_buf, assistant_response, tool_conversation_id, exit_code, started_at, ended_at, cancelled, usage_alert, }; let _ = event_tx.send(RunnerEvent::Finished(result)); Ok(()) } async fn read_plain_stream( mut reader: R, source: StreamSource, tx: mpsc::UnboundedSender, ) where R: tokio::io::AsyncRead + Unpin + Send + 'static, { let mut buf = vec![0_u8; 4096]; loop { match reader.read(&mut buf).await { Ok(0) => { let _ = tx.send(InternalEvent::StreamClosed); break; } Ok(n) => { let text = String::from_utf8_lossy(&buf[..n]).to_string(); let _ = tx.send(InternalEvent::Chunk(source, text)); } Err(err) => { let _ = tx.send(InternalEvent::Chunk( StreamSource::Stderr, format!("stream read error: {err}\n"), )); let _ = tx.send(InternalEvent::StreamClosed); break; } } } } async fn read_structured_stream( reader: R, output_mode: OutputMode, tx: mpsc::UnboundedSender, ) where R: tokio::io::AsyncRead + Unpin + Send + 'static, { let mut lines = BufReader::new(reader).lines(); loop { match lines.next_line().await { Ok(Some(line)) => match parse_structured_line(output_mode, &line) { Ok(parsed) => { if let Some(tool_conversation_id) = parsed.tool_conversation_id { let _ = tx.send(InternalEvent::ToolConversationId(tool_conversation_id)); } if let Some(display_text) = parsed.display_text && !display_text.is_empty() { let _ = tx.send(InternalEvent::Chunk(StreamSource::Stdout, display_text)); } } Err(_) => { let _ = tx.send(InternalEvent::Chunk( StreamSource::Stdout, format!("{line}\n"), )); } }, Ok(None) => { let _ = tx.send(InternalEvent::StreamClosed); break; } Err(err) => { let _ = tx.send(InternalEvent::Chunk( StreamSource::Stderr, format!("stream read error: {err}\n"), )); let _ = tx.send(InternalEvent::StreamClosed); break; } } } } #[derive(Debug)] struct ParsedStructuredLine { tool_conversation_id: Option, display_text: Option, } fn parse_structured_line(output_mode: OutputMode, line: &str) -> Result { let value: Value = serde_json::from_str(line)?; let tool_conversation_id = find_tool_conversation_id(output_mode, &value); let display_text = match output_mode { OutputMode::AgyJson => extract_agy_display_text(&value), OutputMode::CodexJson => extract_codex_display_text(&value), OutputMode::PlainText => None, }; Ok(ParsedStructuredLine { tool_conversation_id, display_text, }) } fn extract_agy_display_text(value: &Value) -> Option { let event_type = value .get("type") .and_then(Value::as_str) .unwrap_or_default(); let role = value .get("role") .and_then(Value::as_str) .unwrap_or_default(); if event_type == "message" && role == "assistant" { return value .get("content") .and_then(Value::as_str) .map(ToString::to_string); } if event_type == "error" { return value .get("message") .and_then(Value::as_str) .map(|message| format!("{message}\n")); } None } fn extract_codex_display_text(value: &Value) -> Option { let event_type = value .get("type") .and_then(Value::as_str) .unwrap_or_default(); if event_type.starts_with("response.output_text.delta") { return value .get("delta") .and_then(Value::as_str) .map(ToString::to_string); } if event_type.starts_with("response.output_text.done") { return value .get("text") .and_then(Value::as_str) .or_else(|| value.get("content").and_then(Value::as_str)) .map(ToString::to_string); } if event_type.contains("error") { return value .get("message") .and_then(Value::as_str) .or_else(|| value.get("content").and_then(Value::as_str)) .map(|message| format!("{message}\n")); } if value.get("role").and_then(Value::as_str) == Some("assistant") { return value .get("content") .and_then(Value::as_str) .map(ToString::to_string); } None } fn find_tool_conversation_id(output_mode: OutputMode, value: &Value) -> Option { let keys = match output_mode { OutputMode::AgyJson => &["session_id", "sessionId"][..], OutputMode::CodexJson => &["thread_id", "threadId", "session_id", "sessionId"][..], OutputMode::PlainText => return None, }; find_string_recursive(value, keys) } fn find_string_recursive(value: &Value, keys: &[&str]) -> Option { match value { Value::Object(map) => { for key in keys { if let Some(candidate) = map.get(*key).and_then(Value::as_str) && !candidate.trim().is_empty() { return Some(candidate.to_string()); } } for nested in map.values() { if let Some(found) = find_string_recursive(nested, keys) { return Some(found); } } None } Value::Array(values) => values .iter() .find_map(|nested| find_string_recursive(nested, keys)), _ => None, } } fn detect_usage_issue(text: &str) -> Option { let lower = text.to_ascii_lowercase(); let lower = lower.trim(); let quota_patterns = [ "insufficient_quota", "usage limit reached", "quota exceeded", "exceeded your current quota", "out of premium requests", "premium requests quota", "request limit reached", "out of usage tokens", "out of tokens", "quota exhausted", "monthly quota", ]; let rate_patterns = [ "rate limit", "too many requests", "please try again later", "temporarily unavailable", "server overloaded", "429", ]; if quota_patterns.iter().any(|pattern| lower.contains(pattern)) { return Some(UsageAlert { kind: UsageAlertKind::QuotaExhausted, evidence: summarize_evidence(text), }); } if rate_patterns.iter().any(|pattern| lower.contains(pattern)) { return Some(UsageAlert { kind: UsageAlertKind::RateLimited, evidence: summarize_evidence(text), }); } None } fn summarize_evidence(text: &str) -> String { text.lines() .map(str::trim) .find(|line| !line.is_empty()) .unwrap_or(text.trim()) .chars() .take(160) .collect() } #[cfg(test)] mod tests { use super::{OutputMode, detect_usage_issue, parse_structured_line, summarize_evidence}; use crate::types::UsageAlertKind; #[test] fn detects_quota_exhaustion() { let alert = detect_usage_issue("Error: usage limit reached for this account") .expect("alert should be detected"); assert_eq!(alert.kind, UsageAlertKind::QuotaExhausted); } #[test] fn detects_rate_limits() { let alert = detect_usage_issue("429 Too many requests").expect("alert should be detected"); assert_eq!(alert.kind, UsageAlertKind::RateLimited); } #[test] fn trims_evidence() { let summary = summarize_evidence(" \n hello world \n second line"); assert_eq!(summary, "hello world"); } #[test] fn parses_agy_session_id() { let parsed = parse_structured_line( OutputMode::AgyJson, r#"{"type":"init","session_id":"agy-session-1"}"#, ) .expect("structured line should parse"); assert_eq!( parsed.tool_conversation_id.as_deref(), Some("agy-session-1") ); } #[test] fn parses_codex_thread_id() { let parsed = parse_structured_line( OutputMode::CodexJson, r#"{"type":"response.output_text.delta","delta":"hi","thread_id":"thread-1"}"#, ) .expect("structured line should parse"); assert_eq!(parsed.tool_conversation_id.as_deref(), Some("thread-1")); assert_eq!(parsed.display_text.as_deref(), Some("hi")); } }