STT-Assistant/src/stt_daemon/assistant/control.py

"""NDJSON control socket for the assistant (clients: TUI and hotkey CLI).

Bidirectional, unlike the STT bus: clients send commands and may `subscribe` as a
viewer to receive the assistant's broadcast events. All outbound traffic for a
client goes through that client's bounded queue and its own writer thread, so a
slow client never blocks the service.
"""

import collections
import itertools
import json
import logging
import os
import socket
import threading
import time
from pathlib import Path

from .. import events
from .service import EXECUTE_BUSY

logger = logging.getLogger(__name__)

CLIENT_QUEUE = 256
POLL_INTERVAL = 0.005
ACCEPT_TIMEOUT = 0.2
READ_CHUNK = 65536


def default_control_path() -> str:
    """Default control socket path under $XDG_RUNTIME_DIR (fallback /tmp)."""
    runtime = os.environ.get("XDG_RUNTIME_DIR") or "/tmp"
    return str(Path(runtime) / "stt-assistant.sock")


class _Client:
    __slots__ = ("conn", "queue", "viewer", "writer", "reader")

    def __init__(self, conn: socket.socket, queue_size: int):
        self.conn = conn
        self.queue: collections.deque[bytes] = collections.deque(maxlen=queue_size)
        self.viewer = False
        self.writer: threading.Thread | None = None
        self.reader: threading.Thread | None = None


