Restore sub-agents from checkpoint with full message history on resume

Previously sub-agents were terminated on Ctrl+C and never properly restored
— only the root agent resumed. This is the full fix.

Architecture change:
- checkpoint.py: Added sub_agent_states field (dict[agent_id -> AgentState
  dump]) saved from _agent_instances at every checkpoint write. Every
  currently-running non-root agent is captured.

- base_agent.py: Replaced fragile _is_root_resume heuristic with an explicit
  is_resumed flag (set via agent config). Works for both root and sub-agents.
  Prevents duplicate task message from being added to restored agents.

- cli.py / tui.py: Added _restore_sub_agents() which, on resume, iterates
  checkpoint sub_agent_states in topological order (parents before children),
  restores each agent's full AgentState, resets blocking flags, clears the
  old sandbox, injects a [SYSTEM - SUB-AGENT RESUMED] message, and spawns
  each agent in a daemon thread — identical to how the root agent is handled.
  Sub-agents are spawned BEFORE execute_scan so root agent can communicate
  with them immediately using their original IDs.

- Root agent's resume message now says "these sub-agents are ALREADY RUNNING
  at IDs [X, Y]" instead of "re-spawn them" — prevents double-spawning.

- agents_graph_actions.py: [SYSTEM - SUB-AGENT RESUMED] filtered from
  inherited context alongside [SYSTEM - SCAN RESUMED] so freshly-spawned
  child agents never see these system markers.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

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:

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):

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

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

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]")
+                    )
                 )
             ]