atlas-iac/dockerfiles/hermes-jetson-stt-server.py

#!/usr/bin/env python3
"""Small OpenAI-compatible Whisper service for the dedicated Jetson."""

from __future__ import annotations

import cgi
import base64
import hashlib
import json
import os
import re
import socket
import struct
import tempfile
import threading
import time
import wave
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

import torch
import whisper


HOST = os.getenv("HERMES_STT_HOST", "0.0.0.0")
PORT = int(os.getenv("HERMES_STT_PORT", "9000"))
MODEL_NAME = os.getenv("HERMES_STT_MODEL", "small")
ROLLING_MODEL_NAME = os.getenv("HERMES_STT_ROLLING_MODEL", "tiny")
CACHE_DIR = Path(os.getenv("HERMES_STT_CACHE", "/cache/whisper"))
MAX_AUDIO_BYTES = 30 * 1024 * 1024
STREAM_SAMPLE_RATE = 16_000
STREAM_SAMPLE_WIDTH = 2
STREAM_CHANNELS = 1
MAX_STREAM_SECONDS = int(os.getenv("HERMES_STT_STREAM_MAX_SECONDS", "90"))
MAX_STREAM_AUDIO_BYTES = min(
    MAX_AUDIO_BYTES,
    STREAM_SAMPLE_RATE * STREAM_SAMPLE_WIDTH * MAX_STREAM_SECONDS,
)
MAX_WS_MESSAGE_BYTES = 256 * 1024
STREAM_IDLE_SECONDS = int(os.getenv("HERMES_STT_STREAM_IDLE_SECONDS", "120"))
VAD_END_SILENCE_MS = int(os.getenv("HERMES_STT_VAD_END_SILENCE_MS", "650"))
VAD_TAIL_MS = int(os.getenv("HERMES_STT_VAD_TAIL_MS", "220"))
ROLLING_MIN_AUDIO_MS = max(
    300,
    min(2_000, int(os.getenv("HERMES_STT_ROLLING_MIN_AUDIO_MS", "1800"))),
)
ROLLING_INTERVAL_MS = max(
    400,
    min(5_000, int(os.getenv("HERMES_STT_ROLLING_INTERVAL_MS", "750"))),
)
ROLLING_WINDOW_MS = max(
    3_000,
    min(30_000, int(os.getenv("HERMES_STT_ROLLING_WINDOW_MS", "12000"))),
)
ROLLING_MAX_RESULT_LAG_MS = max(
    500,
    min(5_000, int(os.getenv("HERMES_STT_ROLLING_MAX_LAG_MS", "2000"))),
)
_WS_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"


class InferenceGate:
    """Give authoritative transcripts priority over disposable partial work."""

    def __init__(self) -> None:
        self._condition = threading.Condition()
        self._busy = False
        self._final_waiters = 0
        self._last_final_wait_ms = 0.0
        self._last_final_decode_ms = 0.0
        self._last_background_decode_ms = 0.0
        self._last_rolling_decode_ms = 0.0
        self._last_speculative_decode_ms = 0.0
        self._background_skips = 0

    def run_final(self, callback):
        """Wait fairly enough that newly arriving partials cannot steal the GPU."""
        wait_started = time.monotonic()
        with self._condition:
            self._final_waiters += 1
            try:
                while self._busy:
                    self._condition.wait()
                self._busy = True
            finally:
                self._final_waiters -= 1
            self._last_final_wait_ms = (time.monotonic() - wait_started) * 1000
        decode_started = time.monotonic()
        try:
            return callback()
        finally:
            with self._condition:
                self._last_final_decode_ms = (time.monotonic() - decode_started) * 1000
                self._busy = False
                self._condition.notify_all()

    def try_background(self, callback, *, kind: str):
        """Run one useful predecode only when no final is active or waiting."""
        if kind not in {"rolling", "speculative"}:
            raise ValueError("unknown background inference kind")
        with self._condition:
            if self._busy or self._final_waiters:
                self._background_skips += 1
                return None
            self._busy = True
        decode_started = time.monotonic()
        try:
            return callback()
        finally:
            with self._condition:
                elapsed_ms = (time.monotonic() - decode_started) * 1000
                self._last_background_decode_ms = elapsed_ms
                if kind == "rolling":
                    self._last_rolling_decode_ms = elapsed_ms
                else:
                    self._last_speculative_decode_ms = elapsed_ms
                self._busy = False
                self._condition.notify_all()

    def snapshot(self) -> dict:
        """Expose bounded timing telemetry for live latency verification."""
        with self._condition:
            return {
                "busy": self._busy,
                "final_waiters": self._final_waiters,
                "last_final_wait_ms": round(self._last_final_wait_ms, 1),
                "last_final_decode_ms": round(self._last_final_decode_ms, 1),
                "last_background_decode_ms": round(self._last_background_decode_ms, 1),
                "last_rolling_decode_ms": round(self._last_rolling_decode_ms, 1),
                "last_speculative_decode_ms": round(
                    self._last_speculative_decode_ms, 1
                ),
                "background_skips": self._background_skips,
            }


