diff --git a/agent/transports/__init__.py b/agent/transports/__init__.py index b606da7fecaea..d7e5e067ba62a 100644 --- a/agent/transports/__init__.py +++ b/agent/transports/__init__.py @@ -6,6 +6,8 @@ result = transport.normalize_response(raw_response) """ +import threading + from agent.transports.types import ( NormalizedResponse, ToolCall, @@ -16,6 +18,7 @@ _REGISTRY: dict = {} _discovered: bool = False +_discover_lock = threading.Lock() def register_transport(api_mode: str, transport_cls: type) -> None: @@ -49,20 +52,28 @@ def get_transport(api_mode: str): def _discover_transports() -> None: """Import all transport modules to trigger auto-registration.""" global _discovered - _discovered = True - try: - import agent.transports.anthropic # noqa: F401 - except ImportError: - pass - try: - import agent.transports.codex # noqa: F401 - except ImportError: - pass - try: - import agent.transports.chat_completions # noqa: F401 - except ImportError: - pass - try: - import agent.transports.bedrock # noqa: F401 - except ImportError: - pass + if _discovered: + return + + with _discover_lock: + if _discovered: + return + + try: + import agent.transports.anthropic # noqa: F401 + except ImportError: + pass + try: + import agent.transports.codex # noqa: F401 + except ImportError: + pass + try: + import agent.transports.chat_completions # noqa: F401 + except ImportError: + pass + try: + import agent.transports.bedrock # noqa: F401 + except ImportError: + pass + + _discovered = True diff --git a/tests/agent/transports/test_transport_discovery_concurrency.py b/tests/agent/transports/test_transport_discovery_concurrency.py new file mode 100644 index 0000000000000..21bd263a7d2a9 --- /dev/null +++ b/tests/agent/transports/test_transport_discovery_concurrency.py @@ -0,0 +1,94 @@ +"""Concurrency regressions for transport discovery.""" + +import builtins +from concurrent.futures import ThreadPoolExecutor +import threading + +import pytest + +import agent.transports as transports + + +class _ObservedLock: + """Expose when the second caller reaches lock acquisition.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self._state_lock = threading.Lock() + self._attempts = 0 + self.second_attempted = threading.Event() + + def __enter__(self) -> "_ObservedLock": + with self._state_lock: + self._attempts += 1 + if self._attempts == 2: + self.second_attempted.set() + self._lock.acquire() + return self + + def __exit__(self, *exc_info: object) -> None: + self._lock.release() + + +def test_discovery_is_published_after_one_serialized_import_sweep(monkeypatch) -> None: + target_imports = ( + "agent.transports.anthropic", + "agent.transports.codex", + "agent.transports.chat_completions", + "agent.transports.bedrock", + ) + observed_lock = _ObservedLock() + first_import_started = threading.Event() + release_first_import = threading.Event() + imported: list[str] = [] + original_import = builtins.__import__ + + def controlled_import(name: str, *args: object, **kwargs: object) -> object: + if name in target_imports: + imported.append(name) + if name == target_imports[0]: + first_import_started.set() + assert release_first_import.wait(timeout=5) + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(transports, "_discovered", False) + monkeypatch.setattr(transports, "_discover_lock", observed_lock) + monkeypatch.setattr(builtins, "__import__", controlled_import) + + pool = ThreadPoolExecutor(max_workers=2) + try: + first = pool.submit(transports._discover_transports) + assert first_import_started.wait(timeout=5) + assert transports._discovered is False + + second = pool.submit(transports._discover_transports) + assert observed_lock.second_attempted.wait(timeout=5) + assert second.done() is False + assert imported == [target_imports[0]] + + release_first_import.set() + first.result(timeout=5) + second.result(timeout=5) + finally: + release_first_import.set() + pool.shutdown(wait=True) + + assert transports._discovered is True + assert imported == list(target_imports) + + +def test_failed_discovery_is_not_published(monkeypatch) -> None: + original_import = builtins.__import__ + + def fail_first_transport(name: str, *args: object, **kwargs: object) -> object: + if name == "agent.transports.anthropic": + raise RuntimeError("transport import failed") + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(transports, "_discovered", False) + monkeypatch.setattr(builtins, "__import__", fail_first_transport) + + with pytest.raises(RuntimeError, match="transport import failed"): + transports._discover_transports() + + assert transports._discovered is False