STT-Assistant/tests/test_control.py

"""Tests for the assistant control socket (Phase 3)."""

import json
import os
import shutil
import socket
import tempfile
import time
from pathlib import Path
from types import SimpleNamespace

import pytest

from stt_daemon.assistant.audit import AuditLog
from stt_daemon.assistant.control import ControlServer
from stt_daemon.assistant.executor import ExecutionResult
from stt_daemon.assistant.service import AssistantService
from stt_daemon.config import Config


@pytest.fixture
def sock_path():
    directory = tempfile.mkdtemp(prefix="sttc-")
    try:
        yield os.path.join(directory, "c.sock")
    finally:
        shutil.rmtree(directory, ignore_errors=True)


class FakeSubscriber:
    def events(self):
        return iter(())

    def close(self):
        pass


class FakeExecutor:
    def __init__(self):
        self.calls = []

    def run(self, command_text, emit):
        self.calls.append(command_text)
        emit("tool_call", tool="get_datetime", args={}, requires_confirmation=False)
        return ExecutionResult(ok=True, result="2026-09-20T12:00:00", tool="get_datetime")


class Client:
    def __init__(self, path):
        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.sock.connect(path)
        self.sock.settimeout(0.2)
        self.buffer = bytearray()
        self.messages: list[dict] = []

    def send(self, **message):
        self.sock.sendall((json.dumps(message) + "\n").encode("utf-8"))

    def read(self, timeout=0.2):
        deadline = time.time() + timeout
        while time.time() < deadline:
            newline = self.buffer.find(b"\n")
            if newline != -1:
                line = bytes(self.buffer[:newline])
                del self.buffer[: newline + 1]
                return json.loads(line.decode("utf-8"))
            try:
                chunk = self.sock.recv(65536)
            except socket.timeout:
                continue
            if not chunk:
                return None
            self.buffer.extend(chunk)
        return None

    def wait_for(self, predicate, timeout=5.0):
        """Return the first message matching `predicate`, accumulating the rest."""
        deadline = time.time() + timeout
        while time.time() < deadline:
            match = next((m for m in self.messages if predicate(m)), None)
            if match is not None:
                return match
            message = self.read()
            if message is not None:
                self.messages.append(message)
        return next((m for m in self.messages if predicate(m)), None)

    def close(self):
        try:
            self.sock.close()
        except OSError:
            pass


@pytest.fixture
def harness(sock_path, tmp_path):
    cfg = Config()
    cfg.assistant.control_socket = sock_path
    cfg.assistant.require_viewer = True

    audit_dir = tmp_path / "audit"
    audit = AuditLog(enabled=True, directory=str(audit_dir))
    executor = FakeExecutor()
    service = AssistantService(cfg, FakeSubscriber(), audit=audit, executor=executor)
    server = ControlServer(cfg, service, path=sock_path)
    service.set_sinks(emit=server.broadcast, relay=server.relay)
    server.start()
    service.start()

    yield SimpleNamespace(
        cfg=cfg,
        service=service,
        server=server,
        audit_dir=audit_dir,
        path=sock_path,
        executor=executor,
    )

    server.close()
    service.stop()


def _transcript(text, ts=None):
    return {
        "v": 1,
        "type": "transcript",
        "seq": 99,
        "ts": time.time() if ts is None else ts,
        "text": text,
    }


def test_subscribe_returns_status_and_counts_viewer(harness):
    client = Client(harness.path)
    client.send(type="subscribe", viewer=True)

    status = client.wait_for(lambda m: m["type"] == "status")
    assert status is not None
    assert status["mode"] == "idle"
    assert status["ready"] is True
    assert harness.server.viewer_count() == 1

    client.close()


def test_command_start_flips_mode_and_broadcasts(harness):
    client = Client(harness.path)
    client.send(type="subscribe", viewer=True)
    assert client.wait_for(lambda m: m["type"] == "status") is not None

    client.send(type="command_start")
    assert client.wait_for(lambda m: m["type"] == "ack" and m["command"] == "command_start")
    assert client.wait_for(lambda m: m["type"] == "mode" and m["mode"] == "command_armed")
    assert harness.service.modes.mode.value == "command_armed"

    client.close()


