"""Single-lane inference queue for serialized LLM work. All LLM calls (chat, command execution, compaction) run through one worker thread, so there are no concurrent requests to the server. """ import collections import logging import threading from dataclasses import dataclass from typing import Callable logger = logging.getLogger(__name__) @dataclass class InferenceJob: """A unit of work for the inference queue.""" name: str # "chat" | "command" | "summary" run: Callable[[], None] # the work to do priority: int = 0 # lower values run first; commands use -1 class InferenceQueue: """Single worker thread owning every LLM inference call.""" def __init__(self, size: int = 8): """Create the queue. Args: size: Max queued jobs (bounded backpressure). """ self.size = size self._queue: collections.deque[InferenceJob] = collections.deque() self._in_flight: InferenceJob | None = None self._lock = threading.Lock() self._cond = threading.Condition(self._lock) self._stop = threading.Event() self._worker: threading.Thread | None = None def start(self) -> None: """Spawn the worker thread.""" if self._worker is not None: return self._worker = threading.Thread( target=self._worker_loop, name="assistant-inference", daemon=True ) self._worker.start() def stop(self) -> None: """Discard queued jobs, let in-flight finish, join the worker.""" with self._cond: self._queue.clear() self._stop.set() self._cond.notify_all() if self._worker is not None: self._worker.join(timeout=1.0) def submit(self, job: InferenceJob) -> bool: """Queue a job for execution. Returns: False only when a priority job (priority < 0) is rejected due to overflow. Normal jobs drop the oldest normal job on overflow and return True. """ with self._cond: if self._stop.is_set(): return False # Capacity includes the in-flight job (if any) and the queue capacity_used = (1 if self._in_flight else 0) + len(self._queue) # Check if at capacity if capacity_used >= self.size: if job.priority < 0 or not self._drop_oldest_normal(): return False # Insert the job: priority jobs at front, normal jobs at back if job.priority < 0: self._queue.appendleft(job) else: self._queue.append(job) self._cond.notify_all() return True def submit_if_room(self, job: InferenceJob, reserve: int = 1) -> bool: """Append a job only if there is room; never drops an existing job. Admission is atomic under the existing condition lock. Capacity counts the in-flight job plus queued jobs, and the job is admitted only while ``capacity_used < size - reserve``. ``reserve`` keeps slots free for higher-priority work, so a ``reserve`` >= ``size`` makes admission always refuse (it never crashes). Unlike ``submit()``, this never evicts an existing job. Returns: True when the job was appended, False when stopped or full. """ with self._cond: if self._stop.is_set(): return False capacity_used = (1 if self._in_flight else 0) + len(self._queue) if capacity_used >= self.size - reserve: return False self._queue.append(job) self._cond.notify_all() return True def drop(self, predicate: Callable[[InferenceJob], bool]) -> int: """Remove queued jobs matching the predicate. Does not touch the in-flight job. Returns the count of jobs removed. """ with self._cond: to_drop = [j for j in self._queue if predicate(j)] for job in to_drop: self._queue.remove(job) return len(to_drop) def depth(self) -> int: """Return the number of queued jobs (not counting in-flight).""" with self._lock: return len(self._queue) def in_flight(self) -> bool: """Return True while a job is executing.""" with self._lock: return self._in_flight is not None def _drop_oldest_normal(self) -> bool: """Remove the oldest normal (priority >= 0) job; False if none exists. Assumes the lock is held. """ for i, job in enumerate(self._queue): if job.priority >= 0: del self._queue[i] return True return False def _worker_loop(self) -> None: """Worker thread: dequeue and execute jobs until stop.""" while True: with self._cond: while not self._queue and not self._stop.is_set(): self._cond.wait() if self._stop.is_set() and not self._queue: break if not self._queue: continue job = self._queue.popleft() self._in_flight = job # Release the lock during execution so a concurrent drop() can proceed. try: job.run() except Exception as e: # noqa: BLE001 - one bad job must not kill the lane logger.error("Inference job %r failed: %s", job.name, e) finally: with self._lock: self._in_flight = None