atlas-iac/testing/tests/test_hermes_tts_language_routing.py

"""Language allow-list contracts for the private Hermes chat TTS voice policy."""

from __future__ import annotations

import importlib.util
import io
import json
import struct
import sys
import wave
from types import SimpleNamespace

import pytest

from testing.tests.test_hermes_chat_support import ROOT

AMY = "en_US-amy-medium"
IRINA = "ru_RU-irina-medium"
CLAUDE = "es_MX-claude-high"


def _load_tts_server(monkeypatch):
    server_path = ROOT / "dockerfiles" / "hermes-jetson-tts-server.py"
    monkeypatch.syspath_prepend(str(server_path.parent))
    spec = importlib.util.spec_from_file_location("hermes_jetson_tts_server", server_path)
    assert spec and spec.loader
    module = importlib.util.module_from_spec(spec)

    class _FakeSessionOptions:
        def __init__(self) -> None:
            self.intra_op_num_threads = None
            self.inter_op_num_threads = None

    fake_onnxruntime = SimpleNamespace(
        SessionOptions=_FakeSessionOptions,
        InferenceSession=lambda *a, **k: SimpleNamespace(),
    )
    fake_piper = SimpleNamespace(
        PiperConfig=SimpleNamespace(from_dict=lambda d: d),
        PiperVoice=lambda **kwargs: SimpleNamespace(**kwargs),
        SynthesisConfig=lambda **kwargs: SimpleNamespace(**kwargs),
    )
    monkeypatch.setitem(sys.modules, "onnxruntime", fake_onnxruntime)
    monkeypatch.setitem(sys.modules, "piper", fake_piper)
    spec.loader.exec_module(module)
    return module


@pytest.fixture
def tts(monkeypatch):
    return _load_tts_server(monkeypatch)


@pytest.mark.parametrize(
    "language,expected",
    [
        ("en", AMY),
        ("en-US", AMY),
        ("en_US", AMY),
        ("EN", AMY),
        ("En-Us", AMY),
        ("ru", IRINA),
        ("ru-RU", IRINA),
        ("ru_RU", IRINA),
        ("RU", IRINA),
        ("es", CLAUDE),
        ("es-MX", CLAUDE),
        ("es_MX", CLAUDE),
        ("es-ES", CLAUDE),
        ("es_ES", CLAUDE),
        ("ES", CLAUDE),
    ],
)
def test_allow_listed_languages_resolve_to_the_approved_voice(tts, language, expected):
    assert tts.resolve_voice_name(language) == expected


@pytest.mark.parametrize(
    "language",
    [
        None,
        "",
        "  ",
        "fr",
        "fr-FR",
        "de-DE",
        "xx",
        "en-GB",
        "es-AR",
        "english",
        123,
        1.5,
        True,
        [],
        {},
        {"lang": "ru"},
        "../../etc/passwd",
        "en_US-amy-medium/../../ru_RU-irina-medium",
        "\x00ru",
        "ru\x00",
    ],
)
def test_unknown_missing_or_malformed_language_falls_back_to_amy(tts, language):
    assert tts.resolve_voice_name(language) == AMY


def test_default_voice_name_matches_the_dockerfile_env_default(tts):
    assert tts.DEFAULT_VOICE_NAME == AMY


def test_deployment_does_not_override_the_image_default() -> None:
    """An old single-voice image and the new image can cross the Flux transition."""
    manifest = (ROOT / "services/hermes/voice-deployment.yaml").read_text()
    assert "HERMES_TTS_VOICE" not in manifest
    assert "piper-multilingual-en-ru-es" in manifest


def test_resolved_voice_is_always_one_of_the_three_baked_names(tts):
    assert frozenset({AMY, IRINA, CLAUDE}) == tts.BAKED_VOICE_NAMES
    fuzz_inputs = [
        "en", "ru", "es", "unknown", "", None, 42, "../../../etc/shadow",
        "en_US-amy-medium\x00; rm -rf /", "RU-ru", "Es-Es", "en-us-extra",
    ]
    for value in fuzz_inputs:
        assert tts.resolve_voice_name(value) in tts.BAKED_VOICE_NAMES


