fix(registry): lock ModeRegistry lazy-load resolution against concurrent first access

instructor/v2/core/registry.py

@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 
+import threading
 from dataclasses import dataclass
 from typing import Callable
 
@@ -80,6 +81,12 @@ def __init__(self) -> None:
         """Initialize empty registry."""
         self._handlers: dict[tuple[Provider, Mode], ModeHandlers] = {}
         self._lazy_loaders: dict[tuple[Provider, Mode], Callable[[], ModeHandlers]] = {}
+        # Guards the lazy-load resolution (check -> pop -> import -> set) in
+        # get_handlers(). Without this, concurrent first-callers for the same
+        # mode_key race: one pops the loader, the others see neither dict
+        # populated yet and raise KeyError. Held for the whole resolution
+        # (not per-key) since lazy-loading only ever runs once per key.
+        self._lazy_load_lock = threading.Lock()
 
     def register(
         self,
@@ -181,12 +188,22 @@ def get_handlers(self, provider: Provider, mode: Mode) -> ModeHandlers:
         if mode_key in self._handlers:
             return self._handlers[mode_key]
 
-        # Try lazy loading
-        if mode_key in self._lazy_loaders:
-            loader = self._lazy_loaders.pop(mode_key)
-            handlers = loader()
-            self._handlers[mode_key] = handlers
-            return handlers
+        # Try lazy loading. Locked because the pop -> import -> set sequence
+        # below is not atomic: without the lock, a thread that loses the race
+        # to pop self._lazy_loaders[mode_key] would find it already gone and
+        # self._handlers[mode_key] not yet set, and raise KeyError even
+        # though the mode genuinely is registered (just still resolving).
+        with self._lazy_load_lock:
+            # Re-check: another thread may have finished loading this key
+            # while we were waiting for the lock.
+            if mode_key in self._handlers:
+                return self._handlers[mode_key]
+
+            if mode_key in self._lazy_loaders:
+                loader = self._lazy_loaders.pop(mode_key)
+                handlers = loader()
+                self._handlers[mode_key] = handlers
+                return handlers
 
         raise KeyError(
             f"Mode {mode_key} is not registered. "

tests/v2/test_registry.py

@@ -116,3 +116,69 @@ def test_registry_invalid_handler_type(provider: Provider, mode: Mode):
     """Test error for invalid handler type."""
     with pytest.raises(ValueError, match="Invalid handler_type"):
         mode_registry.get_handler(provider, mode, "invalid_type")
+
+
+def test_get_handlers_concurrent_first_access_does_not_race():
+    """Regression test for #2422.
+
+    Concurrent first callers for the same lazily-registered (provider, mode)
+    key must all get the same handlers back, never a KeyError, even when the
+    loader is slow (simulating a module import in flight).
+
+    The loader sleeps briefly so every other thread has a chance to reach
+    get_handlers() while the first thread's load is still in progress, this
+    is the exact window in which the unlocked version raced: losing threads
+    would see the lazy loader already popped and self._handlers not yet set,
+    and raise KeyError for a mode that genuinely is registered.
+    """
+    import threading
+    import time
+
+    from instructor.v2.core.registry import ModeHandlers, ModeRegistry
+
+    registry = ModeRegistry()
+    n_threads = 8
+    started = threading.Barrier(n_threads, timeout=5)
+    load_count = 0
+    load_count_lock = threading.Lock()
+
+    def slow_loader() -> ModeHandlers:
+        nonlocal load_count
+        with load_count_lock:
+            load_count += 1
+        # Simulate a slow module import: gives every other thread time to
+        # reach get_handlers() and start contending before this resolves.
+        time.sleep(0.2)
+        return ModeHandlers(
+            request_handler=lambda *_a, **_k: None,
+            reask_handler=lambda *_a, **_k: None,
+            response_parser=lambda *_a, **_k: None,
+        )
+
+    registry.register_lazy(Provider.DEEPSEEK, Mode.TOOLS, slow_loader)
+
+    results: list[object] = [None] * n_threads
+    errors: list[BaseException] = []
+    errors_lock = threading.Lock()
+
+    def worker(idx: int) -> None:
+        # Line up all threads so they call get_handlers() as close to
+        # simultaneously as possible, maximizing contention.
+        started.wait(timeout=5)
+        try:
+            results[idx] = registry.get_handlers(Provider.DEEPSEEK, Mode.TOOLS)
+        except BaseException as exc:  # noqa: BLE001
+            with errors_lock:
+                errors.append(exc)
+
+    threads = [threading.Thread(target=worker, args=(i,)) for i in range(n_threads)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join(timeout=10)
+
+    assert errors == [], f"get_handlers raised for concurrent callers: {errors}"
+    assert all(r is not None for r in results)
+    assert len({id(r) for r in results}) == 1, "all callers should get the same ModeHandlers instance"
+    # The loader must run exactly once, not once per racing thread.
+    assert load_count == 1