boxmaker/crates/bxctl/src/chat.rs

//! `bxctl chat`: a turn client over `loop.sock` and a printer for the events it receives.

use std::io::Write;

use proto::{
    Envelope, ErrorCode, Message, PROTOCOL_VERSION, SessionId, Turn, TurnDone, TurnEvent,
    WireError, read_frame, write_frame,
};
use std::os::unix::net::UnixStream;
use std::path::Path;

#[derive(Debug)]
pub enum ChatError {
    Connect(std::io::Error),
    Frame(proto::FrameError),
    Refused(WireError),
    Protocol(String),
}

impl std::fmt::Display for ChatError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ChatError::Connect(e) => write!(f, "{e}"),
            ChatError::Frame(e) => write!(f, "{e}"),
            ChatError::Refused(w) => write!(f, "{}: {}", code_name(w.code), w.detail),
            ChatError::Protocol(s) => write!(f, "{s}"),
        }
    }
}

impl std::error::Error for ChatError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            ChatError::Connect(e) => Some(e),
            ChatError::Frame(e) => Some(e),
            _ => None,
        }
    }
}

// The ErrorCode name in snake_case, spelled out in words for the owner reading a failure.
fn code_name(code: ErrorCode) -> &'static str {
    match code {
        ErrorCode::BadFrame => "bad frame",
        ErrorCode::BadVersion => "bad version",
        ErrorCode::BadMessage => "bad message",
        ErrorCode::Internal => "internal",
        ErrorCode::SessionFull => "session full",
        ErrorCode::TurnLimit => "turn limit",
        ErrorCode::SessionBusy => "session busy",
        ErrorCode::NoSuchSession => "no such session",
        ErrorCode::SessionExists => "session exists",
        ErrorCode::Inference => "inference",
    }
}

pub fn new_session_id() -> SessionId {
    // The system clock is never before the unix epoch on any machine this runs on.
    let elapsed = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default();
    // "chat-<digits>-<digits>" uses only [a-z0-9-] and stays well under 64 bytes, so this cannot fail.
    let candidate = format!("chat-{}-{}", elapsed.as_secs(), elapsed.subsec_nanos());
    SessionId::new(&candidate).unwrap_or_else(|_| SessionId::new("chat-0-0").unwrap_or_default())
}

/// Sends one turn and reads the reply. `on_event` sees every event frame as it arrives.
pub fn run_turn(
    socket: &Path,
    session: &SessionId,
    content: &str,
    resume: bool,
    on_event: &mut dyn FnMut(&TurnEvent),
) -> Result<TurnDone, ChatError> {
    let mut stream = UnixStream::connect(socket).map_err(ChatError::Connect)?;

    let turn = Turn {
        session: session.clone(),
        content: content.to_string(),
        resume,
    };
    write_frame(
        &mut stream,
        &Envelope {
            v: PROTOCOL_VERSION,
            id: 1,
            r#final: true,
            msg: Message::Turn(turn),
        },
    )
    .map_err(ChatError::Frame)?;

    loop {
        let env = read_frame(&mut stream).map_err(ChatError::Frame)?;
        if env.id != 1 {
            return Err(ChatError::Protocol(
                "expected a frame with id 1".to_string(),
            ));
        }
        match (env.r#final, &env.msg) {
            (false, Message::TurnEvent(e)) => on_event(e),
            (true, Message::TurnDone(d)) => return Ok(d.clone()),
            (true, Message::Error(w)) => return Err(ChatError::Refused(w.clone())),
            _ => return Err(ChatError::Protocol("unexpected final frame".to_string())),
        }
    }
}

pub struct Printer {
    pub show_thinking: bool,
    pub json: bool,
    pub stream_content: bool,
    // Whether a dimmed reasoning block is open and needs closing.
    dimmed_open: bool,
}

impl Printer {
    pub fn new(show_thinking: bool, json: bool) -> Printer {
        Printer {
            show_thinking,
            json,
            stream_content: true,
            dimmed_open: false,
        }
    }

    pub fn event(&mut self, out: &mut dyn Write, event: &TurnEvent) -> std::io::Result<()> {
        if self.json {
            let line = serde_json::to_string(event).map_err(std::io::Error::other)?;
            out.write_all(line.as_bytes())?;
            out.write_all(b"\n")?;
            return out.flush();
        }

        match event {
            // Reasoning opens the dimmed block on the first piece; later pieces just continue it.
            TurnEvent::Reasoning { text } => {
                if self.show_thinking {
                    if !self.dimmed_open {
                        out.write_all(b"\x1b[2m")?;
                        self.dimmed_open = true;
                    }
                    out.write_all(text.as_bytes())?;
                }
            }
            TurnEvent::Content { text } => {
                if self.dimmed_open {
                    out.write_all(b"\x1b[0m\n")?;
                    self.dimmed_open = false;
                }
                if self.stream_content {
                    out.write_all(text.as_bytes())?;
                }
            }
            TurnEvent::ToolCallStarted { name } => {
                self.close_dimmed(out)?;
                writeln!(out, "[tool {name}]")?;
            }
            TurnEvent::ToolResult {
                name,
                class,
                truncated,
            } => {
                self.close_dimmed(out)?;
                if *truncated {
                    writeln!(out, "[{name}: {class:?}, truncated]")?;
                } else {
                    writeln!(out, "[{name}: {class:?}]")?;
                }
            }
            TurnEvent::Waiting { slot_busy } => {
                self.close_dimmed(out)?;
                let state = if *slot_busy { "busy" } else { "idle" };
                writeln!(out, "[waiting: slot {state}]")?;
            }
            TurnEvent::Retrying {
                attempt,
                after_ms,
                error,
            } => {
                self.close_dimmed(out)?;
                writeln!(
                    out,
                    "[retrying: attempt {attempt} in {after_ms} ms: {error}]"
                )?;
            }
            TurnEvent::ThinkingCapped { tokens } => {
                self.close_dimmed(out)?;
                writeln!(out, "[thinking capped at {tokens} tokens]")?;
            }
            TurnEvent::CacheLoss { expected, got } => {
                self.close_dimmed(out)?;
                writeln!(out, "[cache loss: {got} of {expected}]")?;
            }
            TurnEvent::Queued { .. } | TurnEvent::Progress { .. } => {}
        }
        out.flush()
    }

    pub fn end_reasoning(&mut self, out: &mut dyn Write) -> std::io::Result<()> {
        self.close_dimmed(out)
    }

    fn close_dimmed(&mut self, out: &mut dyn Write) -> std::io::Result<()> {
        if self.dimmed_open {
            out.write_all(b"\x1b[0m\n")?;
            self.dimmed_open = false;
        }
        Ok(())
    }
}