def test_client_voice_field_cannot_override_the_language_policy(tts):
    """The POST handler must select the voice from "language" only.

    A malicious or stale "voice" field in a hostile/legacy request must never
    change which baked model answers the request.
    """
    calls: list[str] = []

    class _RecordingVoice:
        def __init__(self, name: str) -> None:
            self.name = name

        def synthesize_wav(self, text, wav_file, syn_config) -> None:
            calls.append(self.name)
            wav_file.setnchannels(1)
            wav_file.setsampwidth(2)
            wav_file.setframerate(16_000)
            wav_file.writeframes(b"\x00\x00")

    class _RecordingHandler(tts.SpeechHandler):
        def __init__(self, payload):
            request = json.dumps(payload).encode("utf-8")
            self.path = "/v1/audio/speech"
            self.headers = {"Content-Length": str(len(request))}
            self.rfile = io.BytesIO(request)
            self.wfile = io.BytesIO()
            self.status = None
            self.response_headers = {}
            self.server = SimpleNamespace(
                voices={
                    AMY: _RecordingVoice(AMY),
                    IRINA: _RecordingVoice(IRINA),
                    CLAUDE: _RecordingVoice(CLAUDE),
                },
                default_voice_name=AMY,
            )

        def send_response(self, status, message=None):
            self.status = status

        def send_header(self, name, value):
            self.response_headers[name] = value

        def end_headers(self):
            return None

    # A payload that supplies an attacker/legacy "voice" value but no
    # language must resolve to the safe default, never the "voice" value.
    handler = _RecordingHandler({"input": "hi", "voice": IRINA})
    handler.do_POST()
    assert handler.status == 200
    assert handler.response_headers["X-TTS-Voice"] == AMY

    # A payload supplying both must still be governed by "language" alone.
    handler = _RecordingHandler({"input": "hi", "voice": CLAUDE, "language": "ru"})
    handler.do_POST()
    assert handler.status == 200
    assert handler.response_headers["X-TTS-Voice"] == IRINA

    assert calls == [AMY, IRINA]


def test_no_client_string_reaches_a_filesystem_path(tts):
    """resolve_voice_name must only ever return a fixed, baked literal.

    This is the property that keeps a client from ever causing the server to
    build a Path out of attacker-controlled text: the return value is always
    a member of the fixed allow-list, regardless of input shape.
    """
    hostile_inputs = [
        "../../../../etc/passwd",
        "/etc/passwd",
        "en_US-amy-medium/../../../etc/passwd",
        "ru_RU-irina-medium\x00.onnx",
        "es_MX-claude-high; cat /etc/shadow",
        "\n\ren",
        "en" + "/" * 200,
        " ",
    ]
    for value in hostile_inputs:
        result = tts.resolve_voice_name(value)
        assert result in tts.BAKED_VOICE_NAMES
        assert "/" not in result
        assert ".." not in result
        assert "\x00" not in result


def test_normalize_language_rejects_non_string_input(tts):
    assert tts.normalize_language(None) is None
    assert tts.normalize_language(123) is None
    assert tts.normalize_language([]) is None
    assert tts.normalize_language("") is None
    assert tts.normalize_language("   ") is None
    assert tts.normalize_language("En_US") == "en-us"


def test_default_voice_name_is_one_of_the_baked_voices(tts):
    assert tts.DEFAULT_VOICE_NAME in tts.BAKED_VOICE_NAMES


@pytest.mark.parametrize(
    "value,expected",
    [
        ("turn-42", "turn-42"),
        (" session_1:chunk.3 ", "session_1:chunk.3"),
        (None, None),
        (42, None),
        ("", None),
        ("turn\r\nX-Injected: true", None),
        ("/not/header-safe", None),
        ("x" * 129, None),
    ],
)
def test_turn_ids_are_safe_to_echo_in_response_headers(tts, value, expected):
    assert tts.normalize_turn_id(value) == expected


