atlas-iac/testing/tests/test_hermes_stt_streaming.py

"""Contract tests for Hermes' bounded, speculative PCM transcription stream."""

from __future__ import annotations

import importlib.util
import io
import struct
import sys
import threading
import time
import wave
from pathlib import Path
from types import SimpleNamespace

import pytest


ROOT = Path(__file__).resolve().parents[2]
SERVER = ROOT / "dockerfiles" / "hermes-jetson-stt-server.py"


def _load_server(monkeypatch):
    """Import the server without requiring the image's CUDA dependencies."""
    spec = importlib.util.spec_from_file_location("hermes_stt_stream_server", SERVER)
    assert spec and spec.loader
    module = importlib.util.module_from_spec(spec)
    monkeypatch.setitem(sys.modules, "cgi", SimpleNamespace())
    monkeypatch.setitem(
        sys.modules,
        "torch",
        SimpleNamespace(
            cuda=SimpleNamespace(is_available=lambda: False, synchronize=lambda: None)
        ),
    )
    monkeypatch.setitem(sys.modules, "whisper", SimpleNamespace())
    spec.loader.exec_module(module)
    return module


class _Connection:
    """Capture server messages without constructing an HTTP handler."""

    def __init__(self) -> None:
        self.messages: list[dict] = []
        self.condition = threading.Condition()

    def send_json(self, payload: dict) -> None:
        with self.condition:
            self.messages.append(payload)
            self.condition.notify_all()

    def wait_for(self, message_type: str, timeout: float = 2.0) -> dict:
        deadline = time.monotonic() + timeout
        with self.condition:
            while time.monotonic() < deadline:
                match = next(
                    (
                        message
                        for message in self.messages
                        if message.get("type") == message_type
                    ),
                    None,
                )
                if match:
                    return match
                self.condition.wait(deadline - time.monotonic())
        raise AssertionError(f"no {message_type} response: {self.messages}")

    def wait_for_revision(self, revision: int, timeout: float = 2.0) -> dict:
        """Wait for one particular rolling partial rather than an older one."""
        deadline = time.monotonic() + timeout
        with self.condition:
            while time.monotonic() < deadline:
                match = next(
                    (
                        message
                        for message in self.messages
                        if message.get("type") == "partial"
                        and message.get("revision") == revision
                    ),
                    None,
                )
                if match:
                    return match
                self.condition.wait(deadline - time.monotonic())
        raise AssertionError(f"no rolling revision {revision}: {self.messages}")


class _Model:
    """Return a deterministic Whisper-shaped result and retain decode options."""

    def __init__(self) -> None:
        self.calls: list[tuple[str, dict]] = []

    def transcribe(self, path: str, **options: object) -> dict:
        self.calls.append((path, options))
        return {
            "text": " hello",
            "language": "en",
            "segments": [
                {
                    "text": " hello",
                    "no_speech_prob": 0.01,
                    "avg_logprob": -0.1,
                }
            ],
        }


class _SequencedModel(_Model):
    """Return evolving text and record how much PCM each decode consumed."""

    def __init__(self, texts: list[str]) -> None:
        super().__init__()
        self.texts = texts
        self.frame_counts: list[int] = []

    def transcribe(self, path: str, **options: object) -> dict:
        with wave.open(path, "rb") as wav_file:
            self.frame_counts.append(wav_file.getnframes())
        text = self.texts[min(len(self.calls), len(self.texts) - 1)]
        self.calls.append((path, options))
        return {
            "text": f" {text}",
            "language": "en",
            "segments": [
                {
                    "text": f" {text}",
                    "no_speech_prob": 0.01,
                    "avg_logprob": -0.1,
                }
            ],
        }


class _BlockingModel(_Model):
    """Hold inference so epoch/cancel behavior can be tested deterministically."""

    def __init__(self) -> None:
        super().__init__()
        self.entered = threading.Event()
        self.release = threading.Event()

    def transcribe(self, path: str, **options: object) -> dict:
        self.entered.set()
        assert self.release.wait(2.0)
        return super().transcribe(path, **options)


def _start(session) -> None:
    session.start(
        {
            "type": "start",
            "turn_id": "turn-10",
            "format": "pcm_s16le",
            "sample_rate": 16_000,
            "language": "auto",
        }
    )


