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

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

strix/interface/cli.py

@@ -13,6 +13,7 @@ from rich.panel import Panel
 from rich.text import Text
 
 from strix.config import load_settings
+from strix.core.inputs import DEFAULT_MAX_TURNS
 from strix.core.runner import run_strix_scan
 from strix.report.state import ReportState, set_global_report_state
 from strix.runtime import session_manager
@@ -184,6 +185,7 @@ async def run_cli(args: Any) -> None:  # noqa: PLR0915
                     local_sources=getattr(args, "local_sources", None) or [],
                     interactive=bool(getattr(args, "interactive", False)),
                     max_budget_usd=getattr(args, "max_budget_usd", None),
+                    max_turns=getattr(args, "max_turns", DEFAULT_MAX_TURNS),
                 )
             finally:
                 stop_updates.set()

strix/interface/main.py

@@ -31,6 +31,7 @@ from strix.config.models import (
     is_known_openai_bare_model,
     is_recommended_or_frontier_model,
 )
+from strix.core.inputs import DEFAULT_MAX_TURNS
 from strix.core.paths import run_dir_for, runtime_state_dir
 from strix.interface.cli import run_cli
 from strix.interface.tui import run_tui
@@ -481,6 +482,16 @@ def _positive_budget(value: str) -> float:
     return budget
 
 
+def _positive_int(value: str) -> int:
+    try:
+        parsed = int(value)
+    except ValueError as exc:
+        raise argparse.ArgumentTypeError(f"invalid int value: {value!r}") from exc
+    if parsed <= 0:
+        raise argparse.ArgumentTypeError("must be an integer greater than 0")
+    return parsed
+
+
 def parse_arguments() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="Strix Multi-Agent Cybersecurity Penetration Testing Tool",
@@ -636,10 +647,27 @@ Examples:
     )
 
     parser.add_argument(
-        "--max-budget-usd",
+        "--max-budget",
+        dest="max_budget_usd",
+        metavar="USD",
         type=_positive_budget,
         default=None,
-        help="Maximum LLM cost in USD (> 0). The scan stops cleanly when this limit is reached.",
+        help=(
+            "Maximum LLM cost in USD (> 0). The scan stops cleanly when this limit is reached. "
+            "Graduated wrap-up warnings are sent to all agents as it is approached."
+        ),
+    )
+
+    parser.add_argument(
+        "--max-turns",
+        dest="max_turns",
+        metavar="N",
+        type=_positive_int,
+        default=DEFAULT_MAX_TURNS,
+        help=(
+            "Maximum turns per agent (> 0, default %(default)s). Each agent is force-stopped "
+            "when it reaches this limit, with graduated wrap-up warnings as it is approached."
+        ),
     )
 
     parser.add_argument(

strix/interface/tui/app.py

@@ -34,6 +34,7 @@ from textual.widgets.tree import TreeNode
 from strix.config import load_settings
 from strix.config.models import is_recommended_or_frontier_model
 from strix.core.hooks import BudgetExceededError
+from strix.core.inputs import DEFAULT_MAX_TURNS
 from strix.core.runner import run_strix_scan
 from strix.interface.tui.live_view import TuiLiveView
 from strix.interface.tui.messages import send_user_message_to_agent
@@ -814,6 +815,7 @@ class StrixTUIApp(App):  # type: ignore[misc]
         self._scan_completed = threading.Event()
         self._scan_error: BaseException | None = None
         self._error_noted_agents: set[str] = set()
+        self._budget_pause_notified = False
 
         self._spinner_frame_index: int = 0
         self._sweep_num_squares: int = 6
@@ -1046,6 +1048,7 @@ class StrixTUIApp(App):  # type: ignore[misc]
                                 self.live_view.record_agent_error(agent_id, error)
                         else:
                             self._error_noted_agents.discard(agent_id)
+                    self._notify_budget_pause(statuses)
 
         if self._scan_loop is None or self._scan_loop.is_closed():
             return
@@ -1057,6 +1060,19 @@ class StrixTUIApp(App):  # type: ignore[misc]
 
         self._agent_graph_sync_future = asyncio.run_coroutine_threadsafe(collect(), self._scan_loop)
 
+    def _notify_budget_pause(self, statuses: dict[str, Any]) -> None:
+        paused = any(status == "budget_paused" for status in statuses.values())
+        if paused and not self._budget_pause_notified:
+            self._budget_pause_notified = True
+            self.notify(
+                "Budget limit reached \u2014 agents paused. Send a message to continue "
+                "(this extends the budget), or ctrl-q to quit.",
+                severity="warning",
+                timeout=15,
+            )
+        elif not paused:
+            self._budget_pause_notified = False
+
     def _update_agent_node(self, agent_id: str, agent_data: dict[str, Any]) -> bool:
         if agent_id not in self.agent_nodes:
             return False
@@ -1069,6 +1085,7 @@ class StrixTUIApp(App):  # type: ignore[misc]
             status_indicators = {
                 "running": "⚪",
                 "waiting": "⏸",
+                "budget_paused": "⏸",
                 "completed": "🟢",
                 "failed": "🔴",
                 "crashed": "🔴",
@@ -1266,10 +1283,17 @@ class StrixTUIApp(App):  # type: ignore[misc]
             self._stop_dot_animation()
             return (text, Text(), False)
 
-        if status == "waiting":
+        if status in {"waiting", "budget_paused"}:
             text = Text()
-            text.append("Send message to resume", style="dim")
-            return (text, Text(), False)
+            keymap = Text()
+            if status == "budget_paused":
+                text.append("Budget limit reached", style="yellow")
+                text.append(" \u00b7 ", style="dim")
+                text.append("Send a message to continue", style="dim")
+                keymap = keymap_styled([("ctrl-q", "quit")])
+            else:
+                text.append("Send message to resume", style="dim")
+            return (text, keymap, False)
 
         if status == "running":
             if self._agent_has_real_activity(agent_id):
@@ -1494,6 +1518,7 @@ class StrixTUIApp(App):  # type: ignore[misc]
                                 coordinator=self.coordinator,
                                 interactive=True,
                                 max_budget_usd=getattr(self.args, "max_budget_usd", None),
+                                max_turns=getattr(self.args, "max_turns", DEFAULT_MAX_TURNS),
                                 event_sink=self._capture_sdk_event,
                             ),
                         )
