"""Dictate service state machine and Unix-domain socket loop.""" from __future__ import annotations from collections.abc import Sequence from concurrent.futures import Executor, Future, ThreadPoolExecutor from dataclasses import dataclass import math import os from queue import Empty, SimpleQueue import selectors import socket import time from typing import Callable, Protocol from .audio import AudioError, pcm_s16le_to_floats from .model import NemoParakeetModel, TranscriptionModel from .protocol import ( MAX_FRAME_BYTES, MAX_TRANSCRIPT_BYTES, PROTOCOL_VERSION, FrameDecoder, ProtocolError, Request, encode_frame, error_event, parse_request, ) from .recorder import Capture, PwRecorder, RecorderError DEFAULT_IDLE_SECONDS = 15 * 60 MAX_CLIENT_CONNECTIONS = 16 MAX_CONNECTION_OUTPUT_BYTES = MAX_FRAME_BYTES LEVEL_EVENT_INTERVAL_SECONDS = 0.1 class Recorder(Protocol): @property def audio_level(self) -> float: ... def start(self) -> None: ... def stop(self) -> Capture: ... def cancel(self) -> None: ... def poll_unexpected(self) -> RecorderError | None: ... @dataclass(frozen=True) class Outbound: client_id: str frame: dict[str, object] @dataclass class _Generation: owner: str generation_id: str request_id: str state: str recorder: Recorder | None = None tombstoned: bool = False last_level_event_at: float | None = None @dataclass(frozen=True) class _Completion: kind: str generation_id: str future: Future[object] class DictateServiceCore: """Single-generation protocol state machine with injected host boundaries.""" def __init__( self, *, model: TranscriptionModel, recorder_factory: Callable[[], Recorder] = PwRecorder, executor: Executor | None = None, clock: Callable[[], float] = time.monotonic, idle_seconds: float = DEFAULT_IDLE_SECONDS, ) -> None: self._model = model self._recorder_factory = recorder_factory self._executor = executor or ThreadPoolExecutor(max_workers=1, thread_name_prefix="dictate") self._owns_executor = executor is None self._clock = clock self._idle_seconds = idle_seconds self._last_activity = clock() self._clients: dict[str, bool] = {} self._active: _Generation | None = None self._model_loaded = bool(getattr(model, "loaded", False)) self._worker_kind: str | None = None self._completions: SimpleQueue[_Completion] = SimpleQueue() self._events: list[Outbound] = [] @property def active_generation_id(self) -> str | None: return self._active.generation_id if self._active else None def accepted(self, client_id: str) -> None: """Track a socket without treating acceptance as protocol activity.""" self._clients.setdefault(client_id, False) def handle_request(self, client_id: str, request: Request) -> None: """Handle one already parsed request and queue its observable events.""" self._last_activity = self._clock() self._clients.setdefault(client_id, False) if request.type == "hello": self._clients[client_id] = True self._emit( client_id, { "version": PROTOCOL_VERSION, "type": "pong", "requestId": request.request_id, "service": "pi-dictate", "ready": self._model_loaded, }, ) return if not self._clients[client_id]: self._request_error(client_id, request, "hello_required", "hello is required first") return if request.type == "ping": self._emit( client_id, { "version": PROTOCOL_VERSION, "type": "pong", "requestId": request.request_id, "service": "pi-dictate", "ready": self._model_loaded, "busy": self._active is not None, }, ) elif request.type == "start": self._start(client_id, request) elif request.type == "stop": self._stop(client_id, request) elif request.type == "cancel": self._cancel(client_id, request) def poll(self) -> None: """Apply completed workers and recorder exits without blocking the socket loop.""" while True: try: completion = self._completions.get_nowait() except Empty: break self._apply_completion(completion) active = self._active if active and active.state == "listening" and active.recorder and self._worker_kind is None: try: failure = active.recorder.poll_unexpected() except RecorderError as error: failure = error if failure is not None: self._async_error(active, failure.code, failure.detail) self._finish_generation(active, notify_idle=True) return self._emit_audio_level(active) def drain_events(self) -> list[Outbound]: events, self._events = self._events, [] return events def disconnect(self, client_id: str) -> None: self._clients.pop(client_id, None) active = self._active if active is None or active.owner != client_id: return active.tombstoned = True if active.state == "listening" and active.recorder is not None: active.state = "stopping" self._submit("cancel", active, active.recorder.cancel) def should_exit(self) -> bool: if self._active is not None or self._worker_kind is not None: return False return self._clock() - self._last_activity >= self._idle_seconds def close(self) -> None: active = self._active if active and active.recorder and active.state == "listening": try: active.recorder.cancel() except RecorderError: pass self._active = None if self._owns_executor: self._executor.shutdown(wait=True, cancel_futures=False) def _start(self, client_id: str, request: Request) -> None: assert request.generation_id is not None if self._active is not None or self._worker_kind is not None: self._request_error(client_id, request, "busy", "another generation is active") return active = _Generation(client_id, request.generation_id, request.request_id, "loading") self._active = active if not self._model_loaded: self._state(active, "loading") self._submit("load", active, self._model.load) return self._begin_recording(active) def _begin_recording(self, active: _Generation) -> None: if active.tombstoned: self._finish_generation(active, notify_idle=False) return try: recorder = self._recorder_factory() recorder.start() except RecorderError as error: self._async_error(active, error.code, error.detail) self._finish_generation(active, notify_idle=True) return except OSError as error: self._async_error(active, "recorder_spawn", str(error)) self._finish_generation(active, notify_idle=True) return active.recorder = recorder active.last_level_event_at = self._clock() self._state(active, "listening") def _emit_audio_level(self, active: _Generation) -> None: now = self._clock() if ( active.last_level_event_at is not None and now - active.last_level_event_at < LEVEL_EVENT_INTERVAL_SECONDS ): return level = active.recorder.audio_level if active.recorder is not None else 0.0 if not math.isfinite(level): level = 0.0 active.last_level_event_at = now self._emit( active.owner, { "version": PROTOCOL_VERSION, "type": "level", "requestId": active.request_id, "generationId": active.generation_id, "level": round(max(0.0, min(1.0, level)), 3), }, ) def _stop(self, client_id: str, request: Request) -> None: active = self._owned_generation(client_id, request) if active is None: return if active.state != "listening" or active.recorder is None: self._request_error(client_id, request, "invalid_state", f"cannot stop while {active.state}") return active.request_id = request.request_id try: failure = active.recorder.poll_unexpected() except RecorderError as error: failure = error if failure is not None: self._async_error(active, failure.code, failure.detail) self._finish_generation(active, notify_idle=True) return active.state = "stopping" self._state(active, "stopping", assign=False) self._submit("stop", active, self._stop_and_convert, active.recorder) def _cancel(self, client_id: str, request: Request) -> None: active = self._owned_generation(client_id, request) if active is None: return if active.tombstoned: self._request_error(client_id, request, "invalid_state", "generation is already cancelled") return active.tombstoned = True active.request_id = request.request_id self._emit( client_id, { "version": PROTOCOL_VERSION, "type": "cancelled", "requestId": active.request_id, "generationId": active.generation_id, }, ) if active.state == "listening" and active.recorder is not None: active.state = "stopping" self._submit("cancel", active, active.recorder.cancel) def _owned_generation(self, client_id: str, request: Request) -> _Generation | None: assert request.generation_id is not None active = self._active if active is None: self._request_error(client_id, request, "invalid_state", "no generation is active") return None if active.owner != client_id: self._request_error(client_id, request, "not_owner", "generation belongs to another client") return None if active.generation_id != request.generation_id: self._request_error( client_id, request, "generation_mismatch", "generation ID does not match the active generation", ) return None return active @staticmethod def _stop_and_convert(recorder: Recorder) -> Sequence[float]: capture = recorder.stop() return pcm_s16le_to_floats(capture.pcm) def _submit( self, kind: str, active: _Generation, function: Callable[..., object], *args: object, ) -> None: self._worker_kind = kind future = self._executor.submit(function, *args) future.add_done_callback( lambda completed, kind=kind, generation_id=active.generation_id: self._completions.put( _Completion(kind, generation_id, completed) ) ) def _apply_completion(self, completion: _Completion) -> None: if self._worker_kind == completion.kind: self._worker_kind = None active = self._active if active is None or active.generation_id != completion.generation_id: return try: result = completion.future.result() except RecorderError as error: if not active.tombstoned: self._async_error(active, error.code, error.detail) self._finish_generation(active, notify_idle=not active.tombstoned) return except AudioError as error: if not active.tombstoned: self._async_error(active, "invalid_audio", str(error)) self._finish_generation(active, notify_idle=not active.tombstoned) return except BaseException as error: if not active.tombstoned: code = "model_failure" if completion.kind in ("load", "infer") else "recorder_failure" self._async_error(active, code, str(error)) self._finish_generation(active, notify_idle=not active.tombstoned) return if completion.kind == "load": self._model_loaded = True if active.tombstoned: self._finish_generation(active, notify_idle=False) else: self._begin_recording(active) elif completion.kind == "stop": if active.tombstoned: self._finish_generation(active, notify_idle=False) return samples = result if not isinstance(samples, Sequence) or isinstance(samples, (str, bytes, bytearray)): self._async_error(active, "invalid_audio", "audio conversion returned invalid samples") self._finish_generation(active, notify_idle=True) return self._state(active, "transcribing") self._submit("infer", active, self._model.transcribe, samples) elif completion.kind == "infer": if active.tombstoned: self._finish_generation(active, notify_idle=False) return if not isinstance(result, str): self._async_error(active, "model_failure", "model returned a non-string transcript") elif len(result.encode("utf-8")) > MAX_TRANSCRIPT_BYTES: self._async_error(active, "transcript_too_large", "transcript exceeds output limit") else: self._emit( active.owner, { "version": PROTOCOL_VERSION, "type": "result", "requestId": active.request_id, "generationId": active.generation_id, "transcript": result, }, ) self._finish_generation(active, notify_idle=True) elif completion.kind == "cancel": self._finish_generation(active, notify_idle=False) def _finish_generation(self, active: _Generation, *, notify_idle: bool) -> None: if self._active is not active: return self._active = None self._last_activity = self._clock() if notify_idle and active.owner in self._clients: self._state(active, "idle", assign=False) def _state(self, active: _Generation, state: str, *, assign: bool = True) -> None: if assign: active.state = state self._emit( active.owner, { "version": PROTOCOL_VERSION, "type": "state", "requestId": active.request_id, "generationId": active.generation_id, "state": state, }, ) def _request_error(self, client_id: str, request: Request, code: str, detail: str) -> None: error = ProtocolError( code, detail, request_id=request.request_id, generation_id=request.generation_id, ) self._emit(client_id, error_event(error)) def _async_error(self, active: _Generation, code: str, detail: str) -> None: error = ProtocolError( code, detail, request_id=active.request_id, generation_id=active.generation_id, ) self._emit(active.owner, error_event(error)) def _emit(self, client_id: str, frame: dict[str, object]) -> None: if client_id in self._clients: self._events.append(Outbound(client_id, frame)) @dataclass class _Connection: socket: socket.socket decoder: FrameDecoder output: bytearray close_after_write: bool = False def inherited_listener(environ: dict[str, str] | None = None) -> socket.socket: """Return systemd's single inherited listener at fd 3.""" environment = os.environ if environ is None else environ try: listen_pid = int(environment.get("LISTEN_PID", "0")) listen_fds = int(environment.get("LISTEN_FDS", "0")) except ValueError as error: raise RuntimeError("invalid systemd socket activation environment") from error if listen_pid != os.getpid() or listen_fds != 1: raise RuntimeError("expected exactly one systemd socket-activated listener") listener = socket.socket(fileno=3) listener.setblocking(False) return listener def serve(listener: socket.socket, core: DictateServiceCore) -> None: """Serve protocol v2 on an explicit listener until the core becomes idle.""" selector = selectors.DefaultSelector() connections: dict[str, _Connection] = {} sequence = 0 listener.setblocking(False) selector.register(listener, selectors.EVENT_READ, None) try: while not core.should_exit(): core.poll() for outbound in core.drain_events(): connection = connections.get(outbound.client_id) if connection is None: continue try: encoded = encode_frame(outbound.frame) except ProtocolError: _close_connection(selector, connections, core, outbound.client_id) continue if len(connection.output) + len(encoded) > MAX_CONNECTION_OUTPUT_BYTES: _close_connection(selector, connections, core, outbound.client_id) continue connection.output.extend(encoded) selector.modify( connection.socket, selectors.EVENT_READ | selectors.EVENT_WRITE, outbound.client_id, ) for key, mask in selector.select(timeout=0.1): if key.data is None: try: client, _address = listener.accept() except BlockingIOError: continue if len(connections) >= MAX_CLIENT_CONNECTIONS: client.close() continue client.setblocking(False) sequence += 1 client_id = f"client-{sequence}" connections[client_id] = _Connection(client, FrameDecoder(), bytearray()) core.accepted(client_id) selector.register(client, selectors.EVENT_READ, client_id) continue client_id = key.data connection = connections.get(client_id) if connection is None: continue if mask & selectors.EVENT_READ: try: data = connection.socket.recv(64 * 1024) except BlockingIOError: data = None except (BrokenPipeError, ConnectionResetError): _close_connection(selector, connections, core, client_id) continue if data == b"": _close_connection(selector, connections, core, client_id) continue if data: try: for frame in connection.decoder.feed(data): core.handle_request(client_id, parse_request(frame)) except ProtocolError as error: connection.output.extend(encode_frame(error_event(error))) if error.close_connection: connection.close_after_write = True if mask & selectors.EVENT_WRITE and connection.output: try: sent = connection.socket.send(connection.output) except BlockingIOError: sent = 0 except (BrokenPipeError, ConnectionResetError): _close_connection(selector, connections, core, client_id) continue del connection.output[:sent] if connection.close_after_write and not connection.output: _close_connection(selector, connections, core, client_id) elif client_id in connections: events = selectors.EVENT_READ if connection.output: events |= selectors.EVENT_WRITE selector.modify(connection.socket, events, client_id) finally: for client_id in list(connections): _close_connection(selector, connections, core, client_id) selector.unregister(listener) listener.close() selector.close() core.close() def _close_connection( selector: selectors.BaseSelector, connections: dict[str, _Connection], core: DictateServiceCore, client_id: str, ) -> None: connection = connections.pop(client_id, None) if connection is None: return core.disconnect(client_id) try: selector.unregister(connection.socket) except KeyError: pass connection.socket.close() def main() -> None: model_path = os.environ.get("PI_DICTATE_MODEL_PATH") if not model_path: raise SystemExit("PI_DICTATE_MODEL_PATH must name an explicit local .nemo file") terms_path = os.environ.get("PI_DICTATE_TERMS_PATH") if not terms_path: raise SystemExit("PI_DICTATE_TERMS_PATH must name the host coding dictionary") idle_value = os.environ.get("PI_DICTATE_IDLE_SECONDS", str(DEFAULT_IDLE_SECONDS)) boost_value = os.environ.get("PI_DICTATE_BOOST_ALPHA", "1.0") try: idle_seconds = float(idle_value) boost_alpha = float(boost_value) except ValueError as error: raise SystemExit("PI_DICTATE_IDLE_SECONDS and PI_DICTATE_BOOST_ALPHA must be numeric") from error if idle_seconds <= 0: raise SystemExit("PI_DICTATE_IDLE_SECONDS must be positive") if boost_alpha <= 0: raise SystemExit("PI_DICTATE_BOOST_ALPHA must be positive") model = NemoParakeetModel(model_path, terms_path=terms_path, boost_alpha=boost_alpha) core = DictateServiceCore(model=model, idle_seconds=idle_seconds) listener = inherited_listener() serve(listener, core) if __name__ == "__main__": main()