class ControlServer:
    """Serves the assistant control protocol over a UNIX socket."""

    def __init__(self, cfg, service, path: str | None = None):
        """Create the server (call `start()` to bind).

        Args:
            cfg: Configuration (uses `[assistant]`).
            service: The `AssistantService` this server drives.
            path: Socket path; defaults to config or `default_control_path()`.
        """
        self.cfg = cfg
        self.service = service
        self.path = path or cfg.assistant.control_socket or default_control_path()
        self._seq = itertools.count(1)
        self._clients: list[_Client] = []
        self._lock = threading.Lock()
        self._stop = threading.Event()
        self._server: socket.socket | None = None
        self._accept: threading.Thread | None = None

    # ------------------------------------------------------------------ lifecycle

    def start(self) -> None:
        """Bind the socket and start accepting clients."""
        self._bind()
        self._accept = threading.Thread(
            target=self._accept_loop, name="control-accept", daemon=True
        )
        self._accept.start()

    def _bind(self) -> None:
        path = Path(self.path)
        path.parent.mkdir(parents=True, exist_ok=True)
        if path.exists():
            path.unlink()
        server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        server.bind(str(path))
        os.chmod(str(path), 0o600)
        server.listen(16)
        server.settimeout(ACCEPT_TIMEOUT)
        self._server = server

    def close(self) -> None:
        """Stop accepting, drop clients, and remove the socket."""
        self._stop.set()
        if self._server is not None:
            try:
                self._server.close()
            except OSError:
                pass
        with self._lock:
            clients = list(self._clients)
            self._clients.clear()
        for client in clients:
            self._shutdown(client)
        if self._accept is not None:
            self._accept.join(timeout=1)
        try:
            Path(self.path).unlink()
        except OSError:
            pass

    # ------------------------------------------------------------------ queries

    def client_count(self) -> int:
        with self._lock:
            return len(self._clients)

    def viewer_count(self) -> int:
        with self._lock:
            return sum(1 for client in self._clients if client.viewer)

    # ------------------------------------------------------------------ outbound

    def broadcast(self, type_: str, **fields) -> None:
        """Send an assistant event to every subscribed viewer."""
        self._enqueue_viewers(events.make_event(type_, next(self._seq), **fields))

    def relay(self, event: dict) -> None:
        """Forward a raw STT event to every subscribed viewer."""
        self._enqueue_viewers(event)

    def _enqueue_viewers(self, event: dict) -> None:
        line = events.encode(event)
        with self._lock:
            clients = [c for c in self._clients if c.viewer]
        for client in clients:
            client.queue.append(line)

    def _send(self, client: _Client, type_: str, **fields) -> None:
        client.queue.append(events.encode(events.make_event(type_, next(self._seq), **fields)))

    # ------------------------------------------------------------------ accept/IO

    def _accept_loop(self) -> None:
        server = self._server
        if server is None:
            return
        while not self._stop.is_set():
            try:
                conn, _ = server.accept()
            except socket.timeout:
                continue
            except OSError:
                break
            client = _Client(conn, CLIENT_QUEUE)
            with self._lock:
                self._clients.append(client)
            client.writer = threading.Thread(
                target=self._writer, args=(client,), name="control-writer", daemon=True
            )
            client.reader = threading.Thread(
                target=self._reader, args=(client,), name="control-reader", daemon=True
            )
            client.writer.start()
            client.reader.start()

    def _writer(self, client: _Client) -> None:
        try:
            while not self._stop.is_set():
                try:
                    line = client.queue.popleft()
                except IndexError:
                    time.sleep(POLL_INTERVAL)
                    continue
                client.conn.sendall(line)
        except OSError:
            pass
        finally:
            self._drop(client)

    def _reader(self, client: _Client) -> None:
        buffer = bytearray()
        try:
            while not self._stop.is_set():
                try:
                    chunk = client.conn.recv(READ_CHUNK)
                except OSError:
                    break
                if not chunk:
                    break
                buffer.extend(chunk)
                while True:
                    newline = buffer.find(b"\n")
                    if newline == -1:
                        break
                    line = bytes(buffer[:newline])
                    del buffer[: newline + 1]
                    if line.strip():
                        self._handle_line(client, line)
        except OSError:
            pass
        finally:
            self._drop(client)

    def _drop(self, client: _Client) -> None:
        last_viewer = False
        with self._lock:
            if client in self._clients:
                self._clients.remove(client)
                # Losing the last viewer ends an unattended chat. Decide under the
                # lock, then notify after releasing it: viewer_detached() emits,
                # and emitting re-enters this non-reentrant lock.
                last_viewer = client.viewer and not any(c.viewer for c in self._clients)
        self._shutdown(client)
        if last_viewer:
            self.service.viewer_detached()

    def _shutdown(self, client: _Client) -> None:
        try:
            client.conn.shutdown(socket.SHUT_RDWR)
        except OSError:
            pass
        try:
            client.conn.close()
        except OSError:
            pass

    # ------------------------------------------------------------------ dispatch

    def _handle_line(self, client: _Client, line: bytes) -> None:
        try:
            message = json.loads(line.decode("utf-8"))
        except (UnicodeDecodeError, json.JSONDecodeError):
            self._send(client, "error", code="bad_message", message="invalid JSON")
            return
        if not isinstance(message, dict) or "type" not in message:
            self._send(
                client, "error", code="bad_message", message="expected an object with a type"
            )
            return
        self._dispatch(client, message)

    def _dispatch(self, client: _Client, message: dict) -> None:
        command = message.get("type")
        service = self.service

        if command == "subscribe":
            client.viewer = bool(message.get("viewer", False))
            self._send_status(client)
        elif command == "status":
            self._send_status(client)
        elif command == "command_start":
            self._ack_or_error(client, command, service.arm(), "cannot start a command now")
        elif command == "command_cancel":
            self._ack_or_error(client, command, service.cancel(), "no command to cancel")
        elif command == "command_edit":
            ok = service.edit(str(message.get("text", "")))
            self._ack_or_error(client, command, ok, "no command to edit")
        elif command == "command_execute":
            self._execute(client, command)
        elif command == "chat_toggle":
            self._ack_or_error(client, command, service.toggle_chat(), "cannot toggle chat now")
        elif command == "chat_send":
            ok = service.submit_chat_text(str(message.get("text", "")))
            self._ack_or_error(client, command, ok, "cannot send chat now")
        elif command == "chat_capture":
            ok = service.set_chat_capture(bool(message.get("on", True)))
            self._ack_or_error(client, command, ok, "cannot set chat capture")
        elif command == "triage_chat":
            ok = service.set_triage_chat(bool(message.get("on", False)))
            self._ack_or_error(client, command, ok, "cannot set triage chat")
        elif command == "clear":
            self._ack_or_error(client, command, service.clear(), "nothing to clear")
        elif command == "speaker_list":
            self._send(client, "speaker_list", speakers=service.speaker_list())
        elif command == "speaker_name":
            speaker_id = str(message.get("speaker_id", ""))
            name = str(message.get("name", ""))
            if not name:
                self._send(client, "error", code="bad_message", message="name is required")
            elif service.speaker_name(speaker_id, name):
                self._send(client, "ack", command="speaker_name", ok=True)
            else:
                self._send(
                    client,
                    "error",
                    code="unknown_speaker",
                    message=f"unknown speaker: {speaker_id}",
                )
        else:
            self._send(client, "error", code="unknown_command", message=f"unknown: {command}")

    def _execute(self, client: _Client, command: str) -> None:
        if self.cfg.assistant.require_viewer and self.viewer_count() == 0:
            self._send(client, "error", code="no_viewer", message="no viewer attached")
            return
        result = self.service.execute()
        if result is EXECUTE_BUSY:
            self._send(client, "error", code="busy", message="inference queue full")
            return
        if result is None:
            self._send(client, "error", code="not_ready", message="command is not ready")
            return
        self._send(client, "ack", command=command, ok=True)

    def _ack_or_error(self, client: _Client, command: str, ok: bool, reason: str) -> None:
        if ok:
            self._send(client, "ack", command=command, ok=True)
        else:
            self._send(client, "error", code="refused", message=reason)

    def _send_status(self, client: _Client) -> None:
        self._send(client, "status", **self.service.status())