@@ -1501,10 +1526,7 @@ class StrixTUIApp(App):  # type: ignore[misc]
                 except (KeyboardInterrupt, asyncio.CancelledError):
                     logger.info("Scan interrupted by user")
                 except BudgetExceededError:
-                    # Defensive: the runner stops the scan cleanly on budget and
-                    # returns, so this normally never propagates. Treat it as a
-                    # graceful stop, not a scan error, if it ever does.
-                    logger.info("Scan stopped: --max-budget-usd limit reached")
+                    logger.info("Scan stopped: --max-budget limit reached")
                 except (ConnectionError, TimeoutError) as e:
                     logging.exception("Network error during scan")
                     self._scan_error = e
@@ -1559,6 +1581,7 @@ class StrixTUIApp(App):  # type: ignore[misc]
         status_indicators = {
             "running": "⚪",
             "waiting": "⏸",
+            "budget_paused": "⏸",
             "completed": "🟢",
             "failed": "🔴",
             "crashed": "🔴",
@@ -1605,6 +1628,7 @@ class StrixTUIApp(App):  # type: ignore[misc]
         status_indicators = {
             "running": "⚪",
             "waiting": "⏸",
+            "budget_paused": "⏸",
             "completed": "🟢",
             "failed": "🔴",
             "crashed": "🔴",
@@ -1729,7 +1753,10 @@ class StrixTUIApp(App):  # type: ignore[misc]
             message=message,
         )
         if not submitted:
-            self.notify("Scan loop is not ready; message was not sent", severity="warning")
+            if self._scan_completed.is_set():
+                self.notify("The scan has ended; message was not sent", severity="warning")
+            else:
+                self.notify("Scan loop is not ready; message was not sent", severity="warning")
             return
 
         self._displayed_events.clear()