Add Client::chat with the first-byte wait, liveness and error mapping

Implemented-By: OpenCode session (model recorded in docs/implementer-log.md)

crates/loopd/src/llama/chat.rs

@@ -0,0 +1,151 @@
+//! 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> {
+        // 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();
+        loop {
+            match events.next_item() {
+                Ok(Some(SseItem::Data(text))) => {
+                    for event in assembler.push(&text)? {
+                        on_event(&event);
+                    }
+                }
+                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(false)
+    }
+}
+
+/// 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,
+            });
+        }
+    }
+}

crates/loopd/src/llama/mod.rs

@@ -1,6 +1,7 @@
 //! The llama inference server's chat-completions types and the request builder.
 
 pub mod assemble;
+pub mod chat;
 pub mod info;
 pub mod request;