Coverage for src/lilbee/server/handlers/sse.py: 100%
185 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-17 10:02 +0000
1"""SSE stream primitives shared by every streaming handler."""
3from __future__ import annotations
5import asyncio
6import json
7import logging
8import threading
9import time
10from collections import deque
11from collections.abc import AsyncGenerator, Callable
12from typing import Any, NamedTuple
14from pydantic import BaseModel
16from lilbee.core.config import cfg
17from lilbee.providers.base import ProviderErrorKind, filter_options
18from lilbee.runtime.progress import (
19 DetailedProgressCallback,
20 EventType,
21 ProgressEvent,
22 SseErrorCode,
23 SseEvent,
24)
25from lilbee.server.chat_completions_api.errors import CompletionsErrorCode
27log = logging.getLogger(__name__)
29# Litestar documents a response's content type from the route decorator, so a
30# streaming route passes this to both the decorator and the Stream it returns.
31SSE_MEDIA_TYPE = "text/event-stream"
33# Machine-readable ``code`` on an SSE error event. Load-time failures use
34# SseErrorCode; failed provider calls reuse ProviderErrorKind directly; the
35# RAG chat stream reuses the wire-layer CompletionsErrorCode for typed
36# dispatch errors (unknown model, no tool support, context overflow) so a
37# single client-facing vocabulary covers both surfaces.
38SseErrorCodeValue = SseErrorCode | ProviderErrorKind | CompletionsErrorCode
41async def frames_with_keepalive(
42 frames: AsyncGenerator[bytes, None] | Any,
43 *,
44 keepalive: bytes,
45 interval_s: float,
46) -> AsyncGenerator[bytes, None]:
47 """Yield *frames*, interleaving *keepalive* whenever the upstream is slow.
49 Keeps clients from tripping their idle-stream timeout during slow
50 first-token latency on local models, which would otherwise fire a retry
51 storm against the chat lock. The in-progress ``__anext__`` task is held
52 across keepalive emissions; only on real completion or error is the task
53 replaced.
54 """
55 iterator = frames.__aiter__()
56 pending = asyncio.ensure_future(iterator.__anext__())
57 try:
58 while True:
59 done, _ = await asyncio.wait({pending}, timeout=interval_s)
60 if pending not in done:
61 yield keepalive
62 continue
63 try:
64 frame = pending.result()
65 except StopAsyncIteration:
66 break
67 yield frame
68 pending = asyncio.ensure_future(iterator.__anext__())
69 finally:
70 if not pending.done():
71 pending.cancel()
72 try:
73 await pending
74 except asyncio.CancelledError:
75 # Expected: the future we just cancelled. By name, not
76 # BaseException, which also swallowed a Ctrl-C landing here.
77 pass
78 except Exception:
79 # The upstream failed as it was torn down. The request is
80 # already unwinding, so record it rather than mask the unwind.
81 log.debug("upstream stream errored during keepalive cleanup", exc_info=True)
84def sse_event(event: str, data: Any) -> str:
85 """Format a single Server-Sent Event string."""
86 return f"event: {event}\ndata: {json.dumps(data)}\n\n"
89def sse_error(
90 message: str, *, code: SseErrorCodeValue | None = None, detail: str | None = None
91) -> str:
92 """Format an SSE error event with optional structured ``code`` / ``detail``."""
93 payload: dict[str, Any] = {"message": message}
94 if code is not None:
95 payload["code"] = code
96 if detail is not None:
97 payload["detail"] = detail
98 return sse_event(SseEvent.ERROR, payload)
101# Phrases that describe an allocation failure. "llama_context" alone used to be
102# here, but that is the prefix llama.cpp stamps on every context-subsystem
103# diagnostic, so n_ctx-over-training-context and KV-cache rejections were all
104# reported as "model too large", pointing at a smaller model when the fix is a
105# config change.
106_OOM_MARKERS = ("failed to load", "free ram", "try a smaller model", "failed to allocate")
107_NOT_INSTALLED_MARKERS = ("is not installed", "is not available", "pull it first")
110def classify_load_error(message: str) -> tuple[SseErrorCode | None, str]:
111 """Return ``(code, user_message)`` for an SSE error event.
113 Maps the llama.cpp out-of-memory diagnostic and the "configured model
114 isn't installed" failure to stable codes. Anything else returns a generic
115 code-less message.
116 """
117 lowered = message.lower()
118 if any(marker in lowered for marker in _OOM_MARKERS):
119 return SseErrorCode.MODEL_TOO_LARGE, "Model too large for available RAM"
120 if any(marker in lowered for marker in _NOT_INSTALLED_MARKERS):
121 return (
122 SseErrorCode.MODEL_NOT_INSTALLED,
123 "Active model isn't installed. Pull it from the catalog.",
124 )
125 return None, "Internal error"
128def sse_done(data: dict[str, Any]) -> str:
129 """Format an SSE done event."""
130 return sse_event(SseEvent.DONE, data)
133def _resolve_generation_options(options: dict[str, Any] | None) -> dict[str, Any] | None:
134 """Merge HTTP-supplied options with config, allowlisting sampling keys only.
136 ``filter_options`` is the validation boundary for untrusted callers: it
137 drops anything outside the sampling allowlist (e.g. injected ``api_base`` /
138 ``api_key``) before the values reach a provider.
139 """
140 return cfg.generation_options(**filter_options(options)) if options else None
143# Cap on buffered SSE events per stream. A bulk sync emits per-file and
144# per-chunk progress far faster than a slow client reads; past this bound the
145# stream sheds the oldest progress event rather than buffering millions.
146SSE_QUEUE_MAX_EVENTS = 1000
148# Poll ceiling once the producer task has finished, for the case where the
149# queue drains empty without the sentinel ever arriving.
150_FINISHED_PRODUCER_POLL_S = 1.0
152# Progress-class event types: high-frequency, safe to coalesce under
153# backpressure. Everything else (done, errors, crawl/setup lifecycle) must land.
154_DROPPABLE_EVENT_TYPES: frozenset[EventType | SseEvent] = frozenset(
155 {
156 EventType.FILE_START,
157 EventType.FILE_DONE,
158 EventType.BATCH_PROGRESS,
159 EventType.EMBED,
160 EventType.EXTRACT,
161 EventType.CRAWL_PAGE,
162 EventType.WIKI_PAGE,
163 EventType.SETUP_PROGRESS,
164 SseEvent.PROGRESS,
165 }
166)
169class _QueuedEvent(NamedTuple):
170 """A queued SSE payload tagged with whether backpressure may shed it."""
172 payload: str | None
173 droppable: bool
176class SseEventQueue(asyncio.Queue[str | None]):
177 """Bounded SSE queue: progress events shed under backpressure, the rest land.
179 ``put_nowait`` (lifecycle events, tokens, the ``None`` sentinel) always
180 enqueues, evicting the oldest progress event first when at capacity.
181 ``put_event_nowait`` enqueues progress-protocol events, dropping the oldest
182 progress event (or the incoming one when the head is not progress) at
183 capacity. ``join()`` semantics are not supported.
184 """
186 _queue: deque[_QueuedEvent]
188 def __init__(self, max_events: int = SSE_QUEUE_MAX_EVENTS) -> None:
189 super().__init__()
190 self._max_events = max_events
191 self._put_droppable = False
192 self.dropped_events = 0
193 # Set once the queue is full of undroppable events, i.e. the consumer
194 # has stopped reading. See put_nowait.
195 self.stalled = False
197 def _put(self, item: str | None) -> None:
198 self._queue.append(_QueuedEvent(item, self._put_droppable))
200 def _get(self) -> str | None:
201 return self._queue.popleft().payload
203 def _evict_oldest_droppable(self) -> bool:
204 """Shed the oldest progress event anywhere in the queue; True when one went.
206 Scans rather than checking only the head. A stream that interleaves
207 tokens with progress puts a non-droppable event at the head almost
208 immediately, and head-only eviction then reported "nothing to shed"
209 while the queue still held progress events it was allowed to drop.
210 """
211 for index, event in enumerate(self._queue):
212 if event.droppable:
213 del self._queue[index]
214 self.dropped_events += 1
215 return True
216 return False
218 def put_nowait(self, item: str | None) -> None:
219 """Enqueue an always-delivered event, evicting old progress when full."""
220 if self.qsize() >= self._max_events and not self._evict_oldest_droppable():
221 # Full of undroppable events (chat and RAG tokens come through
222 # here) and the consumer has taken none in _max_events: a stalled
223 # or gone client. The alternative is unbounded growth.
224 self.stalled = True
225 self._put_droppable = False
226 super().put_nowait(item)
228 def put_event_nowait(self, payload: str, event_type: EventType | SseEvent) -> None:
229 """Enqueue a progress-protocol event, shedding progress when full."""
230 if event_type not in _DROPPABLE_EVENT_TYPES:
231 self.put_nowait(payload)
232 return
233 if self.qsize() >= self._max_events and not self._evict_oldest_droppable():
234 self.dropped_events += 1
235 return
236 self._put_droppable = True
237 try:
238 super().put_nowait(payload)
239 finally:
240 self._put_droppable = False
243class SseStream:
244 """Context object for SSE streaming with cancellation support.
245 Bundles the queue, cancel event, and progress callback that every SSE
246 endpoint needs. Call :meth:`drain` to yield events until the task
247 completes or the client disconnects.
248 """
250 def __init__(self) -> None:
251 self.queue: SseEventQueue = SseEventQueue()
252 self.cancel = threading.Event()
253 self.loop = asyncio.get_running_loop()
254 self.callback: DetailedProgressCallback = self._build_callback()
256 def put_threadsafe(self, item: str | None) -> None:
257 """Enqueue an always-delivered event from a worker thread.
259 ``asyncio.Queue.put_nowait`` is not thread-safe: it wakes a pending
260 getter via ``Future.set_result``, which must run on the loop thread. A
261 producer running under ``run_in_executor`` therefore hands the put back
262 to the loop instead of mutating the queue directly.
263 """
264 self.loop.call_soon_threadsafe(self._put_and_check_stall, item)
266 def _put_and_check_stall(self, item: str | None) -> None:
267 """Enqueue on the loop thread, cancelling the producer if it has stalled.
269 A consumer that has taken nothing in a full queue's worth is gone;
270 cancelling is the same signal a detected disconnect sends.
271 """
272 self.queue.put_nowait(item)
273 if self.queue.stalled and not self.cancel.is_set():
274 log.warning(
275 "SSE consumer stalled with %d undroppable events queued; cancelling the producer.",
276 self.queue.qsize(),
277 )
278 self.cancel.set()
280 def _build_callback(self) -> DetailedProgressCallback:
281 """Create a progress callback that serializes events into the queue.
282 Safe to call from both the event-loop thread and worker threads.
283 """
284 loop = self.loop
285 queue = self.queue
287 def _callback(event_type: EventType, data: ProgressEvent) -> None:
288 serialized = data.model_dump() if isinstance(data, BaseModel) else data
289 payload = f"event: {event_type}\ndata: {json.dumps(serialized)}\n\n"
290 try:
291 running = asyncio.get_running_loop()
292 except RuntimeError:
293 running = None
294 if running is loop:
295 queue.put_event_nowait(payload, event_type)
296 else:
297 loop.call_soon_threadsafe(queue.put_event_nowait, payload, event_type)
299 return _callback
301 async def _flush_pending(self) -> AsyncGenerator[str, None]:
302 """Events left behind the sentinel by a producer that outran the consumer.
304 A fast producer can enqueue its sentinel before its threadsafe progress
305 callbacks run; one loop tick lets them land.
306 """
307 await asyncio.sleep(0)
308 while not self.queue.empty():
309 leftover = self.queue.get_nowait()
310 if leftover is not None:
311 yield leftover
313 def terminal_frame(
314 self,
315 task: asyncio.Task[Any] | asyncio.Future[Any],
316 payload: Callable[[Any], dict[str, Any]],
317 ) -> str | None:
318 """The final SSE frame for a finished producer, or None if there is none.
320 Shared by all five streaming handlers, which used to carry their own
321 copy of this and had drifted: one skipped the cancel check and emitted
322 a done frame to a client that had already disconnected.
323 """
324 if self.cancel.is_set() or not task.done() or task.cancelled():
325 return None
326 exc = task.exception()
327 if exc is not None:
328 return sse_error(str(exc))
329 return sse_done(payload(task.result()))
331 @staticmethod
332 def _drain_waiters(
333 getter: asyncio.Future[str | None],
334 task: asyncio.Task[Any] | asyncio.Future[Any],
335 ) -> set[asyncio.Future[Any]]:
336 """Futures the drain loop waits on. A finished task is left out or it
337 would resolve the wait instantly every pass and spin the loop."""
338 return {getter} if task.done() else {getter, task}
340 @staticmethod
341 def _drain_timeout(task: asyncio.Task[Any] | asyncio.Future[Any]) -> float | None:
342 """How long one drain pass may sleep.
344 While the producer runs the timeout only serves the heartbeat, since
345 the task itself is in the wait set; a disabled heartbeat can wait
346 indefinitely. Once it finishes it leaves the wait set, so this bounds
347 the remaining case: a queue emptying without the sentinel arriving.
348 """
349 if task.done():
350 return _FINISHED_PRODUCER_POLL_S
351 interval = cfg.sse_heartbeat_interval
352 return interval if interval > 0 else None
354 async def drain(
355 self, task: asyncio.Task[Any] | asyncio.Future[Any], label: str
356 ) -> AsyncGenerator[str, None]:
357 """Yield SSE strings until a sentinel arrives; cancel *task* on client disconnect.
359 Emits a ``heartbeat`` event whenever the producer queue stays
360 idle longer than ``cfg.sse_heartbeat_interval`` seconds so
361 clients that enforce a stream-idle timeout don't abort.
363 The pending ``queue.get`` survives across rounds (``asyncio.wait``, not
364 ``wait_for``): cancelling a completed get on the timeout boundary would
365 drop the event it already popped.
367 *task* is in the wait set, so a producer dying without a sentinel wakes
368 the loop directly. That is what lets the timeout be the seconds-scale
369 heartbeat interval instead of a 10Hz tick.
370 """
371 last_yielded = time.monotonic()
372 getter: asyncio.Future[str | None] | None = None
373 try:
374 while True:
375 if getter is None:
376 getter = asyncio.ensure_future(self.queue.get())
377 done, _ = await asyncio.wait(
378 self._drain_waiters(getter, task),
379 timeout=self._drain_timeout(task),
380 return_when=asyncio.FIRST_COMPLETED,
381 )
382 if getter not in done:
383 now = time.monotonic()
384 heartbeat_interval = cfg.sse_heartbeat_interval
385 if heartbeat_interval > 0 and now - last_yielded >= heartbeat_interval:
386 last_yielded = now
387 yield sse_event(SseEvent.HEARTBEAT, {"ts": time.time()})
388 # Fallback for producers that die without a sentinel.
389 if task.done() and self.queue.empty():
390 getter.cancel()
391 break
392 continue
393 item = getter.result()
394 getter = None
395 if item is None:
396 async for leftover in self._flush_pending():
397 last_yielded = time.monotonic()
398 yield leftover
399 break
400 last_yielded = time.monotonic()
401 yield item
402 except (asyncio.CancelledError, GeneratorExit):
403 log.info("%s cancelled by client", label)
404 self.cancel.set()
405 task.cancel()
406 if getter is not None:
407 getter.cancel()