def _speech(samples: int = 1_600) -> bytes:
    return struct.pack(f"<{samples}h", *([12_000] * samples))


def _silence(samples: int = 12_000) -> bytes:
    return bytes(samples * 2)


def _configure_fast_rolling(module, monkeypatch) -> None:
    """Make rolling thresholds deterministic without adding test sleeps."""
    monkeypatch.setattr(module, "ROLLING_MIN_AUDIO_MS", 100)
    monkeypatch.setattr(module, "ROLLING_INTERVAL_MS", 100)
    monkeypatch.setattr(module, "ROLLING_MAX_RESULT_LAG_MS", 5_000)


def _wait_rolling_idle(session, timeout: float = 1.0) -> None:
    """Wait until an inference worker has applied or discarded its result."""
    deadline = time.monotonic() + timeout
    with session._lock:
        while session._rolling_inflight and time.monotonic() < deadline:
            session._lock.wait(deadline - time.monotonic())
        assert not session._rolling_inflight


def test_websocket_handshake_matches_rfc_example(monkeypatch):
    module = _load_server(monkeypatch)

    assert (
        module.WebSocketConnection.accept_value("dGhlIHNhbXBsZSBub25jZQ==")
        == "s3pPLMBiTxaQ9kYGzzhZRbK+xOo="
    )


def test_websocket_read_exact_accumulates_partial_socket_reads(monkeypatch):
    module = _load_server(monkeypatch)
    payload = b'{"type":"resume"}'
    mask = b"mask"
    masked = bytes(value ^ mask[index % 4] for index, value in enumerate(payload))
    frame = bytes([0x81, 0x80 | len(payload)]) + mask + masked

    class ShortReader(io.BytesIO):
        def read(self, size: int = -1) -> bytes:
            return super().read(1 if size > 0 else size)

    handler = SimpleNamespace(rfile=ShortReader(frame), wfile=io.BytesIO())
    connection = module.WebSocketConnection(handler)

    assert connection.receive() == (0x1, payload)


def test_stream_speculation_is_reused_for_unchanged_commit(monkeypatch):
    module = _load_server(monkeypatch)
    connection = _Connection()
    model = _Model()
    session = module.StreamingTranscription(connection, model)
    _start(session)

    session.append(_silence(samples=32_000))
    session.append(_speech(samples=20_800))
    assert not session._rolling_inflight
    assert model.calls == []
    session.append(_silence())
    partial = connection.wait_for("partial")
    assert partial == {
        "type": "partial",
        "turn_id": "turn-10",
        "transcript": "hello",
        "language": "en",
        "speculative": True,
    }

    session.commit()
    final = connection.wait_for("final")
    assert final["transcript"] == "hello"
    assert final["language"] == "en"
    assert len(model.calls) == 1
    assert model.calls[0][1]["language"] is None
    assert model.calls[0][1]["condition_on_previous_text"] is False


def test_active_speech_emits_only_one_disposable_rolling_partial(monkeypatch):
    module = _load_server(monkeypatch)
    connection = _Connection()
    model = _SequencedModel(["please open", "please open my calendar"])
    session = module.StreamingTranscription(connection, model)
    _start(session)

    session.append(_speech(samples=32_000))
    first = connection.wait_for_revision(1)
    assert first == {
        "type": "partial",
        "turn_id": "turn-10",
        "transcript": "please open",
        "stable_transcript": "",
        "language": "en",
        "speculative": True,
        "rolling": True,
        "revision": 1,
        "epoch": 0,
        "window_start_ms": 0,
    }

    # More PCM in the same epoch must not launch disposable GPU work again.
    session._rolling_last_started_at = 0.0
    session.append(_speech(samples=12_000))
    time.sleep(0.02)
    assert len(model.calls) == 1
    assert not any(message.get("revision") == 2 for message in connection.messages)


def test_rolling_window_is_bounded_but_final_decode_uses_full_utterance(monkeypatch):
    module = _load_server(monkeypatch)
    _configure_fast_rolling(module, monkeypatch)
    monkeypatch.setattr(module, "ROLLING_WINDOW_MS", 300)
    connection = _Connection()
    model = _SequencedModel(["partial", "complete request"])
    session = module.StreamingTranscription(connection, model)
    _start(session)

    utterance = _speech(samples=16_000)
    session.append(utterance)
    partial = connection.wait_for_revision(1)
    assert partial["window_start_ms"] == 700
    assert model.frame_counts == [4_800]

    session.commit()
    assert connection.wait_for("final")["transcript"] == "complete request"
    assert model.frame_counts == [4_800, 16_000]


