From 1ce43d1b94240a3fd4b9f45cd3c0376b4feb36c9 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 20 Aug 2026 20:24:08 -0700 Subject: [PATCH] perf: take heavy imports off the startup path and pre-warm them in the background (#1141) Co-authored-by: Ahmed Allam --- pyproject.toml | 11 +++++++ strix/core/execution.py | 20 ++++++++++-- strix/interface/main.py | 4 +++ strix/interface/utils.py | 7 ++-- strix/llm/compaction.py | 19 +++++++++-- strix/llm/context_budget.py | 6 ++-- strix/llm/warmup.py | 55 ++++++++++++++++++++++++++++++++ strix/runtime/caido_bootstrap.py | 10 ++++-- strix/tools/proxy/caido_api.py | 26 +++++++++------ tests/test_context_budget.py | 6 ++-- 10 files changed, 139 insertions(+), 25 deletions(-) create mode 100644 strix/llm/warmup.py diff --git a/pyproject.toml b/pyproject.toml index 77be738f..195a9646 100644 --- a/pyproject.toml +++ b/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"] diff --git a/strix/core/execution.py b/strix/core/execution.py index bd99e7c3..9f2674d2 100644 --- a/strix/core/execution.py +++ b/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( diff --git a/strix/interface/main.py b/strix/interface/main.py index 06966f4c..45e114b5 100644 --- a/strix/interface/main.py +++ b/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() diff --git a/strix/interface/utils.py b/strix/interface/utils.py index 6789abe0..faab1772 100644 --- a/strix/interface/utils.py +++ b/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: diff --git a/strix/llm/compaction.py b/strix/llm/compaction.py index e40caf6a..ecdb9a83 100644 --- a/strix/llm/compaction.py +++ b/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 diff --git a/strix/llm/context_budget.py b/strix/llm/context_budget.py index b7589a9e..baa02c4b 100644 --- a/strix/llm/context_budget.py +++ b/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")) diff --git a/strix/llm/warmup.py b/strix/llm/warmup.py new file mode 100644 index 00000000..98da959d --- /dev/null +++ b/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 diff --git a/strix/runtime/caido_bootstrap.py b/strix/runtime/caido_bootstrap.py index 0b9ad5b1..8398350b 100644 --- a/strix/runtime/caido_bootstrap.py +++ b/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) diff --git a/strix/tools/proxy/caido_api.py b/strix/tools/proxy/caido_api.py index 6cfee56c..af81b500 100644 --- a/strix/tools/proxy/caido_api.py +++ b/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( diff --git a/tests/test_context_budget.py b/tests/test_context_budget.py index a9a48e0c..8c06fa49 100644 --- a/tests/test_context_budget.py +++ b/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