def test_command_execute_refused_without_viewer(harness):
    client = Client(harness.path)  # connected but not subscribed as a viewer

    client.send(type="command_execute")
    error = client.wait_for(lambda m: m["type"] == "error")
    assert error is not None
    assert error["code"] == "no_viewer"
    assert harness.executor.calls == []

    client.close()


def test_command_execute_runs_executor_and_writes_audit(harness):
    client = Client(harness.path)
    client.send(type="subscribe", viewer=True)
    assert client.wait_for(lambda m: m["type"] == "status") is not None

    client.send(type="command_start")
    assert client.wait_for(lambda m: m["type"] == "ack" and m["command"] == "command_start")

    service = harness.service
    service.handle_event({"v": 1, "type": "speech_stop", "seq": 1, "ts": time.time(), "seg_id": 1})
    service.handle_event(
        {"v": 1, "type": "segment_done", "seq": 2, "ts": time.time(), "seg_id": 1, "emitted": True}
    )
    service.handle_event(_transcript("what time is it"))
    assert service.ready() is True

    client.send(type="command_execute")
    assert client.wait_for(lambda m: m["type"] == "ack" and m["command"] == "command_execute")

    result = client.wait_for(lambda m: m["type"] == "tool_result")
    assert result is not None
    assert result["ok"] is True
    assert result["tool"] == "get_datetime"
    assert harness.executor.calls == ["what time is it"]

    files = list(harness.audit_dir.glob("*.jsonl"))
    assert len(files) == 1
    record = json.loads(files[0].read_text().splitlines()[0])
    assert record["tool"] == "get_datetime"
    assert record["ok"] is True
    assert record["text"] == "what time is it"

    client.close()


def test_command_edit_updates_buffer(harness):
    client = Client(harness.path)
    client.send(type="subscribe", viewer=True)
    assert client.wait_for(lambda m: m["type"] == "status") is not None

    client.send(type="command_start")
    assert client.wait_for(lambda m: m["type"] == "ack")

    client.send(type="command_edit", text="corrected text")
    assert client.wait_for(lambda m: m["type"] == "command_text" and m["text"] == "corrected text")
    assert harness.service.command_text == "corrected text"
    assert harness.service.command_edited is True

    client.close()


def test_chat_toggle_broadcasts_mode(harness):
    client = Client(harness.path)
    client.send(type="subscribe", viewer=True)
    assert client.wait_for(lambda m: m["type"] == "status") is not None

    client.send(type="chat_toggle")
    assert client.wait_for(lambda m: m["type"] == "mode" and m["mode"] == "chat")

    client.close()


def test_chat_send_acks_and_broadcasts_user_turn(harness):
    client = Client(harness.path)
    client.send(type="subscribe", viewer=True)
    assert client.wait_for(lambda m: m["type"] == "status") is not None

    client.send(type="chat_toggle")
    assert client.wait_for(lambda m: m["type"] == "mode" and m["mode"] == "chat")

    client.send(type="chat_send", text="hello there")
    assert client.wait_for(lambda m: m["type"] == "ack" and m["command"] == "chat_send")
    assert client.wait_for(
        lambda m: m["type"] == "chat_message"
        and m.get("role") == "user"
        and m.get("text") == "hello there"
    ) is not None
    assert harness.service.chat_history == ["hello there"]

    client.close()


def test_chat_send_refused_outside_chat(harness):
    client = Client(harness.path)
    client.send(type="subscribe", viewer=True)
    assert client.wait_for(lambda m: m["type"] == "status") is not None

    client.send(type="chat_send", text="hello")
    error = client.wait_for(lambda m: m["type"] == "error" and m["code"] == "refused")
    assert error is not None
    assert harness.service.chat_history == []

    client.close()


