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;
 

crates/loopd/tests/chat.rs

@@ -0,0 +1,294 @@
+//! Tests for one request: events, waits, liveness and errors. Do not edit.
+//!
+//! Limits come from `support::test_config`: poll 40 ms, busy wait 400 ms, load wait 300 ms,
+//! idle grace 200 ms, liveness 150 ms.
+
+mod support;
+
+use loopd::llama::{ChatEvent, ChatMessage, ChatRequest, Client, FinishReason, InferError};
+use std::time::{Duration, Instant};
+use support::{FakeServer, Reply};
+
+const CHAT: &str = "/v1/chat/completions";
+const IDLE: &str = r#"[{"id":0,"is_processing":false},{"id":1,"is_processing":false}]"#;
+const SLOT0_BUSY: &str = r#"[{"id":0,"is_processing":true},{"id":1,"is_processing":false}]"#;
+const SLOT1_BUSY: &str = r#"[{"id":0,"is_processing":false},{"id":1,"is_processing":true}]"#;
+
+fn request() -> ChatRequest {
+    ChatRequest {
+        slot: 0,
+        messages: vec![ChatMessage::User {
+            content: "hi".to_string(),
+        }],
+        tools: vec![],
+        thinking: false,
+    }
+}
+
+fn run(
+    server: &FakeServer,
+) -> (
+    Result<loopd::llama::Completion, InferError>,
+    Vec<ChatEvent>,
+    Duration,
+) {
+    let client = Client::new(support::test_config(&server.socket));
+    let mut events = Vec::new();
+    let started = Instant::now();
+    let result = client.chat(&request(), &mut |e| events.push(e.clone()));
+    (result, events, started.elapsed())
+}
+
+fn waiting(events: &[ChatEvent]) -> Vec<bool> {
+    events
+        .iter()
+        .filter_map(|e| match e {
+            ChatEvent::Waiting { slot_busy } => Some(*slot_busy),
+            _ => None,
+        })
+        .collect()
+}
+
+#[test]
+fn a_recorded_completion_comes_back_whole() {
+    let server = FakeServer::start();
+    server.route(CHAT, vec![Reply::fixture("tool_call")]);
+    let (result, events, _) = run(&server);
+    let done = result.unwrap();
+    let want = support::expected("tool_call");
+    assert_eq!(done.finish_reason, FinishReason::ToolCalls);
+    assert_eq!(done.content.as_deref(), want["content"].as_str());
+    assert_eq!(done.tool_calls[0].arguments, r#"{"path":"/etc/hostname"}"#);
+    assert_eq!(done.timings.prompt_n, 312);
+    assert!(!done.thinking_capped);
+    assert!(
+        waiting(&events).is_empty(),
+        "the answer came at once, so nothing was polled"
+    );
+    assert!(matches!(
+        events.first(),
+        Some(ChatEvent::Progress { total: 312, .. })
+    ));
+
+    let sent = &server.requests_to(CHAT)[0];
+    assert_eq!(sent.method, "POST");
+    assert_eq!(sent.json()["messages"][0]["content"], "hi");
+    assert_eq!(sent.json()["id_slot"], 0);
+}
+
+#[test]
+fn the_result_does_not_depend_on_how_the_bytes_arrive() {
+    for piece in [1, 17, 4096] {
+        let server = FakeServer::start();
+        server.route(CHAT, vec![Reply::fixture("thinking").trickle(piece, 0)]);
+        let done = run(&server).0.unwrap();
+        let want = support::expected("thinking");
+        assert_eq!(
+            done.content.as_deref(),
+            want["content"].as_str(),
+            "pieces of {piece}"
+        );
+        assert_eq!(
+            done.reasoning_content.as_deref(),
+            want["reasoning_content"].as_str()
+        );
+        assert_eq!(done.reasoning_tokens, 49);
+    }
+}
+
+#[test]
+fn a_busy_slot_is_waited_out() {
+    let server = FakeServer::start();
+    server.route(CHAT, vec![Reply::fixture("plain").head_delay(250)]);
+    server.route("/slots", vec![Reply::json(200, SLOT0_BUSY)]);
+    let (result, events, took) = run(&server);
+    assert!(result.is_ok(), "{result:?}");
+    let polls = waiting(&events);
+    assert!(polls.len() >= 4, "250 ms at 40 ms per poll: {polls:?}");
+    assert!(polls.iter().all(|busy| *busy));
+    assert!(took >= Duration::from_millis(250));
+    assert!(server.requests_to("/slots").len() >= 4);
+    assert_eq!(
+        server.requests_to("/slots")[0].target,
+        "/slots?model=test-model"
+    );
+}
+
+#[test]
+fn a_slot_that_stays_busy_is_a_wait_timeout() {
+    let server = FakeServer::start();
+    server.route(CHAT, vec![Reply::fixture("plain").head_delay(5_000)]);
+    server.route("/slots", vec![Reply::json(200, SLOT0_BUSY)]);
+    let (result, events, took) = run(&server);
+    assert!(matches!(result, Err(InferError::WaitTimeout)), "{result:?}");
+    assert!(
+        took >= Duration::from_millis(400) && took < Duration::from_millis(1500),
+        "{took:?}"
+    );
+    assert!(waiting(&events).iter().all(|busy| *busy));
+}
+
+#[test]
+fn only_the_requests_own_slot_counts_as_busy() {
+    // Slot 1 is busy, slot 0 (ours) is idle and silent: that is a stall, not a queue.
+    let server = FakeServer::start();
+    server.route(CHAT, vec![Reply::fixture("plain").head_delay(5_000)]);
+    server.route("/slots", vec![Reply::json(200, SLOT1_BUSY)]);
+    let (result, events, took) = run(&server);
+    assert!(matches!(result, Err(InferError::Stalled)), "{result:?}");
+    assert!(
+        took >= Duration::from_millis(200) && took < Duration::from_millis(1000),
+        "{took:?}"
+    );
+    assert!(waiting(&events).iter().all(|busy| !*busy));
+}
+
+#[test]
+fn an_unavailable_server_is_a_load_timeout() {
+    for slots_reply in [
+        Reply::json(503, r#"{"error":"loading"}"#),
+        Reply::raw("").cut_after(0),
+    ] {
+        let server = FakeServer::start();
+        server.route(CHAT, vec![Reply::fixture("plain").head_delay(5_000)]);
+        server.route("/slots", vec![slots_reply]);
+        let (result, _, took) = run(&server);
+        assert!(matches!(result, Err(InferError::LoadTimeout)), "{result:?}");
+        assert!(
+            took >= Duration::from_millis(300) && took < Duration::from_millis(1500),
+            "{took:?}"
+        );
+    }
+}
+
+#[test]
+fn the_wait_clocks_restart_when_the_state_changes() {
+    // Busy for about 5 polls, then idle: the idle grace starts counting from there, so the
+    // request outlives 200 ms of busy plus most of the 200 ms grace and then completes.
+    let server = FakeServer::start();
+    server.route(CHAT, vec![Reply::fixture("plain").head_delay(330)]);
+    let busy = Reply::json(200, SLOT0_BUSY);
+    server.route(
+        "/slots",
+        vec![
+            busy.clone(),
+            busy.clone(),
+            busy.clone(),
+            busy.clone(),
+            busy,
+            Reply::json(200, IDLE),
+        ],
+    );
+    let (result, events, _) = run(&server);
+    assert!(result.is_ok(), "{result:?}");
+    let polls = waiting(&events);
+    assert!(
+        polls.starts_with(&[true, true, true, true, true]),
+        "{polls:?}"
+    );
+    assert_eq!(polls.last(), Some(&false), "{polls:?}");
+}
+
+#[test]
+fn silence_in_mid_stream_is_a_stall() {
+    let server = FakeServer::start();
+    let reply = Reply::fixture("thinking");
+    let cut = reply.offset_after_events(10);
+    server.route(CHAT, vec![reply.hang_after(cut, 5_000)]);
+    let (result, events, took) = run(&server);
+    assert!(matches!(result, Err(InferError::Stalled)), "{result:?}");
+    assert!(
+        took >= Duration::from_millis(150) && took < Duration::from_millis(1000),
+        "{took:?}"
+    );
+    assert!(
+        events.iter().any(|e| matches!(e, ChatEvent::Reasoning(_))),
+        "events before the stall were delivered"
+    );
+}
+
+#[test]
+fn a_slow_but_steady_stream_is_not_a_stall() {
+    // 60 ms between pieces is well inside the 150 ms liveness limit, however long the whole takes.
+    let server = FakeServer::start();
+    let size = support::fixture_bytes("http", "turn1.http").len();
+    server.route(CHAT, vec![Reply::fixture("turn1").trickle(size / 8, 60)]);
+    let (result, _, took) = run(&server);
+    assert_eq!(result.unwrap().content.as_deref(), Some("Blue"));
+    assert!(
+        took >= Duration::from_millis(400),
+        "the stream took longer than the liveness limit: {took:?}"
+    );
+}
+
+#[test]
+fn a_stream_that_is_cut_is_closed_early_wherever_the_cut_falls() {
+    let reply = Reply::fixture("plain");
+    let whole = support::fixture_bytes("http", "plain.http").len();
+    let after_five = reply.offset_after_events(5);
+    // After the head only; in the middle of an event; between events; just before [DONE].
+    for cut in [300, after_five + 20, after_five, whole - 30] {
+        let server = FakeServer::start();
+        server.route(CHAT, vec![reply.clone().cut_after(cut)]);
+        let (result, _, _) = run(&server);
+        assert!(
+            matches!(result, Err(InferError::StreamClosedEarly)),
+            "cut at {cut}: {result:?}"
+        );
+    }
+    // Cut before the head is complete.
+    let server = FakeServer::start();
+    server.route(CHAT, vec![reply.cut_after(40)]);
+    assert!(matches!(run(&server).0, Err(InferError::StreamClosedEarly)));
+}
+
+#[test]
+fn error_statuses_keep_their_body() {
+    let server = FakeServer::start();
+    server.route(CHAT, vec![Reply::fixture("bad_request")]);
+    match run(&server).0 {
+        Err(InferError::Http { status: 400, body }) => {
+            assert!(body.contains("Expected 'messages'"), "{body}")
+        }
+        other => panic!("{other:?}"),
+    }
+    let server = FakeServer::start();
+    server.route(
+        CHAT,
+        vec![Reply::json(503, r#"{"error":{"message":"Loading model"}}"#)],
+    );
+    assert!(matches!(
+        run(&server).0,
+        Err(InferError::Http { status: 503, .. })
+    ));
+}
+
+#[test]
+fn garbage_in_the_stream_is_a_protocol_error() {
+    let server = FakeServer::start();
+    let raw = "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\r\ndata: {not json}\n\n";
+    server.route(CHAT, vec![Reply::raw(raw)]);
+    assert!(matches!(run(&server).0, Err(InferError::Protocol(_))));
+
+    let server = FakeServer::start();
+    let huge = format!(
+        "HTTP/1.1 200 OK\r\n\r\ndata: {}\n\n",
+        "x".repeat(2 * 1024 * 1024)
+    );
+    server.route(CHAT, vec![Reply::raw(huge)]);
+    assert!(
+        matches!(run(&server).0, Err(InferError::Protocol(_))),
+        "a line over 1 MiB"
+    );
+}
+
+#[test]
+fn no_server_is_a_connect_error() {
+    let server = FakeServer::start();
+    let mut cfg = support::test_config(&server.socket);
+    cfg.infer.socket = std::env::temp_dir().join("loopd-no-such-socket.sock");
+    let result = Client::new(cfg).chat(&request(), &mut |_| {});
+    assert!(matches!(result, Err(InferError::Connect(_))), "{result:?}");
+    let e: Box<dyn std::error::Error> = Box::new(InferError::Stalled);
+    assert!(!e.to_string().is_empty());
+}