From fd5366f61229cee72d34c9015502b367a2347a21 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 19 Mar 2026 08:26:17 +0100 Subject: [PATCH] fix: save and restore sub-agent state in checkpoints MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously checkpoints only saved tracer.chat_messages and tracer.vulnerability_reports, leaving tracer.agents and tracer.tool_executions empty on resume — so all sub-agents (both in-progress and completed) were invisible after resuming. Changes: - checkpoint.py: add tracer_agents, tracer_tool_executions, tracer_next_execution_id fields to CheckpointModel; populate them in CheckpointManager.save() from the live tracer - cli.py: on resume, restore agents dict, tool_executions dict, and advance _next_execution_id to avoid ID collisions - tui.py: same restore logic so TUI sidebar shows all previous agents and their tool results immediately on resume Co-Authored-By: Claude Sonnet 4.6 --- strix/interface/cli.py | 9 +++++++++ strix/interface/tui.py | 10 ++++++++++ strix/telemetry/checkpoint.py | 20 ++++++++++++++++++++ 3 files changed, 39 insertions(+) diff --git a/strix/interface/cli.py b/strix/interface/cli.py index f370b919..16d27da3 100644 --- a/strix/interface/cli.py +++ b/strix/interface/cli.py @@ -204,9 +204,18 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915 tracer.set_scan_config(scan_config) # Added for Resume Feature — pre-populate tracer so stats/vulns are correct + # Also restores sub-agent registry and tool executions so the full run is visible if is_resuming and checkpoint_data: tracer.chat_messages.extend(checkpoint_data.tracer_chat_messages) tracer.vulnerability_reports.extend(checkpoint_data.tracer_vulnerability_reports) + # Restore every agent (root + sub-agents) with their last-known status + tracer.agents.update(checkpoint_data.tracer_agents) + # Restore tool execution records; keys were serialised as str, restore as int + for k, v in checkpoint_data.tracer_tool_executions.items(): + tracer.tool_executions[int(k)] = v + # Advance the ID counter so new executions don't overwrite saved ones + if checkpoint_data.tracer_next_execution_id > tracer._next_execution_id: + tracer._next_execution_id = checkpoint_data.tracer_next_execution_id # Added for Resume Feature — show resume banner + replay previous output if is_resuming and checkpoint_data: diff --git a/strix/interface/tui.py b/strix/interface/tui.py index b20055c0..d99ec3c4 100644 --- a/strix/interface/tui.py +++ b/strix/interface/tui.py @@ -711,10 +711,20 @@ class StrixTUIApp(App): # type: ignore[misc] # Added for Resume Feature — pre-populate tracer with checkpoint data so # stats and findings reflect the full scan history including past sessions. + # Also restores sub-agents and their tool executions so the TUI sidebar + # shows every agent (completed or in-progress) from the previous session. _cp = getattr(args, "_checkpoint_data", None) if _cp and getattr(args, "resume_from_checkpoint", False): self.tracer.chat_messages.extend(_cp.tracer_chat_messages) self.tracer.vulnerability_reports.extend(_cp.tracer_vulnerability_reports) + # Restore full agent registry (root + all sub-agents) + self.tracer.agents.update(_cp.tracer_agents) + # Restore tool execution records (keys were serialised as str) + for k, v in _cp.tracer_tool_executions.items(): + self.tracer.tool_executions[int(k)] = v + # Advance execution ID counter to avoid collisions + if _cp.tracer_next_execution_id > self.tracer._next_execution_id: + self.tracer._next_execution_id = _cp.tracer_next_execution_id self.agent_nodes: dict[str, TreeNode] = {} diff --git a/strix/telemetry/checkpoint.py b/strix/telemetry/checkpoint.py index a3725b98..217da0ae 100644 --- a/strix/telemetry/checkpoint.py +++ b/strix/telemetry/checkpoint.py @@ -41,6 +41,13 @@ class CheckpointModel(BaseModel): tracer_chat_messages: list[dict[str, Any]] = Field(default_factory=list) tracer_vulnerability_reports: list[dict[str, Any]] = Field(default_factory=list) + # Added for sub-agent persistence: full agent registry and tool execution log + # tracer.agents → every agent (root + sub) with status, name, task, parent_id + # tracer.tool_executions → every tool call result across all agents + tracer_agents: dict[str, Any] = Field(default_factory=dict) + tracer_tool_executions: dict[str, Any] = Field(default_factory=dict) # key is str(int) + tracer_next_execution_id: int = 1 # restore the ID counter so new IDs don't collide + # Original scan configuration (passed to execute_scan) scan_config: dict[str, Any] = Field(default_factory=dict) @@ -100,11 +107,21 @@ class CheckpointManager: tracer_chat_messages: list[dict[str, Any]] = [] tracer_vulnerability_reports: list[dict[str, Any]] = [] + tracer_agents: dict[str, Any] = {} + tracer_tool_executions: dict[str, Any] = {} + tracer_next_execution_id: int = 1 if tracer: tracer_chat_messages = list(getattr(tracer, "chat_messages", [])) tracer_vulnerability_reports = list( getattr(tracer, "vulnerability_reports", []) ) + # Added for sub-agent persistence — capture full agent registry + # and all tool execution records so sub-agents are fully restored + tracer_agents = dict(getattr(tracer, "agents", {})) + # tool_executions keys are ints; serialise as strings for JSON + raw_execs = getattr(tracer, "tool_executions", {}) + tracer_tool_executions = {str(k): v for k, v in raw_execs.items()} + tracer_next_execution_id = getattr(tracer, "_next_execution_id", 1) checkpoint = CheckpointModel( run_name=self.run_name, @@ -114,6 +131,9 @@ class CheckpointManager: agent_state=state_dict, tracer_chat_messages=tracer_chat_messages, tracer_vulnerability_reports=tracer_vulnerability_reports, + tracer_agents=tracer_agents, + tracer_tool_executions=tracer_tool_executions, + tracer_next_execution_id=tracer_next_execution_id, scan_config=scan_config, )