fix(stt): join a decoder migration that started during the probe

src/kiro_crew/stt/decoder.py

@@ -560,6 +560,19 @@ def _remove_retired() -> None:
             logger.debug("Could not remove retired decoder %s", name, exc_info=True)
 
 
+async def _join_migration(task: asyncio.Task[Path | None]) -> Path | None:
+    """Wait for the migration *task* another caller started; ``None`` when it
+    runs on a different event loop, which this caller cannot await."""
+    try:
+        same_loop = task.get_loop() is asyncio.get_running_loop()
+    except RuntimeError:
+        same_loop = False
+    if not same_loop:
+        return None
+    # Shielded: one caller giving up must not cancel the shared transfer.
+    return await asyncio.shield(task)
+
+
 async def migrate_retired_decoder() -> Path | None:
     """Fetch the current pin for an install whose store has only a retired one.
 
@@ -569,19 +582,17 @@ async def migrate_retired_decoder() -> Path | None:
     by :meth:`DecoderStore.ensure` once the current pin is installed.
     """
     global _migration_task
-    task = _migration_task
-    if task is not None:
-        try:
-            same_loop = task.get_loop() is asyncio.get_running_loop()
-        except RuntimeError:
-            same_loop = False
-        if not same_loop:
-            return None
-        # Shielded: one caller giving up must not cancel the shared transfer.
-        return await asyncio.shield(task)
+    if _migration_task is not None:
+        return await _join_migration(_migration_task)
     if bundle_carries_decoder() or artifact_for() is None:
         return None
-    if not await asyncio.to_thread(has_retired_decoder):
+    has_retired = await asyncio.to_thread(has_retired_decoder)
+    # Another caller can start the migration while this one awaits the probe, and
+    # its transfer removes the retired file, so a probe that ran after it finds
+    # none: join that transfer rather than report the decoder missing.
+    if _migration_task is not None:
+        return await _join_migration(_migration_task)
+    if not has_retired:
         return None
     # Started at most once per process: on a host that cannot reach the release,
     # every later transcode would otherwise wait out another failed fetch. The

test/test_stt_decoder.py

@@ -19,6 +19,7 @@
 import hashlib
 import os
 import stat
+import threading
 from pathlib import Path
 
 import pytest
@@ -605,6 +606,54 @@ async def test_concurrent_callers_share_one_transfer(self, monkeypatch, tmp_path
         assert first == second == tmp_path / artifact.filename
         assert len(calls) == 1
 
+    @pytest.mark.asyncio
+    async def test_a_caller_whose_probe_sees_the_migrated_store_joins_the_transfer(
+        self, monkeypatch, tmp_path
+    ):
+        """The second caller's retired-file probe runs only after the first caller's
+        migration installed the pin and removed the retired file, so it finds none.
+        A migration is under way all the same, and that caller still gets its path."""
+        artifact = _artifact()
+        _pin(monkeypatch, tmp_path, artifact)
+        _allow_download(monkeypatch)
+        (tmp_path / decoder.RETIRED_FILENAMES[1]).write_bytes(b"old decoder")
+        calls: list[dict] = []
+        _serve(monkeypatch, _PAYLOAD, calls)
+
+        migrated = threading.Event()
+        probes = 0
+        probes_lock = threading.Lock()
+        real_probe = decoder.has_retired_decoder
+
+        def _probe() -> bool:
+            nonlocal probes
+            with probes_lock:
+                probes += 1
+                later = probes > 1
+            if later:
+                assert migrated.wait(timeout=10), "the first migration never finished"
+            return real_probe()
+
+        shared_store = decoder.store()
+        real_ensure = shared_store.ensure
+
+        async def _ensure() -> Path | None:
+            try:
+                return await real_ensure()
+            finally:
+                migrated.set()
+
+        monkeypatch.setattr(decoder, "has_retired_decoder", _probe)
+        monkeypatch.setattr(shared_store, "ensure", _ensure)
+
+        first, second = await asyncio.wait_for(
+            asyncio.gather(decoder.migrate_retired_decoder(), decoder.migrate_retired_decoder()),
+            timeout=30,
+        )
+
+        assert first == second == tmp_path / artifact.filename
+        assert len(calls) == 1
+
     @pytest.mark.asyncio
     async def test_no_retired_file_means_no_network(self, monkeypatch, tmp_path):
         _pin(monkeypatch, tmp_path, _artifact())