diff --git a/strix/agents/base_agent.py b/strix/agents/base_agent.py index 2644000e..b1909e9f 100644 --- a/strix/agents/base_agent.py +++ b/strix/agents/base_agent.py @@ -82,6 +82,9 @@ class BaseAgent(metaclass=AgentMeta): self._checkpoint_manager = config.get("checkpoint_manager") self._scan_config: dict[str, Any] = config.get("scan_config", {}) self._target_hash: str = config.get("target_hash", "") + # True when this agent (root OR sub) is being restored from a checkpoint. + # Prevents _initialize_sandbox_and_state from adding a duplicate task message. + self._is_resumed: bool = bool(config.get("is_resumed", False)) with contextlib.suppress(Exception): self.llm.set_agent_identity(self.state.agent_name, self.state.agent_id) @@ -382,14 +385,10 @@ class BaseAgent(metaclass=AgentMeta): if not self.state.task: self.state.task = task - # Added for Resume Feature: only skip the task message when this is the - # ROOT agent being resumed (parent_id is None AND messages already has - # history from the checkpoint). - # Sub-agents can have pre-loaded context messages and still need their task - # message added — the old `if not self.state.messages` guard broke them. - # Original behavior is 100% unchanged for all non-resume paths. - _is_root_resume = (self.state.parent_id is None and bool(self.state.messages)) - if not _is_root_resume: + # Skip adding the task message when this agent (root or sub) is being + # restored from a checkpoint — the full message history including the + # task is already present. Fresh agents always get the task message. + if not self._is_resumed: self.state.add_message("user", task) async def _process_iteration(self, tracer: Optional["Tracer"]) -> bool | None: diff --git a/strix/interface/cli.py b/strix/interface/cli.py index aa249755..9a654664 100644 --- a/strix/interface/cli.py +++ b/strix/interface/cli.py @@ -97,51 +97,162 @@ def _replay_previous_output( console.print() -def _build_resume_context_message(state: Any, checkpoint_data: Any) -> None: - """Inject a user message telling the LLM it was interrupted and must continue. +def _restore_sub_agents(checkpoint_data: Any, llm_config: Any) -> list[str]: + """Spawn previously-running sub-agents from checkpoint with their full history. - Added for Resume Feature — prevents the model from calling finish_scan or - agent_finish just because the history ends abruptly, and explicitly lists - which sub-agents were alive so the LLM knows what to re-spawn. + Each sub-agent is restored exactly like the root agent: its saved AgentState + (messages, iteration, task) is loaded, blocking flags reset, sandbox cleared, + and a resume-context message injected. The agent is then started in a daemon + thread with is_resumed=True so it continues from where it left off. + + Returns a list of the agent_ids that were successfully restored. + """ + import asyncio + import threading + + from strix.agents.StrixAgent import StrixAgent + from strix.agents.state import AgentState + from strix.tools.agents_graph import agents_graph_actions + + sub_agent_states: dict[str, Any] = checkpoint_data.sub_agent_states or {} + if not sub_agent_states: + return [] + + # Topological sort so parents are created (and registered in _agent_graph) + # before their children. + def _depth(aid: str, _memo: dict = {}) -> int: # noqa: B006 + if aid in _memo: + return _memo[aid] + parent = sub_agent_states.get(aid, {}).get("parent_id") + _memo[aid] = 0 if (parent is None or parent not in sub_agent_states) else 1 + _depth(parent, _memo) + return _memo[aid] + + ordered_ids = sorted(sub_agent_states.keys(), key=_depth) + + restored_ids: list[str] = [] + for agent_id in ordered_ids: + state_dict = sub_agent_states[agent_id] + try: + state = AgentState.model_validate(state_dict) + + # Reset any blocking flags captured at the moment of interruption + state.waiting_for_input = False + state.waiting_start_time = None + state.stop_requested = False + state.completed = False + state.llm_failed = False + + # Clear old sandbox — it no longer exists + state.sandbox_id = None + state.sandbox_token = None + state.sandbox_info = None + + # Give a fresh iteration budget from the resume point + state.max_iterations = state.iteration + checkpoint_data.original_max_iterations + state.max_iterations_warning_sent = False + + # Inject a sub-agent resume message so the LLM knows to continue + state.add_message( + "user", + f"[SYSTEM - SUB-AGENT RESUMED]\n" + f"You were interrupted at iteration {state.iteration}. " + f"Your sandbox has been reset and a fresh one will be created. " + f"Review your conversation history above and continue your task " + f"from where you left off. " + f"Call agent_finish only when your task is genuinely complete.", + ) + + agent_cfg: dict[str, Any] = { + "llm_config": llm_config, + "state": state, + "is_resumed": True, + } + agent = StrixAgent(agent_cfg) + + def _run(a: Any = agent, s: Any = state) -> None: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + loop.run_until_complete(a.agent_loop(s.task)) + except Exception: # noqa: BLE001 + pass + finally: + loop.close() + + t = threading.Thread( + target=_run, + daemon=True, + name=f"ResumedAgent-{state.agent_name}-{agent_id[:8]}", + ) + t.start() + agents_graph_actions._running_agents[agent_id] = t + restored_ids.append(agent_id) + + except Exception: # noqa: BLE001 + pass + + return restored_ids + + +def _build_resume_context_message( + state: Any, + checkpoint_data: Any, + restored_ids: "list[str] | None" = None, +) -> None: + """Inject a user message telling the root agent what happened and what is live. + + If sub-agents were already restored automatically, the message tells the root + agent their IDs so it can communicate with them directly. Otherwise it lists + what was running so the root agent knows what to re-spawn. """ iteration = checkpoint_data.iteration + sub_agent_states: dict[str, Any] = checkpoint_data.sub_agent_states or {} - # Collect sub-agents that were NOT completed at the time of interruption - # so the LLM knows exactly what to re-create. - dead_sub_agents = [] - for agent_id, node in (checkpoint_data.tracer_agents or {}).items(): - if node.get("parent_id") is None: - continue # skip root agent - status = node.get("status", "unknown") - if status not in ("completed", "finished", "stopped", "error", "failed"): - dead_sub_agents.append({ - "name": node.get("name", "sub-agent"), - "task": (node.get("task") or "")[:300], - "status": status, - }) - - sub_agent_section = "" - if dead_sub_agents: + if restored_ids: lines = [ - "\n\nThe following sub-agents were ACTIVE at the time of interruption. " - "They no longer exist — their agent IDs are completely invalid. " - "Re-spawn each one if their work is not yet reflected in the findings above:" + "\n\nThe following sub-agents have been AUTOMATICALLY RESTORED and are " + "already running. Communicate with them using their original IDs:" ] - for sa in dead_sub_agents: - lines.append(f" • {sa['name']} (was doing: {sa['task']})") + for aid in restored_ids: + sd = sub_agent_states.get(aid, {}) + name = sd.get("agent_name", "sub-agent") + task = (sd.get("task") or "")[:200] + lines.append(f" • {name} (ID: {aid}) — task: {task}") + lines.append("\nDo NOT re-spawn these agents — they are already active.") sub_agent_section = "\n".join(lines) + else: + # No saved sub-agent states — fall back to listing from tracer_agents + dead: list[dict[str, Any]] = [] + for aid, node in (checkpoint_data.tracer_agents or {}).items(): + if node.get("parent_id") is None: + continue + status = node.get("status", "unknown") + if status not in ("completed", "finished", "stopped", "error", "failed"): + dead.append({ + "name": node.get("name", "sub-agent"), + "task": (node.get("task") or "")[:300], + }) + if dead: + lines = [ + "\n\nThe following sub-agents were active at interruption but could " + "not be automatically restored. Re-spawn them if their work is incomplete:" + ] + for sa in dead: + lines.append(f" • {sa['name']} (task: {sa['task']})") + sub_agent_section = "\n".join(lines) + else: + sub_agent_section = "" msg = ( f"[SYSTEM - SCAN RESUMED]\n" f"This penetration test was interrupted at iteration {iteration}. " - f"ALL previous sub-agents have been terminated and their agent IDs no longer exist in the graph. " f"A fresh sandbox will be created automatically.\n\n" - f"CRITICAL: Do NOT attempt to send_message_to_agent or interact with ANY agent ID " - f"that appears in the conversation history above — every one of those IDs is now dead. " - f"Call view_agent_graph to see the current graph (only you, the root agent, exist now)." + f"CRITICAL: Any agent IDs that appear in the conversation history ABOVE " + f"this message are from the old session and are DEAD — do not interact " + f"with them. Only the agents listed below are currently alive." f"{sub_agent_section}\n\n" - f"Review the history to understand what was done, then CONTINUE the penetration test. " - f"Re-spawn sub-agents for any incomplete work. " + f"Review the conversation history to understand what has already been done, " + f"then CONTINUE the penetration test. " f"Do NOT call finish_scan unless all testing is genuinely complete." ) state.add_message("user", msg) @@ -181,12 +292,7 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915 resumed_state.stop_requested = False resumed_state.completed = False resumed_state.llm_failed = False - - # Inject a resume-context message so the LLM does NOT call finish_scan - # or agent_finish just because the history ended abruptly. - # Without this the model sees a dangling tool call (sub-agent that was - # killed mid-execution) and may decide the task is complete or broken. - _build_resume_context_message(resumed_state, checkpoint_data) + # Resume message is injected AFTER sub-agents are restored (below) start_text = Text() if is_resuming: @@ -261,6 +367,7 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915 # Added for Resume Feature — pass restored state into the agent config if resumed_state is not None: agent_config["state"] = resumed_state + agent_config["is_resumed"] = True tracer = Tracer(args.run_name) tracer.set_scan_config(scan_config) @@ -382,6 +489,13 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915 try: agent = StrixAgent(agent_config) _agent_ref.append(agent) # expose to interrupt handler + + # Restore sub-agents THEN inject resume message so the root + # agent knows exactly which sub-agents are already running. + if is_resuming and checkpoint_data: + restored_ids = _restore_sub_agents(checkpoint_data, llm_config) + _build_resume_context_message(agent.state, checkpoint_data, restored_ids) + result = await agent.execute_scan(scan_config) if isinstance(result, dict) and not result.get("success", True): diff --git a/strix/interface/tui.py b/strix/interface/tui.py index 0981f9dd..50e1a784 100644 --- a/strix/interface/tui.py +++ b/strix/interface/tui.py @@ -41,53 +41,136 @@ from strix.telemetry.tracer import Tracer, set_global_tracer logger = logging.getLogger(__name__) -def _inject_resume_context_message(state: Any, checkpoint_data: Any) -> None: - """Inject a user message telling the LLM it was interrupted and must continue. +def _restore_sub_agents_tui(checkpoint_data: Any, llm_config: Any) -> list[str]: + """Restore previously-running sub-agents from checkpoint (TUI variant).""" + import asyncio + import threading - Added for Resume Feature — prevents the model from calling finish_scan or - agent_finish just because the history ends abruptly, and explicitly lists - which sub-agents were alive so the LLM knows what to re-spawn. - """ + from strix.agents.StrixAgent import StrixAgent + from strix.agents.state import AgentState + from strix.tools.agents_graph import agents_graph_actions + + sub_agent_states: dict[str, Any] = checkpoint_data.sub_agent_states or {} + if not sub_agent_states: + return [] + + def _depth(aid: str, _memo: dict = {}) -> int: # noqa: B006 + if aid in _memo: + return _memo[aid] + parent = sub_agent_states.get(aid, {}).get("parent_id") + _memo[aid] = 0 if (parent is None or parent not in sub_agent_states) else 1 + _depth(parent, _memo) + return _memo[aid] + + restored_ids: list[str] = [] + for agent_id in sorted(sub_agent_states.keys(), key=_depth): + state_dict = sub_agent_states[agent_id] + try: + state = AgentState.model_validate(state_dict) + state.waiting_for_input = False + state.waiting_start_time = None + state.stop_requested = False + state.completed = False + state.llm_failed = False + state.sandbox_id = None + state.sandbox_token = None + state.sandbox_info = None + state.max_iterations = state.iteration + checkpoint_data.original_max_iterations + state.max_iterations_warning_sent = False + state.add_message( + "user", + f"[SYSTEM - SUB-AGENT RESUMED]\n" + f"You were interrupted at iteration {state.iteration}. " + f"Your sandbox has been reset and a fresh one will be created. " + f"Review your conversation history above and continue your task " + f"from where you left off. " + f"Call agent_finish only when your task is genuinely complete.", + ) + agent_cfg: dict[str, Any] = { + "llm_config": llm_config, + "state": state, + "is_resumed": True, + } + agent = StrixAgent(agent_cfg) + + def _run(a: Any = agent, s: Any = state) -> None: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + loop.run_until_complete(a.agent_loop(s.task)) + except Exception: # noqa: BLE001 + pass + finally: + loop.close() + + t = threading.Thread( + target=_run, + daemon=True, + name=f"ResumedAgent-{state.agent_name}-{agent_id[:8]}", + ) + t.start() + agents_graph_actions._running_agents[agent_id] = t + restored_ids.append(agent_id) + except Exception: # noqa: BLE001 + pass + + return restored_ids + + +def _build_root_resume_message( + state: Any, + checkpoint_data: Any, + restored_ids: "list[str] | None" = None, +) -> None: + """Inject the root agent's resume context message after sub-agents are known.""" iteration = checkpoint_data.iteration + sub_agent_states: dict[str, Any] = checkpoint_data.sub_agent_states or {} - dead_sub_agents = [] - for agent_id, node in (checkpoint_data.tracer_agents or {}).items(): - if node.get("parent_id") is None: - continue - status = node.get("status", "unknown") - if status not in ("completed", "finished", "stopped", "error", "failed"): - dead_sub_agents.append({ - "name": node.get("name", "sub-agent"), - "task": (node.get("task") or "")[:300], - }) - - sub_agent_section = "" - if dead_sub_agents: + if restored_ids: lines = [ - "\n\nThe following sub-agents were ACTIVE at the time of interruption. " - "They no longer exist — their agent IDs are completely invalid. " - "Re-spawn each one if their work is not yet reflected in the findings above:" + "\n\nThe following sub-agents have been AUTOMATICALLY RESTORED and are " + "already running. Communicate with them using their original IDs:" ] - for sa in dead_sub_agents: - lines.append(f" • {sa['name']} (was doing: {sa['task']})") + for aid in restored_ids: + sd = sub_agent_states.get(aid, {}) + name = sd.get("agent_name", "sub-agent") + task = (sd.get("task") or "")[:200] + lines.append(f" \u2022 {name} (ID: {aid}) \u2014 task: {task}") + lines.append("\nDo NOT re-spawn these agents \u2014 they are already active.") sub_agent_section = "\n".join(lines) + else: + dead: list[dict[str, Any]] = [] + for aid, node in (checkpoint_data.tracer_agents or {}).items(): + if node.get("parent_id") is None: + continue + status = node.get("status", "unknown") + if status not in ("completed", "finished", "stopped", "error", "failed"): + dead.append({"name": node.get("name", "sub-agent"), + "task": (node.get("task") or "")[:300]}) + if dead: + lines = [ + "\n\nThe following sub-agents were active at interruption but could " + "not be automatically restored. Re-spawn them if their work is incomplete:" + ] + for sa in dead: + lines.append(f" \u2022 {sa['name']} (task: {sa['task']})") + sub_agent_section = "\n".join(lines) + else: + sub_agent_section = "" msg = ( f"[SYSTEM - SCAN RESUMED]\n" f"This penetration test was interrupted at iteration {iteration}. " - f"ALL previous sub-agents have been terminated and their agent IDs no longer exist in the graph. " f"A fresh sandbox will be created automatically.\n\n" - f"CRITICAL: Do NOT attempt to send_message_to_agent or interact with ANY agent ID " - f"that appears in the conversation history above — every one of those IDs is now dead. " - f"Call view_agent_graph to see the current graph (only you, the root agent, exist now)." + f"CRITICAL: Any agent IDs that appear in the conversation history ABOVE " + f"this message are from the old session and are DEAD. " + f"Only the agents listed below are currently alive." f"{sub_agent_section}\n\n" - f"Review the history to understand what was done, then CONTINUE the penetration test. " - f"Re-spawn sub-agents for any incomplete work. " + f"Review the conversation history to understand what has already been done, " + f"then CONTINUE the penetration test. " f"Do NOT call finish_scan unless all testing is genuinely complete." ) state.add_message("user", msg) - def get_package_version() -> str: try: return pkg_version("strix-agent") @@ -838,10 +921,9 @@ class StrixTUIApp(App): # type: ignore[misc] resumed_state.stop_requested = False resumed_state.completed = False resumed_state.llm_failed = False - # Inject resume-context message so the LLM does NOT call finish_scan - # or agent_finish just because the history ended abruptly. - _inject_resume_context_message(resumed_state, _cp) + # Resume message injected in _start_scan_thread after sub-agents restored config["state"] = resumed_state + config["is_resumed"] = True _mgr = getattr(args, "_checkpoint_manager", None) if _mgr: @@ -1588,6 +1670,14 @@ class StrixTUIApp(App): # type: ignore[misc] agent = StrixAgent(self.agent_config) self._current_agent = agent # expose for checkpoint on interrupt + # Restore sub-agents THEN inject the root resume message so + # the root agent knows exactly which sub-agents are running. + _cp = getattr(self.args, "_checkpoint_data", None) + if _cp and self.agent_config.get("is_resumed"): + _llm = self.agent_config.get("llm_config") + _restored = _restore_sub_agents_tui(_cp, _llm) + _build_root_resume_message(agent.state, _cp, _restored) + if not self._scan_stop_event.is_set(): loop.run_until_complete(agent.execute_scan(self.scan_config)) diff --git a/strix/telemetry/checkpoint.py b/strix/telemetry/checkpoint.py index 217da0ae..c8f76b50 100644 --- a/strix/telemetry/checkpoint.py +++ b/strix/telemetry/checkpoint.py @@ -51,6 +51,10 @@ class CheckpointModel(BaseModel): # Original scan configuration (passed to execute_scan) scan_config: dict[str, Any] = Field(default_factory=dict) + # Full AgentState dumps for every non-root, non-completed sub-agent. + # Keyed by agent_id so they can be restored with their original IDs. + sub_agent_states: dict[str, dict[str, Any]] = Field(default_factory=dict) + def compute_target_hash(targets_info: list[dict[str, Any]]) -> str: """Return a short stable hash of the target list for checkpoint validation. @@ -123,6 +127,21 @@ class CheckpointManager: tracer_tool_executions = {str(k): v for k, v in raw_execs.items()} tracer_next_execution_id = getattr(tracer, "_next_execution_id", 1) + # Capture full AgentState for every running sub-agent so they + # can be restored with their complete message history on resume. + sub_agent_states: dict[str, Any] = {} + try: + from strix.tools.agents_graph import agents_graph_actions + + for sid, inst in list(agents_graph_actions._agent_instances.items()): + s = getattr(inst, "state", None) + if s is not None and s.parent_id is not None: + sub_agent_states[sid] = ( + s.model_dump() if hasattr(s, "model_dump") else {} + ) + except Exception: # noqa: BLE001 + pass + checkpoint = CheckpointModel( run_name=self.run_name, target_hash=target_hash, @@ -135,6 +154,7 @@ class CheckpointManager: tracer_tool_executions=tracer_tool_executions, tracer_next_execution_id=tracer_next_execution_id, scan_config=scan_config, + sub_agent_states=sub_agent_states, ) # Atomic write: .tmp → rename diff --git a/strix/tools/agents_graph/agents_graph_actions.py b/strix/tools/agents_graph/agents_graph_actions.py index 02dff117..a5f5d51f 100644 --- a/strix/tools/agents_graph/agents_graph_actions.py +++ b/strix/tools/agents_graph/agents_graph_actions.py @@ -271,7 +271,10 @@ def create_agent( for msg in agent_state.get_conversation_history() if not ( isinstance(msg.get("content"), str) - and msg["content"].startswith("[SYSTEM - SCAN RESUMED]") + and ( + msg["content"].startswith("[SYSTEM - SCAN RESUMED]") + or msg["content"].startswith("[SYSTEM - SUB-AGENT RESUMED]") + ) ) ]