Coverage for src/lilbee/server/handlers/sse.py: 100%
185 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +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 EventType.ANALYZE,
165 SseEvent.PROGRESS,
166 }
167)
170class _QueuedEvent(NamedTuple):
171 """A queued SSE payload tagged with whether backpressure may shed it."""
173 payload: str | None
174 droppable: bool
177class SseEventQueue(asyncio.Queue[str | None]):
178 """Bounded SSE queue: progress events shed under backpressure, the rest land.
180 ``put_nowait`` (lifecycle events, tokens, the ``None`` sentinel) always
181 enqueues, evicting the oldest progress event first when at capacity.
182 ``put_event_nowait`` enqueues progress-protocol events, dropping the oldest
183 progress event (or the incoming one when the head is not progress) at
184 capacity. ``join()`` semantics are not supported.
185 """
187 _queue: deque[_QueuedEvent]
189 def __init__(self, max_events: int = SSE_QUEUE_MAX_EVENTS) -> None:
190 super().__init__()
191 self._max_events = max_events
192 self._put_droppable = False
193 self.dropped_events = 0
194 # Set once the queue is full of undroppable events, i.e. the consumer
195 # has stopped reading. See put_nowait.
196 self.stalled = False
198 def _put(self, item: str | None) -> None:
199 self._queue.append(_QueuedEvent(item, self._put_droppable))
201 def _get(self) -> str | None:
202 return self._queue.popleft().payload
204 def _evict_oldest_droppable(self) -> bool:
205 """Shed the oldest progress event anywhere in the queue; True when one went.
207 Scans rather than checking only the head. A stream that interleaves
208 tokens with progress puts a non-droppable event at the head almost
209 immediately, and head-only eviction then reported "nothing to shed"
210 while the queue still held progress events it was allowed to drop.
211 """
212 for index, event in enumerate(self._queue):
213 if event.droppable:
214 del self._queue[index]
215 self.dropped_events += 1
216 return True
217 return False
219 def put_nowait(self, item: str | None) -> None:
220 """Enqueue an always-delivered event, evicting old progress when full."""
221 if self.qsize() >= self._max_events and not self._evict_oldest_droppable():
222 # Full of undroppable events (chat and RAG tokens come through
223 # here) and the consumer has taken none in _max_events: a stalled
224 # or gone client. The alternative is unbounded growth.
225 self.stalled = True
226 self._put_droppable = False
227 super().put_nowait(item)
229 def put_event_nowait(self, payload: str, event_type: EventType | SseEvent) -> None:
230 """Enqueue a progress-protocol event, shedding progress when full."""
231 if event_type not in _DROPPABLE_EVENT_TYPES:
232 self.put_nowait(payload)
233 return
234 if self.qsize() >= self._max_events and not self._evict_oldest_droppable():
235 self.dropped_events += 1
236 return
237 self._put_droppable = True
238 try:
239 super().put_nowait(payload)
240 finally:
241 self._put_droppable = False
244class SseStream:
245 """Context object for SSE streaming with cancellation support.
246 Bundles the queue, cancel event, and progress callback that every SSE
247 endpoint needs. Call :meth:`drain` to yield events until the task
248 completes or the client disconnects.
249 """
251 def __init__(self) -> None:
252 self.queue: SseEventQueue = SseEventQueue()
253 self.cancel = threading.Event()
254 self.loop = asyncio.get_running_loop()
255 self.callback: DetailedProgressCallback = self._build_callback()
257 def put_threadsafe(self, item: str | None) -> None:
258 """Enqueue an always-delivered event from a worker thread.
260 ``asyncio.Queue.put_nowait`` is not thread-safe: it wakes a pending
261 getter via ``Future.set_result``, which must run on the loop thread. A
262 producer running under ``run_in_executor`` therefore hands the put back
263 to the loop instead of mutating the queue directly.
264 """
265 self.loop.call_soon_threadsafe(self._put_and_check_stall, item)
267 def _put_and_check_stall(self, item: str | None) -> None:
268 """Enqueue on the loop thread, cancelling the producer if it has stalled.
270 A consumer that has taken nothing in a full queue's worth is gone;
271 cancelling is the same signal a detected disconnect sends.
272 """
273 self.queue.put_nowait(item)
274 if self.queue.stalled and not self.cancel.is_set():
275 log.warning(
276 "SSE consumer stalled with %d undroppable events queued; cancelling the producer.",
277 self.queue.qsize(),
278 )
279 self.cancel.set()
281 def _build_callback(self) -> DetailedProgressCallback:
282 """Create a progress callback that serializes events into the queue.
283 Safe to call from both the event-loop thread and worker threads.
284 """
285 loop = self.loop
286 queue = self.queue
288 def _callback(event_type: EventType, data: ProgressEvent) -> None:
289 serialized = data.model_dump() if isinstance(data, BaseModel) else data
290 payload = f"event: {event_type}\ndata: {json.dumps(serialized)}\n\n"
291 try:
292 running = asyncio.get_running_loop()
293 except RuntimeError:
294 running = None
295 if running is loop:
296 queue.put_event_nowait(payload, event_type)
297 else:
298 loop.call_soon_threadsafe(queue.put_event_nowait, payload, event_type)
300 return _callback
302 async def _flush_pending(self) -> AsyncGenerator[str, None]:
303 """Events left behind the sentinel by a producer that outran the consumer.
305 A fast producer can enqueue its sentinel before its threadsafe progress
306 callbacks run; one loop tick lets them land.
307 """
308 await asyncio.sleep(0)
309 while not self.queue.empty():
310 leftover = self.queue.get_nowait()
311 if leftover is not None:
312 yield leftover
314 def terminal_frame(
315 self,
316 task: asyncio.Task[Any] | asyncio.Future[Any],
317 payload: Callable[[Any], dict[str, Any]],
318 ) -> str | None:
319 """The final SSE frame for a finished producer, or None if there is none.
321 Shared by all five streaming handlers, which used to carry their own
322 copy of this and had drifted: one skipped the cancel check and emitted
323 a done frame to a client that had already disconnected.
324 """
325 if self.cancel.is_set() or not task.done() or task.cancelled():
326 return None
327 exc = task.exception()
328 if exc is not None:
329 return sse_error(str(exc))
330 return sse_done(payload(task.result()))
332 @staticmethod
333 def _drain_waiters(
334 getter: asyncio.Future[str | None],
335 task: asyncio.Task[Any] | asyncio.Future[Any],
336 ) -> set[asyncio.Future[Any]]:
337 """Futures the drain loop waits on. A finished task is left out or it
338 would resolve the wait instantly every pass and spin the loop."""
339 return {getter} if task.done() else {getter, task}
341 @staticmethod
342 def _drain_timeout(task: asyncio.Task[Any] | asyncio.Future[Any]) -> float | None:
343 """How long one drain pass may sleep.
345 While the producer runs the timeout only serves the heartbeat, since
346 the task itself is in the wait set; a disabled heartbeat can wait
347 indefinitely. Once it finishes it leaves the wait set, so this bounds
348 the remaining case: a queue emptying without the sentinel arriving.
349 """
350 if task.done():
351 return _FINISHED_PRODUCER_POLL_S
352 interval = cfg.sse_heartbeat_interval
353 return interval if interval > 0 else None
355 async def drain(
356 self, task: asyncio.Task[Any] | asyncio.Future[Any], label: str
357 ) -> AsyncGenerator[str, None]:
358 """Yield SSE strings until a sentinel arrives; cancel *task* on client disconnect.
360 Emits a ``heartbeat`` event whenever the producer queue stays
361 idle longer than ``cfg.sse_heartbeat_interval`` seconds so
362 clients that enforce a stream-idle timeout don't abort.
364 The pending ``queue.get`` survives across rounds (``asyncio.wait``, not
365 ``wait_for``): cancelling a completed get on the timeout boundary would
366 drop the event it already popped.
368 *task* is in the wait set, so a producer dying without a sentinel wakes
369 the loop directly. That is what lets the timeout be the seconds-scale
370 heartbeat interval instead of a 10Hz tick.
371 """
372 last_yielded = time.monotonic()
373 getter: asyncio.Future[str | None] | None = None
374 try:
375 while True:
376 if getter is None:
377 getter = asyncio.ensure_future(self.queue.get())
378 done, _ = await asyncio.wait(
379 self._drain_waiters(getter, task),
380 timeout=self._drain_timeout(task),
381 return_when=asyncio.FIRST_COMPLETED,
382 )
383 if getter not in done:
384 now = time.monotonic()
385 heartbeat_interval = cfg.sse_heartbeat_interval
386 if heartbeat_interval > 0 and now - last_yielded >= heartbeat_interval:
387 last_yielded = now
388 yield sse_event(SseEvent.HEARTBEAT, {"ts": time.time()})
389 # Fallback for producers that die without a sentinel.
390 if task.done() and self.queue.empty():
391 getter.cancel()
392 break
393 continue
394 item = getter.result()
395 getter = None
396 if item is None:
397 async for leftover in self._flush_pending():
398 last_yielded = time.monotonic()
399 yield leftover
400 break
401 last_yielded = time.monotonic()
402 yield item
403 except (asyncio.CancelledError, GeneratorExit):
404 log.info("%s cancelled by client", label)
405 self.cancel.set()
406 task.cancel()
407 if getter is not None:
408 getter.cancel()