boxmaker/crates/loopd/src/llama/chat.rs

//! One chat request: the wait for the first byte, the liveness of the stream, and the errors.
//!
//! There is no total deadline anywhere. The only question while waiting for the head is whether
//! bytes are still arriving; silence before the first byte means the slot is busy or the server is
//! loading, and only `slots` can say which. Once bytes flow, the only question is whether they keep
//! arriving, measured by the read timeout, never a wall-clock total.

use std::time::Duration;

use super::assemble::Assembler;
use super::{ChatEvent, ChatRequest, Client, Completion, InferError};
use crate::http::{Connection, Head, HttpError, Request, read_capped};
use crate::llama::info::{MAX_BODY, error_text, map_http};
use crate::llama::request::build_body;
use crate::sse::{Events, SseError, SseItem};

pub const MAX_SSE_LINE: usize = 1024 * 1024;

const CHAT: &str = "/v1/chat/completions";

/// What the wait for the head is doing when it times out.
#[derive(Clone, Copy, PartialEq)]
enum WaitState {
    /// The head has started to arrive; no poll is needed.
    Idle,
    /// The request's own slot is processing.
    Busy,
    /// The server could not answer `/slots`, so it is loading or restarting.
    Unavailable,
}

impl Client {
    pub fn chat(
        &self,
        req: &ChatRequest,
        on_event: &mut dyn FnMut(&ChatEvent),
    ) -> Result<Completion, InferError> {
        // The gate holds one request per slot from before it is sent until chat returns; a second
        // request on the same slot queues here, bounded, and is told (Queued) or refused (Busy).
        let _permit = self
            .gate
            .acquire(req.slot, self.cfg.limits.queue_len, &mut |ahead| {
                on_event(&ChatEvent::Queued { ahead })
            })
            .map_err(|_| InferError::Busy)?;

        // 1. Send.
        let body_text =
            build_body(&self.cfg, req).map_err(|e| InferError::Protocol(e.to_string()))?;
        let mut conn = Connection::open(&self.cfg.infer.socket).map_err(map_http)?;
        conn.send(&Request {
            method: "POST",
            path: CHAT,
            body: Some(body_text.as_bytes()),
        })
        .map_err(map_http)?;

        // 2. Wait for the head.
        let head = wait_for_head(&mut conn, self, req, on_event)?;

        // 3. Status.
        conn.set_read_timeout(Duration::from_millis(self.cfg.limits.liveness_ms))
            .map_err(map_http)?;
        if head.status != 200 {
            let mut body = conn.body(&head).map_err(map_http)?;
            let bytes = read_capped(&mut body, MAX_BODY).map_err(map_http)?;
            return Err(InferError::Http {
                status: head.status,
                body: error_text(&bytes),
            });
        }

        // 4. Stream.
        let mut body = conn.body(&head).map_err(map_http)?;
        let mut events = Events::new(&mut body, MAX_SSE_LINE);
        let mut assembler = Assembler::new();
        // Where the cap fired, if it has: the overrun allowance is measured from here.
        let mut cap_at: Option<u64> = None;
        loop {
            match events.next_item() {
                Ok(Some(SseItem::Data(text))) => {
                    for event in assembler.push(&text)? {
                        on_event(&event);
                    }
                    if assembler.in_reasoning() {
                        let tokens = assembler.reasoning_tokens();
                        match cap_at {
                            None if tokens >= self.cfg.limits.thinking_cap => {
                                let Some(id) = assembler.id() else {
                                    return Err(InferError::Protocol(
                                        "a reasoning chunk carried no completion id".to_string(),
                                    ));
                                };
                                match self.end_reasoning(id) {
                                    Ok(true) => {
                                        cap_at = Some(tokens);
                                        on_event(&ChatEvent::ThinkingCapped { tokens });
                                    }
                                    Ok(false) | Err(_) => {
                                        return Err(InferError::ThinkingOverrun);
                                    }
                                }
                            }
                            Some(at) if tokens >= at + self.cfg.limits.thinking_overrun => {
                                return Err(InferError::ThinkingOverrun);
                            }
                            _ => {}
                        }
                    }
                }
                Ok(Some(SseItem::Done)) | Ok(None) => break,
                Err(SseError::Timeout) => return Err(InferError::Stalled),
                Err(SseError::Truncated) => return Err(InferError::StreamClosedEarly),
                Err(e) => return Err(InferError::Protocol(e.to_string())),
            }
        }

        // 5. Finish. A stream that ended without a finish_reason is a server that died between two
        // events.
        assembler.finish(cap_at.is_some())
    }
}

/// Reads the head of the chat response, waiting out a busy, unavailable or idle slot as it goes.
fn wait_for_head(
    conn: &mut Connection,
    client: &Client,
    req: &ChatRequest,
    on_event: &mut dyn FnMut(&ChatEvent),
) -> Result<Head, InferError> {
    let poll = Duration::from_millis(client.cfg.limits.poll_ms);
    conn.set_read_timeout(poll).map_err(map_http)?;
    let mut prev: Option<WaitState> = None;
    let mut since: Duration = Duration::ZERO;
    loop {
        match conn.read_head() {
            Ok(head) => return Ok(head),
            Err(HttpError::Timeout) => {}
            Err(e) => return Err(map_http(e)),
        }
        let received = conn.received_any();
        let (state, busy) = if received {
            (WaitState::Idle, false)
        } else {
            match client.slots() {
                Ok(slots) => {
                    let busy = slots
                        .iter()
                        .find(|s| s.id == req.slot)
                        .map(|s| s.is_processing)
                        .unwrap_or(false);
                    (
                        if busy {
                            WaitState::Busy
                        } else {
                            WaitState::Idle
                        },
                        busy,
                    )
                }
                Err(_) => (WaitState::Unavailable, false),
            }
        };
        // A poll was made exactly when the head had not started to arrive.
        if !received {
            on_event(&ChatEvent::Waiting { slot_busy: busy });
        }
        since = if prev == Some(state) {
            since + poll
        } else {
            poll
        };
        prev = Some(state);
        let wait_for = match state {
            WaitState::Busy => client.cfg.limits.busy_wait_ms,
            WaitState::Unavailable => client.cfg.limits.load_wait_ms,
            WaitState::Idle => client.cfg.limits.idle_grace_ms,
        };
        if since >= Duration::from_millis(wait_for) {
            return Err(match state {
                WaitState::Busy => InferError::WaitTimeout,
                WaitState::Unavailable => InferError::LoadTimeout,
                WaitState::Idle => InferError::Stalled,
            });
        }
    }
}