Add props, slots, tokenize and the cache-loss check to the llama client

Implemented-By: OpenCode session (model recorded in docs/implementer-log.md)
2026-09-18 03:39:05 -07:00
parent 0fa481df29
commit 3663a1e4d5
4 changed files with 333 additions and 0 deletions
@@ -0,0 +1,145 @@
+//! Small non-streaming server calls and the cache-loss check.
+//!
+//! The responses are the server's format, so unknown fields are ignored; we read only the fields we
+//! need and treat anything else as a shape error.
+
+use std::time::Duration;
+
+use super::{Client, InferError, Timings};
+use crate::http::{Connection, HttpError, Request, read_capped};
+
+pub const MAX_BODY: usize = 4 * 1024 * 1024; // bodies that are not streams
+pub const MAX_ERROR_BODY: usize = 4 * 1024; // how much of an error body is kept
+pub const CACHE_TOLERANCE: u64 = 64;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Props {
+    pub chat_template: String,
+    pub n_ctx: u64,
+    pub total_slots: u32,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Deserialize)]
+pub struct SlotInfo {
+    pub id: u32,
+    pub is_processing: bool,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum CacheOutcome {
+    Hit,
+    Loss { expected: u64, got: u64 },
+}
+
+pub fn cache_outcome(previous: &Timings, current: &Timings) -> CacheOutcome {
+    let expected = previous
+        .cache_n
+        .saturating_add(previous.prompt_n)
+        .saturating_add(previous.predicted_n);
+    if current.cache_n.saturating_add(CACHE_TOLERANCE) >= expected {
+        CacheOutcome::Hit
+    } else {
+        CacheOutcome::Loss {
+            expected,
+            got: current.cache_n,
+        }
+    }
+}
+
+pub(crate) fn map_http(e: HttpError) -> InferError {
+    match e {
+        HttpError::Connect(e) => InferError::Connect(e),
+        HttpError::Timeout => InferError::Stalled,
+        HttpError::Closed => InferError::StreamClosedEarly,
+        other => InferError::Protocol(other.to_string()),
+    }
+}
+
+pub(crate) fn error_text(bytes: &[u8]) -> String {
+    let text = String::from_utf8_lossy(bytes).into_owned();
+    if text.len() <= MAX_ERROR_BODY {
+        return text;
+    }
+    let mut end = MAX_ERROR_BODY;
+    while end > 0 && !text.is_char_boundary(end) {
+        end -= 1;
+    }
+    text[..end].to_string()
+}
+
+impl Client {
+    pub(crate) fn call(
+        &self,
+        method: &str,
+        path: &str,
+        body: Option<&[u8]>,
+    ) -> Result<Vec<u8>, InferError> {
+        let mut conn = Connection::open(&self.cfg.infer.socket).map_err(map_http)?;
+        conn.set_read_timeout(Duration::from_millis(self.cfg.limits.liveness_ms))
+            .map_err(map_http)?;
+        conn.send(&Request { method, path, body })
+            .map_err(map_http)?;
+        let head = conn.read_head().map_err(map_http)?;
+        let mut body = conn.body(&head).map_err(map_http)?;
+        let bytes = read_capped(&mut body, MAX_BODY).map_err(map_http)?;
+        if head.status != 200 {
+            return Err(InferError::Http {
+                status: head.status,
+                body: error_text(&bytes),
+            });
+        }
+        Ok(bytes)
+    }
+
+    pub fn props(&self) -> Result<Props, InferError> {
+        let path = format!("/props?model={}", self.cfg.infer.model);
+        let bytes = self.call("GET", &path, None)?;
+        let value: serde_json::Value = serde_json::from_slice(&bytes)
+            .map_err(|_| InferError::Protocol("props response is not JSON".to_string()))?;
+        let chat_template = value
+            .get("chat_template")
+            .and_then(|v| v.as_str())
+            .ok_or_else(|| {
+                InferError::Protocol("props response has no chat_template".to_string())
+            })?;
+        let n_ctx = value
+            .get("default_generation_settings")
+            .and_then(|v| v.get("n_ctx"))
+            .and_then(|v| v.as_u64())
+            .ok_or_else(|| InferError::Protocol("props response has no n_ctx".to_string()))?;
+        let total_slots = value
+            .get("total_slots")
+            .and_then(|v| v.as_u64())
+            .ok_or_else(|| InferError::Protocol("props response has no total_slots".to_string()))?;
+        Ok(Props {
+            chat_template: chat_template.to_string(),
+            n_ctx,
+            total_slots: u32::try_from(total_slots)
+                .map_err(|_| InferError::Protocol("props total_slots is too large".to_string()))?,
+        })
+    }
+
+    pub fn slots(&self) -> Result<Vec<SlotInfo>, InferError> {
+        let path = format!("/slots?model={}", self.cfg.infer.model);
+        let bytes = self.call("GET", &path, None)?;
+        let slots: Vec<SlotInfo> = serde_json::from_slice(&bytes)
+            .map_err(|_| InferError::Protocol("slots response is not a list".to_string()))?;
+        Ok(slots)
+    }
+
+    pub fn tokenize(&self, text: &str) -> Result<usize, InferError> {
+        let body = serde_json::to_string(&serde_json::json!({
+            "model": self.cfg.infer.model,
+            "content": text,
+        }))
+        .map_err(|_| InferError::Protocol("could not encode tokenize request".to_string()))?;
+        let bytes = self.call("POST", "/tokenize", Some(body.as_bytes()))?;
+        let value: serde_json::Value = serde_json::from_slice(&bytes)
+            .map_err(|_| InferError::Protocol("tokenize response is not JSON".to_string()))?;
+        value
+            .get("tokens")
+            .and_then(|v| v.as_array())
+            .map(Vec::len)
+            .ok_or_else(|| InferError::Protocol("tokenize response has no tokens".to_string()))
+    }
+}
@@ -1,6 +1,7 @@
 //! The llama inference server's chat-completions types and the request builder.

 pub mod assemble;
+pub mod info;
 pub mod request;

 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -0,0 +1,186 @@
+//! Tests for the small server calls and the cache-loss check. Do not edit.
+
+mod support;
+
+use loopd::llama::info::{CacheOutcome, SlotInfo, cache_outcome};
+use loopd::llama::{Client, InferError, Timings};
+use support::{FakeServer, Reply};
+
+fn client(server: &FakeServer) -> Client {
+    Client::new(support::test_config(&server.socket))
+}
+
+fn t(cache_n: u64, prompt_n: u64, predicted_n: u64) -> Timings {
+    Timings {
+        cache_n,
+        prompt_n,
+        predicted_n,
+    }
+}
+
+#[test]
+fn props_slots_and_tokenize_read_the_recorded_responses() {
+    let server = FakeServer::start();
+    server.route("/props", vec![Reply::fixture("props")]);
+    server.route("/slots", vec![Reply::fixture("slots")]);
+    server.route("/tokenize", vec![Reply::fixture("tokenize")]);
+    let c = client(&server);
+
+    let props = c.props().unwrap();
+    assert_eq!(props.n_ctx, 131_072);
+    assert_eq!(props.total_slots, 2);
+    assert_eq!(
+        proto::sha256(props.chat_template.as_bytes())
+            .unwrap()
+            .to_hex(),
+        "f55f52930aa8bf44ab5cb85f99370fcc3c56e9a85640b812086d5330bce5d86b"
+    );
+
+    let want = vec![
+        SlotInfo {
+            id: 0,
+            is_processing: false,
+        },
+        SlotInfo {
+            id: 1,
+            is_processing: false,
+        },
+    ];
+    assert_eq!(c.slots().unwrap(), want);
+
+    assert_eq!(c.tokenize("The box is made.").unwrap(), 5);
+
+    // In router mode the model goes in the query for GET and in the body for POST.
+    assert_eq!(
+        server.requests_to("/props")[0].target,
+        "/props?model=test-model"
+    );
+    assert_eq!(
+        server.requests_to("/slots")[0].target,
+        "/slots?model=test-model"
+    );
+    let sent = server.requests_to("/tokenize")[0].json();
+    assert_eq!(
+        sent,
+        serde_json::json!({"model": "test-model", "content": "The box is made."})
+    );
+}
+
+#[test]
+fn an_error_status_keeps_its_body() {
+    let server = FakeServer::start();
+    server.route("/props", vec![Reply::fixture("unknown_model")]);
+    match client(&server).props() {
+        Err(InferError::Http { status: 400, body }) => {
+            assert!(body.contains("no-such-model"), "{body}")
+        }
+        other => panic!("{other:?}"),
+    }
+}
+
+#[test]
+fn a_long_error_body_is_cut_at_4_kib_on_a_character_boundary() {
+    let server = FakeServer::start();
+    // 3 bytes of padding, then 2-byte characters, so that byte 4096 falls inside a character.
+    let long = format!("abc{}", "\u{e9}".repeat(5000));
+    server.route("/props", vec![Reply::json(500, &long)]);
+    match client(&server).props() {
+        Err(InferError::Http { status: 500, body }) => {
+            assert!(
+                body.len() <= 4096 && body.len() >= 4094,
+                "{} bytes",
+                body.len()
+            );
+            assert!(body.starts_with("abc\u{e9}"));
+        }
+        other => panic!("{other:?}"),
+    }
+}
+
+#[test]
+fn unexpected_bodies_are_protocol_errors() {
+    let server = FakeServer::start();
+    server.route("/props", vec![Reply::json(200, r#"{"chat_template": 5}"#)]);
+    server.route("/slots", vec![Reply::json(200, "not json")]);
+    server.route("/tokenize", vec![Reply::json(200, r#"{"tokens": "many"}"#)]);
+    let c = client(&server);
+    assert!(matches!(c.props(), Err(InferError::Protocol(_))));
+    assert!(matches!(c.slots(), Err(InferError::Protocol(_))));
+    assert!(matches!(c.tokenize("x"), Err(InferError::Protocol(_))));
+}
+
+#[test]
+fn a_dead_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");
+    assert!(matches!(
+        Client::new(cfg).props(),
+        Err(InferError::Connect(_))
+    ));
+}
+
+/// Numbers from docs/inference-contract.md.
+#[test]
+fn cache_outcome_matches_what_was_measured() {
+    // (b) a normal turn: 539 processed + 53 generated, then 591 reused.
+    assert_eq!(
+        cache_outcome(&t(0, 539, 53), &t(591, 29, 61)),
+        CacheOutcome::Hit
+    );
+    // (k) a natural end to thinking: 64 + 104 = 168 left, 167 reused.
+    assert_eq!(
+        cache_outcome(&t(0, 64, 104), &t(167, 22, 2)),
+        CacheOutcome::Hit
+    );
+    // (k) a forced end to thinking: 89 + 717 = 806 left, only 85 reused.
+    assert_eq!(
+        cache_outcome(&t(0, 89, 717), &t(85, 743, 2)),
+        CacheOutcome::Loss {
+            expected: 806,
+            got: 85
+        }
+    );
+    // (h) the tool list changed: everything re-read.
+    assert_eq!(
+        cache_outcome(&t(685, 27, 55), &t(23, 920, 126)),
+        CacheOutcome::Loss {
+            expected: 767,
+            got: 23
+        }
+    );
+    // The recorded two-turn exchange used by the self-test.
+    assert_eq!(
+        cache_outcome(&t(15, 29, 2), &t(45, 30, 2)),
+        CacheOutcome::Hit
+    );
+}
+
+#[test]
+fn cache_outcome_tolerates_64_tokens_and_no_more() {
+    let previous = t(1000, 100, 50); // 1150 left in the slot
+    assert_eq!(cache_outcome(&previous, &t(1150, 5, 5)), CacheOutcome::Hit);
+    assert_eq!(
+        cache_outcome(&previous, &t(1086, 5, 5)),
+        CacheOutcome::Hit,
+        "64 short"
+    );
+    assert_eq!(
+        cache_outcome(&previous, &t(1085, 5, 5)),
+        CacheOutcome::Loss {
+            expected: 1150,
+            got: 1085
+        },
+        "65 short"
+    );
+    assert_eq!(
+        cache_outcome(&t(0, 0, 0), &t(0, 10, 1)),
+        CacheOutcome::Hit,
+        "a first request"
+    );
+    assert_eq!(
+        cache_outcome(&previous, &t(5000, 5, 5)),
+        CacheOutcome::Hit,
+        "more than expected"
+    );
+}