@pytest.mark.parametrize(
    "text,expected",
    [
        ("A short clause,", 90),
        ("A continued thought:", 110),
        ("Is this natural?", 170),
        ("This is natural.", 190),
        ("Thinking...", 220),
        ("First paragraph.\n\n", 280),
        ("A provisional phrase", 70),
        ('She said, "done."', 190),
    ],
)
def test_trailing_pause_budget_tracks_punctuation(tts, text, expected):
    assert tts.trailing_pause_ms(text) == expected


def test_sentence_pause_targets_keep_paragraph_and_final_clause_cadence(tts):
    assert tts.sentence_pause_targets("First.\n\nSecond? A tail,") == [280, 170, 90]


def _pcm(*runs: tuple[int, int]) -> bytes:
    return b"".join(struct.pack(f"<{count}h", *([value] * count)) for value, count in runs)


def test_pcm_tail_trim_preserves_voice_and_keeps_period_pause(tts):
    # A deliberately low-amplitude final phoneme remains above the very
    # conservative silence threshold and must remain intact.
    pcm = _pcm((1200, 50), (40, 50), (0, 600))
    normalized = tts.trim_trailing_pcm_silence(pcm, sample_rate=1000, pause_ms=190)

    assert bytes(normalized[:200]) == pcm[:200]
    assert len(normalized) == (100 + 190) * 2
    assert bytes(normalized[-190 * 2 :]) == b"\x00" * (190 * 2)


def test_pcm_tail_trim_leaves_short_and_all_silent_audio_unchanged(tts):
    near_target = _pcm((900, 100), (0, 205))
    all_silent = _pcm((0, 800))

    assert bytes(tts.trim_trailing_pcm_silence(near_target, 1000, 190)) == near_target
    assert bytes(tts.trim_trailing_pcm_silence(all_silent, 1000, 190)) == all_silent


def test_pcm_tail_trim_rejects_malformed_or_unknown_format_without_mutation(tts):
    malformed = b"\x01\x02\x03"
    assert bytes(tts.trim_trailing_pcm_silence(malformed, 22_050, 190)) == malformed
    assert bytes(tts.trim_trailing_pcm_silence(b"\x01\x00", 0, 190)) == b"\x01\x00"


def test_compatibility_wav_uses_the_same_period_pause_budget(tts):
    source = io.BytesIO()
    with wave.open(source, "wb") as wav_file:
        wav_file.setnchannels(1)
        wav_file.setsampwidth(2)
        wav_file.setframerate(1000)
        wav_file.writeframes(_pcm((1000, 100), (0, 600)))

    normalized = tts.normalize_wav_trailing_pause(source.getvalue(), "Finished.")
    with wave.open(io.BytesIO(normalized), "rb") as wav_file:
        assert wav_file.getparams().nchannels == 1
        assert wav_file.getparams().sampwidth == 2
        assert wav_file.getparams().framerate == 1000
        assert wav_file.getnframes() == 100 + 190


def _decode_chunked_response(data: bytes) -> bytes:
    decoded = bytearray()
    cursor = 0
    while True:
        line_end = data.index(b"\r\n", cursor)
        size = int(data[cursor:line_end], 16)
        cursor = line_end + 2
        if size == 0:
            assert data[cursor:] == b"\r\n"
            break
        decoded.extend(data[cursor : cursor + size])
        cursor += size
        assert data[cursor : cursor + 2] == b"\r\n"
        cursor += 2
    return bytes(decoded)


