Add the thinking cap to Client::chat

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

crates/loopd/src/llama/chat.rs

@@ -65,12 +65,39 @@ impl Client {
         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),
@@ -81,7 +108,7 @@ impl Client {
 
         // 5. Finish. A stream that ended without a finish_reason is a server that died between two
         // events.
-        assembler.finish(false)
+        assembler.finish(cap_at.is_some())
     }
 }
 

crates/loopd/src/llama/info.rs

@@ -142,4 +142,24 @@ impl Client {
             .map(Vec::len)
             .ok_or_else(|| InferError::Protocol("tokenize response has no tokens".to_string()))
     }
+
+    pub(crate) fn end_reasoning(&self, completion_id: &str) -> Result<bool, InferError> {
+        let body = serde_json::to_string(&serde_json::json!({
+            "id": completion_id,
+            "action": "reasoning_end",
+            "model": self.cfg.infer.model,
+        }))
+        .map_err(|_| InferError::Protocol("could not encode reasoning_end request".to_string()))?;
+        let bytes = self.call(
+            "POST",
+            "/v1/chat/completions/control",
+            Some(body.as_bytes()),
+        )?;
+        let value: serde_json::Value = serde_json::from_slice(&bytes)
+            .map_err(|_| InferError::Protocol("control response is not JSON".to_string()))?;
+        value
+            .get("success")
+            .and_then(|v| v.as_bool())
+            .ok_or_else(|| InferError::Protocol("control response has no success".to_string()))
+    }
 }

crates/loopd/tests/cap.rs

@@ -0,0 +1,168 @@
+//! Tests for the thinking cap. Do not edit.
+//!
+//! The "capped" recording has 63 reasoning chunks (63 reasoning tokens) and then an answer. The
+//! fake server plays it regardless of the control call, so the cap and the overrun allowance
+//! decide what the client makes of it.
+
+mod support;
+
+use loopd::llama::{ChatEvent, ChatMessage, ChatRequest, Client, Completion, InferError};
+use support::{FakeServer, Reply};
+
+const CHAT: &str = "/v1/chat/completions";
+const CONTROL: &str = "/v1/chat/completions/control";
+
+fn run(
+    server: &FakeServer,
+    cap: u64,
+    overrun: u64,
+) -> (Result<Completion, InferError>, Vec<ChatEvent>) {
+    let mut cfg = support::test_config(&server.socket);
+    cfg.limits.thinking_cap = cap;
+    cfg.limits.thinking_overrun = overrun;
+    let req = ChatRequest {
+        slot: 0,
+        messages: vec![ChatMessage::User {
+            content: "think".to_string(),
+        }],
+        tools: vec![],
+        thinking: true,
+    };
+    let mut events = Vec::new();
+    let result = Client::new(cfg).chat(&req, &mut |e| events.push(e.clone()));
+    (result, events)
+}
+
+fn capped_events(events: &[ChatEvent]) -> Vec<u64> {
+    events
+        .iter()
+        .filter_map(|e| match e {
+            ChatEvent::ThinkingCapped { tokens } => Some(*tokens),
+            _ => None,
+        })
+        .collect()
+}
+
+#[test]
+fn under_the_cap_nothing_happens() {
+    let server = FakeServer::start();
+    server.route(CHAT, vec![Reply::fixture("capped")]);
+    let (result, events) = run(&server, 64, 10);
+    let done = result.unwrap();
+    assert!(!done.thinking_capped);
+    assert_eq!(done.reasoning_tokens, 63);
+    assert!(capped_events(&events).is_empty());
+    assert!(
+        server.requests_to(CONTROL).is_empty(),
+        "63 tokens is under a cap of 64"
+    );
+}
+
+#[test]
+fn at_the_cap_the_control_call_is_sent_once() {
+    let server = FakeServer::start();
+    server.route(CHAT, vec![Reply::fixture("capped")]);
+    server.route(CONTROL, vec![Reply::fixture("control")]);
+    let (result, events) = run(&server, 20, 100);
+    let done = result.unwrap();
+    assert!(done.thinking_capped);
+    assert_eq!(
+        done.reasoning_tokens, 63,
+        "the count keeps running after the cap"
+    );
+    assert_eq!(
+        done.content.as_deref(),
+        support::expected("capped")["content"].as_str()
+    );
+    assert_eq!(
+        capped_events(&events),
+        vec![20],
+        "one event, at the token count that hit the cap"
+    );
+
+    let control = server.requests_to(CONTROL);
+    assert_eq!(
+        control.len(),
+        1,
+        "the control call is made once, not on every later chunk"
+    );
+    let id = support::expected("capped")["id"].clone();
+    assert_eq!(
+        control[0].json(),
+        serde_json::json!({"id": id, "action": "reasoning_end", "model": "test-model"})
+    );
+
+    // The event sits between the reasoning that hit the cap and what came after.
+    let at = events
+        .iter()
+        .position(|e| matches!(e, ChatEvent::ThinkingCapped { .. }))
+        .unwrap();
+    let before = events[..at]
+        .iter()
+        .filter(|e| matches!(e, ChatEvent::Reasoning(_)))
+        .count();
+    assert_eq!(before, 20);
+}
+
+#[test]
+fn thinking_on_past_the_allowance_is_an_overrun() {
+    let server = FakeServer::start();
+    server.route(CHAT, vec![Reply::fixture("capped")]);
+    server.route(CONTROL, vec![Reply::fixture("control")]);
+    // Capped at 20; 63 reasoning tokens is more than 20 + 30.
+    let (result, events) = run(&server, 20, 30);
+    assert!(
+        matches!(result, Err(InferError::ThinkingOverrun)),
+        "{result:?}"
+    );
+    assert_eq!(capped_events(&events), vec![20]);
+    let reasoning = events
+        .iter()
+        .filter(|e| matches!(e, ChatEvent::Reasoning(_)))
+        .count();
+    assert_eq!(reasoning, 50, "the client stops reading at 20 + 30 tokens");
+}
+
+#[test]
+fn the_allowance_is_exact() {
+    let server = FakeServer::start();
+    server.route(CHAT, vec![Reply::fixture("capped")]);
+    server.route(CONTROL, vec![Reply::fixture("control")]);
+    // 63 is not 20 + 44, so this is allowed; with 43 it would be an overrun at token 63.
+    assert!(run(&server, 20, 44).0.is_ok());
+    assert!(matches!(
+        run(&server, 20, 43).0,
+        Err(InferError::ThinkingOverrun)
+    ));
+}
+
+#[test]
+fn a_control_call_that_fails_is_an_overrun() {
+    for reply in [
+        Reply::json(200, r#"{"success":false,"message":"no such completion"}"#),
+        Reply::json(500, r#"{"error":"boom"}"#),
+        Reply::json(200, "not json"),
+    ] {
+        let server = FakeServer::start();
+        server.route(CHAT, vec![Reply::fixture("capped")]);
+        server.route(CONTROL, vec![reply]);
+        let (result, events) = run(&server, 20, 100);
+        assert!(
+            matches!(result, Err(InferError::ThinkingOverrun)),
+            "{result:?}"
+        );
+        assert!(
+            capped_events(&events).is_empty(),
+            "the cap did not take effect, so no event"
+        );
+    }
+}
+
+#[test]
+fn a_completion_without_reasoning_is_never_capped() {
+    let server = FakeServer::start();
+    server.route(CHAT, vec![Reply::fixture("plain")]);
+    let (result, _) = run(&server, 1, 0);
+    assert!(!result.unwrap().thinking_capped);
+    assert!(server.requests_to(CONTROL).is_empty());
+}