INFERENCE_GATE = InferenceGate()
STARTUP_WARMUP_MS: float | None = None
STARTUP_ROLLING_WARMUP_MS: float | None = None


def _repetitive_token(token: str) -> bool:
    """Identify long periodic Whisper hallucinations caused by steady noise."""
    letters = "".join(re.findall(r"[a-z]+", token.lower()))
    if len(letters) < 10:
        return False
    for period in range(1, 5):
        pattern = letters[:period]
        matches = sum(
            character == pattern[index % period]
            for index, character in enumerate(letters)
        )
        if matches / len(letters) >= 0.86:
            return True
    return (
        max(letters.count(character) for character in set(letters)) / len(letters)
        >= 0.78
    )


def _clean_transcript(result: dict) -> str:
    """Drop noise-only segments and repetitive tokens while retaining speech."""
    segments = result.get("segments")
    if not isinstance(segments, list):
        segments = [{"text": result.get("text") or ""}]
    kept: list[str] = []
    for segment in segments:
        if not isinstance(segment, dict):
            continue
        text = str(segment.get("text") or "").strip()
        if not text:
            continue
        no_speech = float(segment.get("no_speech_prob") or 0.0)
        average_logprob = float(segment.get("avg_logprob") or 0.0)
        if no_speech >= 0.55 and average_logprob <= -0.55:
            continue
        words = [word for word in text.split() if not _repetitive_token(word)]
        if words:
            kept.append(" ".join(words))
    return " ".join(kept).strip()


def _detected_language(result: object) -> str:
    """Return the bare ISO-639 code Whisper decoded with, or nothing at all.

    ``whisper.transcribe`` reports the language it auto-detected (or the one it
    was told to use) as a plain lowercase token such as ``en``/``ru``/``yue``.
    Anything that is not that exact shape is dropped rather than guessed at, so
    a surprising model result can never become a downstream voice selector.
    """
    if not isinstance(result, dict):
        return ""
    value = result.get("language")
    if not isinstance(value, str):
        return ""
    code = value.strip().lower()
    if not 2 <= len(code) <= 3 or not code.isascii() or not code.isalpha():
        return ""
    return code


def _transcription_payload(result: dict) -> dict:
    """Build the transcription contract: text plus the model's own language."""
    return {
        "text": _clean_transcript(result),
        "model": MODEL_NAME,
        "language": _detected_language(result),
    }


def _decode_path(model: object, path: str, language: str) -> dict:
    """Decode one path with the quality settings shared by partial and final work."""
    result = model.transcribe(
        path,
        language=None if language in {"", "auto"} else language,
        task="transcribe",
        fp16=torch.cuda.is_available(),
        condition_on_previous_text=False,
        temperature=0,
        compression_ratio_threshold=2.0,
        logprob_threshold=-0.8,
        no_speech_threshold=0.5,
        verbose=False,
    )
    return _transcription_payload(result)


def _transcribe_path(model: object, path: str, language: str) -> dict:
    """Run the one canonical full-utterance Whisper decode configuration."""
    return INFERENCE_GATE.run_final(lambda: _decode_path(model, path, language))


def _write_pcm_wav(pcm: bytes) -> str:
    """Write an aligned browser PCM snapshot to a temporary Whisper input."""
    with tempfile.NamedTemporaryFile(
        prefix="atlas-stt-stream-", suffix=".wav", delete=False
    ) as temp:
        temp_path = temp.name
    with wave.open(temp_path, "wb") as wav_file:
        wav_file.setnchannels(STREAM_CHANNELS)
        wav_file.setsampwidth(STREAM_SAMPLE_WIDTH)
        wav_file.setframerate(STREAM_SAMPLE_RATE)
        wav_file.writeframes(pcm)
    return temp_path


def _transcribe_pcm(model: object, pcm: bytes, language: str) -> dict:
    """Decode one complete 16-kHz mono PCM snapshot through the canonical path."""
    temp_path = ""
    try:
        temp_path = _write_pcm_wav(pcm)
        return _transcribe_path(model, temp_path, language)
    finally:
        if temp_path:
            try:
                os.unlink(temp_path)
            except OSError:
                pass


def _warm_model(model: object) -> float:
    """Pay Whisper's lazy CUDA initialization cost before readiness."""
    temp_path = _write_pcm_wav(bytes(_pcm_bytes_for_ms(1_000)))
    started = time.monotonic()
    try:
        _decode_path(model, temp_path, "auto")
        torch.cuda.synchronize()
    finally:
        try:
            os.unlink(temp_path)
        except OSError:
            pass
    return (time.monotonic() - started) * 1000


