971 lines
39 KiB
Python
971 lines
39 KiB
Python
"""Dictate service state machine and Unix-domain socket loop."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
from concurrent.futures import Executor, Future
|
|
from dataclasses import dataclass
|
|
import math
|
|
import os
|
|
from queue import Empty, Queue, SimpleQueue
|
|
import selectors
|
|
import socket
|
|
import threading
|
|
import time
|
|
from typing import Callable, Protocol
|
|
|
|
from .audio import AudioError, pcm_s16le_to_floats
|
|
from .backend import ProviderError, TranscriptionBackend, TranscriptionTurn, backend_from_config
|
|
from .config import ConfigError, load_config
|
|
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
|
|
|
|
|
|
def _looks_like_turn(value: object) -> bool:
|
|
"""Keep injected test backends source-compatible while validating turns."""
|
|
return all(callable(getattr(value, name, None)) for name in ("enqueue_pcm", "finish", "poll_error", "cancel", "close"))
|
|
|
|
|
|
def _turn_cleanup_complete(turn: object) -> bool:
|
|
return bool(getattr(turn, "cleanup_complete", True))
|
|
|
|
|
|
def _recorder_cleanup_complete(recorder: object) -> bool:
|
|
return bool(getattr(recorder, "cleanup_complete", True))
|
|
|
|
|
|
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
|
|
recorder_cleanup_started: bool = False
|
|
turn: TranscriptionTurn | None = None
|
|
tombstoned: bool = False
|
|
failure_reported: bool = False
|
|
cleanup_pending: bool = False
|
|
cleanup_notify_idle: bool = False
|
|
last_level_event_at: float | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _Completion:
|
|
kind: str
|
|
generation_id: str
|
|
future: Future[object]
|
|
|
|
|
|
class _DaemonExecutor(Executor):
|
|
"""Single-worker executor that cannot hold interpreter shutdown hostage."""
|
|
|
|
def __init__(self) -> None:
|
|
self._queue: Queue[tuple[Future[object], Callable[..., object], tuple[object, ...], dict[str, object]] | None] = Queue()
|
|
self._lock = threading.Lock()
|
|
self._shutdown = False
|
|
self._worker = threading.Thread(target=self._run, daemon=True, name="dictate-worker")
|
|
self._worker.start()
|
|
|
|
def submit(self, fn: Callable[..., object], /, *args: object, **kwargs: object) -> Future[object]:
|
|
future: Future[object] = Future()
|
|
with self._lock:
|
|
if self._shutdown:
|
|
raise RuntimeError("cannot schedule new futures after shutdown")
|
|
self._queue.put((future, fn, args, kwargs))
|
|
return future
|
|
|
|
def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None:
|
|
with self._lock:
|
|
if not self._shutdown:
|
|
self._shutdown = True
|
|
if cancel_futures:
|
|
while True:
|
|
try:
|
|
pending = self._queue.get_nowait()
|
|
except Empty:
|
|
break
|
|
if pending is not None:
|
|
pending[0].cancel()
|
|
self._queue.put(None)
|
|
if wait:
|
|
self._worker.join()
|
|
|
|
def _run(self) -> None:
|
|
while True:
|
|
work = self._queue.get()
|
|
if work is None:
|
|
return
|
|
future, fn, args, kwargs = work
|
|
if not future.set_running_or_notify_cancel():
|
|
continue
|
|
try:
|
|
result = fn(*args, **kwargs)
|
|
except BaseException as error:
|
|
future.set_exception(error)
|
|
else:
|
|
future.set_result(result)
|
|
|
|
|
|
class DictateServiceCore:
|
|
"""Single-generation protocol state machine with injected host boundaries."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
model: TranscriptionModel | None = None,
|
|
backend: TranscriptionBackend | None = None,
|
|
recorder_factory: Callable[..., Recorder] = PwRecorder,
|
|
executor: Executor | None = None,
|
|
clock: Callable[[], float] = time.monotonic,
|
|
idle_seconds: float = DEFAULT_IDLE_SECONDS,
|
|
) -> None:
|
|
if model is None and backend is None:
|
|
raise ValueError("either model or backend is required")
|
|
if model is not None and backend is not None:
|
|
raise ValueError("model and backend are mutually exclusive")
|
|
self._model = model
|
|
self._backend = backend
|
|
self._recorder_factory = recorder_factory
|
|
self._executor = executor or _DaemonExecutor()
|
|
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)) if model is not None else 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 _ready(self) -> bool:
|
|
if self._backend is None:
|
|
return self._model_loaded
|
|
return bool(getattr(self._backend, "ready", False))
|
|
|
|
def _diagnostics(self) -> dict[str, str]:
|
|
descriptor = getattr(self._backend, "descriptor", None)
|
|
if descriptor is None:
|
|
return {}
|
|
# Descriptors are fixed, bounded values supplied by the host backend,
|
|
# never paths, endpoint URLs, or credential state.
|
|
return {"backend": str(descriptor.backend)[:64], "model": str(descriptor.model)[:128]}
|
|
|
|
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._ready(),
|
|
**self._diagnostics(),
|
|
},
|
|
)
|
|
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._ready(),
|
|
"busy": self._active is not None,
|
|
**self._diagnostics(),
|
|
},
|
|
)
|
|
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.cleanup_pending and self._worker_kind is None:
|
|
if active.recorder is not None and not _recorder_cleanup_complete(active.recorder):
|
|
self._submit("cleanup", active, active.recorder.poll_unexpected)
|
|
return
|
|
self._finish_generation(active, notify_idle=active.cleanup_notify_idle)
|
|
if self._active is active:
|
|
return
|
|
active = self._active
|
|
if active and active.tombstoned and self._worker_kind is None:
|
|
if active.recorder is not None and not active.recorder_cleanup_started:
|
|
active.recorder_cleanup_started = True
|
|
active.state = "stopping"
|
|
self._submit("cancel", active, active.recorder.cancel)
|
|
else:
|
|
self._finish_generation(active, notify_idle=False)
|
|
return
|
|
if active and active.state == "listening" and active.recorder and self._worker_kind is None:
|
|
if active.turn is not None:
|
|
failure_from_turn = active.turn.poll_error()
|
|
if failure_from_turn is not None:
|
|
active.failure_reported = True
|
|
self._async_error(active, failure_from_turn.code, failure_from_turn.detail)
|
|
self._abort_active(active)
|
|
return
|
|
try:
|
|
failure = active.recorder.poll_unexpected()
|
|
except RecorderError as error:
|
|
failure = error
|
|
if failure is not None:
|
|
active.failure_reported = True
|
|
self._async_error(active, failure.code, failure.detail)
|
|
self._abort_active(active)
|
|
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
|
|
self._cancel_preparation(active)
|
|
if active.turn is not None:
|
|
self._cancel_turn(active)
|
|
if active.recorder is not None and self._worker_kind is None and not active.recorder_cleanup_started:
|
|
active.recorder_cleanup_started = True
|
|
active.state = "stopping"
|
|
self._submit("cancel", active, active.recorder.cancel)
|
|
elif self._worker_kind is None:
|
|
self._finish_generation(active, notify_idle=False)
|
|
|
|
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:
|
|
self._cancel_preparation(active)
|
|
if active.turn is not None:
|
|
self._cancel_turn(active)
|
|
if active.recorder is not None:
|
|
try:
|
|
active.recorder.cancel()
|
|
except Exception:
|
|
pass
|
|
if active.turn is not None:
|
|
deadline = time.monotonic() + 2.0
|
|
while not _turn_cleanup_complete(active.turn) and time.monotonic() < deadline:
|
|
try:
|
|
active.turn.close()
|
|
except Exception:
|
|
break
|
|
time.sleep(0.01)
|
|
self._active = None
|
|
if self._owns_executor:
|
|
# A recorder/provider worker may be stuck in native or network I/O.
|
|
# The owned worker is daemonized and isn't registered with Python's
|
|
# ThreadPoolExecutor exit hook, so this bounded shutdown also permits
|
|
# normal interpreter exit.
|
|
self._executor.shutdown(wait=False, cancel_futures=True)
|
|
|
|
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 self._backend is not None:
|
|
begin_prepare = getattr(self._backend, "begin_prepare", None)
|
|
if callable(begin_prepare):
|
|
begin_prepare()
|
|
self._state(active, "loading")
|
|
self._submit("prepare", active, self._backend.start_turn)
|
|
return
|
|
assert self._model is not None
|
|
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, turn: TranscriptionTurn | None = None) -> None:
|
|
if active.tombstoned:
|
|
self._finish_generation(active, notify_idle=False)
|
|
return
|
|
try:
|
|
active.turn = turn
|
|
if turn is None:
|
|
recorder = self._recorder_factory()
|
|
else:
|
|
# PwRecorder accepts these options. Keep injected legacy test
|
|
# factories source-compatible while requiring a sink for cloud.
|
|
sink = None if turn.retain_pcm else turn
|
|
try:
|
|
recorder = self._recorder_factory(
|
|
sample_rate_hz=turn.sample_rate_hz,
|
|
pcm_sink=sink,
|
|
retain_pcm=turn.retain_pcm,
|
|
)
|
|
except TypeError:
|
|
try:
|
|
recorder = self._recorder_factory(
|
|
sample_rate_hz=turn.sample_rate_hz, sink=sink
|
|
)
|
|
except TypeError:
|
|
try:
|
|
recorder = self._recorder_factory(turn.sample_rate_hz, sink)
|
|
except TypeError:
|
|
if not turn.retain_pcm:
|
|
raise RecorderError(
|
|
"recorder_config",
|
|
"realtime recorder factory must accept a PCM sink",
|
|
)
|
|
recorder = self._recorder_factory()
|
|
active.recorder = recorder
|
|
recorder.start()
|
|
except RecorderError as error:
|
|
self._async_error(active, error.code, error.detail)
|
|
active.failure_reported = True
|
|
self._abort_active(active)
|
|
return
|
|
except OSError as error:
|
|
self._async_error(active, "recorder_spawn", str(error))
|
|
active.failure_reported = True
|
|
self._abort_active(active)
|
|
return
|
|
except Exception as error:
|
|
self._async_error(active, "recorder_failure", str(error))
|
|
active.failure_reported = True
|
|
self._abort_active(active)
|
|
return
|
|
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)
|
|
active.failure_reported = True
|
|
self._abort_active(active)
|
|
return
|
|
active.state = "stopping"
|
|
self._state(active, "stopping", assign=False)
|
|
if active.turn is None:
|
|
self._submit("stop", active, self._stop_and_convert, active.recorder)
|
|
else:
|
|
self._submit("stop", active, self._stop_capture, 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
|
|
# Emit the delivery tombstone before potentially blocking provider or
|
|
# recorder cleanup. The client must learn cancellation promptly while
|
|
# the host remains busy reaping the old generation.
|
|
self._emit(
|
|
client_id,
|
|
{
|
|
"version": PROTOCOL_VERSION,
|
|
"type": "cancelled",
|
|
"requestId": active.request_id,
|
|
"generationId": active.generation_id,
|
|
},
|
|
)
|
|
self._cancel_preparation(active)
|
|
if active.turn is not None:
|
|
self._cancel_turn(active)
|
|
if active.recorder is not None and self._worker_kind is None and not active.recorder_cleanup_started:
|
|
active.recorder_cleanup_started = True
|
|
active.state = "stopping"
|
|
self._submit("cancel", active, active.recorder.cancel)
|
|
elif active.recorder is None and self._worker_kind is None:
|
|
self._finish_generation(active, notify_idle=False)
|
|
|
|
def _cancel_preparation(self, active: _Generation) -> None:
|
|
if self._worker_kind != "prepare" or self._backend is None:
|
|
return
|
|
cancel_prepare = getattr(self._backend, "cancel_prepare", None)
|
|
if callable(cancel_prepare):
|
|
try:
|
|
cancel_prepare()
|
|
except Exception:
|
|
pass
|
|
|
|
@staticmethod
|
|
def _cancel_turn(active: _Generation) -> None:
|
|
if active.turn is None:
|
|
return
|
|
try:
|
|
active.turn.cancel()
|
|
except Exception:
|
|
pass
|
|
|
|
def _abort_active(self, active: _Generation) -> None:
|
|
"""Tombstone a failed generation and always reap recorder/provider work."""
|
|
active.tombstoned = True
|
|
self._cancel_preparation(active)
|
|
if active.turn is not None:
|
|
self._cancel_turn(active)
|
|
if active.recorder is not None and getattr(active.recorder, "cleanup_complete", False) is True:
|
|
self._finish_generation(active, notify_idle=False)
|
|
elif active.recorder is not None and self._worker_kind is None and not active.recorder_cleanup_started:
|
|
active.recorder_cleanup_started = True
|
|
active.state = "stopping"
|
|
self._submit("cancel", active, active.recorder.cancel)
|
|
elif active.recorder is None and self._worker_kind is None:
|
|
self._finish_generation(active, notify_idle=False)
|
|
|
|
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)
|
|
|
|
@staticmethod
|
|
def _stop_capture(recorder: Recorder) -> Capture:
|
|
return recorder.stop()
|
|
|
|
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 and not active.failure_reported:
|
|
self._async_error(active, error.code, error.detail)
|
|
active.failure_reported = True
|
|
self._abort_active(active)
|
|
return
|
|
except AudioError as error:
|
|
if not active.tombstoned and not active.failure_reported:
|
|
self._async_error(active, "invalid_audio", str(error))
|
|
active.failure_reported = True
|
|
self._abort_active(active)
|
|
return
|
|
except ProviderError as error:
|
|
if not active.tombstoned and not active.failure_reported:
|
|
self._async_error(active, error.code, error.detail)
|
|
active.failure_reported = True
|
|
self._abort_active(active)
|
|
return
|
|
except BaseException as error:
|
|
if not active.tombstoned and not active.failure_reported:
|
|
code = "model_failure" if completion.kind in ("load", "prepare", "infer") else "recorder_failure"
|
|
self._async_error(active, code, str(error))
|
|
active.failure_reported = True
|
|
self._abort_active(active)
|
|
return
|
|
|
|
if completion.kind == "prepare":
|
|
if active.tombstoned:
|
|
if _looks_like_turn(result):
|
|
active.turn = result
|
|
try:
|
|
result.cancel()
|
|
except Exception:
|
|
pass
|
|
self._finish_generation(active, notify_idle=False)
|
|
elif _looks_like_turn(result):
|
|
self._begin_recording(active, result)
|
|
else:
|
|
self._async_error(active, "backend_failure", "backend returned an invalid transcription turn")
|
|
active.failure_reported = True
|
|
self._abort_active(active)
|
|
elif 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._abort_active(active)
|
|
return
|
|
if active.turn is not None:
|
|
if not isinstance(result, Capture):
|
|
self._async_error(active, "invalid_audio", "recorder returned invalid capture")
|
|
active.failure_reported = True
|
|
self._abort_active(active)
|
|
return
|
|
self._state(active, "transcribing")
|
|
self._submit("infer", active, active.turn.finish, result)
|
|
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")
|
|
active.failure_reported = True
|
|
self._abort_active(active)
|
|
return
|
|
self._state(active, "transcribing")
|
|
assert self._model is not None
|
|
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")
|
|
else:
|
|
try:
|
|
transcript_bytes = len(result.encode("utf-8"))
|
|
except UnicodeEncodeError:
|
|
self._async_error(active, "model_failure", "model returned invalid Unicode")
|
|
self._finish_generation(active, notify_idle=True)
|
|
return
|
|
if transcript_bytes > 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=not active.tombstoned)
|
|
elif completion.kind == "cleanup":
|
|
self._finish_generation(active, notify_idle=active.cleanup_notify_idle)
|
|
|
|
def _finish_generation(self, active: _Generation, *, notify_idle: bool) -> None:
|
|
if self._active is not active:
|
|
return
|
|
active.cleanup_pending = True
|
|
active.cleanup_notify_idle = active.cleanup_notify_idle or notify_idle
|
|
if active.recorder is not None and not _recorder_cleanup_complete(active.recorder):
|
|
# Cleanup retries run through the executor so a stubborn child never
|
|
# blocks the socket loop. The pending generation remains exclusive.
|
|
return
|
|
if active.turn is not None:
|
|
try:
|
|
active.turn.close()
|
|
except Exception:
|
|
pass
|
|
# A provider close can be blocked by a wedged transport. Keep the
|
|
# tombstoned generation busy until both provider workers have
|
|
# actually exited, so no thread or transport crosses generations.
|
|
if not _turn_cleanup_complete(active.turn):
|
|
return
|
|
active.cleanup_pending = False
|
|
self._active = None
|
|
self._last_activity = self._clock()
|
|
if (notify_idle or active.failure_reported) 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,
|
|
*,
|
|
expected_path: str | os.PathLike[str] | None = None,
|
|
) -> socket.socket:
|
|
"""Return systemd's fd 3 only when it is the expected Unix stream socket."""
|
|
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 | None = None
|
|
try:
|
|
listener = socket.socket(fileno=3)
|
|
if listener.family != socket.AF_UNIX:
|
|
raise RuntimeError("inherited listener must be AF_UNIX")
|
|
socket_type = listener.type & getattr(socket, "SOCK_TYPE_MASK", 0xF)
|
|
if socket_type != socket.SOCK_STREAM:
|
|
raise RuntimeError("inherited listener must be SOCK_STREAM")
|
|
try:
|
|
accepting = listener.getsockopt(socket.SOL_SOCKET, socket.SO_ACCEPTCONN)
|
|
except OSError as error:
|
|
raise RuntimeError("could not validate inherited listener state") from error
|
|
if accepting != 1:
|
|
raise RuntimeError("inherited listener is not listening")
|
|
actual = listener.getsockname()
|
|
if isinstance(actual, bytes):
|
|
actual = actual.decode("utf-8", errors="strict")
|
|
if not isinstance(actual, str):
|
|
raise RuntimeError("inherited listener has an invalid socket path")
|
|
expected = expected_path
|
|
if expected is None:
|
|
runtime_dir = environment.get("XDG_RUNTIME_DIR", f"/run/user/{os.getuid()}")
|
|
expected = os.path.join(runtime_dir, "pi-dictate.sock")
|
|
if os.path.abspath(actual) != os.path.abspath(os.fspath(expected)):
|
|
raise RuntimeError("inherited listener is not the expected pi-dictate socket")
|
|
listener.setblocking(False)
|
|
return listener
|
|
except Exception:
|
|
if listener is not None:
|
|
listener.close()
|
|
raise
|
|
|
|
|
|
def _queue_connection_frame(
|
|
selector: selectors.BaseSelector,
|
|
connections: dict[str, _Connection],
|
|
core: DictateServiceCore,
|
|
client_id: str,
|
|
frame: dict[str, object],
|
|
) -> bool:
|
|
"""Encode and append one frame without exceeding the per-client bound."""
|
|
connection = connections.get(client_id)
|
|
if connection is None:
|
|
return False
|
|
try:
|
|
encoded = encode_frame(frame)
|
|
except ProtocolError:
|
|
_close_connection(selector, connections, core, client_id)
|
|
return False
|
|
if len(connection.output) + len(encoded) > MAX_CONNECTION_OUTPUT_BYTES:
|
|
_close_connection(selector, connections, core, client_id)
|
|
return False
|
|
connection.output.extend(encoded)
|
|
selector.modify(
|
|
connection.socket,
|
|
selectors.EVENT_READ | selectors.EVENT_WRITE,
|
|
client_id,
|
|
)
|
|
return True
|
|
|
|
|
|
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
|
|
_queue_connection_frame(
|
|
selector, connections, core, outbound.client_id, outbound.frame
|
|
)
|
|
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:
|
|
if _queue_connection_frame(
|
|
selector,
|
|
connections,
|
|
core,
|
|
client_id,
|
|
error_event(error),
|
|
) and error.close_connection:
|
|
connection.close_after_write = True
|
|
# Queueing a protocol error can close this connection when its
|
|
# output bound is exceeded. The selector event was created for
|
|
# the old socket, so do not send to or modify that stale socket.
|
|
if connections.get(client_id) is not connection:
|
|
continue
|
|
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:
|
|
# Installed services use the strict host config. The environment fallback
|
|
# keeps direct development invocations and existing local deployments
|
|
# compatible until install.sh has written config.json.
|
|
config_path = os.environ.get("PI_DICTATE_CONFIG_PATH")
|
|
default_config = os.path.expanduser("~/.config/pi-dictate/config.json")
|
|
if config_path or os.path.isfile(default_config):
|
|
try:
|
|
config = load_config(config_path)
|
|
network_mode = os.environ.get("PI_DICTATE_NETWORK_MODE")
|
|
expected_network_mode = "offline" if config.backend == "local" else "cloud"
|
|
if network_mode is not None and network_mode != expected_network_mode:
|
|
raise ConfigError("configured backend does not match service network mode")
|
|
backend = backend_from_config(config)
|
|
except (ConfigError, ValueError) as error:
|
|
raise SystemExit(f"Invalid pi-dictate configuration: {error}") from error
|
|
core = DictateServiceCore(backend=backend, idle_seconds=config.idle_seconds)
|
|
else:
|
|
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 = int(idle_value)
|
|
if str(idle_seconds) != idle_value:
|
|
raise ValueError
|
|
boost_alpha = float(boost_value)
|
|
except ValueError as error:
|
|
raise SystemExit("PI_DICTATE_IDLE_SECONDS must be an integer; boost alpha must be numeric") from error
|
|
if not 60 <= idle_seconds <= 86_400:
|
|
raise SystemExit("PI_DICTATE_IDLE_SECONDS must be between 60 and 86400")
|
|
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()
|