def test_chat_capture_command_sets_and_broadcasts(harness):
    client = Client(harness.path)
    client.send(type="subscribe", viewer=True)
    status = client.wait_for(lambda m: m["type"] == "status")
    assert status is not None
    assert status["chat_capture"] is False

    client.send(type="chat_capture", on=True)
    assert client.wait_for(lambda m: m["type"] == "ack" and m["command"] == "chat_capture")
    assert harness.service.chat_capture is True
    assert client.wait_for(
        lambda m: m["type"] == "chat_capture" and m["on"] is True
    ) is not None

    client.close()


def test_triage_chat_command_sets_and_broadcasts(harness):
    client = Client(harness.path)
    client.send(type="subscribe", viewer=True)
    status = client.wait_for(lambda m: m["type"] == "status")
    assert status is not None
    assert status["triage_chat"] is False

    client.send(type="triage_chat", on=True)
    assert client.wait_for(lambda m: m["type"] == "ack" and m["command"] == "triage_chat")
    assert harness.service.triage_chat is True
    assert client.wait_for(
        lambda m: m["type"] == "triage_chat" and m["on"] is True
    ) is not None

    client.close()


def test_last_viewer_disconnect_ends_chat(harness):
    client = Client(harness.path)
    client.send(type="subscribe", viewer=True)
    assert client.wait_for(lambda m: m["type"] == "status") is not None

    client.send(type="chat_toggle")
    assert client.wait_for(lambda m: m["type"] == "mode" and m["mode"] == "chat")
    assert harness.service.modes.mode.value == "chat"

    client.close()

    deadline = time.time() + 5.0
    while time.time() < deadline and harness.service.modes.mode.value == "chat":
        time.sleep(0.02)
    assert harness.service.modes.mode.value == "idle"


def test_non_last_viewer_disconnect_keeps_chat(harness):
    first = Client(harness.path)
    first.send(type="subscribe", viewer=True)
    assert first.wait_for(lambda m: m["type"] == "status") is not None

    second = Client(harness.path)
    second.send(type="subscribe", viewer=True)
    assert second.wait_for(lambda m: m["type"] == "status") is not None

    first.send(type="chat_toggle")
    assert first.wait_for(lambda m: m["type"] == "mode" and m["mode"] == "chat")

    second.close()
    time.sleep(0.4)
    assert harness.service.modes.mode.value == "chat"

    first.close()


def test_status_reply_contains_fields(harness):
    client = Client(harness.path)
    client.send(type="status")

    status = client.wait_for(lambda m: m["type"] == "status")
    assert status is not None
    for field in (
        "mode",
        "ready",
        "stt_ok",
        "phase",
        "command_text",
        "speech_open",
        "chat_turns",
        "chat_capture",
        "triage_chat",
    ):
        assert field in status

    client.close()


def test_unknown_command_is_rejected(harness):
    client = Client(harness.path)
    client.send(type="bogus")

    error = client.wait_for(lambda m: m["type"] == "error")
    assert error is not None
    assert error["code"] == "unknown_command"

    client.close()


def test_clear_broadcasts_cleared_event(harness):
    client = Client(harness.path)
    client.send(type="subscribe", viewer=True)
    assert client.wait_for(lambda m: m["type"] == "status") is not None

    client.send(type="chat_toggle")
    assert client.wait_for(lambda m: m["type"] == "mode" and m["mode"] == "chat")

    service = harness.service
    service.set_chat_capture(True)
    service.handle_event(_transcript("hello"))

    assert len(service.chat_history) == 1

    client.send(type="clear")
    assert client.wait_for(lambda m: m["type"] == "ack" and m["command"] == "clear")
    assert client.wait_for(lambda m: m["type"] == "cleared") is not None

    assert service.chat_history == []
    assert service._pending_chat == []

    client.close()


def test_control_socket_is_private_and_removed(tmp_path):
    cfg = Config()
    path = str(Path(tempfile.mkdtemp(prefix="sttc-")) / "c.sock")
    service = AssistantService(cfg, FakeSubscriber())
    server = ControlServer(cfg, service, path=path)
    server.start()

    assert os.stat(path).st_mode & 0o777 == 0o600

    server.close()
    assert not os.path.exists(path)
    service.stop()