def _transcribe_pcm_rolling(model: object, pcm: bytes, language: str) -> dict | None:
    """Decode a rolling window only when the single GPU worker is immediately free.

    Partial work is disposable. It must never form an inference backlog ahead of
    a final utterance or another user's request, so a busy model means this
    speech epoch is skipped rather than retried.
    """
    temp_path = ""
    try:
        temp_path = _write_pcm_wav(pcm)
        return INFERENCE_GATE.try_background(
            lambda: _decode_path(model, temp_path, language),
            kind="rolling",
        )
    finally:
        if temp_path:
            try:
                os.unlink(temp_path)
            except OSError:
                pass


def _transcribe_pcm_speculative(
    model: object, pcm: bytes, language: str
) -> dict | None:
    """Predecode one EOS snapshot only when the model is currently idle."""
    temp_path = ""
    try:
        temp_path = _write_pcm_wav(pcm)
        return INFERENCE_GATE.try_background(
            lambda: _decode_path(model, temp_path, language),
            kind="speculative",
        )
    finally:
        if temp_path:
            try:
                os.unlink(temp_path)
            except OSError:
                pass


def _pcm_bytes_for_ms(milliseconds: int) -> int:
    """Convert milliseconds to an aligned byte count for the stream format."""
    return (
        STREAM_SAMPLE_RATE
        * STREAM_SAMPLE_WIDTH
        * milliseconds
        // 1000
        // STREAM_SAMPLE_WIDTH
        * STREAM_SAMPLE_WIDTH
    )


def _token_identity(token: str) -> str:
    """Normalize punctuation/case drift before comparing consecutive decodes."""
    return re.sub(r"[^\w']+", "", token.casefold(), flags=re.UNICODE)


def _stable_token_prefix(previous: str, current: str) -> str:
    """Return only tokens repeated in the same order by consecutive decodes."""
    old_tokens = previous.split()
    new_tokens = current.split()
    agreed = 0
    for old, new in zip(old_tokens, new_tokens):
        if not _token_identity(old) or _token_identity(old) != _token_identity(new):
            break
        agreed += 1
    return " ".join(new_tokens[:agreed])


class WebSocketError(Exception):
    """Protocol error carrying an RFC 6455 close code safe to expose."""

    def __init__(self, message: str, close_code: int = 1002) -> None:
        super().__init__(message)
        self.close_code = close_code


class WebSocketConnection:
    """Minimal server-side RFC 6455 framing for the private streaming endpoint."""

    def __init__(self, handler: BaseHTTPRequestHandler) -> None:
        self.handler = handler
        self._send_lock = threading.Lock()
        self._closed = False

    @staticmethod
    def accept_value(key: str) -> str:
        """Return the RFC 6455 handshake response for a client nonce."""
        digest = hashlib.sha1((key + _WS_GUID).encode("ascii")).digest()
        return base64.b64encode(digest).decode("ascii")

    def _read_exact(self, size: int) -> bytes:
        data = bytearray()
        while len(data) < size:
            chunk = self.handler.rfile.read(size - len(data))
            if not chunk:
                raise EOFError("websocket peer disconnected")
            data.extend(chunk)
        return bytes(data)

    def receive(self) -> tuple[int, bytes]:
        """Read and reassemble one client message, answering ping inline."""
        opcode = 0
        payload = bytearray()
        while True:
            header = self._read_exact(2)
            first, second = header
            final = bool(first & 0x80)
            reserved = first & 0x70
            frame_opcode = first & 0x0F
            masked = bool(second & 0x80)
            length = second & 0x7F
            if reserved or not masked:
                raise WebSocketError("invalid websocket frame")
            if length == 126:
                length = struct.unpack("!H", self._read_exact(2))[0]
            elif length == 127:
                length = struct.unpack("!Q", self._read_exact(8))[0]
            if frame_opcode >= 0x8 and (not final or length > 125):
                raise WebSocketError("invalid websocket control frame")
            if (
                length > MAX_WS_MESSAGE_BYTES
                or len(payload) + length > MAX_WS_MESSAGE_BYTES
            ):
                raise WebSocketError("websocket message is too large", 1009)
            mask = self._read_exact(4)
            frame = bytearray(self._read_exact(length))
            for index in range(length):
                frame[index] ^= mask[index % 4]

            if frame_opcode == 0x8:
                raise EOFError("websocket peer closed")
            if frame_opcode == 0x9:
                self.send(bytes(frame), opcode=0xA)
                continue
            if frame_opcode == 0xA:
                continue
            if frame_opcode in {0x1, 0x2}:
                if opcode:
                    raise WebSocketError("nested websocket message")
                opcode = frame_opcode
            elif frame_opcode != 0x0 or not opcode:
                raise WebSocketError("unexpected websocket continuation")
            payload.extend(frame)
            if final:
                return opcode, bytes(payload)

    def send(self, payload: bytes, opcode: int) -> None:
        """Send one unmasked server frame while serializing background writers."""
        with self._send_lock:
            if self._closed:
                return
            length = len(payload)
            header = bytearray([0x80 | opcode])
            if length < 126:
                header.append(length)
            elif length <= 0xFFFF:
                header.extend((126,))
                header.extend(struct.pack("!H", length))
            else:
                header.extend((127,))
                header.extend(struct.pack("!Q", length))
            self.handler.wfile.write(bytes(header) + payload)
            self.handler.wfile.flush()

    def send_json(self, payload: dict) -> None:
        """Send a compact UTF-8 JSON message."""
        self.send(json.dumps(payload, separators=(",", ":")).encode("utf-8"), 0x1)

    def close(self, code: int = 1000, reason: str = "") -> None:
        """Close once; callers may safely race disconnect and worker completion."""
        reason_bytes = reason.encode("utf-8")[:123]
        with self._send_lock:
            if self._closed:
                return
            try:
                payload = struct.pack("!H", code) + reason_bytes
                self.handler.wfile.write(bytes([0x88, len(payload)]) + payload)
                self.handler.wfile.flush()
            except (BrokenPipeError, ConnectionError, OSError):
                pass
            self._closed = True


