feat(runtime): graduated wrap-up warnings, budget reserve, and interactive budget pause/continue (#893)

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

strix/core/agents.py

@@ -14,13 +14,15 @@ from strix.core.sessions import session_write_lock
 
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from agents.items import TResponseInputItem
     from agents.memory import Session
 
 
 logger = logging.getLogger(__name__)
 
-Status = Literal["running", "waiting", "completed", "stopped", "crashed", "failed"]
+Status = Literal["running", "waiting", "completed", "stopped", "crashed", "failed", "budget_paused"]
 
 
 @dataclass(slots=True)
@@ -47,6 +49,9 @@ class AgentCoordinator:
         self._snapshot_path: Path | None = None
         self.is_shutting_down = False
         self._budget_stopped = False
+        self._reserve_stopped = False
+        self._budget_paused = False
+        self._extend_budget: Callable[[], None] | None = None
 
     def set_snapshot_path(self, path: Path) -> None:
         self._snapshot_path = path
@@ -65,6 +70,71 @@ class AgentCoordinator:
             for runtime in self.runtimes.values():
                 runtime.wake.set()
 
+    @property
+    def reserve_stopped(self) -> bool:
+        return self._reserve_stopped
+
+    @property
+    def budget_paused(self) -> bool:
+        return self._budget_paused
+
+    def set_budget_extender(self, extend: Callable[[], None]) -> None:
+        self._extend_budget = extend
+
+    async def pause_for_budget(self, agent_id: str) -> None:
+        async with self._lock:
+            self._budget_paused = True
+        await self.set_status(agent_id, "budget_paused")
+
+    async def resume_from_budget_pause(self, *, exclude: str | None = None) -> None:
+        async with self._lock:
+            if not self._budget_paused:
+                return
+            self._budget_paused = False
+            paused = [aid for aid, status in self.statuses.items() if status == "budget_paused"]
+        if self._extend_budget is not None:
+            self._extend_budget()
+        for aid in paused:
+            await self.set_status(aid, "waiting")
+            if aid != exclude:
+                await self.send(
+                    aid,
+                    {
+                        "from": "system",
+                        "type": "budget_extended",
+                        "content": (
+                            "[Budget] The user extended the scan budget \u2014 continue your "
+                            "current task."
+                        ),
+                    },
+                )
+
+    async def reset_budget_stops(
+        self,
+        *,
+        budget_stopped: bool,
+        reserve_stopped: bool,
+        budget_paused: bool = False,
+    ) -> None:
+        async with self._lock:
+            self._budget_stopped = budget_stopped
+            self._reserve_stopped = reserve_stopped
+            if not budget_paused:
+                self._budget_paused = False
+                for aid, status in self.statuses.items():
+                    if status == "budget_paused":
+                        self.statuses[aid] = "waiting"
+        await self._maybe_snapshot()
+
+    async def claim_reserve_notification(self) -> str | None:
+        async with self._lock:
+            if self._reserve_stopped:
+                return None
+            self._reserve_stopped = True
+            for runtime in self.runtimes.values():
+                runtime.wake.set()
+            return next((aid for aid, parent in self.parent_of.items() if parent is None), None)
+
     async def register(
         self,
         agent_id: str,
@@ -132,6 +202,8 @@ class AgentCoordinator:
 
     async def send(self, target_agent_id: str, message: dict[str, Any]) -> bool:
         """Deliver a user/peer message by appending it to the target SDK session."""
+        if message.get("from") == "user" and self._budget_paused:
+            await self.resume_from_budget_pause(exclude=target_agent_id)
         async with self._lock:
             if target_agent_id not in self.statuses:
                 logger.debug("agent.send dropped unknown target=%s", target_agent_id)
@@ -166,7 +238,8 @@ class AgentCoordinator:
     async def wait_for_message(self, agent_id: str) -> None:
         while True:
             async with self._lock:
-                if self._budget_stopped or self.pending_counts.get(agent_id, 0) > 0:
+                reserve_exit = self._reserve_stopped and self.parent_of.get(agent_id) is not None
+                if self._budget_stopped or reserve_exit or self.pending_counts.get(agent_id, 0) > 0:
                     return
                 wake = self.runtimes.setdefault(agent_id, AgentRuntime()).wake
                 wake.clear()
@@ -300,6 +373,9 @@ class AgentCoordinator:
                 "metadata": {aid: dict(md) for aid, md in self.metadata.items()},
                 "pending_counts": dict(self.pending_counts),
                 "errors": dict(self.errors),
+                "budget_stopped": self._budget_stopped,
+                "reserve_stopped": self._reserve_stopped,
+                "budget_paused": self._budget_paused,
             }
 
     async def restore(self, snap: dict[str, Any]) -> None:
@@ -310,6 +386,9 @@ class AgentCoordinator:
             self.metadata = {aid: dict(md) for aid, md in snap.get("metadata", {}).items()}
             self.pending_counts = dict(snap.get("pending_counts", {}))
             self.errors = dict(snap.get("errors", {}))
+            self._budget_stopped = bool(snap.get("budget_stopped", False))
+            self._reserve_stopped = bool(snap.get("reserve_stopped", False))
+            self._budget_paused = bool(snap.get("budget_paused", False))
             for aid in self.statuses:
                 self.runtimes.setdefault(aid, AgentRuntime())
 

strix/core/execution.py

@@ -21,7 +21,11 @@ from openai import (
     RateLimitError,
 )
 
-from strix.core.hooks import BudgetExceededError
+from strix.core.hooks import (
+    BudgetExceededError,
+    BudgetPausedError,
+    SubagentBudgetReservedError,
+)
 from strix.core.inputs import child_initial_input
 from strix.core.sessions import (
     enforce_image_budget,
@@ -134,21 +138,34 @@ async def run_agent_loop(
     )
     result: RunResultBase | None = None
 
+    budget_stopped = coordinator.budget_stopped
+    reserve_stopped = coordinator.reserve_stopped
+    if budget_stopped:
+        await coordinator.set_status(agent_id, "stopped")
+        raise BudgetExceededError("scan budget reached")
+    if reserve_stopped and context.get("parent_id") is not None:
+        await coordinator.set_status(agent_id, "stopped")
+        raise SubagentBudgetReservedError("scan reached the sub-agent budget reserve")
+
+    if reserve_stopped and start_parked and interactive and context.get("parent_id") is None:
+        await coordinator.send(agent_id, _reserve_notice())
+
     if not (start_parked and interactive):
         if interactive:
-            result = await _run_cycle(
-                agent,
-                coordinator,
-                agent_id,
-                input_data=initial_input,
-                run_config=run_config,
-                context=context,
-                max_turns=max_turns,
-                session=session,
-                interactive=interactive,
-                event_sink=event_sink,
-                hooks=hooks,
-            )
+            with contextlib.suppress(BudgetPausedError):
+                result = await _run_cycle(
+                    agent,
+                    coordinator,
+                    agent_id,
+                    input_data=initial_input,
+                    run_config=run_config,
+                    context=context,
+                    max_turns=max_turns,
+                    session=session,
+                    interactive=interactive,
+                    event_sink=event_sink,
+                    hooks=hooks,
+                )
         else:
             result = await _run_noninteractive_until_lifecycle(
                 agent,
@@ -176,20 +193,25 @@ async def run_agent_loop(
             await coordinator.set_status(agent_id, "stopped")
             raise BudgetExceededError("scan budget reached")
 
+        if coordinator.reserve_stopped and context.get("parent_id") is not None:
+            await coordinator.set_status(agent_id, "stopped")
+            raise SubagentBudgetReservedError("scan reached the sub-agent budget reserve")
+
         await coordinator.consume_pending(agent_id)
-        result = await _run_cycle(
-            agent,
-            coordinator,
-            agent_id,
-            input_data=[],
-            run_config=run_config,
-            context=context,
-            max_turns=max_turns,
-            session=session,
-            interactive=interactive,
-            event_sink=event_sink,
-            hooks=hooks,
-        )
+        with contextlib.suppress(BudgetPausedError):
+            result = await _run_cycle(
+                agent,
+                coordinator,
+                agent_id,
+                input_data=[],
+                run_config=run_config,
+                context=context,
+                max_turns=max_turns,
+                session=session,
+                interactive=interactive,
+                event_sink=event_sink,
+                hooks=hooks,
+            )
 
 
 async def spawn_child_agent(
@@ -361,6 +383,10 @@ async def _run_noninteractive_until_lifecycle(
             await coordinator.set_status(agent_id, "stopped")
             raise BudgetExceededError("scan budget reached")
 
+        if coordinator.reserve_stopped and context.get("parent_id") is not None:
+            await coordinator.set_status(agent_id, "stopped")
+            raise SubagentBudgetReservedError("scan reached the sub-agent budget reserve")
+
         result = await _run_cycle(
             agent,
             coordinator,
@@ -456,9 +482,7 @@ async def _run_cycle(  # noqa: PLR0912, PLR0915
                                 logger.exception("stream event sink failed for %s", agent_id)
                     if stream.run_loop_exception is not None:
                         raise stream.run_loop_exception
-                except BudgetExceededError:
-                    # A RuntimeError subclass: re-raise explicitly so it is never
-                    # mistaken for the LiteLLM "after shutdown" race below.
+                except (BudgetExceededError, BudgetPausedError, SubagentBudgetReservedError):
                     raise
                 except RuntimeError as stream_exc:
                     if "after shutdown" not in str(stream_exc):
@@ -477,6 +501,15 @@ async def _run_cycle(  # noqa: PLR0912, PLR0915
                     )
             finally:
                 await coordinator.detach_stream(agent_id, stream)
+        except BudgetPausedError as exc:
+            logger.info("agent %s paused at the scan budget limit: %s", agent_id, exc)
+            await coordinator.pause_for_budget(agent_id)
+            raise
+        except SubagentBudgetReservedError as exc:
+            logger.info("sub-agent %s stopped at the budget reserve: %s", agent_id, exc)
+            await coordinator.set_status(agent_id, "stopped")
+            await _notify_root_on_budget_reserve(coordinator)
+            raise
         except BudgetExceededError as exc:
             logger.info(
                 "agent %s reached the scan budget limit; stopping the scan: %s", agent_id, exc
@@ -639,6 +672,28 @@ async def _notify_parent_on_crash(
     )
 
 
+def _reserve_notice() -> dict[str, Any]:
+    return {
+        "from": "system",
+        "type": "budget_reserve_stop",
+        "priority": "high",
+        "content": (
+            "[Budget reserve] The scan has reached the sub-agent budget reserve: every "
+            "sub-agent is being force-stopped as soon as its in-flight turn completes, and "
+            "none will send a completion report. Their confirmed vulnerabilities are "
+            "already filed as they were found. Do not wait on any sub-agents and do not "
+            "spawn new ones — wrap up now and call finish_scan."
+        ),
+    }
+
+
+async def _notify_root_on_budget_reserve(coordinator: AgentCoordinator) -> None:
+    root = await coordinator.claim_reserve_notification()
+    if root is None:
+        return
+    await coordinator.send(root, _reserve_notice())
+
+
 async def _start_child_runner(
     *,
     parent_ctx: dict[str, Any],
@@ -690,6 +745,8 @@ async def _start_child_runner(
             )
         except BudgetExceededError:
             logger.info("child %s stopped after reaching the scan budget limit", child_id)
+        except SubagentBudgetReservedError:
+            logger.info("child %s stopped at the sub-agent budget reserve", child_id)
 
     task_handle = asyncio.create_task(_child_loop(), name=f"agent-{name}-{child_id}")
     await coordinator.attach_runtime(child_id, task=task_handle)

strix/core/hooks.py

@@ -14,26 +14,210 @@ from strix.report.state import get_global_report_state
 if TYPE_CHECKING:
     from agents import RunContextWrapper
     from agents.agent import Agent
-    from agents.items import ModelResponse
+    from agents.items import ModelResponse, TResponseInputItem
 
 
 logger = logging.getLogger(__name__)
 
 
+_STAGE_LABELS: tuple[str, ...] = ("NOTICE", "URGENT", "CRITICAL")
+_TURN_WARN_BANDS: tuple[float, ...] = (0.70, 0.85, 0.95)
+_ROOT_BUDGET_WARN_BANDS: tuple[float, ...] = (0.70, 0.85, 0.95)
+_SUBAGENT_BUDGET_WARN_BANDS: tuple[float, ...] = (0.75, 0.80, 0.85)
+_SUBAGENT_BUDGET_RESERVE = 0.90
+
+
 class BudgetExceededError(RuntimeError):
     """Raised when the accumulated LLM cost reaches the configured budget."""
 
 
-class ReportUsageHooks(RunHooks[dict[str, Any]]):
-    """Persist SDK-native usage after every model response."""
+class SubagentBudgetReservedError(RuntimeError):
+    """Raised to stop a single sub-agent once the reserve threshold is crossed."""
 
-    def __init__(self, *, model: str, max_budget_usd: float | None = None) -> None:
+
+class BudgetPausedError(RuntimeError):
+    """Raised to park one agent when an interactive scan reaches its budget."""
+
+
+def recomputed_budget_flags(
+    cost: float,
+    max_budget_usd: float | None,
+    *,
+    interactive: bool,
+) -> tuple[bool, bool]:
+    """Return the (budget_stopped, reserve_stopped) flags a resumed scan should carry."""
+    if max_budget_usd is None:
+        return False, False
+    if interactive:
+        return False, False
+    budget_stopped = cost >= max_budget_usd
+    reserve_stopped = cost >= max_budget_usd * _SUBAGENT_BUDGET_RESERVE
+    return budget_stopped, reserve_stopped
+
+
+def _crossed_stage(fraction: float, bands: tuple[float, ...]) -> int | None:
+    crossed: int | None = None
+    for index, band in enumerate(bands):
+        if fraction >= band:
+            crossed = index
+    return crossed
+
+
+_ROOT_DIRECTIVES: tuple[str, ...] = (
+    (
+        "As the root agent, begin planning your wind-down of the whole scan: avoid "
+        "starting large new lines of investigation, and keep your required objectives on "
+        "track so you can call finish_scan comfortably before the limit."
+    ),
+    (
+        "As the root agent, prioritize wrapping up the whole scan now: stop opening new "
+        "lines of investigation, close out only what is essential, and move toward calling "
+        "finish_scan to compile and deliver the final report."
+    ),
+    (
+        "As the root agent, STOP all other work on the whole scan and finish immediately: "
+        "secure your findings and call finish_scan now — anything left unfinished when the "
+        "limit is hit is discarded."
+    ),
+)
+_SUBAGENT_DIRECTIVES: tuple[str, ...] = (
+    (
+        "As a sub-agent, begin planning your wind-down: avoid starting large new subtasks, "
+        "and if you are close to a confirmed, validated vulnerability, drive it to a result "
+        "you can report."
+    ),
+    (
+        "As a sub-agent, prioritize wrapping up your task now: report any confirmed, "
+        "validated vulnerability, finish work that is nearly done rather than starting "
+        "anything new, and prepare to call agent_finish."
+    ),
+    (
+        "As a sub-agent, STOP all other work and finish immediately: report any confirmed "
+        "vulnerability right now and call agent_finish to hand your results back to your "
+        "parent before you are cut off."
+    ),
+)
+
+
+def _wrapup_directive(context: RunContextWrapper[dict[str, Any]], stage: int) -> str:
+    is_root = context.context.get("parent_id") is None
+    directives = _ROOT_DIRECTIVES if is_root else _SUBAGENT_DIRECTIVES
+    return directives[stage]
+
+
+def _urgency(stage: int) -> str:
+    return _STAGE_LABELS[stage]
+
+
+class ReportUsageHooks(RunHooks[dict[str, Any]]):
+    """Persist SDK-native usage and warn/stop as turn and cost budgets are consumed."""
+
+    def __init__(
+        self,
+        *,
+        model: str,
+        max_budget_usd: float | None = None,
+        max_turns: int | None = None,
+        interactive: bool = False,
+    ) -> None:
         if max_budget_usd is not None and (
             not math.isfinite(max_budget_usd) or max_budget_usd <= 0
         ):
             raise ValueError("max_budget_usd must be a finite number greater than 0")
+        if max_turns is not None and max_turns <= 0:
+            raise ValueError("max_turns must be a positive integer")
         self._model = model
         self._max_budget_usd = max_budget_usd
+        self._budget_increment = max_budget_usd
+        self._max_turns = max_turns
+        self._interactive = interactive
+
+    def extend_budget(self) -> None:
+        if self._max_budget_usd is None or self._budget_increment is None:
+            return
+        self._max_budget_usd += self._budget_increment
+
+    async def on_llm_start(
+        self,
+        context: RunContextWrapper[dict[str, Any]],
+        agent: Agent[dict[str, Any]],  # noqa: ARG002
+        system_prompt: str | None,  # noqa: ARG002
+        input_items: list[TResponseInputItem],
+    ) -> None:
+        try:
+            self._maybe_warn_turns(context, input_items)
+            self._maybe_warn_budget(context, input_items)
+        except Exception:
+            logger.exception("budget/turn warning injection failed")
+
+    def _maybe_warn_turns(
+        self,
+        context: RunContextWrapper[dict[str, Any]],
+        input_items: list[TResponseInputItem],
+    ) -> None:
+        if not self._max_turns:
+            return
+        usage = getattr(context, "usage", None)
+        requests = getattr(usage, "requests", None)
+        if not isinstance(requests, int):
+            return
+        turns_used = requests + 1
+        stage = _crossed_stage(turns_used / self._max_turns, _TURN_WARN_BANDS)
+        if stage is None:
+            return
+        remaining = max(self._max_turns - turns_used, 0)
+        pct = round(100 * turns_used / self._max_turns)
+        content = (
+            f"[{_urgency(stage)}] Turn budget: {turns_used}/{self._max_turns} used ({pct}%). "
+            f"About {remaining} turn(s) remain before this agent is force-stopped and any "
+            f"in-progress work is discarded. {_wrapup_directive(context, stage)}"
+        )
+        input_items.append({"role": "user", "content": content})
+
+    def _maybe_warn_budget(
+        self,
+        context: RunContextWrapper[dict[str, Any]],
+        input_items: list[TResponseInputItem],
+    ) -> None:
+        if self._max_budget_usd is None:
+            return
+        report_state = get_global_report_state()
+        if report_state is None:
+            return
+        cost = report_state.get_total_llm_cost()
+        is_root = context.context.get("parent_id") is None
+        if self._interactive:
+            bands = _ROOT_BUDGET_WARN_BANDS
+        else:
+            bands = _ROOT_BUDGET_WARN_BANDS if is_root else _SUBAGENT_BUDGET_WARN_BANDS
+        stage = _crossed_stage(cost / self._max_budget_usd, bands)
+        if stage is None:
+            return
+        pct = round(100 * cost / self._max_budget_usd)
+        reserve_pct = round(_SUBAGENT_BUDGET_RESERVE * 100)
+        if self._interactive:
+            content = (
+                f"[{_urgency(stage)}] Scan cost budget: ${cost:.2f}/${self._max_budget_usd:.2f} "
+                f"spent ({pct}%). This budget is shared across every agent in the scan; when it "
+                "is reached all agents are paused until the user chooses to continue. "
+                f"{_wrapup_directive(context, stage)}"
+            )
+        elif is_root:
+            content = (
+                f"[{_urgency(stage)}] Scan cost budget: ${cost:.2f}/${self._max_budget_usd:.2f} "
+                f"spent ({pct}%). This budget is shared across every agent in the scan; when it "
+                "is reached the whole scan is stopped immediately, and sub-agents are stopped at "
+                f"{reserve_pct}% to reserve the remainder for your final report. "
+                f"{_wrapup_directive(context, stage)}"
+            )
+        else:
+            content = (
+                f"[{_urgency(stage)}] Scan cost budget: ${cost:.2f}/${self._max_budget_usd:.2f} "
+                f"spent ({pct}%). This budget is shared across every agent in the scan; "
+                f"sub-agents are stopped at {reserve_pct}% to leave the remainder for the root "
+                f"agent's final report. {_wrapup_directive(context, stage)}"
+            )
+        input_items.append({"role": "user", "content": content})
 
     async def on_llm_end(
         self,
@@ -66,6 +250,21 @@ class ReportUsageHooks(RunHooks[dict[str, Any]]):
         if self._max_budget_usd is not None:
             cost = report_state.get_total_llm_cost()
             if cost >= self._max_budget_usd:
+                if self._interactive:
+                    raise BudgetPausedError(
+                        f"Scan budget of ${self._max_budget_usd:.2f} reached "
+                        f"(spent ${cost:.4f}); pausing until the user continues"
+                    )
                 raise BudgetExceededError(
                     f"Token budget of ${self._max_budget_usd:.2f} exceeded (spent ${cost:.4f})"
                 )
+            is_root = ctx.get("parent_id") is None
+            if not self._interactive and not is_root:
+                reserve_limit = self._max_budget_usd * _SUBAGENT_BUDGET_RESERVE
+                if cost >= reserve_limit:
+                    raise SubagentBudgetReservedError(
+                        f"Sub-agent budget reserve reached: spent ${cost:.4f} of "
+                        f"${self._max_budget_usd:.2f} "
+                        f"(>= {round(_SUBAGENT_BUDGET_RESERVE * 100)}% reserve); stopping this "
+                        "sub-agent so the root agent can finish the scan."
+                    )

strix/core/runner.py

@@ -31,7 +31,7 @@ from strix.core.execution import (
 from strix.core.execution import (
     spawn_child_agent as start_child_agent,
 )
-from strix.core.hooks import BudgetExceededError, ReportUsageHooks
+from strix.core.hooks import BudgetExceededError, ReportUsageHooks, recomputed_budget_flags
 from strix.core.inputs import (
     DEFAULT_MAX_TURNS,
     build_root_task,
@@ -40,6 +40,7 @@ from strix.core.inputs import (
 )
 from strix.core.paths import run_dir_for, runtime_state_dir
 from strix.core.sessions import open_agent_session
+from strix.report.state import get_global_report_state
 from strix.runtime import session_manager
 from strix.telemetry.logging import set_scan_id, setup_scan_logging
 from strix.tools.output_store import (
@@ -185,6 +186,18 @@ async def run_strix_scan(
                 f"Cannot resume scan {scan_id}: missing SDK session database at {agents_db}",
             )
         await coordinator.restore(snap)
+        report_state = get_global_report_state()
+        if report_state is not None:
+            budget_stopped, reserve_stopped = recomputed_budget_flags(
+                report_state.get_total_llm_cost(),
+                max_budget_usd,
+                interactive=interactive,
+            )
+            await coordinator.reset_budget_stops(
+                budget_stopped=budget_stopped,
+                reserve_stopped=reserve_stopped,
+                budget_paused=interactive and coordinator.budget_paused,
+            )
         for aid, parent in coordinator.parent_of.items():
             if parent is None:
                 root_id = aid
@@ -245,7 +258,14 @@ async def run_strix_scan(
             sandbox=SandboxRunConfig(client=bundle["client"], session=bundle["session"]),
             trace_include_sensitive_data=False,
         )
-        hooks = ReportUsageHooks(model=resolved_model, max_budget_usd=max_budget_usd)
+        hooks = ReportUsageHooks(
+            model=resolved_model,
+            max_budget_usd=max_budget_usd,
+            max_turns=max_turns,
+            interactive=interactive,
+        )
+        if interactive:
+            coordinator.set_budget_extender(hooks.extend_budget)
 
         scope_context = build_scope_context(scan_config)
         root_context = _merge_root_prompt_context(scope_context, extra_system_prompt_context)