def test_cancel_drops_inflight_rolling_result(monkeypatch):
    module = _load_server(monkeypatch)
    _configure_fast_rolling(module, monkeypatch)
    connection = _Connection()
    model = _BlockingModel()
    session = module.StreamingTranscription(connection, model)
    _start(session)

    session.append(_speech())
    assert model.entered.wait(1.0)
    session.cancel()
    model.release.set()
    _wait_rolling_idle(session)

    assert connection.messages == []


def test_resume_epoch_drops_old_rolling_result(monkeypatch):
    module = _load_server(monkeypatch)
    _configure_fast_rolling(module, monkeypatch)
    connection = _Connection()
    model = _BlockingModel()
    session = module.StreamingTranscription(connection, model)
    _start(session)

    session.append(_speech())
    assert model.entered.wait(1.0)
    session.resume()
    model.release.set()
    _wait_rolling_idle(session)

    assert connection.messages == []
    assert session._epoch == 1


def test_busy_gpu_skips_partial_instead_of_queuing_work(monkeypatch):
    module = _load_server(monkeypatch)
    _configure_fast_rolling(module, monkeypatch)
    connection = _Connection()
    model = _Model()
    session = module.StreamingTranscription(connection, model)
    _start(session)

    release = threading.Event()
    entered = threading.Event()

    def occupy_gpu():
        module.INFERENCE_GATE.try_background(
            lambda: (entered.set(), release.wait(2.0)),
            kind="rolling",
        )

    worker = threading.Thread(target=occupy_gpu)
    worker.start()
    assert entered.wait(1.0)
    try:
        session.append(_speech())
        _wait_rolling_idle(session)
        assert model.calls == []
        assert connection.messages == []
    finally:
        release.set()
        worker.join(timeout=1.0)

    # A skipped epoch is never retried, so it cannot race the final later.
    session._rolling_last_started_at = 0.0
    session.append(_speech())
    time.sleep(0.02)
    assert connection.messages == []
    assert model.calls == []


def test_waiting_final_prevents_new_rolling_work_from_stealing_gpu(monkeypatch):
    module = _load_server(monkeypatch)
    first_entered = threading.Event()
    release_first = threading.Event()
    final_ran = threading.Event()

    first = threading.Thread(
        target=lambda: module.INFERENCE_GATE.try_background(
            lambda: (first_entered.set(), release_first.wait(2.0)),
            kind="rolling",
        )
    )
    first.start()
    assert first_entered.wait(1.0)
    final = threading.Thread(
        target=lambda: module.INFERENCE_GATE.run_final(final_ran.set)
    )
    final.start()
    deadline = time.monotonic() + 1.0
    while module.INFERENCE_GATE._final_waiters < 1 and time.monotonic() < deadline:
        time.sleep(0.005)
    assert module.INFERENCE_GATE._final_waiters == 1
    assert module.INFERENCE_GATE.try_background(lambda: "ran", kind="rolling") is None
    release_first.set()
    first.join(timeout=1.0)
    final.join(timeout=1.0)
    assert final_ran.is_set()
    telemetry = module.INFERENCE_GATE.snapshot()
    assert telemetry["background_skips"] == 1
    assert telemetry["last_final_wait_ms"] >= 0
    assert telemetry["last_final_decode_ms"] >= 0


def test_resume_invalidates_speculative_snapshot(monkeypatch):
    module = _load_server(monkeypatch)
    connection = _Connection()
    model = _Model()
    session = module.StreamingTranscription(connection, model)
    _start(session)

    session.append(_speech())
    session.append(_silence())
    connection.wait_for("partial")
    session.resume()
    session.append(_speech(800))
    session.commit()

    assert len(model.calls) == 2
    assert connection.wait_for("final")["turn_id"] == "turn-10"