def _pcm_rms(pcm: bytes) -> float:
    """Return normalized RMS for aligned signed-16-bit little-endian PCM."""
    sample_count = len(pcm) // 2
    if not sample_count:
        return 0.0
    samples = struct.unpack(f"<{sample_count}h", pcm[: sample_count * 2])
    square_mean = sum(sample * sample for sample in samples) / sample_count
    return (square_mean**0.5) / 32768.0


class StreamingTranscription:
    """Own one turn's bounded PCM, VAD state and speculative Whisper work."""

    def __init__(
        self,
        connection: WebSocketConnection,
        model: object,
        rolling_model: object | None = None,
    ) -> None:
        self.connection = connection
        self.model = model
        self.rolling_model = rolling_model if rolling_model is not None else model
        self.turn_id = ""
        self.language = "auto"
        self._pcm = bytearray()
        self._lock = threading.Condition()
        self._started = False
        self._closed = False
        self._committing = False
        self._heard_speech = False
        self._at_eos = False
        self._silence_bytes = 0
        self._last_speech_byte = 0
        self._noise_floor = 0.004
        self._client_active = False
        self._epoch = 0
        self._pending_key = ""
        self._cached_key = ""
        self._cached_payload: dict | None = None
        self._rolling_inflight = False
        self._rolling_last_started_at = 0.0
        self._rolling_speech_bytes = 0
        self._rolling_last_speech_bytes = 0
        self._rolling_previous = ""
        self._rolling_previous_window_start = 0
        self._rolling_revision = 0
        self._rolling_attempted = False

    def _response(self, message_type: str, **values: object) -> dict:
        return {"type": message_type, "turn_id": self.turn_id, **values}

    def error(self, message: str, *, fatal: bool = False) -> None:
        """Report a turn-scoped error without leaking implementation details."""
        self.connection.send_json(self._response("error", error=message, fatal=fatal))

    def start(self, message: dict) -> None:
        """Validate and initialize the only accepted browser audio contract."""
        turn_id = str(message.get("turn_id") or "")
        language = str(message.get("language") or "auto").strip().lower()
        if self._started:
            raise WebSocketError("stream is already started", 1008)
        if not re.fullmatch(r"[A-Za-z0-9_.:-]{1,128}", turn_id):
            raise WebSocketError("invalid turn_id", 1008)
        if message.get("format") != "pcm_s16le":
            raise WebSocketError("pcm_s16le audio is required", 1003)
        try:
            sample_rate = int(message.get("sample_rate") or 0)
        except (TypeError, ValueError) as exc:
            raise WebSocketError("invalid sample rate", 1003) from exc
        if sample_rate != STREAM_SAMPLE_RATE:
            raise WebSocketError("16000 Hz audio is required", 1003)
        if language != "auto" and not re.fullmatch(r"[a-z]{2,3}", language):
            raise WebSocketError("invalid language", 1008)
        self.turn_id = turn_id
        self.language = language
        self._started = True

    def _snapshot_locked(self) -> tuple[bytes, str]:
        """Keep the full utterance and discard only bounded post-speech silence.

        Browser ``resume`` is authoritative: once sent, every following PCM
        frame is preserved until the next speculative EOS marker, even when a
        quiet syllable falls below the server's conservative energy threshold.
        """
        cutoff = len(self._pcm)
        if self._heard_speech and self._last_speech_byte:
            tail_bytes = STREAM_SAMPLE_RATE * STREAM_SAMPLE_WIDTH * VAD_TAIL_MS // 1000
            cutoff = min(cutoff, self._last_speech_byte + tail_bytes)
        pcm = bytes(self._pcm[:cutoff])
        key = f"{self._epoch}:{hashlib.sha256(pcm).hexdigest()}"
        return pcm, key

    def _rolling_snapshot_locked(self) -> tuple[bytes, int]:
        """Return a bounded tail window and its absolute start in milliseconds."""
        window_bytes = _pcm_bytes_for_ms(ROLLING_WINDOW_MS)
        start = max(0, len(self._pcm) - window_bytes)
        start -= start % STREAM_SAMPLE_WIDTH
        pcm = bytes(self._pcm[start:])
        start_ms = start * 1000 // (STREAM_SAMPLE_RATE * STREAM_SAMPLE_WIDTH)
        return pcm, start_ms

    def _reset_rolling_history_locked(self, *, rearm: bool) -> None:
        """Discard provisional agreement after the client's speech epoch changes."""
        self._rolling_previous = ""
        self._rolling_previous_window_start = 0
        # Only speech in the new epoch may satisfy the rolling threshold.
        self._rolling_speech_bytes = 0
        self._rolling_last_speech_bytes = 0
        self._rolling_last_started_at = 0.0
        if rearm:
            self._rolling_attempted = False

    def _begin_rolling_locked(self, now: float) -> tuple[bytes, int, int, int] | None:
        """Reserve one due rolling decode without ever queuing a second one."""
        total_bytes = len(self._pcm)
        minimum_bytes = _pcm_bytes_for_ms(ROLLING_MIN_AUDIO_MS)
        interval_bytes = _pcm_bytes_for_ms(ROLLING_INTERVAL_MS)
        required_new_bytes = (
            interval_bytes if self._rolling_last_speech_bytes else minimum_bytes
        )
        if (
            self._rolling_inflight
            or self._rolling_attempted
            or self._rolling_speech_bytes < minimum_bytes
            or self._rolling_speech_bytes - self._rolling_last_speech_bytes
            < required_new_bytes
            or (
                self._rolling_last_started_at
                and (now - self._rolling_last_started_at) * 1000 < ROLLING_INTERVAL_MS
            )
        ):
            return None
        pcm, window_start_ms = self._rolling_snapshot_locked()
        # Reserve the epoch before spawning. A busy gate, decode error or stale
        # result must not create a retry loop that can collide with the final.
        self._rolling_attempted = True
        self._rolling_inflight = True
        self._rolling_last_started_at = now
        self._rolling_last_speech_bytes = self._rolling_speech_bytes
        return pcm, window_start_ms, total_bytes, self._epoch

    def _start_rolling(
        self,
        pcm: bytes,
        window_start_ms: int,
        audio_bytes: int,
        epoch: int,
    ) -> None:
        """Run one disposable partial decode and publish only a current result."""

        def worker() -> None:
            payload: dict | None = None
            try:
                payload = _transcribe_pcm_rolling(
                    self.rolling_model,
                    pcm,
                    self.language,
                )
            except Exception as exc:
                # Partial inference is best-effort. The canonical commit path
                # remains authoritative and reports its own actionable errors.
                print(f"[stt] rolling transcription skipped: {exc}", flush=True)

            message: dict | None = None
            with self._lock:
                self._rolling_inflight = False
                lag_bytes = max(0, len(self._pcm) - audio_bytes)
                max_lag_bytes = _pcm_bytes_for_ms(ROLLING_MAX_RESULT_LAG_MS)
                current = (
                    payload is not None
                    and not self._closed
                    and not self._committing
                    and not self._at_eos
                    and self._epoch == epoch
                    and lag_bytes <= max_lag_bytes
                )
                if current:
                    transcript = str(payload.get("text") or "")
                    stable = ""
                    if self._rolling_previous_window_start == window_start_ms:
                        stable = _stable_token_prefix(
                            self._rolling_previous,
                            transcript,
                        )
                    self._rolling_previous = transcript
                    self._rolling_previous_window_start = window_start_ms
                    self._rolling_revision += 1
                    message = self._response(
                        "partial",
                        transcript=transcript,
                        stable_transcript=stable,
                        language=payload.get("language") or "",
                        speculative=True,
                        rolling=True,
                        revision=self._rolling_revision,
                        epoch=epoch,
                        window_start_ms=window_start_ms,
                    )
                self._lock.notify_all()
            if message is not None:
                try:
                    self.connection.send_json(message)
                except (BrokenPipeError, ConnectionError, OSError):
                    self.cancel()

        threading.Thread(
            target=worker,
            name=f"stt-rolling-{self.turn_id}",
            daemon=True,
        ).start()

    def append(self, pcm: bytes) -> None:
        """Append PCM while overlapping rolling understanding with active speech."""
        if not self._started:
            raise WebSocketError("start must precede audio", 1008)
        if not pcm or len(pcm) % STREAM_SAMPLE_WIDTH:
            raise WebSocketError("unaligned PCM audio", 1003)
        auto_speculate = False
        rolling: tuple[bytes, int, int, int] | None = None
        with self._lock:
            if self._closed or self._committing:
                raise WebSocketError("turn is no longer accepting audio", 1008)
            if len(self._pcm) + len(pcm) > MAX_STREAM_AUDIO_BYTES:
                raise WebSocketError("stream audio is too long", 1009)
            self._pcm.extend(pcm)
            rms = _pcm_rms(pcm)
            threshold = max(0.012, self._noise_floor * 2.5 + 0.004)
            speech = self._client_active or rms >= threshold
            if not self._heard_speech and not speech:
                self._noise_floor = self._noise_floor * 0.96 + rms * 0.04
            if speech:
                if self._at_eos:
                    self._epoch += 1
                    self._cached_key = ""
                    self._cached_payload = None
                    self._reset_rolling_history_locked(rearm=True)
                self._heard_speech = True
                self._at_eos = False
                self._silence_bytes = 0
                self._last_speech_byte = len(self._pcm)
                self._rolling_speech_bytes += len(pcm)
                rolling = self._begin_rolling_locked(time.monotonic())
            elif self._heard_speech:
                self._silence_bytes += len(pcm)
                required = (
                    STREAM_SAMPLE_RATE
                    * STREAM_SAMPLE_WIDTH
                    * VAD_END_SILENCE_MS
                    // 1000
                )
                if self._silence_bytes >= required and not self._at_eos:
                    self._at_eos = True
                    self._reset_rolling_history_locked(rearm=False)
                    auto_speculate = True
        if auto_speculate:
            self.speculate()
        elif rolling is not None:
            self._start_rolling(*rolling)

    def resume(self) -> None:
        """Invalidate an EOS snapshot when client-side VAD hears resumed speech."""
        with self._lock:
            if not self._started or self._closed:
                return
            self._at_eos = False
            self._silence_bytes = 0
            self._epoch += 1
            self._client_active = True
            self._cached_key = ""
            self._cached_payload = None
            self._reset_rolling_history_locked(rearm=True)

    def speculate(self) -> None:
        """Decode a stable snapshot once; only an unchanged turn may consume it."""
        with self._lock:
            if not self._started or self._closed or self._committing or not self._pcm:
                return
            self._client_active = False
            self._at_eos = True
            pcm, key = self._snapshot_locked()
            if not pcm or key in {self._pending_key, self._cached_key}:
                return
            self._pending_key = key

        def worker() -> None:
            try:
                payload = _transcribe_pcm_speculative(self.model, pcm, self.language)
            except Exception as exc:
                print(f"[stt] speculative transcription failed: {exc}", flush=True)
                with self._lock:
                    if self._pending_key == key:
                        self._pending_key = ""
                        self._lock.notify_all()
                if not self._closed and not self._committing:
                    self.error("speculative transcription failed")
                return

            if payload is None:
                with self._lock:
                    if self._pending_key == key:
                        self._pending_key = ""
                    self._lock.notify_all()
                return

            should_send = False
            with self._lock:
                if self._pending_key == key:
                    self._pending_key = ""
                current_pcm, current_key = self._snapshot_locked()
                del current_pcm
                if not self._closed and current_key == key:
                    self._cached_key = key
                    self._cached_payload = payload
                    should_send = not self._committing
                self._lock.notify_all()
            if should_send:
                try:
                    self.connection.send_json(
                        self._response(
                            "partial",
                            transcript=payload["text"],
                            language=payload["language"],
                            speculative=True,
                        )
                    )
                except (BrokenPipeError, ConnectionError, OSError):
                    self.cancel()

        threading.Thread(
            target=worker,
            name=f"stt-speculate-{self.turn_id}",
            daemon=True,
        ).start()

    def commit(self) -> None:
        """Emit one final result from the complete stable utterance snapshot."""
        with self._lock:
            if not self._started or self._closed or self._committing:
                raise WebSocketError("turn cannot be committed", 1008)
            if not self._pcm:
                raise WebSocketError("turn has no audio", 1008)
            self._committing = True
            pcm, key = self._snapshot_locked()
            deadline = time.monotonic() + STREAM_IDLE_SECONDS
            while self._pending_key == key and not self._cached_payload:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    break
                self._lock.wait(timeout=remaining)
            payload = self._cached_payload if self._cached_key == key else None

        try:
            if payload is None:
                payload = _transcribe_pcm(self.model, pcm, self.language)
            self.connection.send_json(
                self._response(
                    "final",
                    transcript=payload["text"],
                    language=payload["language"],
                    model=payload["model"],
                )
            )
        except Exception as exc:
            print(f"[stt] streaming transcription failed: {exc}", flush=True)
            self.error("transcription failed", fatal=True)
        finally:
            with self._lock:
                self._closed = True
                self._lock.notify_all()

    def cancel(self) -> None:
        """Invalidate queued work and prevent background writers after disconnect."""
        with self._lock:
            self._closed = True
            self._cached_payload = None
            self._cached_key = ""
            self._lock.notify_all()


