PicoBot/src/client/oneshot.rs

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<dyn std::error::Error>;

#[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<crate::providers::Usage>,
    error: Option<String>,
}

pub fn read_run_prompt(parts: Vec<String>) -> Result<String, DynError> {
    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<dyn std::error::Error>> {
    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::<WsOutbound>(&text)? {
                    WsOutbound::SessionEstablished { session_id, .. } => break session_id,
                    WsOutbound::Error { code, message } => {
                        return Err::<RunOutput, DynError>(
                            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<String, (String, ToolStatus)> = HashMap::new();
        loop {
            match receiver.next().await {
                Some(Ok(Message::Text(text))) => match serde_json::from_str::<WsOutbound>(&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<S>(sender: &mut S)
where
    S: futures_util::Sink<Message> + 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<dyn std::error::Error>> {
    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::<IpAddr>()
                .is_ok_and(|address| address.is_loopback())
    });
    Ok((url.to_string(), local_gateway))
}

fn read_prompt_from(reader: &mut impl Read) -> Result<String, Box<dyn std::error::Error>> {
    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<String, Box<dyn std::error::Error>> {
    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::<Vec<_>>()
        .join("\n\n")
}

fn report_progress(
    snapshot: &TurnSnapshot,
    last_phase: &mut Option<TurnPhase>,
    tool_states: &mut HashMap<String, (String, ToolStatus)>,
) {
    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(&current) {
            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<dyn std::error::Error>> {
    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);
    }
}