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::{BufRead, Write};
use std::os::unix::net::UnixStream;
use std::path::Path;

use proto::{
    DenyReason, Envelope, ErrorCode, Message, PROTOCOL_VERSION, SessionId, Timestamp, Turn,
    TurnDone, TurnEvent, WireError, read_frame, write_frame,
};

use crate::admin;
use crate::escape::{escape_json_text, escape_model_text};

#[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}"),
            // serde quotes the offending text of a bad frame, decoded: escape it.
            ChatError::Frame(e) => write!(f, "{}", escape_json_text(&e.to_string())),
            // The detail may carry the inference server's body: escape it like model text.
            ChatError::Refused(w) => {
                write!(f, "{}: {}", code_name(w.code), escape_json_text(&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.
pub 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",
        ErrorCode::Forbidden => "forbidden",
        ErrorCode::NoSuchApproval => "no such approval",
    }
}

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(escape_model_text(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(escape_model_text(text).as_bytes())?;
                }
            }
            TurnEvent::ToolCallStarted { name } => {
                self.close_dimmed(out)?;
                let name = escape_model_text(name);
                writeln!(out, "[tool {name}]")?;
            }
            TurnEvent::ToolResult {
                name,
                class,
                truncated,
            } => {
                self.close_dimmed(out)?;
                let name = escape_model_text(name);
                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: {}]",
                    escape_json_text(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 { .. } => {}
            TurnEvent::ApprovalPending { .. } => {
                self.close_dimmed(out)?;
            }
            TurnEvent::ToolDenied { name, reason } => {
                self.close_dimmed(out)?;
                let shown = escape_model_text(name);
                writeln!(out, "[denied {shown}: {}]", admin::reason_name(*reason))?;
                match reason {
                    DenyReason::GrantsInvalid => {
                        writeln!(out, "see docs/runbook.md#grants-invalid")?;
                    }
                    DenyReason::AuditUnavailable => {
                        writeln!(out, "see docs/runbook.md#audit-unavailable")?;
                    }
                    DenyReason::StateUnreadable => {
                        writeln!(out, "see docs/runbook.md#broker-state-damaged")?;
                    }
                    DenyReason::NoGrant
                    | DenyReason::GrantExpired
                    | DenyReason::TaintTooHigh
                    | DenyReason::DeniedByGrant
                    | DenyReason::ApprovalRefused
                    | DenyReason::ApprovalExpired
                    | DenyReason::InvalidArguments => {}
                }
            }
        }
        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(())
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OnPending {
    EventOnly,
    Show,
    Ask,
}

pub struct Approvals<'a> {
    pub admin_socket: &'a Path,
    pub on_pending: OnPending,
}

pub struct TurnIo<'a> {
    pub printer: &'a mut Printer,
    pub input: &'a mut dyn BufRead,
    pub out: &'a mut dyn Write,
}

pub fn handle_pending(
    admin_socket: &Path,
    approval: u64,
    ask: bool,
    now: Timestamp,
    input: &mut dyn BufRead,
    out: &mut dyn Write,
) -> std::io::Result<()> {
    let items = match admin::list(admin_socket) {
        Ok(items) => items,
        Err(e) => {
            writeln!(out, "approval {approval}: cannot ask brokerd: {e}")?;
            return Ok(());
        }
    };
    let Some(item) = items.iter().find(|item| item.approval == approval) else {
        writeln!(out, "approval {approval} is no longer pending")?;
        return Ok(());
    };
    out.write_all(b"\x1b[0m")?;
    admin::write_block(out, item, now)?;
    if !ask {
        return Ok(());
    }
    write!(out, "type {approval} to approve, anything else refuses: ")?;
    out.flush()?;
    let mut line = String::new();
    input.read_line(&mut line)?;
    if line.ends_with('\n') {
        line.pop();
    }
    if line.ends_with('\r') {
        line.pop();
    }
    let result = if line == approval.to_string() {
        admin::cmd_approve(admin_socket, approval, out)
    } else {
        admin::cmd_refuse(admin_socket, approval, None, out)
    };
    match result {
        Ok(_) => Ok(()),
        // The output itself failed: writing another line to it would fail the same way.
        Err(admin::AdminError::Io(e)) => Err(e),
        Err(e) => {
            writeln!(out, "approval {approval}: {e}")?;
            Ok(())
        }
    }
}

/// The outer error is a failed write to `io.out`; the inner one is the turn's.
pub fn stream_turn(
    socket: &Path,
    session: &SessionId,
    text: &str,
    resume: bool,
    approvals: &Approvals<'_>,
    io: &mut TurnIo<'_>,
) -> std::io::Result<Result<TurnDone, ChatError>> {
    let mut err = None;
    let outcome = run_turn(socket, session, text, resume, &mut |event| {
        if err.is_some() {
            return;
        }
        if let Err(e) = io.printer.event(io.out, event) {
            err = Some(e);
            return;
        }
        if let TurnEvent::ApprovalPending { approval, .. } = event
            && approvals.on_pending != OnPending::EventOnly
        {
            let ask = approvals.on_pending == OnPending::Ask;
            if let Err(e) = handle_pending(
                approvals.admin_socket,
                *approval,
                ask,
                Timestamp::now(),
                io.input,
                io.out,
            ) {
                err = Some(e);
            }
        }
    });
    if let Some(e) = err {
        return Err(e);
    }
    io.printer.end_reasoning(io.out)?;
    Ok(outcome)
}