def _json(handler: BaseHTTPRequestHandler, status: int, payload: dict) -> None:
    body = json.dumps(payload).encode("utf-8")
    handler.send_response(status)
    handler.send_header("Content-Type", "application/json")
    handler.send_header("Content-Length", str(len(body)))
    handler.send_header("Cache-Control", "no-store")
    handler.end_headers()
    handler.wfile.write(body)


class SpeechHandler(BaseHTTPRequestHandler):
    """Serve health and transcription without exposing a general runtime."""

    server_version = "AtlasWhisper/1"

    def log_message(self, message: str, *args: object) -> None:
        print(f"[stt] {self.address_string()} {message % args}", flush=True)

    def do_GET(self) -> None:
        path = self.path.split("?", 1)[0]
        if path == "/v1/audio/transcriptions/stream":
            self._serve_stream()
            return
        if path != "/health":
            _json(self, 404, {"error": "not found"})
            return
        _json(
            self,
            200,
            {
                "ok": True,
                "model": MODEL_NAME,
                "rolling_model": ROLLING_MODEL_NAME,
                "device": "cuda" if torch.cuda.is_available() else "cpu",
                "inference": INFERENCE_GATE.snapshot(),
                "startup": {
                    "warmed": (
                        STARTUP_WARMUP_MS is not None
                        and STARTUP_ROLLING_WARMUP_MS is not None
                    ),
                    "warmup_ms": round(STARTUP_WARMUP_MS or 0.0, 1),
                    "rolling_warmup_ms": round(
                        STARTUP_ROLLING_WARMUP_MS or 0.0,
                        1,
                    ),
                },
                "streaming": {
                    "enabled": True,
                    "path": "/v1/audio/transcriptions/stream",
                    "format": "pcm_s16le",
                    "sample_rate": STREAM_SAMPLE_RATE,
                    "max_seconds": MAX_STREAM_SECONDS,
                    "server_vad": True,
                    "speculative": True,
                    "rolling": {
                        "enabled": True,
                        "minimum_audio_ms": ROLLING_MIN_AUDIO_MS,
                        "interval_ms": ROLLING_INTERVAL_MS,
                        "window_ms": ROLLING_WINDOW_MS,
                        "stable_prefix": True,
                    },
                },
            },
        )

    def _serve_stream(self) -> None:
        """Upgrade one request and process a single, isolated voice turn."""
        upgrade = self.headers.get("Upgrade", "").strip().lower()
        connection_tokens = {
            token.strip().lower()
            for token in self.headers.get("Connection", "").split(",")
        }
        key = self.headers.get("Sec-WebSocket-Key", "").strip()
        version = self.headers.get("Sec-WebSocket-Version", "").strip()
        if upgrade != "websocket" or "upgrade" not in connection_tokens or not key:
            _json(self, 426, {"error": "websocket upgrade is required"})
            return
        if version != "13":
            self.send_response(426)
            self.send_header("Sec-WebSocket-Version", "13")
            self.send_header("Content-Length", "0")
            self.end_headers()
            return
        try:
            decoded_key = base64.b64decode(key, validate=True)
        except (ValueError, TypeError):
            decoded_key = b""
        if len(decoded_key) != 16:
            _json(self, 400, {"error": "invalid websocket key"})
            return

        self.send_response(101, "Switching Protocols")
        self.send_header("Upgrade", "websocket")
        self.send_header("Connection", "Upgrade")
        self.send_header("Sec-WebSocket-Accept", WebSocketConnection.accept_value(key))
        self.end_headers()
        self.wfile.flush()
        self.close_connection = True
        self.connection.settimeout(STREAM_IDLE_SECONDS)
        websocket = WebSocketConnection(self)
        session = StreamingTranscription(
            websocket,
            self.server.model,  # type: ignore[attr-defined]
            self.server.rolling_model,  # type: ignore[attr-defined]
        )
        close_code = 1000
        close_reason = ""
        try:
            while True:
                opcode, payload = websocket.receive()
                if opcode == 0x2:
                    session.append(payload)
                    continue
                try:
                    message = json.loads(payload.decode("utf-8"))
                except (UnicodeDecodeError, json.JSONDecodeError):
                    raise WebSocketError("control message must be JSON", 1007)
                if not isinstance(message, dict):
                    raise WebSocketError("control message must be an object", 1007)
                message_type = str(message.get("type") or "")
                if message_type == "start":
                    session.start(message)
                    continue
                if message.get("turn_id") != session.turn_id:
                    raise WebSocketError("turn_id does not match", 1008)
                if message_type == "speculate":
                    session.speculate()
                elif message_type == "resume":
                    session.resume()
                elif message_type == "commit":
                    session.commit()
                    break
                elif message_type == "cancel":
                    break
                else:
                    raise WebSocketError("unknown stream control", 1008)
        except WebSocketError as exc:
            close_code = exc.close_code
            close_reason = str(exc)
            try:
                session.error(str(exc), fatal=True)
            except (BrokenPipeError, ConnectionError, OSError):
                pass
        except (EOFError, BrokenPipeError, ConnectionError, socket.timeout, OSError):
            close_reason = "peer disconnected"
        finally:
            session.cancel()
            websocket.close(close_code, close_reason)

    def do_POST(self) -> None:
        if self.path != "/v1/audio/transcriptions":
            _json(self, 404, {"error": "not found"})
            return
        content_length = int(self.headers.get("Content-Length", "0") or "0")
        if content_length <= 0 or content_length > MAX_AUDIO_BYTES:
            _json(self, 413, {"error": "audio payload is missing or too large"})
            return

        content_type = self.headers.get("Content-Type", "")
        if not content_type.lower().startswith("multipart/form-data"):
            _json(self, 400, {"error": "multipart/form-data is required"})
            return

        form = cgi.FieldStorage(
            fp=self.rfile,
            headers=self.headers,
            environ={
                "REQUEST_METHOD": "POST",
                "CONTENT_TYPE": content_type,
                "CONTENT_LENGTH": str(content_length),
            },
        )
        audio = form["file"] if "file" in form else None
        if audio is None or not getattr(audio, "file", None):
            _json(self, 400, {"error": "file is required"})
            return

        suffix = Path(getattr(audio, "filename", "audio.wav") or "audio.wav").suffix
        suffix = (
            suffix if suffix in {".wav", ".webm", ".ogg", ".mp3", ".m4a"} else ".wav"
        )
        language = str(form.getfirst("language", "auto") or "auto").strip().lower()
        temp_path = ""
        try:
            with tempfile.NamedTemporaryFile(
                prefix="atlas-stt-", suffix=suffix, delete=False
            ) as temp:
                temp_path = temp.name
                while True:
                    chunk = audio.file.read(1024 * 1024)
                    if not chunk:
                        break
                    temp.write(chunk)

            payload = _transcribe_path(
                self.server.model,  # type: ignore[attr-defined]
                temp_path,
                language,
            )
            _json(self, 200, payload)
        except Exception as exc:
            print(f"[stt] transcription failed: {exc}", flush=True)
            _json(self, 500, {"error": "transcription failed"})
        finally:
            if temp_path:
                try:
                    os.unlink(temp_path)
                except OSError:
                    pass


