perf: take heavy imports off the startup path and pre-warm them in the background (#1141)

Co-authored-by: Ahmed Allam <ahmed39652003@gmail.com>

pyproject.toml

@@ -270,6 +270,10 @@ ignore = [
 "strix/tools/thinking/tool.py" = ["TC002"]
 "strix/tools/web_search/tool.py" = ["TC002"]
 "strix/tools/proxy/tools.py" = ["TC002", "PLR0911"]
+# The generated Caido GraphQL schema is slow to import, so the SDK is imported
+# on first proxy call instead of at module scope (keeps it off the launch path).
+"strix/tools/proxy/caido_api.py" = ["PLC0415"]
+"strix/runtime/caido_bootstrap.py" = ["PLC0415"]
 "strix/tools/agents_graph/tools.py" = ["TC002"]
 "strix/agents/factory.py" = ["TC002"]
 # Entry point: ``Path`` is used at runtime by the typing of the
@@ -280,6 +284,13 @@ ignore = [
 # a runtime ``Callable`` annotation on ``vulnerability_found_callback``.
 "strix/report/state.py" = ["TC003", "PLR0912", "PLR0915", "E501", "PERF401", "PLC0415"]
 "strix/report/usage.py" = ["PLC0415"]
+# LiteLLM and the Docker SDK are imported on first use, not at module scope:
+# both cost seconds to import and neither is needed until a model call is made
+# (or, for Docker, unless the Docker runtime backend is in use).
+"strix/core/execution.py" = ["PLC0415"]
+"strix/report/pricing.py" = ["PLC0415"]
+"strix/llm/compaction.py" = ["PLC0415"]
+"strix/llm/context_budget.py" = ["PLC0415"]
 # Lazy import of strix.config.models avoids a circular dependency between the
 # report pipeline and the config layer.
 "strix/report/dedupe.py" = ["PLC0415"]

strix/core/execution.py

@@ -7,13 +7,12 @@ import contextlib
 import logging
 import uuid
 from collections.abc import Callable
+from functools import cache
 from typing import TYPE_CHECKING, Any, cast
 
-import litellm
 from agents import RunConfig, Runner
 from agents.exceptions import AgentsException, MaxTurnsExceeded, UserError
 from agents.sandbox.errors import ExecTransportError
-from docker import errors as docker_errors  # type: ignore[import-untyped, unused-ignore]
 from openai import (
     APIConnectionError,
     APIError,
@@ -56,6 +55,19 @@ _INPUT_REJECTION_CODES = frozenset({400, 404, 422})
 _MAX_COMPACTIONS_PER_CYCLE = 2
 
 
+@cache
+def _teardown_sandbox_errors() -> tuple[type[BaseException], ...]:
+    """Sandbox-gone errors, tolerated during shutdown.
+
+    The Docker SDK is imported here rather than at module scope: it is only
+    reachable with the Docker runtime backend, and importing it eagerly puts it
+    on every launch's critical path.
+    """
+    from docker import errors as docker_errors  # type: ignore[import-untyped, unused-ignore]
+
+    return (ExecTransportError, docker_errors.NotFound)
+
+
 class ProviderRefusalError(AgentsException):
     """Raised when a provider returns a structured refusal instead of an exception."""
 
@@ -126,6 +138,8 @@ def _is_transient_model_error(exc: BaseException) -> bool:
         return True
     code = _model_error_status_code(exc)
     if code is not None:
+        import litellm
+
         return bool(litellm._should_retry(code))
     return isinstance(exc, APIError)
 
@@ -692,7 +706,7 @@ async def _run_cycle(  # noqa: PLR0912, PLR0915
                         "Ignoring LiteLLM end-of-stream shutdown race for %s",
                         agent_id,
                     )
-                except (ExecTransportError, docker_errors.NotFound):
+                except _teardown_sandbox_errors():
                     if not coordinator.is_shutting_down:
                         raise
                     logger.warning(

strix/interface/main.py

@@ -431,6 +431,10 @@ def main() -> None:
 
         sys.exit(run_auth(sys.argv[2:]))
 
+    from strix.llm.warmup import start_import_warmup
+
+    start_import_warmup()
+
     args = parse_arguments()
 
     start_background_check()

strix/interface/utils.py

@@ -13,9 +13,7 @@ from pathlib import Path
 from typing import Any
 from urllib.parse import parse_qs, urlparse
 
-import docker
 import requests
-from docker.errors import DockerException, ImageNotFound
 from rich.console import Console
 from rich.panel import Panel
 from rich.text import Text
@@ -1599,6 +1597,9 @@ def clone_repository(repo_url: str, run_name: str, dest_name: str | None = None)
 
 
 def check_docker_connection() -> Any:
+    import docker
+    from docker.errors import DockerException
+
     try:
         return docker.from_env()
     except DockerException:
@@ -1624,6 +1625,8 @@ def check_docker_connection() -> Any:
 
 
 def image_exists(client: Any, image_name: str) -> bool:
+    from docker.errors import ImageNotFound
+
     try:
         client.images.get(image_name)
     except ImageNotFound:

strix/llm/compaction.py

@@ -10,11 +10,11 @@ pairing so the trimmed history is still valid provider input.
 from __future__ import annotations
 
 import logging
+from functools import cache
 from typing import TYPE_CHECKING, Any
 
 from agents.model_settings import ModelSettings
 from agents.models.interface import ModelTracing
-from litellm.exceptions import BadRequestError, ContextWindowExceededError
 from openai.types.responses import ResponseOutputMessage, ResponseOutputText
 
 from strix.config import load_settings
@@ -63,6 +63,18 @@ _OVERFLOW_MARKERS = (
 )
 
 
+@cache
+def _overflow_error_types() -> tuple[type[BaseException], type[BaseException]]:
+    """``(ContextWindowExceededError, BadRequestError)``, imported on first use.
+
+    LiteLLM costs seconds to import, and nothing needs it until a model call is
+    actually made, so it stays off the launch path.
+    """
+    from litellm.exceptions import BadRequestError, ContextWindowExceededError
+
+    return ContextWindowExceededError, BadRequestError
+
+
 def is_context_overflow(exc: BaseException) -> bool:
     """Whether ``exc`` is a model context-window-overflow error.
 
@@ -70,9 +82,10 @@ def is_context_overflow(exc: BaseException) -> bool:
     OpenRouter branch raises a plain BadRequestError, so for that we fall back to
     matching the provider message.
     """
-    if isinstance(exc, ContextWindowExceededError):
+    context_window_exceeded, bad_request = _overflow_error_types()
+    if isinstance(exc, context_window_exceeded):
         return True
-    if isinstance(exc, BadRequestError):
+    if isinstance(exc, bad_request):
         msg = str(exc).lower()
         if any(x in msg for x in _OVERFLOW_EXCLUSIONS):
             return False

strix/llm/context_budget.py

@@ -8,8 +8,6 @@ import logging
 from functools import lru_cache
 from typing import Any
 
-import litellm
-
 from strix.config import load_settings
 
 
@@ -38,6 +36,8 @@ def _lookup_key(model: str) -> str:
 
 def _safe_get_model_info(model: str) -> dict[str, Any] | None:
     try:
+        import litellm
+
         return dict(litellm.get_model_info(model))
     except Exception:  # noqa: BLE001 - unmapped models raise; caller falls back.
         return None
@@ -82,6 +82,8 @@ def count_tokens(model: str, text: str) -> int:
     if not text:
         return 0
     try:
+        import litellm
+
         return int(litellm.token_counter(model=_lookup_key(model), text=text))
     except Exception:  # noqa: BLE001 - tokenizer may be unavailable for some models.
         return len(text.encode("utf-8"))

strix/llm/warmup.py

@@ -0,0 +1,55 @@
+"""Background pre-import of the heavy scan dependencies.
+
+The scan engine's import graph (the agents SDK, OpenAI client, LiteLLM, the
+Caido SDK, the Docker SDK) costs seconds to import cold, but none of it is
+needed until a scan actually starts. Importing it on a daemon thread at CLI
+entry overlaps that cost with the I/O-bound startup work that always precedes
+a scan (argument parsing, Docker checks, image pull, TUI setup), so by the
+time the scan begins the modules are already in ``sys.modules``. Any thread
+that needs one of them before the warm-up finishes just blocks on the normal
+import lock, so behaviour is unchanged either way.
+"""
+
+from __future__ import annotations
+
+import importlib
+import logging
+import threading
+
+
+logger = logging.getLogger(__name__)
+
+WARMUP_MODULES = (
+    "strix.core.runner",
+    "litellm",
+    "caido_sdk_client",
+    "docker",
+)
+
+_lock = threading.Lock()
+_thread: threading.Thread | None = None
+
+
+def _warm(modules: tuple[str, ...]) -> None:
+    for name in modules:
+        try:
+            importlib.import_module(name)
+        except Exception:  # noqa: BLE001 - a failed warm-up must never fail the run.
+            logger.debug("Import warm-up for %r failed", name, exc_info=True)
+
+
+def start_import_warmup(modules: tuple[str, ...] = WARMUP_MODULES) -> threading.Thread:
+    """Start importing the heavy scan dependencies in the background, once.
+
+    ``modules`` lets embedders that never touch some backends (e.g. a cloud
+    runtime that has no local Docker) warm a narrower set.
+    """
+    global _thread  # noqa: PLW0603
+    with _lock:
+        if _thread is not None:
+            return _thread
+        _thread = threading.Thread(
+            target=_warm, args=(modules,), name="strix-import-warmup", daemon=True
+        )
+        _thread.start()
+        return _thread

strix/runtime/caido_bootstrap.py

@@ -15,12 +15,10 @@ import json
 import logging
 from typing import TYPE_CHECKING
 
-from caido_sdk_client import Client, TokenAuthOptions
-from caido_sdk_client.types import CreateProjectOptions
-
 
 if TYPE_CHECKING:
     from agents.sandbox.session import BaseSandboxSession
+    from caido_sdk_client import Client
 
 
 logger = logging.getLogger(__name__)
@@ -87,6 +85,12 @@ async def bootstrap_caido(
     container_url: str,
 ) -> Client:
     """Connect to the in-container Caido sidecar and select a fresh project."""
+    # The Caido SDK (and its generated GraphQL schema) is slow to import and is
+    # only needed once a sandbox is actually being bootstrapped, so it is
+    # imported here rather than at module scope.
+    from caido_sdk_client import Client, TokenAuthOptions
+    from caido_sdk_client.types import CreateProjectOptions
+
     logger.info("Bootstrapping Caido client (host=%s, container=%s)", host_url, container_url)
 
     access_token = await _login_as_guest(session, container_url=container_url)

strix/tools/proxy/caido_api.py

@@ -10,20 +10,16 @@ import urllib.request
 from typing import TYPE_CHECKING, Any, Literal
 from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
 
-from caido_sdk_client import Client, TokenAuthOptions
-from caido_sdk_client.types import (
-    ConnectionInfoInput,
-    CreateScopeOptions,
-    ReplaySendOptions,
-    RequestGetOptions,
-    UpdateScopeOptions,
-)
-
 
+# The generated Caido GraphQL schema module is slow to import and is only needed
+# once a proxy tool actually runs, so the SDK is imported on first use rather
+# than at module scope, which would put it on every launch's critical path.
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable
 
+    from caido_sdk_client import Client
     from caido_sdk_client import Client as CaidoClient
+    from caido_sdk_client.types import ConnectionInfoInput
 
 
 RequestPart = Literal["request", "response"]
@@ -85,6 +81,8 @@ def _login_as_guest() -> str:
 
 
 async def _new_client() -> Client:
+    from caido_sdk_client import Client, TokenAuthOptions
+
     token = await asyncio.to_thread(_login_as_guest)
     client = Client(caido_url(), auth=TokenAuthOptions(token=token))
     await client.connect()
@@ -163,6 +161,8 @@ async def get_request_with_client(
     # Passing False for either causes pydantic validation to fail with
     # "Field required" on the missing raw field. Always request both —
     # the caller picks which one to surface via ``part``.
+    from caido_sdk_client.types import RequestGetOptions
+
     opts = RequestGetOptions(request_raw=True, response_raw=True)
     return await client.request.get(request_id, opts)
 
@@ -206,6 +206,8 @@ def build_raw_request(
     if body:
         final_headers["Content-Length"] = str(len(body.encode("utf-8")))
 
+    from caido_sdk_client.types import ConnectionInfoInput
+
     lines = [f"{method.upper()} {path} HTTP/1.1"]
     lines.extend(f"{k}: {v}" for k, v in final_headers.items())
     raw = ("\r\n".join(lines) + "\r\n\r\n" + body).encode("utf-8")
@@ -334,6 +336,8 @@ async def replay_send_raw(
     raw: bytes,
     connection: ConnectionInfoInput,
 ) -> dict[str, Any]:
+    from caido_sdk_client.types import ReplaySendOptions
+
     started = time.time()
     # Create an empty replay session, then dispatch via ``send()``.
     # Passing ``CreateReplaySessionFromRaw`` here would also seed a stored
@@ -391,6 +395,8 @@ async def scope_create(
     allowlist: list[str] | None = None,
     denylist: list[str] | None = None,
 ) -> Any:
+    from caido_sdk_client.types import CreateScopeOptions
+
     return await client.scope.create(
         CreateScopeOptions(
             name=name,
@@ -408,6 +414,8 @@ async def scope_update(
     allowlist: list[str] | None = None,
     denylist: list[str] | None = None,
 ) -> Any:
+    from caido_sdk_client.types import UpdateScopeOptions
+
     return await client.scope.update(
         scope_id,
         UpdateScopeOptions(

tests/test_context_budget.py

@@ -31,7 +31,7 @@ def test_context_window_chatgpt_prefix_skips_provider_auth(
         calls.append(model)
         return {"max_input_tokens": 1_050_000, "max_output_tokens": 128_000}
 
-    monkeypatch.setattr("strix.llm.context_budget.litellm.get_model_info", _model_info)
+    monkeypatch.setattr("litellm.get_model_info", _model_info)
     try:
         assert context_budget.context_window("chatgpt/gpt-5.6-luna") == 1_050_000
         assert calls == ["gpt-5.6-luna"]
@@ -45,7 +45,7 @@ def test_context_window_unmapped_uses_fallback(monkeypatch: pytest.MonkeyPatch)
     def _raise(_model: str) -> dict[str, int]:
         raise ValueError("This model isn't mapped yet.")
 
-    monkeypatch.setattr("strix.llm.context_budget.litellm.get_model_info", _raise)
+    monkeypatch.setattr("litellm.get_model_info", _raise)
     expected = load_settings().context.fallback_context_tokens
     assert context_budget.context_window("totally-made-up-model") == expected
     context_budget._model_info.cache_clear()
@@ -55,7 +55,7 @@ def test_count_tokens_fallback_on_error(monkeypatch: pytest.MonkeyPatch) -> None
     def _raise(**_kwargs: object) -> int:
         raise RuntimeError("no tokenizer")
 
-    monkeypatch.setattr("strix.llm.context_budget.litellm.token_counter", _raise)
+    monkeypatch.setattr("litellm.token_counter", _raise)
     # Falls back to UTF-8 byte length (upper bound on tokens).
     assert context_budget.count_tokens("weird-model", "x" * 400) == 400
     assert context_budget.count_tokens("weird-model", "😀" * 10) == 40