diff --git a/strix/core/runner.py b/strix/core/runner.py index 0ebcc48b4..20cdaf538 100644 --- a/strix/core/runner.py +++ b/strix/core/runner.py @@ -210,6 +210,15 @@ async def run_strix_scan( scan_mode = str(scan_config.get("scan_mode") or "deep") is_whitebox = any(t.get("type") == "local_code" for t in targets) skills = list(scan_config.get("skills") or []) + + # Run registered session-setup hooks now that the sandbox is ready and + # the target is materialised, but before the agent's first turn — the + # window an addon needs to prepare in-sandbox state (e.g. a code-graph + # index). Runs once per materialized session (a resume reuses the cached + # bundle, so setups don't re-run). Best-effort: a hook failure is logged, + # never fatal. + await session_manager.run_session_setups(bundle, scan_config) + root_task = build_root_task(scan_config) model_settings = make_model_settings( settings.llm.reasoning_effort, diff --git a/strix/runtime/session_manager.py b/strix/runtime/session_manager.py index 5aa714238..64726314c 100644 --- a/strix/runtime/session_manager.py +++ b/strix/runtime/session_manager.py @@ -4,6 +4,7 @@ import logging import shutil +from collections.abc import Awaitable, Callable from pathlib import Path from typing import Any @@ -175,3 +176,82 @@ async def cleanup(scan_id: str) -> None: "cleanup(%s): client.delete raised; container may need manual reaping", scan_id, ) + + +# --- session-setup hooks ----------------------------------------------------- +# +# Extension point mirroring ``backends.register_backend`` / +# ``agents.factory.register_agent_tools``: a callback registered here runs once +# per scan, right after the sandbox session is ready and BEFORE the agent's +# first turn, with the live session and the scan config. This is the seam an +# addon needs to prepare in-sandbox state the agent will use — e.g. building a +# code-graph index against the freshly-materialised target and copying it out. +# +# A ``@function_tool`` runs runner-side and only receives a ``RunContextWrapper`` +# (not the session), so setup that must ``session.exec`` inside the container +# has no other place to hook. Callbacks are awaited in registration order; an +# exception in one is logged and swallowed so an optional setup step can never +# break the scan. + +SessionSetupCallback = Callable[["Any", dict[str, Any]], Awaitable[None]] + +_SESSION_SETUP_CALLBACKS: list[SessionSetupCallback] = [] + + +def register_session_setup(callback: SessionSetupCallback) -> None: + """Register a coroutine ``callback(session, scan_config)`` to run once per + scan after the sandbox session is ready. Duplicate registrations are ignored + so repeated imports don't double-run a setup step.""" + if callback not in _SESSION_SETUP_CALLBACKS: + _SESSION_SETUP_CALLBACKS.append(callback) + logger.info( + "Registered session-setup callback: %s", + getattr(callback, "__name__", callback), + ) + + +def registered_session_setups() -> tuple[SessionSetupCallback, ...]: + """Return the currently registered session-setup callbacks.""" + return tuple(_SESSION_SETUP_CALLBACKS) + + +_SETUP_DONE_KEY = "_session_setups_done" + + +async def run_session_setups(bundle: dict[str, Any], scan_config: dict[str, Any]) -> None: + """Invoke every registered session-setup callback ONCE per materialized + sandbox session. + + ``bundle`` is the session bundle from :func:`create_or_reuse`; it's cached + and REUSED across resumes of the same ``scan_id`` in one process (the + cred-boundary iteration loop resumes in-process). Running setups again on a + reused session would re-execute non-idempotent hooks against an + already-prepared sandbox (start a service twice, re-append config, or fail + mid-way and — since failures are swallowed — leave the resumed agent on + invalid state). So we stamp the bundle once setups have run and skip on + reuse. A fresh session (new bundle) gets a fresh run. + + On skipping and staleness: the reused bundle is the SAME container with the + SAME materialised source, so setup-derived state (e.g. a code-graph index) + is not stale — nothing changed between in-process iterations. A cross-push + PR resume runs in a NEW process against a freshly materialised sandbox, so + it gets a new bundle and setups re-run — the index is rebuilt against the + new code. So skipping here never serves a stale artifact. + + Callback failures are logged and swallowed — setup is best-effort and must + not fail the scan. + """ + if bundle.get(_SETUP_DONE_KEY): + logger.debug("session setups already ran for this session; skipping on reuse") + return + bundle[_SETUP_DONE_KEY] = True # stamp first: a partial run must not retry-loop + session = bundle.get("session") + for callback in _SESSION_SETUP_CALLBACKS: + try: + await callback(session, scan_config) + except Exception: # noqa: BLE001 — an optional setup step must not break the scan + logger.warning( + "session-setup callback %s raised; continuing", + getattr(callback, "__name__", callback), + exc_info=True, + ) diff --git a/tests/test_session_setup_hooks.py b/tests/test_session_setup_hooks.py new file mode 100644 index 000000000..232830c10 --- /dev/null +++ b/tests/test_session_setup_hooks.py @@ -0,0 +1,124 @@ +"""Tests for the session-setup hook registry (register_session_setup).""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from strix.runtime import session_manager + + +@pytest.fixture(autouse=True) +def _clear_registry() -> Any: + """Each test starts with an empty registry and restores it after.""" + saved = list(session_manager._SESSION_SETUP_CALLBACKS) + session_manager._SESSION_SETUP_CALLBACKS.clear() + yield + session_manager._SESSION_SETUP_CALLBACKS[:] = saved + + +def _bundle(session: Any = None) -> dict[str, Any]: + """A minimal session bundle, as create_or_reuse returns + caches.""" + return {"session": session if session is not None else object()} + + +async def test_callbacks_run_in_registration_order() -> None: + calls: list[str] = [] + + async def a(_session: Any, _cfg: dict[str, Any]) -> None: + calls.append("a") + + async def b(_session: Any, _cfg: dict[str, Any]) -> None: + calls.append("b") + + session_manager.register_session_setup(a) + session_manager.register_session_setup(b) + await session_manager.run_session_setups(_bundle(), scan_config={}) + + assert calls == ["a", "b"] + + +async def test_callback_receives_session_and_config() -> None: + seen: dict[str, Any] = {} + sentinel_session = object() + + async def grab(session: Any, cfg: dict[str, Any]) -> None: + seen["session"] = session + seen["cfg"] = cfg + + session_manager.register_session_setup(grab) + cfg = {"targets": [{"type": "local_code"}]} + await session_manager.run_session_setups(_bundle(sentinel_session), scan_config=cfg) + + assert seen["session"] is sentinel_session + assert seen["cfg"] is cfg + + +async def test_setups_run_once_per_session_on_reuse() -> None: + """The bundle is cached + reused when a scan resumes in-process; setups must + run ONCE per materialized session, not again on reuse (non-idempotent hooks + would double-apply). A fresh bundle runs them again.""" + calls: list[str] = [] + + async def setup(_session: Any, _cfg: dict[str, Any]) -> None: + calls.append("run") + + session_manager.register_session_setup(setup) + + bundle = _bundle() + await session_manager.run_session_setups(bundle, scan_config={}) + await session_manager.run_session_setups(bundle, scan_config={}) # resume: same bundle + assert calls == ["run"], "setup must not re-run on a reused session" + + # A different materialized session (new bundle) runs setups afresh. + await session_manager.run_session_setups(_bundle(), scan_config={}) + assert calls == ["run", "run"] + + +async def test_setup_marked_done_even_when_callback_fails() -> None: + """A partial/failed setup still marks the session done, so a resume doesn't + retry-loop a broken non-idempotent hook.""" + calls: list[str] = [] + + async def boom(_session: Any, _cfg: dict[str, Any]) -> None: + calls.append("boom") + raise RuntimeError("half-applied") + + session_manager.register_session_setup(boom) + bundle = _bundle() + await session_manager.run_session_setups(bundle, scan_config={}) + await session_manager.run_session_setups(bundle, scan_config={}) + assert calls == ["boom"] + + +async def test_duplicate_registration_ignored() -> None: + async def once(_session: Any, _cfg: dict[str, Any]) -> None: + pass + + session_manager.register_session_setup(once) + session_manager.register_session_setup(once) + + assert session_manager.registered_session_setups() == (once,) + + +async def test_failing_callback_is_swallowed_and_others_still_run() -> None: + calls: list[str] = [] + + async def boom(_session: Any, _cfg: dict[str, Any]) -> None: + calls.append("boom") + raise RuntimeError("setup blew up") + + async def after(_session: Any, _cfg: dict[str, Any]) -> None: + calls.append("after") + + session_manager.register_session_setup(boom) + session_manager.register_session_setup(after) + # Must not raise — a best-effort setup step can't fail the scan. + await session_manager.run_session_setups(_bundle(), scan_config={}) + + assert calls == ["boom", "after"] + + +async def test_no_callbacks_is_a_noop() -> None: + await session_manager.run_session_setups(_bundle(), scan_config={}) # no raise