def test_stream_endpoint_returns_progressive_pcm_and_preserves_voice_policy(tts):
    events: list[str] = []
    first_pcm = b"\x01\x00" * (tts.STREAM_WRITE_BYTES // 2 + 3)
    second_pcm = b"\x02\x00\x03\x00"

    class _StreamingVoice:
        config = SimpleNamespace(sample_rate=22_050)

        def synthesize(self, text, syn_config):
            assert text == "Hola mundo. Otra frase."
            assert syn_config.length_scale == pytest.approx(0.8)
            assert "headers" in events
            events.append("synthesis")
            yield SimpleNamespace(
                sample_rate=22_050,
                sample_width=2,
                sample_channels=1,
                audio_int16_bytes=first_pcm,
            )
            yield SimpleNamespace(
                sample_rate=22_050,
                sample_width=2,
                sample_channels=1,
                audio_int16_bytes=second_pcm,
            )

    class _StreamingHandler(tts.SpeechHandler):
        def __init__(self):
            payload = json.dumps(
                {
                    "text": "Hola mundo. Otra frase.",
                    "language": "es",
                    "speed": 1.25,
                    "turn_id": "turn-9:chunk.1",
                }
            ).encode("utf-8")
            self.path = "/v1/audio/speech/stream"
            self.headers = {"Content-Length": str(len(payload))}
            self.rfile = io.BytesIO(payload)
            self.wfile = io.BytesIO()
            self.status = None
            self.response_headers = {}
            self.close_connection = False
            self.server = SimpleNamespace(
                voices={CLAUDE: _StreamingVoice()},
                default_voice_name=AMY,
            )

        def send_response(self, status, message=None):
            self.status = status

        def send_header(self, name, value):
            self.response_headers[name] = value

        def end_headers(self):
            events.append("headers")

    handler = _StreamingHandler()
    handler.do_POST()

    assert handler.status == 200
    assert handler.response_headers["Transfer-Encoding"] == "chunked"
    assert handler.response_headers["X-Audio-Format"] == "pcm_s16le"
    assert handler.response_headers["X-Audio-Sample-Rate"] == "22050"
    assert handler.response_headers["X-Audio-Channels"] == "1"
    assert handler.response_headers["X-Audio-Sample-Width"] == "2"
    assert handler.response_headers["X-TTS-Voice"] == CLAUDE
    assert handler.response_headers["X-TTS-Turn-ID"] == "turn-9:chunk.1"
    assert events == ["headers", "synthesis"]
    assert _decode_chunked_response(handler.wfile.getvalue()) == first_pcm + second_pcm


def test_cached_thinking_cue_is_lock_free_and_ignores_client_text(tts):
    cached = SimpleNamespace(
        cue_id="thinking",
        language="ru",
        voice_name=IRINA,
        sample_rate=22_050,
        pcm=b"\x01\x00\x02\x00",
    )

    class _CachedHandler(tts.SpeechHandler):
        def __init__(self):
            payload = json.dumps(
                {
                    "text": "attacker-controlled text is ignored",
                    "language": "ru",
                    "cue_id": "thinking",
                    "turn_id": "turn-9:thinking-cue:1",
                }
            ).encode()
            self.path = "/v1/audio/speech/stream"
            self.headers = {"Content-Length": str(len(payload))}
            self.rfile = io.BytesIO(payload)
            self.wfile = io.BytesIO()
            self.status = None
            self.response_headers = {}
            self.close_connection = False
            self.server = SimpleNamespace(
                voices={}, default_voice_name=AMY, cue_cache={("ru", "thinking"): cached}
            )

        def send_response(self, status, message=None):
            self.status = status

        def send_header(self, name, value):
            self.response_headers[name] = value

        def end_headers(self):
            return None

    handler = _CachedHandler()
    handler.do_POST()

    assert handler.status == 200
    assert handler.wfile.getvalue() == cached.pcm
    assert handler.response_headers["X-TTS-Voice"] == IRINA
    assert handler.response_headers["X-TTS-Cue-ID"] == "thinking"
    assert handler.response_headers["X-TTS-Cache"] == "HIT"
    assert handler.response_headers["X-TTS-Turn-ID"] == "turn-9:thinking-cue:1"


def test_unknown_cached_thinking_cue_fails_closed(tts):
    with pytest.raises(ValueError, match="unknown thinking cue"):
        tts.resolve_cached_cue(
            {"language": "en", "cue_id": "invented"},
            {},
        )


def test_thinking_cue_cache_precomputes_all_localized_voice_pairs(tts):
    class _CueVoice:
        config = SimpleNamespace(sample_rate=22_050)

        def synthesize(self, text, syn_config):
            assert text
            assert syn_config.length_scale == 1.0
            yield SimpleNamespace(
                sample_rate=22_050,
                sample_width=2,
                sample_channels=1,
                audio_int16_bytes=b"\x01\x00\x02\x00",
            )

    voices = {name: _CueVoice() for name in (AMY, IRINA, CLAUDE)}
    cache = tts.build_cue_cache(
        voices,
        tts.SynthesisConfig,
        lambda pcm, _rate, _pause: memoryview(pcm),
        lambda _text: 190,
    )

    assert len(cache) == 12
    assert {key[0] for key in cache} == {"en", "ru", "es"}
    assert {cue.voice_name for (language, _), cue in cache.items() if language == "en"} == {AMY}
    assert {cue.voice_name for (language, _), cue in cache.items() if language == "ru"} == {IRINA}
    assert {cue.voice_name for (language, _), cue in cache.items() if language == "es"} == {CLAUDE}
    assert all(cue.pcm and len(cue.pcm) % 2 == 0 for cue in cache.values())


def test_thinking_cue_cache_fails_startup_on_missing_or_malformed_audio(tts):
    class _EmptyVoice:
        config = SimpleNamespace(sample_rate=22_050)

        def synthesize(self, _text, _syn_config):
            return iter(())

    voices = {name: _EmptyVoice() for name in (AMY, IRINA, CLAUDE)}
    with pytest.raises(RuntimeError, match="returned no audio"):
        tts.build_cue_cache(
            voices,
            tts.SynthesisConfig,
            lambda pcm, _rate, _pause: memoryview(pcm),
            lambda _text: 190,
        )


def test_stream_endpoint_trims_each_sentence_with_its_own_pause_budget(tts):
    sentence_pcm = _pcm((1000, 100), (0, 600))

    class _StreamingVoice:
        config = SimpleNamespace(sample_rate=1000)

        def synthesize(self, text, syn_config):
            assert text == "First. Final clause,"
            for _ in range(2):
                yield SimpleNamespace(
                    sample_rate=1000,
                    sample_width=2,
                    sample_channels=1,
                    audio_int16_bytes=sentence_pcm,
                )

    class _StreamingHandler(tts.SpeechHandler):
        def __init__(self):
            self.wfile = io.BytesIO()
            self.response_headers = {}
            self.close_connection = False

        def send_response(self, status, message=None):
            return None

        def send_header(self, name, value):
            self.response_headers[name] = value

        def end_headers(self):
            return None

    handler = _StreamingHandler()
    handler._stream_pcm("First. Final clause,", 1.0, AMY, _StreamingVoice(), "turn-2")
    decoded = _decode_chunked_response(handler.wfile.getvalue())

    # 100 ms of voiced audio plus 190 ms after the period and 90 ms after the
    # final clause. The audio itself remains byte-for-byte identical.
    assert len(decoded) == ((100 + 190) + (100 + 90)) * 2
    assert decoded[: 100 * 2] == sentence_pcm[: 100 * 2]
    second_start = (100 + 190) * 2
    assert decoded[second_start : second_start + 100 * 2] == sentence_pcm[: 100 * 2]


def test_stream_disconnect_stops_before_synthesizing_another_sentence(tts):
    generated: list[int] = []

    class _StreamingVoice:
        config = SimpleNamespace(sample_rate=16_000)

        def synthesize(self, text, syn_config):
            for sentence in (1, 2):
                generated.append(sentence)
                yield SimpleNamespace(
                    sample_rate=16_000,
                    sample_width=2,
                    sample_channels=1,
                    audio_int16_bytes=b"\x00\x00",
                )

    class _DisconnectedHandler(tts.SpeechHandler):
        def __init__(self):
            self.wfile = io.BytesIO()
            self.response_headers = {}
            self.close_connection = False

        def send_response(self, status, message=None):
            return None

        def send_header(self, name, value):
            self.response_headers[name] = value

        def end_headers(self):
            return None

        def _write_stream_chunk(self, audio):
            self.close_connection = True
            return False

    handler = _DisconnectedHandler()
    handler._stream_pcm("hello", 1.0, AMY, _StreamingVoice(), "turn-1")

    assert generated == [1]
    assert handler.close_connection is True


def test_load_voices_fails_closed_when_a_baked_model_is_missing(tts, tmp_path):
    with pytest.raises(RuntimeError, match="baked Piper voice is missing"):
        tts.load_voices(tmp_path, threads=1)