[Bugfix] Invalidate stale ref_audio cache on local file modification (#4873) (#5670)

Signed-off-by: pranavthakur0-0 <pranavthakurdew@gmail.com>
Co-authored-by: Zhou Taichang <tzhouam@connect.ust.hk>

tests/entrypoints/openai_api/test_serving_speech.py

@@ -1237,10 +1237,364 @@ class TestTTSMethods:
         assert second[1] == 24000
         assert first[0] is second[0]
         assert first[0][0] == pytest.approx(float(wav[0]), abs=1e-4)
+        cache_key = first[2]
         assert speech_server._get_resolved_ref_audio_artifact_key(
-            ref_audio
+            cache_key
         ) == speech_server._make_ref_audio_artifact_cache_key(np.asarray(first[0], dtype=np.float32), 24000)
 
+    # ── ref-audio cache key tests (static helper) ──
+
+    def test_local_file_cache_key_invalidation(self):
+        """Modified file must produce a different cache key (mtime change)."""
+        import pathlib
+        import tempfile
+        import time as _time
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            test_file = pathlib.Path(tmp_dir) / "test_audio.wav"
+            test_file.write_text("dummy audio data 1")
+            file_uri = f"file://{test_file}"
+            key1 = OmniOpenAIServingSpeech._get_ref_audio_cache_key(
+                file_uri,
+                allowed_local_media_path=tmp_dir,
+            )
+            _time.sleep(0.01)
+            test_file.write_text("dummy audio data 2")
+            key2 = OmniOpenAIServingSpeech._get_ref_audio_cache_key(
+                file_uri,
+                allowed_local_media_path=tmp_dir,
+            )
+            assert key1 != key2, "Cache key should change when file is modified!"
+
+    def test_remote_url_cache_key(self):
+        """Remote URLs must produce a stable, repeatable cache key."""
+        url = "https://example.com/audio.wav"
+        key1 = OmniOpenAIServingSpeech._get_ref_audio_cache_key(url)
+        key2 = OmniOpenAIServingSpeech._get_ref_audio_cache_key(url)
+        assert key1 == key2, "Cache key for URLs should remain constant!"
+
+    def test_missing_file_cache_key(self, caplog):
+        """A missing file falls back to the URI-string hash and emits a log
+        so the stale-cache risk is visible."""
+        file_uri = "file:///path/to/nonexistent/file.wav"
+        target_logger = logging.getLogger("vllm_omni.entrypoints.openai.serving_speech")
+        target_logger.addHandler(caplog.handler)
+        prev_level = target_logger.level
+        target_logger.setLevel(logging.DEBUG)
+        try:
+            key = OmniOpenAIServingSpeech._get_ref_audio_cache_key(
+                file_uri,
+                allowed_local_media_path="/path/to/nonexistent",
+            )
+        finally:
+            target_logger.removeHandler(caplog.handler)
+            target_logger.setLevel(prev_level)
+        expected_key = hashlib.sha1(file_uri.encode("utf-8")).hexdigest()
+        assert key == expected_key, "Missing files should fallback to the string hash."
+        assert any("stale cache" in r.getMessage() for r in caplog.records), (
+            "A log message about stale cache must be emitted when os.stat fails"
+        )
+
+    def test_percent_encoded_file_uri(self):
+        """Percent-encoded file:// URIs must be decoded so os.stat hits the real file."""
+        import pathlib
+        import tempfile
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            test_file = pathlib.Path(tmp_dir) / "my audio.wav"
+            test_file.write_text("dummy audio data")
+            encoded_uri = f"file://{str(test_file).replace(' ', '%20')}"
+            key = OmniOpenAIServingSpeech._get_ref_audio_cache_key(
+                encoded_uri,
+                allowed_local_media_path=tmp_dir,
+            )
+            string_only_key = hashlib.sha1(encoded_uri.encode("utf-8")).hexdigest()
+            assert key != string_only_key, (
+                "Percent-encoded URI should be decoded and stat'd, not fall back to string-only cache key!"
+            )
+
+    def test_no_allowlist_skips_stat_for_local_paths(self):
+        """When allowed_local_media_path is None, local paths must NOT be stat'd
+        — the key falls back to the string-only hash."""
+        import pathlib
+        import tempfile
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            test_file = pathlib.Path(tmp_dir) / "test.wav"
+            test_file.write_text("data")
+            file_uri = f"file://{test_file}"
+            key = OmniOpenAIServingSpeech._get_ref_audio_cache_key(
+                file_uri,
+                allowed_local_media_path=None,
+            )
+            # With no allowlist the stat is skipped, so the key is just sha1(uri)
+            expected = hashlib.sha1(file_uri.encode("utf-8")).hexdigest()
+            assert key == expected, "No allowlist should skip stat and use string-only key"
+
+    def test_empty_string_allowlist_skips_stat_for_local_paths(self):
+        """vLLM's default allowed_local_media_path is '' not None; that must
+        also skip stat rather than treating cwd as the allowlist."""
+        import pathlib
+        import tempfile
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            test_file = pathlib.Path(tmp_dir) / "test.wav"
+            test_file.write_text("data")
+            file_uri = f"file://{test_file}"
+            key = OmniOpenAIServingSpeech._get_ref_audio_cache_key(
+                file_uri,
+                allowed_local_media_path="",
+            )
+            expected = hashlib.sha1(file_uri.encode("utf-8")).hexdigest()
+            assert key == expected, "Empty-string allowlist should skip stat and use string-only key"
+
+    def test_uppercase_file_uri_uses_local_metadata(self):
+        """FILE:// must be treated as a local file, not a string-only locator."""
+        import pathlib
+        import tempfile
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            test_file = pathlib.Path(tmp_dir) / "test.wav"
+            test_file.write_text("data")
+            upper_uri = f"FILE://{test_file}"
+            key = OmniOpenAIServingSpeech._get_ref_audio_cache_key(
+                upper_uri,
+                allowed_local_media_path=tmp_dir,
+            )
+            string_only = hashlib.sha1(upper_uri.encode("utf-8")).hexdigest()
+            assert key != string_only
+
+    def test_uppercase_http_uri_is_remote(self, speech_server):
+        url = "HTTP://example.com/audio.wav"
+        key = OmniOpenAIServingSpeech._get_ref_audio_cache_key(url)
+        assert key == hashlib.sha1(url.encode("utf-8")).hexdigest()
+        assert speech_server._validate_ref_audio_format(url) is None
+        assert speech_server._validate_ref_audio_format("FILE:///tmp/a.wav") is None
+        assert speech_server._validate_ref_audio_format("/tmp/a.wav") is not None
+
+    @pytest.mark.asyncio
+    async def test_resolve_ref_audio_retries_when_file_changes_during_fetch(self, speech_server, mocker):
+        import pathlib
+        import tempfile
+
+        wav = np.linspace(-0.5, 0.5, 48000, dtype=np.float32)
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            path = pathlib.Path(tmp_dir) / "spk.wav"
+            path.write_bytes(b"v1")
+            uri = f"file://{path}"
+            speech_server.model_config.allowed_local_media_path = tmp_dir
+            speech_server.model_config.allowed_media_domains = None
+            speech_server._diffusion_mode = False
+
+            fetches = {"n": 0}
+
+            async def fake_fetch(_ref):
+                fetches["n"] += 1
+                if fetches["n"] == 1:
+                    path.write_bytes(b"v2-longer-content")
+                return wav, 24000
+
+            mock_connector = mocker.MagicMock()
+            mock_connector.fetch_audio_async = fake_fetch
+            mocker.patch(
+                "vllm_omni.entrypoints.openai.serving_speech.MediaConnector",
+                return_value=mock_connector,
+            )
+
+            wav_list, sr, cache_key = await speech_server._resolve_ref_audio(uri)
+            expected_key = OmniOpenAIServingSpeech._get_ref_audio_cache_key(uri, allowed_local_media_path=tmp_dir)
+            assert sr == 24000
+            assert len(wav_list) == 48000
+            assert fetches["n"] == 2
+            assert cache_key == expected_key
+            assert cache_key in speech_server._ref_audio_resolve_cache
+
+    def test_conditioning_cache_salt_changes_with_ref_audio_cache_key(self):
+        """Prefix-cache salt must fold the content-aware resolve key so a
+        same-path file rewrite cannot reuse KV from the previous reference."""
+        from vllm_omni.entrypoints.openai.serving_speech import _conditioning_cache_salt
+
+        req = OpenAICreateSpeechRequest(input="hello", ref_audio="file:///data/spk.wav")
+        salt_a = _conditioning_cache_salt(req, {"ref_audio_cache_key": "key_aaa"})
+        salt_b = _conditioning_cache_salt(req, {"ref_audio_cache_key": "key_bbb"})
+        salt_locator_only = _conditioning_cache_salt(req, {})
+        assert salt_a != salt_b
+        assert salt_a != salt_locator_only
+
+    @pytest.mark.asyncio
+    async def test_higgs_v3_cache_salt_changes_with_ref_audio_cache_key(self, speech_server, mocker):
+        """Higgs v3 voice clone uses placeholder prompt ids + prefix caching.
+        The salt must move when the resolve key moves, or a same-path rewrite
+        reuses KV from the previous reference."""
+        from unittest.mock import AsyncMock, MagicMock
+
+        adapter = MagicMock()
+        adapter.build_prompt.return_value = [1, 1, 1, 1]
+        speech_server._resolve_higgs_audio_v3_adapter = AsyncMock(return_value=adapter)
+        codes = torch.arange(8, dtype=torch.long).view(8, 1)
+        speech_server._resolve_higgs_audio_v3_ref_codes = AsyncMock(return_value=(codes, False, False))
+        speech_server._get_resolved_ref_audio_artifact_key = MagicMock(return_value="art")
+
+        req = OpenAICreateSpeechRequest(
+            input="hello",
+            ref_audio="file:///data/spk.wav",
+            ref_text="transcript",
+        )
+        speech_server._resolve_ref_audio = AsyncMock(return_value=([0.1] * 48000, 24000, "key_aaa"))
+        prompt_a = await speech_server._build_higgs_audio_v3_params(req)
+        speech_server._resolve_ref_audio = AsyncMock(return_value=([0.9] * 48000, 24000, "key_bbb"))
+        prompt_b = await speech_server._build_higgs_audio_v3_params(req)
+
+        assert prompt_a["prompt_token_ids"] == prompt_b["prompt_token_ids"]
+        assert prompt_a["cache_salt"] != prompt_b["cache_salt"]
+        assert prompt_a["additional_information"]["ref_audio_cache_key"] == "key_aaa"
+        assert prompt_b["additional_information"]["ref_audio_cache_key"] == "key_bbb"
+
+    # ── encode_reference_codes: speaker cache invalidation ──
+
+    @pytest.mark.asyncio
+    async def test_encode_reference_codes_reencodes_on_key_change(self):
+        """Changing the resolve_cache_key must trigger a fresh encode, not serve
+        the stale cached tensor.  This is the user-visible fix for delay-family
+        voice clones that were reproducing the old speaker."""
+        from unittest.mock import MagicMock
+
+        from vllm_omni.model_executor.models.moss_tts.reference_encoder import encode_reference_codes
+        from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache
+
+        speaker_cache = SpeakerEmbeddingCache(max_bytes=64 * 1024 * 1024)
+
+        call_count = 0
+        codes_v1 = torch.arange(10, dtype=torch.int64)
+        codes_v2 = torch.arange(10, 20, dtype=torch.int64)
+
+        # Stub processor that tracks call count and returns different codes
+        processor = MagicMock()
+
+        def fake_encode(wavs, sampling_rate, n_vq):
+            nonlocal call_count
+            call_count += 1
+            return [codes_v1 if call_count == 1 else codes_v2]
+
+        processor.encode_audios_from_wav = fake_encode
+
+        # resolve_ref_audio returns different cache keys on successive calls
+        resolve_keys = iter(["key_aaa", "key_bbb"])
+
+        async def mock_resolve(ref_str):
+            return ([0.1, 0.2, 0.3], 24000, next(resolve_keys))
+
+        # First call: cold miss, encodes and stores
+        result1 = await encode_reference_codes(
+            "data:audio/wav;base64,fake",
+            processor=processor,
+            resolve_ref_audio=mock_resolve,
+            speaker_cache=speaker_cache,
+            variant="tts",
+            n_vq=32,
+            sr_target=24000,
+        )
+        assert call_count == 1
+        assert torch.equal(result1, codes_v1)
+
+        # Second call: different resolve key → must re-encode, not serve cached
+        result2 = await encode_reference_codes(
+            "data:audio/wav;base64,fake",
+            processor=processor,
+            resolve_ref_audio=mock_resolve,
+            speaker_cache=speaker_cache,
+            variant="tts",
+            n_vq=32,
+            sr_target=24000,
+        )
+        assert call_count == 2, "Different resolve key should trigger re-encode"
+        assert torch.equal(result2, codes_v2)
+
+    @pytest.mark.asyncio
+    async def test_encode_reference_codes_named_voice_skips_resolve_on_hit(self):
+        """Named-voice cache hit must NOT call resolve_ref_audio."""
+        from unittest.mock import MagicMock
+
+        from vllm_omni.model_executor.models.moss_tts.reference_encoder import encode_reference_codes
+        from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache
+
+        speaker_cache = SpeakerEmbeddingCache(max_bytes=64 * 1024 * 1024)
+
+        processor = MagicMock()
+        processor.encode_audios_from_wav = lambda wavs, sampling_rate, n_vq: [torch.zeros(5)]
+
+        resolve_called = 0
+
+        async def mock_resolve(ref_str):
+            nonlocal resolve_called
+            resolve_called += 1
+            return ([0.1], 24000, "some_key")
+
+        # First call: cold miss, resolves + encodes
+        await encode_reference_codes(
+            "data:audio/wav;base64,fake",
+            processor=processor,
+            resolve_ref_audio=mock_resolve,
+            speaker_cache=speaker_cache,
+            variant="tts",
+            n_vq=32,
+            sr_target=24000,
+            voice_name="alice",
+            voice_created_at=100,
+        )
+        assert resolve_called == 1
+
+        # Second call: warm hit, must NOT resolve
+        await encode_reference_codes(
+            "data:audio/wav;base64,fake",
+            processor=processor,
+            resolve_ref_audio=mock_resolve,
+            speaker_cache=speaker_cache,
+            variant="tts",
+            n_vq=32,
+            sr_target=24000,
+            voice_name="alice",
+            voice_created_at=100,
+        )
+        assert resolve_called == 1, "Named-voice cache hit should skip resolve_ref_audio"
+
+    @pytest.mark.asyncio
+    async def test_ttsd_second_reference_does_not_use_named_voice(self, speech_server, mocker):
+        """Uploaded voice names speaker 1 only; speaker 2 must stay anonymous."""
+        speech_server._moss_variant = "ttsd"
+        speech_server.uploaded_speakers = {"alice": {}}
+        mocker.patch.object(speech_server, "_voice_created_at", return_value=42)
+
+        proc = mocker.MagicMock()
+        proc.model_config.n_vq = 8
+        proc.model_config.sampling_rate = 24000
+        proc.build_user_message.return_value = "msg"
+        proc.return_value = {"input_ids": [torch.zeros((4, 9), dtype=torch.int64)]}
+        mocker.patch.object(speech_server, "_get_moss_processor", return_value=proc)
+
+        seen: list[tuple[str, str | None]] = []
+
+        async def fake_encode(ref_str, **kwargs):
+            seen.append((ref_str, kwargs.get("voice_name")))
+            return torch.zeros(3, dtype=torch.int64)
+
+        mocker.patch(
+            "vllm_omni.model_executor.models.moss_tts.reference_encoder.encode_reference_codes",
+            fake_encode,
+        )
+
+        req = OpenAICreateSpeechRequest(
+            input="hello",
+            voice="alice",
+            ref_audio="data:audio/wav;base64,aaa",
+            ref_audio_2="data:audio/wav;base64,bbb",
+        )
+        await speech_server._build_moss_tts_params(req, has_inline_ref_audio=False)
+        assert seen == [
+            ("data:audio/wav;base64,aaa", "alice"),
+            ("data:audio/wav;base64,bbb", None),
+        ]
+
     def test_precomputed_qwen3_voice_infers_base_without_ref_audio(self, speech_server):
         """Precomputed Qwen3 voices are reusable by name without per-request ref_audio."""
         speech_server._tts_model_type = "qwen3_tts"
@@ -1896,7 +2250,7 @@ class TestTTSMethods:
         mocker.patch.object(
             speech_server,
             "_resolve_ref_audio",
-            new=mocker.AsyncMock(return_value=ref_audio_data),
+            new=mocker.AsyncMock(return_value=(*ref_audio_data, "fake_cache_key")),
         )
         mock_prompt = mocker.patch.object(
             speech_server,
@@ -2013,7 +2367,7 @@ class TestTTSMethods:
         mocker.patch.object(
             speech_server,
             "_resolve_ref_audio",
-            new=mocker.AsyncMock(return_value=ref_audio_data),
+            new=mocker.AsyncMock(return_value=(*ref_audio_data, "fake_cache_key")),
         )
         mock_prompt = mocker.patch.object(
             speech_server,
@@ -4382,7 +4736,7 @@ class TestTTSAsyncOffloading:
     ):
         """Base explicit true should reach the model prompt additional_information."""
         qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None)
-        qwen3_tts_server._resolve_ref_audio = mocker.AsyncMock(return_value=([0.0] * 48000, 24000))
+        qwen3_tts_server._resolve_ref_audio = mocker.AsyncMock(return_value=([0.0] * 48000, 24000, "fake_cache_key"))
         qwen3_tts_server._get_resolved_ref_audio_artifact_key = mocker.MagicMock(return_value=None)
         qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512)