def test_final_snapshot_preserves_resumed_audio_and_only_trims_tail(monkeypatch):
    module = _load_server(monkeypatch)
    session = module.StreamingTranscription(_Connection(), _Model())
    _start(session)
    voiced = _speech()
    session.append(voiced)
    session.append(_silence())

    with session._lock:
        trimmed, _ = session._snapshot_locked()
    tail_bytes = (
        module.STREAM_SAMPLE_RATE
        * module.STREAM_SAMPLE_WIDTH
        * module.VAD_TAIL_MS
        // 1000
    )
    assert len(trimmed) == len(voiced) + tail_bytes

    session.resume()
    quiet_resumed_speech = struct.pack("<800h", *([100] * 800))
    session.append(quiet_resumed_speech)
    with session._lock:
        resumed, _ = session._snapshot_locked()
    assert resumed.endswith(quiet_resumed_speech)


def test_stream_contract_rejects_wrong_format_and_oversize(monkeypatch):
    module = _load_server(monkeypatch)
    session = module.StreamingTranscription(_Connection(), _Model())
    with pytest.raises(module.WebSocketError, match="pcm_s16le"):
        session.start(
            {
                "turn_id": "turn-10",
                "format": "webm",
                "sample_rate": 16_000,
            }
        )

    malformed_rate = module.StreamingTranscription(_Connection(), _Model())
    with pytest.raises(module.WebSocketError, match="invalid sample rate"):
        malformed_rate.start(
            {
                "turn_id": "turn-10",
                "format": "pcm_s16le",
                "sample_rate": "sixteen kilohertz",
            }
        )

    session = module.StreamingTranscription(_Connection(), _Model())
    _start(session)
    monkeypatch.setattr(module, "MAX_STREAM_AUDIO_BYTES", 4)
    with pytest.raises(module.WebSocketError, match="too long"):
        session.append(bytes(6))


class _RecordingModel:
    """Record the exact decode options each path passes to Whisper."""

    def transcribe(self, path: str, **options: object) -> dict:
        self.options = options
        return {
            "text": " ok",
            "language": "en",
            "segments": [{"text": " ok", "no_speech_prob": 0.0, "avg_logprob": -0.1}],
        }


def test_final_decode_is_accurate_and_primed_while_rolling_stays_fast_and_clean(
    monkeypatch,
):
    """The committed/speculative path uses beam search + a temperature ladder +
    an initial_prompt; the disposable rolling preview stays greedy and clean."""
    module = _load_server(monkeypatch)
    wav = module._write_pcm_wav(_speech(samples=1_600))

    final_model = _RecordingModel()
    module._decode_final(final_model, wav, "auto")
    assert final_model.options["beam_size"] == module.FINAL_BEAM_SIZE >= 2
    assert final_model.options["temperature"] == module.FINAL_TEMPERATURE
    assert final_model.options["temperature"][0] == 0.0
    assert len(final_model.options["temperature"]) > 1
    assert final_model.options["initial_prompt"] == module.INITIAL_PROMPT
    assert final_model.options["condition_on_previous_text"] is False

    rolling_model = _RecordingModel()
    module._decode_rolling(rolling_model, wav, "auto")
    assert rolling_model.options["beam_size"] == 1
    assert rolling_model.options["temperature"] == 0
    assert "initial_prompt" not in rolling_model.options
    assert rolling_model.options["condition_on_previous_text"] is False

    Path(wav).unlink()


def test_initial_prompt_is_env_overridable_and_default_primes_places(monkeypatch):
    module = _load_server(monkeypatch)
    # The default primes accented, capitalized place names across EN/ES/RU.
    assert "Córdoba" in module.DEFAULT_INITIAL_PROMPT
    assert "Москва" in module.DEFAULT_INITIAL_PROMPT
    assert "Amy" in module.DEFAULT_INITIAL_PROMPT
    # Domain acronyms are primed so they are not misheard as common words
    # (e.g. "CUI" transcribed as "cue").
    assert "CUI" in module.DEFAULT_INITIAL_PROMPT
    assert "DoD" in module.DEFAULT_INITIAL_PROMPT
    assert module.INITIAL_PROMPT == module.DEFAULT_INITIAL_PROMPT

    monkeypatch.setenv("HERMES_STT_INITIAL_PROMPT", "Custom prime.")
    reloaded = _load_server(monkeypatch)
    assert reloaded.INITIAL_PROMPT == "Custom prime."