From 7d3adf8e48a48a1b5ffecebe9101ef7a91750ba7 Mon Sep 17 00:00:00 2001 From: Arthur Date: Mon, 31 Aug 2026 12:49:13 +0200 Subject: [PATCH 1/2] fix(ai): stop giving up on a model that is still working The AI client waited for a whole response with nothing on the wire, so the only question it could ask was whether too much time had passed. A long document's prefill looks exactly like a stuck server from there, so filing a big scan on a small Mac got cancelled mid-flight and every minute already spent on it was thrown away. It now streams and times the silence instead. While an answer keeps arriving there is no deadline at all, however slow the machine. A server that actually stops is caught in 90 seconds rather than 300. --- lib/stack/ai/client.py | 133 ++++++++++++++++---- tests/framework/test_ai_client.py | 194 +++++++++++++++++++++++++----- tests/integration/openai_stub.py | 32 ++--- 3 files changed, 289 insertions(+), 70 deletions(-) diff --git a/lib/stack/ai/client.py b/lib/stack/ai/client.py index 52bd5700..389e3e54 100644 --- a/lib/stack/ai/client.py +++ b/lib/stack/ai/client.py @@ -21,10 +21,12 @@ from __future__ import annotations +import asyncio import base64 import os from dataclasses import dataclass +import httpx import openai from loguru import logger from openai import AsyncOpenAI @@ -148,6 +150,35 @@ class LLMImage: ) +# ── Liveness, not deadlines ────────────────────────────────────────────── +# +# A local model on a small Mac is slow, not broken. A 30K-token prefill on +# a scanned document can run for minutes before the first token, and a +# wall-clock deadline cannot tell that apart from a wedged server: it just +# hangs up on work that was progressing fine, throwing away everything +# spent so far. That is what a non-streaming call forces, because not one +# byte comes back until generation is complete. +# +# So we stream and time the *gaps* instead. Two different questions, two +# different budgets: +# +# FIRST_TOKEN_TIMEOUT — how long we allow silence before anything comes +# back. This covers prefill, which scales with input size and with how +# hard the machine is swapping, so it has to be generous. +# STALL_TIMEOUT — how long we allow silence *between* pieces of an answer +# already underway. Once tokens flow, gaps are short on any machine, so +# a long one means the server stopped rather than slowed. Measured +# against oMLX, a warm 9B on an idle Mac still went quiet for ~4s +# mid-answer (it batches tokens), so this needs real headroom, not a +# tight bound. +# +# The consequence that matters: while output is arriving we never give up, +# no matter how long the document takes. A genuinely hung endpoint is +# caught sooner than the old deadline caught it. +FIRST_TOKEN_TIMEOUT = 900.0 +STALL_TIMEOUT = 90.0 + + class LLM: """An OpenAI-compatible chat client scoped to one stacklet/bot. @@ -158,9 +189,13 @@ class LLM: """ def __init__(self, client: AsyncOpenAI, *, namespace: str | None = None, - capabilities: ModelCapabilities | None = None): + capabilities: ModelCapabilities | None = None, + first_token_timeout: float = FIRST_TOKEN_TIMEOUT, + stall_timeout: float = STALL_TIMEOUT): self._client = client self.namespace = namespace + self.first_token_timeout = first_token_timeout + self.stall_timeout = stall_timeout # Default to in-memory; bots inject a disk-backed cache so the # vision probe survives container restarts. self.capabilities = capabilities or ModelCapabilities() @@ -190,24 +225,17 @@ def from_env(cls, *, namespace: str | None = None, "No AI endpoint configured — set up AI with 'stack up ai'" ) key = os.environ.get("OPENAI_KEY", "") or "not-needed" - # This is a read timeout, and the call is not streamed, so it - # covers the entire silent wait while the model reads a document - # and generates an answer. Prefill alone, measured against local - # oMLX (Qwen3.5-9B-MLX-4bit) on a dense 14-page PDF: - # - # 1,000 tokens -> 7.6s 15,000 tokens -> 58.4s - # 3,000 tokens -> 20.4s 28,400 tokens -> 163.6s - # 7,500 tokens -> 36.7s - # - # That machine had RAM to spare. A 16GB Mac running photos and - # documents alongside is several times slower, so 120s and then - # 300s were both cancelling healthy work partway through a big - # scan and discarding every minute already spent on it. 900s - # covers ~158K tokens here, past any context this serves, and - # still bounds a genuinely dead endpoint. + # No read timeout on purpose: `complete` streams and enforces + # liveness per chunk (see FIRST_TOKEN_TIMEOUT / STALL_TIMEOUT). + # A deadline down here could only ever fire on work that was still + # producing output, which is the bug being fixed. Raising the + # ceiling (120s, then 300s, then 900s) only moved the point at + # which a big enough document got cancelled; timing the silence + # instead removes it. Connect stays short: an endpoint that will + # not accept a socket is down, and that answer arrives at once. client = AsyncOpenAI( base_url=url, api_key=key, max_retries=max_retries, - timeout=900.0, + timeout=httpx.Timeout(None, connect=10.0), ) return cls(client, namespace=namespace, capabilities=capabilities) @@ -242,11 +270,26 @@ async def complete(self, role: str, prompt: str, *, kwargs["temperature"] = temperature try: - resp = await self._client.chat.completions.create( - model=model, - messages=[{"role": "user", "content": content}], - **kwargs, + # The first-token budget has to cover this call, not just the + # chunks after it: a server that withholds its response headers + # until the first token spends the whole prefill in here, where + # no per-chunk timeout can see it. With no read timeout on the + # SDK, that would otherwise be an unbounded wait. + stream = await asyncio.wait_for( + self._client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": content}], + stream=True, + **kwargs, + ), + timeout=self.first_token_timeout, ) + return await self._read_stream(stream, model) + except asyncio.TimeoutError as e: + raise LLMTimeoutError( + f"{model} — went silent before answering " + f"for {self.first_token_timeout:.0f}s" + ) from e # Order matters: APITimeoutError < APIConnectionError, and # Authentication/NotFound < APIStatusError — catch specific first. except openai.APITimeoutError as e: @@ -260,7 +303,53 @@ async def complete(self, role: str, prompt: str, *, except openai.APIStatusError as e: raise LLMUnavailableError(f"HTTP {e.status_code}: {str(e)[:200]}") from e - return resp.choices[0].message.content or "" + async def _read_stream(self, stream, model: str) -> str: + """Reassemble a streamed completion, giving up only on silence. + + Callers still get one string back; streaming is an implementation + detail they should not have to care about. What it buys is the + progress signal: every chunk that arrives proves the server is + still working, so the budget for the next one resets. + + The switch from the first-token budget to the stall budget waits + for actual *content*, not merely the first chunk. Some servers open + with a role-only preamble before prefill has finished, and treating + that as "generation started" would drop us onto the tight budget + with the long silent part still ahead. + """ + parts: list[str] = [] + chunks = stream.__aiter__() + while True: + budget = self.stall_timeout if parts else self.first_token_timeout + try: + chunk = await asyncio.wait_for(chunks.__anext__(), timeout=budget) + except StopAsyncIteration: + break + except asyncio.TimeoutError as e: + await self._close_stream(stream) + where = "mid-answer" if parts else "before answering" + raise LLMTimeoutError( + f"{model} — went silent {where} for {budget:.0f}s" + ) from e + # A chunk with no choices (a usage-only trailer) still counts as + # the server being alive, it just carries no text. + if chunk.choices: + delta = chunk.choices[0].delta + if delta is not None and delta.content: + parts.append(delta.content) + return "".join(parts) + + @staticmethod + async def _close_stream(stream) -> None: + """Release an abandoned stream's connection, best effort. + + We only get here on a timeout, where the useful error is the + timeout itself; a failure to hang up cleanly must not replace it. + """ + try: + await stream.close() + except Exception as e: + logger.debug("[llm] closing abandoned stream failed: {}", e) async def has_vision(self, *, role: str = "classifier", model_override: str | None = None) -> bool: diff --git a/tests/framework/test_ai_client.py b/tests/framework/test_ai_client.py index 3e80f72e..9692f2c5 100644 --- a/tests/framework/test_ai_client.py +++ b/tests/framework/test_ai_client.py @@ -14,6 +14,7 @@ import base64 import json import sys +import time from pathlib import Path import httpx @@ -49,7 +50,9 @@ def stub_resolver(monkeypatch): def _make_llm(httpserver: HTTPServer, *, timeout: float = 5.0, - capabilities: ModelCapabilities | None = None) -> LLM: + capabilities: ModelCapabilities | None = None, + first_token_timeout: float = 5.0, + stall_timeout: float = 5.0) -> LLM: """Build an LLM pointed at the local httpserver mock. `max_retries=0` keeps tests fast and lets a 401/404 surface on the @@ -60,21 +63,162 @@ def _make_llm(httpserver: HTTPServer, *, timeout: float = 5.0, max_retries=0, timeout=timeout, ) - return LLM(client, namespace="test-bot", capabilities=capabilities) + return LLM(client, namespace="test-bot", capabilities=capabilities, + first_token_timeout=first_token_timeout, + stall_timeout=stall_timeout) -def _completion_payload(content: str, *, model: str = "test-model") -> dict: - return { +def _chunk(delta: dict, *, model: str = "test-model") -> str: + """One `chat.completion.chunk` as an SSE event.""" + body = { "id": "cmpl-test", - "object": "chat.completion", + "object": "chat.completion.chunk", "model": model, - "choices": [{ - "index": 0, - "message": {"role": "assistant", "content": content}, - "finish_reason": "stop", - }], - "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, + "choices": [{"index": 0, "delta": delta, "finish_reason": None}], } + return f"data: {json.dumps(body)}\n\n" + + +def _sse_body(*contents: str, model: str = "test-model") -> str: + """A complete streamed answer: one chunk per piece, then DONE.""" + events = [_chunk({"content": c}, model=model) for c in contents] + return "".join(events) + "data: [DONE]\n\n" + + +def _sse_response(*contents: str, model: str = "test-model"): + """The whole answer at once. `complete` still returns it joined.""" + from werkzeug.wrappers import Response + return Response(_sse_body(*contents, model=model), + content_type="text/event-stream") + + +def _sse_trickle(*contents: str, gap: float = 0.0, lead: float = 0.0, + tail: float = 0.0, preamble: bool = False): + """A stream paced like a real one, so silence can be tested. + + `lead` is the quiet before anything arrives (prefill). `gap` is the + quiet between pieces. `tail` is a hang after the last piece but before + DONE. `preamble` opens with a role-only chunk carrying no content, the + way some servers announce themselves before generating. + """ + from werkzeug.wrappers import Response + + def generate(): + if preamble: + yield _chunk({"role": "assistant"}) + if lead: + time.sleep(lead) + for i, c in enumerate(contents): + if i and gap: + time.sleep(gap) + yield _chunk({"content": c}) + if tail: + time.sleep(tail) + yield "data: [DONE]\n\n" + + return Response(generate(), content_type="text/event-stream") + + +# ── complete(): liveness ───────────────────────────────────────────────── + +class TestSlowIsNotBroken: + """The distinction the old wall-clock timeout could not make. + + A local model on a small Mac is slow, not broken. The client used to + wait for a complete response with no bytes on the wire, so the only + question it could ask was "has too much time passed", and the answer + was the same for a 30K-token prefill that was progressing fine and a + server that had wedged. It cancelled healthy work mid-flight and threw + away everything spent on it. + + Streaming turns that into a question worth asking: has it gone *quiet*. + """ + + async def test_an_answer_that_keeps_arriving_is_never_cut_off( + self, httpserver: HTTPServer): + """Total time far exceeds the stall budget; no single gap does. + + This is the property that matters for a slow machine: as long as + output keeps coming, there is no deadline at all. + """ + httpserver.expect_request( + "/v1/chat/completions", method="POST", + ).respond_with_handler( + lambda _r: _sse_trickle("slow ", "but ", "still ", "working", gap=0.3), + ) + + llm = _make_llm(httpserver, stall_timeout=0.6, first_token_timeout=2.0) + result = await llm.complete("classifier", "long document") + assert result == "slow but still working" + await llm.aclose() + + async def test_a_long_silent_prefill_is_allowed(self, httpserver: HTTPServer): + """Nothing comes back while the model reads a long document. That + silence is the work, not a fault, so it gets its own budget.""" + httpserver.expect_request( + "/v1/chat/completions", method="POST", + ).respond_with_handler(lambda _r: _sse_trickle("done", lead=0.8)) + + llm = _make_llm(httpserver, first_token_timeout=3.0, stall_timeout=0.3) + assert await llm.complete("classifier", "scan") == "done" + await llm.aclose() + + async def test_a_role_only_preamble_does_not_start_the_stall_clock( + self, httpserver: HTTPServer): + """Some servers announce themselves before generating. Counting + that as "generation started" would drop us onto the tight budget + with the whole silent prefill still ahead.""" + httpserver.expect_request( + "/v1/chat/completions", method="POST", + ).respond_with_handler( + lambda _r: _sse_trickle("here", preamble=True, lead=0.8), + ) + + llm = _make_llm(httpserver, first_token_timeout=3.0, stall_timeout=0.3) + assert await llm.complete("classifier", "scan") == "here" + await llm.aclose() + + +class TestSilenceIsBroken: + """The other half: a server that stops producing must still be caught, + and caught sooner than the old deadline caught it.""" + + async def test_silence_before_any_answer_gives_up(self, httpserver: HTTPServer): + httpserver.expect_request( + "/v1/chat/completions", method="POST", + ).respond_with_handler(lambda _r: _sse_trickle("never", lead=2.0)) + + llm = _make_llm(httpserver, first_token_timeout=0.3, stall_timeout=0.3) + with pytest.raises(LLMTimeoutError) as excinfo: + await llm.complete("classifier", "hi") + assert "before answering" in str(excinfo.value) + await llm.aclose() + + async def test_silence_partway_through_an_answer_gives_up( + self, httpserver: HTTPServer): + """Generation started and then stopped. The tight budget applies + here precisely because tokens were already flowing.""" + httpserver.expect_request( + "/v1/chat/completions", method="POST", + ).respond_with_handler(lambda _r: _sse_trickle("started", "stalled", gap=2.0)) + + llm = _make_llm(httpserver, first_token_timeout=3.0, stall_timeout=0.3) + with pytest.raises(LLMTimeoutError) as excinfo: + await llm.complete("classifier", "hi") + assert "mid-answer" in str(excinfo.value) + await llm.aclose() + + +class TestReassembly: + async def test_deltas_are_joined_into_one_answer(self, httpserver: HTTPServer): + """Streaming is an implementation detail; callers get a string.""" + httpserver.expect_request( + "/v1/chat/completions", method="POST", + ).respond_with_handler(lambda _r: _sse_response("Kwik-", "E-", "Mart")) + + llm = _make_llm(httpserver) + assert await llm.complete("classifier", "shop") == "Kwik-E-Mart" + await llm.aclose() # ── complete(): happy path ─────────────────────────────────────────────── @@ -83,7 +227,7 @@ class TestComplete: async def test_returns_assistant_text(self, httpserver: HTTPServer): httpserver.expect_request( "/v1/chat/completions", method="POST", - ).respond_with_json(_completion_payload("hello world")) + ).respond_with_handler(lambda _r: _sse_response("hello world")) llm = _make_llm(httpserver) result = await llm.complete("classifier", "say hi") @@ -98,11 +242,7 @@ async def test_sends_resolved_model(self, httpserver: HTTPServer, monkeypatch): def handler(request): captured["body"] = json.loads(request.get_data().decode()) - from werkzeug.wrappers import Response - return Response( - json.dumps(_completion_payload("ok")), - content_type="application/json", - ) + return _sse_response("ok") httpserver.expect_request( "/v1/chat/completions", method="POST", @@ -118,11 +258,7 @@ async def test_json_mode_sets_response_format(self, httpserver: HTTPServer): def handler(request): captured["body"] = json.loads(request.get_data().decode()) - from werkzeug.wrappers import Response - return Response( - json.dumps(_completion_payload('{"k":1}')), - content_type="application/json", - ) + return _sse_response('{"k":1}') httpserver.expect_request( "/v1/chat/completions", method="POST", @@ -138,11 +274,7 @@ async def test_images_sent_as_data_url(self, httpserver: HTTPServer): def handler(request): captured["body"] = json.loads(request.get_data().decode()) - from werkzeug.wrappers import Response - return Response( - json.dumps(_completion_payload("ok")), - content_type="application/json", - ) + return _sse_response("ok") httpserver.expect_request( "/v1/chat/completions", method="POST", @@ -164,11 +296,7 @@ async def test_model_override_bypasses_resolver(self, httpserver: HTTPServer): def handler(request): captured["body"] = json.loads(request.get_data().decode()) - from werkzeug.wrappers import Response - return Response( - json.dumps(_completion_payload("ok")), - content_type="application/json", - ) + return _sse_response("ok") httpserver.expect_request( "/v1/chat/completions", method="POST", @@ -322,7 +450,7 @@ async def test_probe_success_caches_true(self, httpserver: HTTPServer, tmp_path) cap = ModelCapabilities(path=tmp_path / "caps.json") httpserver.expect_request( "/v1/chat/completions", method="POST", - ).respond_with_json(_completion_payload("ok")) + ).respond_with_handler(lambda _r: _sse_response("ok")) llm = _make_llm(httpserver, capabilities=cap) assert await llm.has_vision() is True diff --git a/tests/integration/openai_stub.py b/tests/integration/openai_stub.py index c1afc38d..601383d1 100644 --- a/tests/integration/openai_stub.py +++ b/tests/integration/openai_stub.py @@ -53,19 +53,27 @@ _KINDS = ("classify", "reformat", "rewrite", "synthesize") -def _chat_completion(content: str, model: str = "test-model") -> dict: - """OpenAI chat.completion response envelope.""" - return { +def _chat_completion(content: str, model: str = "test-model") -> Response: + """A streamed chat completion, the way real backends answer. + + The client streams every call so it can tell a slow model from a stopped + one, so a stub that replies with a single JSON body is no longer + speaking the same protocol as the thing it stands in for. Sent as one + content chunk: the stub's job is to be a believable endpoint, not to + reproduce token-by-token pacing (the client's own tests cover pacing). + """ + chunk = { "id": "cmpl-test", - "object": "chat.completion", + "object": "chat.completion.chunk", "model": model, "choices": [{ "index": 0, - "message": {"role": "assistant", "content": content}, - "finish_reason": "stop", + "delta": {"role": "assistant", "content": content}, + "finish_reason": None, }], - "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, } + body = f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n" + return Response(body, content_type="text/event-stream") def _prompt_of(body: dict) -> str: @@ -125,10 +133,7 @@ def _dispatch(self, request: Request) -> Response: body = {} prompt = _prompt_of(body) if prompt.lstrip().startswith(_VISION_PROBE_MARKER): - return Response( - json.dumps(_chat_completion("ok")), - content_type="application/json", - ) + return _chat_completion("ok") head = " ".join(prompt.split())[:120] kind = _route(prompt) with self._lock: @@ -142,10 +147,7 @@ def _dispatch(self, request: Request) -> Response: ) return Response(f"openai stub: {kind} queue empty", status=500) content = queue.popleft() - return Response( - json.dumps(_chat_completion(content)), - content_type="application/json", - ) + return _chat_completion(content) # ── queueing ───────────────────────────────────────────────────────── From e1c2f28e85862045c1bd5a1e0374214c9bd830cf Mon Sep 17 00:00:00 2001 From: Arthur Date: Mon, 31 Aug 2026 14:04:53 +0200 Subject: [PATCH 2/2] docs(ai): record the measured prefill curve behind the timeout budgets --- lib/stack/ai/client.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/lib/stack/ai/client.py b/lib/stack/ai/client.py index 389e3e54..3117d7a3 100644 --- a/lib/stack/ai/client.py +++ b/lib/stack/ai/client.py @@ -167,10 +167,21 @@ class LLMImage: # hard the machine is swapping, so it has to be generous. # STALL_TIMEOUT — how long we allow silence *between* pieces of an answer # already underway. Once tokens flow, gaps are short on any machine, so -# a long one means the server stopped rather than slowed. Measured -# against oMLX, a warm 9B on an idle Mac still went quiet for ~4s -# mid-answer (it batches tokens), so this needs real headroom, not a -# tight bound. +# a long one means the server stopped rather than slowed. +# +# Measured against local oMLX (Qwen3.5-9B-MLX-4bit, dense 14-page PDF), to +# first token: +# +# 1,000 tokens -> 7.6s 15,000 tokens -> 58.4s +# 3,000 tokens -> 20.4s 28,400 tokens -> 163.6s +# 7,500 tokens -> 36.7s +# +# Roughly linear at ~5.7ms/token, and that is a machine with RAM to spare; +# a 16GB Mac running photos and documents alongside is several times +# slower. Gaps *after* the first token were 0.07-0.13s. Hence the lopsided +# pair: the first-token budget must swallow minutes of nothing, while the +# stall budget can be tight enough to notice a dead server and still sit +# orders of magnitude above any real gap. # # The consequence that matters: while output is arriving we never give up, # no matter how long the document takes. A genuinely hung endpoint is