feat(runtime): session-setup hook (register_session_setup)

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,

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,
+            )

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