use super::{WsInbound, WsOutbound, load_auth_token}; use crate::config::get_user_config_dir; use crate::gateway::auth::ADMIN_TOKEN_HEADER; use crate::session::{ToolStatus, TurnBlock, TurnPhase, TurnSnapshot, TurnStatus}; use futures_util::{SinkExt, StreamExt}; use serde::Serialize; use std::collections::HashMap; use std::io::{self, IsTerminal, Read, Write}; use std::net::IpAddr; use std::time::Duration; use tokio_tungstenite::connect_async; use tokio_tungstenite::tungstenite::{ Message, client::IntoClientRequest, http::{HeaderValue, header}, }; const MAX_RUN_PROMPT_BYTES: usize = 1024 * 1024; type DynError = Box; #[derive(Debug, Clone, Copy)] pub struct RunOptions { pub timeout: Duration, pub json: bool, pub verbose: bool, } #[derive(Debug, Serialize)] struct RunOutput { session_id: String, turn_id: String, status: TurnStatus, content: String, usage: Option, error: Option, } pub fn read_run_prompt(parts: Vec) -> Result { if !parts.is_empty() { return validate_prompt(parts.join(" ")); } if io::stdin().is_terminal() { return Err("provide a prompt as arguments or pipe it on stdin".into()); } let stdin = io::stdin(); let mut locked = stdin.lock(); read_prompt_from(&mut locked) } pub async fn run_once( gateway_url: &str, prompt: String, options: RunOptions, ) -> Result<(), Box> { if options.timeout.is_zero() { return Err("run timeout must be greater than zero".into()); } let prompt = validate_prompt(prompt)?; let client_id = format!("run-{}", uuid::Uuid::new_v4().simple()); let (connect_url, local_gateway) = websocket_url(gateway_url, &client_id)?; let admin_token = local_gateway .then(|| std::fs::read_to_string(get_user_config_dir().join("web_admin_token")).ok()) .flatten() .map(|token| token.trim().to_string()) .filter(|token| !token.is_empty()); let bearer_token = admin_token.is_none().then(load_auth_token).flatten(); let mut request = connect_url.into_client_request()?; if let Some(token) = &admin_token { let mut value = HeaderValue::from_str(token)?; value.set_sensitive(true); request.headers_mut().insert(ADMIN_TOKEN_HEADER, value); } else if let Some(token) = &bearer_token { let mut value = HeaderValue::from_str(&format!("Bearer {token}"))?; value.set_sensitive(true); request.headers_mut().insert(header::AUTHORIZATION, value); } let (stream, _) = connect_async(request).await.map_err(|error| { if local_gateway && admin_token.is_none() { format!( "gateway connection failed: {error}; local admin token is unavailable at {}", get_user_config_dir().join("web_admin_token").display() ) } else { format!( "gateway connection failed: {error}. Remote gateways require an existing paired CLI token" ) } })?; let (mut sender, mut receiver) = stream.split(); let operation = async { let session_id = loop { match receiver.next().await { Some(Ok(Message::Text(text))) => match serde_json::from_str::(&text)? { WsOutbound::SessionEstablished { session_id, .. } => break session_id, WsOutbound::Error { code, message } => { return Err::( format!("gateway error {code}: {message}").into(), ); } _ => {} }, Some(Ok(Message::Close(_))) | None => { return Err("gateway closed before establishing a session".into()); } Some(Err(error)) => return Err(error.into()), _ => {} } }; let input = WsInbound::UserInput { content: prompt, upload_ids: Vec::new(), channel: None, chat_id: None, sender_id: None, }; sender .send(Message::Text(serde_json::to_string(&input)?.into())) .await?; let mut turn_id = None; let mut last_phase = None; let mut tool_states: HashMap = HashMap::new(); loop { match receiver.next().await { Some(Ok(Message::Text(text))) => match serde_json::from_str::(&text)? { WsOutbound::TurnUpdated { snapshot } if snapshot.session_id == session_id && turn_id.as_ref().is_none_or(|id| id == &snapshot.id.0) => { turn_id.get_or_insert_with(|| snapshot.id.0.clone()); if options.verbose { report_progress(&snapshot, &mut last_phase, &mut tool_states); } if snapshot.status != TurnStatus::Running { break Ok(output_from_snapshot(snapshot)); } } WsOutbound::Error { code, message } => { break Err(format!("gateway error {code}: {message}").into()); } _ => {} }, Some(Ok(Message::Close(_))) | None => { break Err("gateway closed before the run completed".into()); } Some(Err(error)) => break Err(error.into()), _ => {} } } }; let output = tokio::select! { result = tokio::time::timeout(options.timeout, operation) => { match result { Ok(result) => result?, Err(_) => { send_stop(&mut sender).await; return Err(format!("run timed out after {} seconds", options.timeout.as_secs()).into()); } } } signal = tokio::signal::ctrl_c() => { send_stop(&mut sender).await; signal?; return Err("run cancelled".into()); } }; render_output(&output, options.json)?; match output.status { TurnStatus::Completed => Ok(()), TurnStatus::Cancelled => Err(output .error .unwrap_or_else(|| "run cancelled".to_string()) .into()), TurnStatus::Failed => Err(output .error .unwrap_or_else(|| "run failed".to_string()) .into()), TurnStatus::Running => Err("gateway returned a non-terminal run result".into()), } } async fn send_stop(sender: &mut S) where S: futures_util::Sink + Unpin, { let stop = WsInbound::UserInput { content: "/stop".to_string(), upload_ids: Vec::new(), channel: None, chat_id: None, sender_id: None, }; if let Ok(text) = serde_json::to_string(&stop) { let _ = sender.send(Message::Text(text.into())).await; let _ = sender.flush().await; } } fn websocket_url( gateway_url: &str, client_id: &str, ) -> Result<(String, bool), Box> { let mut url = reqwest::Url::parse(gateway_url)?; let scheme = match url.scheme() { "ws" => "ws", "wss" => "wss", "http" => "ws", "https" => "wss", other => return Err(format!("unsupported gateway URL scheme: {other}").into()), }; url.set_scheme(scheme) .map_err(|_| "failed to set gateway URL scheme")?; if url.path().is_empty() || url.path() == "/" { url.set_path("/ws"); } url.query_pairs_mut().append_pair("client_id", client_id); let local_gateway = url.host_str().is_some_and(|host| { let host = host .strip_prefix('[') .and_then(|value| value.strip_suffix(']')) .unwrap_or(host); host.eq_ignore_ascii_case("localhost") || host .parse::() .is_ok_and(|address| address.is_loopback()) }); Ok((url.to_string(), local_gateway)) } fn read_prompt_from(reader: &mut impl Read) -> Result> { let mut bytes = Vec::new(); reader .take((MAX_RUN_PROMPT_BYTES + 1) as u64) .read_to_end(&mut bytes)?; if bytes.len() > MAX_RUN_PROMPT_BYTES { return Err(format!("prompt exceeds {MAX_RUN_PROMPT_BYTES} bytes").into()); } let prompt = String::from_utf8(bytes)?; validate_prompt(prompt.trim_end_matches(['\r', '\n']).to_string()) } fn validate_prompt(prompt: String) -> Result> { if prompt.len() > MAX_RUN_PROMPT_BYTES { return Err(format!("prompt exceeds {MAX_RUN_PROMPT_BYTES} bytes").into()); } if prompt.trim().is_empty() { return Err("prompt is empty".into()); } Ok(prompt) } fn output_from_snapshot(snapshot: TurnSnapshot) -> RunOutput { RunOutput { session_id: snapshot.session_id, turn_id: snapshot.id.0, status: snapshot.status, content: assistant_text(&snapshot.blocks), usage: snapshot.usage, error: snapshot.error, } } fn assistant_text(blocks: &[TurnBlock]) -> String { blocks .iter() .filter_map(|block| match block { TurnBlock::Assistant { text, .. } if !text.is_empty() => Some(text.as_str()), _ => None, }) .collect::>() .join("\n\n") } fn report_progress( snapshot: &TurnSnapshot, last_phase: &mut Option, tool_states: &mut HashMap, ) { if last_phase.as_ref() != Some(&snapshot.phase) { eprintln!("[phase: {}]", phase_name(snapshot.phase)); *last_phase = Some(snapshot.phase); } for block in &snapshot.blocks { let TurnBlock::Tool { id, name, status, .. } = block else { continue; }; let current = (name.clone(), *status); if tool_states.get(id) != Some(¤t) { eprintln!("[tool: {name}: {}]", tool_status_name(*status)); tool_states.insert(id.clone(), current); } } } fn phase_name(phase: TurnPhase) -> &'static str { match phase { TurnPhase::Queued => "queued", TurnPhase::Reasoning => "reasoning", TurnPhase::Responding => "responding", TurnPhase::Acting => "acting", TurnPhase::Finalizing => "finalizing", } } fn tool_status_name(status: ToolStatus) -> &'static str { match status { ToolStatus::Running => "running", ToolStatus::Completed => "completed", ToolStatus::Failed => "failed", } } fn render_output(output: &RunOutput, json: bool) -> Result<(), Box> { if json { println!("{}", serde_json::to_string(output)?); } else if output.status == TurnStatus::Completed { print!("{}", output.content); if !output.content.ends_with('\n') { println!(); } io::stdout().flush()?; } Ok(()) } #[cfg(test)] mod tests { use super::*; use crate::session::{BlockId, TurnId}; #[test] fn positional_prompt_parts_are_joined() { assert_eq!( validate_prompt(["hello", "world"].join(" ")).unwrap(), "hello world" ); } #[test] fn stdin_prompt_preserves_lines_and_trims_terminal_newline() { let mut input = "first\nsecond\n".as_bytes(); assert_eq!(read_prompt_from(&mut input).unwrap(), "first\nsecond"); } #[test] fn empty_and_oversized_prompts_are_rejected() { assert!(validate_prompt(" \n".to_string()).is_err()); assert!(validate_prompt("x".repeat(MAX_RUN_PROMPT_BYTES + 1)).is_err()); } #[test] fn assistant_blocks_are_joined_without_reasoning_or_tools() { let blocks = vec![ TurnBlock::Reasoning { id: BlockId("reasoning".to_string()), iteration: 0, text: "hidden".to_string(), }, TurnBlock::Assistant { id: BlockId("answer-1".to_string()), iteration: 0, text: "hello".to_string(), }, TurnBlock::Tool { id: "tool".to_string(), iteration: 0, name: "bash".to_string(), arguments: serde_json::json!({}), status: ToolStatus::Completed, preview: None, }, TurnBlock::Assistant { id: BlockId("answer-2".to_string()), iteration: 1, text: "world".to_string(), }, ]; assert_eq!(assistant_text(&blocks), "hello\n\nworld"); } #[test] fn loopback_gate_uses_the_url_host() { assert!( websocket_url("ws://127.0.0.1:19876/ws", "run-id") .unwrap() .1 ); assert!(websocket_url("ws://[::1]:19876/ws", "run-id").unwrap().1); assert!(websocket_url("http://localhost:19876", "run-id").unwrap().1); assert!( !websocket_url("wss://gateway.example/ws", "run-id") .unwrap() .1 ); } #[test] fn terminal_snapshot_becomes_script_output() { let output = output_from_snapshot(TurnSnapshot { id: TurnId("turn".to_string()), session_id: "session".to_string(), message_id: "message".to_string(), revision: 1, status: TurnStatus::Completed, phase: TurnPhase::Finalizing, blocks: vec![TurnBlock::Assistant { id: BlockId("answer".to_string()), iteration: 0, text: "done".to_string(), }], usage: None, error: None, }); assert_eq!(output.turn_id, "turn"); assert_eq!(output.content, "done"); assert_eq!(output.status, TurnStatus::Completed); } }