From 13140e403216edb0c2607940ce58a258c96031f6 Mon Sep 17 00:00:00 2001 From: Sean Turner Date: Thu, 16 Jul 2026 08:02:37 +0100 Subject: [PATCH 1/2] feat(runtime): session-setup hook (register_session_setup) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an extension point that runs a callback 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: from strix.runtime.session_manager import register_session_setup async def prepare(session, scan_config): await session.exec(...) # set up in-sandbox state the agent will use register_session_setup(prepare) Mirrors backends.register_backend / agents.factory.register_agent_tools. The gap it fills: a @function_tool runs runner-side and only receives a RunContextWrapper, never the sandbox session — so setup that must session.exec inside the container (build an index, warm a cache, install a probe) has no seam to hook. run_strix_scan invokes the callbacks after create_or_reuse and before building the root agent. Best-effort by contract: callbacks are awaited in registration order, and an exception in one is logged and swallowed so an optional setup step can never fail the scan. Duplicate registrations are ignored. Motivating consumer: a SCIP code-graph addon that indexes the freshly materialised target in-sandbox and copies the index out for its query tools — but the hook is generic (cache warming, sidecar probes, etc.). Tests: registration order, session/config passthrough, dedup, failure-is- swallowed, empty-registry no-op. --- strix/core/runner.py | 7 +++ strix/runtime/session_manager.py | 53 ++++++++++++++++++++ tests/test_session_setup_hooks.py | 82 +++++++++++++++++++++++++++++++ 3 files changed, 142 insertions(+) create mode 100644 tests/test_session_setup_hooks.py diff --git a/strix/core/runner.py b/strix/core/runner.py index 0ebcc48b4..bbc104412 100644 --- a/strix/core/runner.py +++ b/strix/core/runner.py @@ -210,6 +210,13 @@ 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). Best-effort: a hook failure is logged, never fatal. + await session_manager.run_session_setups(bundle["session"], 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..f877815f0 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,55 @@ 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) + + +async def run_session_setups(session: Any, scan_config: dict[str, Any]) -> None: + """Invoke every registered session-setup callback. A callback failure is + logged and swallowed — session setup is best-effort and must not fail the + scan.""" + 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..7a4cd514c --- /dev/null +++ b/tests/test_session_setup_hooks.py @@ -0,0 +1,82 @@ +"""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 + + +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(session=object(), 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(session=sentinel_session, scan_config=cfg) + + assert seen["session"] is sentinel_session + assert seen["cfg"] is cfg + + +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(session=object(), scan_config={}) + + assert calls == ["boom", "after"] + + +async def test_no_callbacks_is_a_noop() -> None: + await session_manager.run_session_setups(session=object(), scan_config={}) # no raise From 33eb270916197e0bddb506c5dfe3a6f37e3b4640 Mon Sep 17 00:00:00 2001 From: Sean Turner Date: Thu, 16 Jul 2026 10:16:46 +0100 Subject: [PATCH 2/2] fix(runtime): run session setups once per materialized session (Greptile #781) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit create_or_reuse caches + reuses the session bundle across in-process resumes (the cred-boundary iteration loop). run_session_setups was called unconditionally, so a resume re-ran every setup callback against an already-prepared sandbox — double-applying non-idempotent hooks (start a service twice, re-append config) or failing mid-way and, since failures are swallowed, leaving the resumed agent on invalid state. Take the bundle (not the raw session) and stamp it once setups have run; skip on reuse. Stamp BEFORE running so a partial/failed setup doesn't retry-loop. A fresh session (new bundle) — including a cross-push PR resume in a new process — gets a fresh run, so setup-derived state is never stale: the reused bundle is the same container with the same materialised source. Tests: run-once-on-reuse, fresh-bundle-reruns, marked-done-even-on-failure. --- strix/core/runner.py | 6 ++-- strix/runtime/session_manager.py | 35 +++++++++++++++++++--- tests/test_session_setup_hooks.py | 50 ++++++++++++++++++++++++++++--- 3 files changed, 81 insertions(+), 10 deletions(-) diff --git a/strix/core/runner.py b/strix/core/runner.py index bbc104412..20cdaf538 100644 --- a/strix/core/runner.py +++ b/strix/core/runner.py @@ -214,8 +214,10 @@ async def run_strix_scan( # 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). Best-effort: a hook failure is logged, never fatal. - await session_manager.run_session_setups(bundle["session"], scan_config) + # 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( diff --git a/strix/runtime/session_manager.py b/strix/runtime/session_manager.py index f877815f0..64726314c 100644 --- a/strix/runtime/session_manager.py +++ b/strix/runtime/session_manager.py @@ -215,10 +215,37 @@ def registered_session_setups() -> tuple[SessionSetupCallback, ...]: return tuple(_SESSION_SETUP_CALLBACKS) -async def run_session_setups(session: Any, scan_config: dict[str, Any]) -> None: - """Invoke every registered session-setup callback. A callback failure is - logged and swallowed — session setup is best-effort and must not fail the - scan.""" +_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) diff --git a/tests/test_session_setup_hooks.py b/tests/test_session_setup_hooks.py index 7a4cd514c..232830c10 100644 --- a/tests/test_session_setup_hooks.py +++ b/tests/test_session_setup_hooks.py @@ -18,6 +18,11 @@ def _clear_registry() -> Any: 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] = [] @@ -29,7 +34,7 @@ async def b(_session: Any, _cfg: dict[str, Any]) -> None: session_manager.register_session_setup(a) session_manager.register_session_setup(b) - await session_manager.run_session_setups(session=object(), scan_config={}) + await session_manager.run_session_setups(_bundle(), scan_config={}) assert calls == ["a", "b"] @@ -44,12 +49,49 @@ async def grab(session: Any, cfg: dict[str, Any]) -> None: session_manager.register_session_setup(grab) cfg = {"targets": [{"type": "local_code"}]} - await session_manager.run_session_setups(session=sentinel_session, scan_config=cfg) + 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 @@ -73,10 +115,10 @@ async def after(_session: Any, _cfg: dict[str, Any]) -> None: 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(session=object(), scan_config={}) + 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(session=object(), scan_config={}) # no raise + await session_manager.run_session_setups(_bundle(), scan_config={}) # no raise