def main() -> None:
    """Warm Whisper once, then serve concurrent clients through one GPU lock."""
    global STARTUP_ROLLING_WARMUP_MS, STARTUP_WARMUP_MS
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required for the Atlas Whisper service")
    print(f"[stt] loading Whisper {MODEL_NAME} for final CUDA", flush=True)
    model = whisper.load_model(MODEL_NAME, device="cuda", download_root=str(CACHE_DIR))
    print(f"[stt] loading Whisper {ROLLING_MODEL_NAME} for rolling CUDA", flush=True)
    rolling_model = whisper.load_model(
        ROLLING_MODEL_NAME,
        device="cuda",
        download_root=str(CACHE_DIR),
    )
    STARTUP_ROLLING_WARMUP_MS = _warm_model(rolling_model)
    print(
        "[stt] rolling CUDA warm-up completed in " f"{STARTUP_ROLLING_WARMUP_MS:.1f}ms",
        flush=True,
    )
    # Warm the authoritative model last so the first real final never inherits
    # lazy CUDA work from either model's startup path.
    STARTUP_WARMUP_MS = _warm_model(model)
    print(
        f"[stt] final CUDA warm-up completed in {STARTUP_WARMUP_MS:.1f}ms", flush=True
    )
    server = ThreadingHTTPServer((HOST, PORT), SpeechHandler)
    server.model = model  # type: ignore[attr-defined]
    server.rolling_model = rolling_model  # type: ignore[attr-defined]
    print(f"[stt] ready on {HOST}:{PORT}", flush=True)
    server.serve_forever(poll_interval=0.25)


if __name__ == "__main__":
    main()