deepstrike_core/runtime/
snapshot.rs

1//! W2-2: First-class `KernelSnapshot` — live kernel state serialization.
2//!
3//! Enables:
4//! - Crash recovery (save state, restart, restore)
5//! - Agent migration (move running agent between hosts)
6//! - Checkpointing for long-running tasks
7//! - Cross-session state persistence
8//!
9//! The snapshot captures the essential kernel state needed to resume execution:
10//! - TaskTable (all TCBs with budget, wait state, proc info)
11//! - Context summary (messages, task state, signals)
12//! - Turn counter and budget totals
13//!
14//! NOT captured (recreated on restore):
15//! - SignalRouter dedup set (cleared on restore)
16//! - Governance pipeline closures (recreated from policy data)
17//! - MilestoneTracker state (recreated from config)
18//! - HandleTable (recreated from context history)
19
20use serde::{Deserialize, Serialize};
21
22use crate::scheduler::tcb::{TaskId, TaskState, Tcb, TaskTable};
23use crate::types::agent::{AgentIsolation, AgentRole, ContextInheritance};
24use crate::types::message::Message;
25
26/// Serializable snapshot of a TCB (Task Control Block).
27#[derive(Debug, Clone, Serialize, Deserialize)]
28pub struct TcbSnapshot {
29    pub id: TaskId,
30    pub parent: Option<TaskId>,
31    pub state: TaskState,
32    pub turns: u32,
33    pub total_tokens: u64,
34    pub started_at_ms: Option<u64>,
35    pub max_tokens: u32,
36    pub max_turns: u32,
37    pub max_total_tokens: u64,
38    pub max_wall_ms: Option<u64>,
39    pub wait_reason: Option<String>,  // Serialized as string label
40    pub wait_children: Option<Vec<TaskId>>,  // For SubAgentJoin
41    pub deferred_until: Option<u64>,
42    pub caps: Vec<TaskId>,
43    pub proc: Option<ProcInfoSnapshot>,
44}
45
46/// Snapshot of ProcInfo (sub-agent identity).
47#[derive(Debug, Clone, Serialize, Deserialize)]
48pub struct ProcInfoSnapshot {
49    pub parent_session_id: TaskId,
50    pub role: AgentRole,
51    pub isolation: AgentIsolation,
52    pub context_inheritance: ContextInheritance,
53    pub result: Option<ResultSnapshot>,
54}
55
56/// Snapshot of SubAgentResult.
57#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct ResultSnapshot {
59    pub termination: String,
60}
61
62/// K1: per-entry knowledge identity, parallel to `knowledge_messages` by index. Absent/short
63/// vectors (old snapshots) restore as unkeyed/unpinned — graceful, additive.
64#[derive(Debug, Clone, Serialize, Deserialize, Default)]
65pub struct KnowledgeEntryMeta {
66    #[serde(default, skip_serializing_if = "Option::is_none")]
67    pub key: Option<String>,
68    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
69    pub pinned: bool,
70}
71
72/// Snapshot of context state (simplified representation).
73#[derive(Debug, Clone, Serialize, Deserialize, Default)]
74pub struct ContextSnapshot {
75    pub system_messages: Vec<Message>,
76    pub knowledge_messages: Vec<Message>,
77    /// K1: identity metadata for `knowledge_messages` (index-parallel). `#[serde(default)]` keeps
78    /// pre-K1 snapshots loadable; pending upserts/eviction marks are deliberately NOT snapshotted
79    /// (graceful reset, same philosophy as `frozen_history_len`).
80    #[serde(default, skip_serializing_if = "Vec::is_empty")]
81    pub knowledge_entries_meta: Vec<KnowledgeEntryMeta>,
82    pub task_goal: Option<String>,
83    pub task_plan: Option<String>,
84    pub task_progress: Option<String>,
85    pub task_open_steps: Vec<String>,
86    /// Durable user directives — preserved across snapshot/restore like goal/plan.
87    #[serde(default)]
88    pub task_directives: Vec<String>,
89    pub history_messages: Vec<Message>,
90    pub signals: Vec<String>,
91    pub max_tokens: u32,
92    pub sprint: u32,
93}
94
95impl ContextSnapshot {
96    /// Create a snapshot from context manager state.
97    pub fn from_context(ctx: &crate::context::manager::ContextManager) -> Self {
98        // Convert plan steps to JSON string representation
99        let task_plan = if ctx.partitions.task_state.plan.is_empty() {
100            None
101        } else {
102            serde_json::to_string(&ctx.partitions.task_state.plan).ok()
103        };
104
105        Self {
106            system_messages: ctx.partitions.system.messages.clone(),
107            knowledge_messages: ctx.partitions.knowledge.messages().cloned().collect(),
108            knowledge_entries_meta: ctx
109                .partitions
110                .knowledge
111                .entries
112                .iter()
113                .map(|e| KnowledgeEntryMeta {
114                    key: e.key.as_ref().map(|k| k.to_string()),
115                    pinned: e.pinned,
116                })
117                .collect(),
118            task_goal: Some(ctx.partitions.task_state.goal.clone()),
119            task_plan,
120            task_progress: Some(ctx.partitions.task_state.progress.clone()),
121            task_open_steps: ctx.partitions.task_state.open_steps(),
122            task_directives: ctx.partitions.task_state.directives.clone(),
123            history_messages: ctx.partitions.history.messages.clone(),
124            signals: ctx.partitions.signals.clone(),
125            max_tokens: ctx.max_tokens,
126            sprint: ctx.sprint,
127        }
128    }
129}
130
131/// Full kernel snapshot.
132#[derive(Debug, Clone, Serialize, Deserialize)]
133pub struct KernelSnapshot {
134    pub turn: u32,
135    pub total_tokens: u64,
136    pub tasks: Vec<TcbSnapshot>,
137    pub context: ContextSnapshot,
138    pub run_spec: Option<String>,  // JSON-encoded AgentRunSpec
139}
140
141impl KernelSnapshot {
142    /// Create a snapshot from kernel state components.
143    pub fn from_state(
144        turn: u32,
145        total_tokens: u64,
146        tasks: &TaskTable,
147        context: &ContextSnapshot,
148        run_spec: Option<&crate::AgentRunSpec>,
149    ) -> Self {
150        Self {
151            turn,
152            total_tokens,
153            tasks: tasks.all().iter().map(TcbSnapshot::from).collect(),
154            context: context.clone(),
155            run_spec: run_spec.and_then(|s| serde_json::to_string(s).ok()),
156        }
157    }
158
159    /// Convert back to AgentRunSpec if present.
160    pub fn run_spec(&self) -> Option<crate::AgentRunSpec> {
161        self.run_spec.as_ref().and_then(|s| serde_json::from_str(s).ok())
162    }
163
164    /// W2-2: Convert to OsSnapshot-compatible process records for verification with
165    /// `rebuild_os_snapshot_from_events`. This enables cross-validation between the
166    /// live state snapshot and the event-log-derived audit view.
167    pub fn to_os_process_records(&self) -> Vec<crate::runtime::replay::ProcessRecord> {
168        use crate::runtime::replay::ProcessRecord;
169        let mut records = Vec::new();
170        for tcb_snap in &self.tasks {
171            // Only include tasks with ProcInfo (sub-agents), not the root task
172            if let Some(proc) = &tcb_snap.proc {
173                records.push(ProcessRecord {
174                    turn: tcb_snap.turns,
175                    agent_id: tcb_snap.id.to_string(),
176                    parent_session_id: proc.parent_session_id.to_string(),
177                    state: tcb_snap.state.label().to_string(),
178                });
179            }
180        }
181        records
182    }
183
184    /// W2-2: Restore TCB from snapshot. Returns None if the snapshot data is invalid.
185    pub fn restore_tcb(&self, snapshot: &TcbSnapshot) -> Option<Tcb> {
186        use crate::scheduler::policy::SchedulerBudget;
187
188        // Reconstruct BudgetLedger limits
189        let limits = SchedulerBudget {
190            max_tokens: snapshot.max_tokens,
191            max_turns: snapshot.max_turns,
192            max_total_tokens: snapshot.max_total_tokens,
193            max_wall_ms: snapshot.max_wall_ms,
194        };
195
196        // Reconstruct wait reason from label
197        let wait = snapshot.wait_reason.as_ref().and_then(|label| match label.as_str() {
198            "approval" => Some(crate::scheduler::tcb::WaitReason::Approval),
199            "sub_agent_join" => snapshot.wait_children.as_ref().map(|children| {
200                crate::scheduler::tcb::WaitReason::SubAgentJoin(
201                    children.iter().map(|id| id.clone().into()).collect()
202                )
203            }),
204            "tool" => Some(crate::scheduler::tcb::WaitReason::Tool),
205            "milestone" => Some(crate::scheduler::tcb::WaitReason::Milestone),
206            "signal" => Some(crate::scheduler::tcb::WaitReason::Signal),
207            "external" => Some(crate::scheduler::tcb::WaitReason::External),
208            _ => None,
209        });
210
211        // Reconstruct ProcInfo if present
212        let proc = snapshot.proc.as_ref().and_then(|p| {
213            let result = p.result.as_ref().and_then(|r| {
214                // Parse termination string back to TerminationReason
215                match r.termination.as_str() {
216                    "\"Completed\"" | "Completed" => Some(crate::types::result::SubAgentResult {
217                        agent_id: snapshot.id.clone(),
218                        result: crate::types::result::LoopResult {
219                            termination: crate::types::result::TerminationReason::Completed,
220                            final_message: None,
221                            turns_used: 0,
222                            total_tokens_used: 0,
223                            loop_continue: None,
224                            classify_branch: None,
225                            tournament_winner: None,
226                        },
227                    }),
228                    _ => None,
229                }
230            });
231
232            Some(crate::scheduler::tcb::ProcInfo {
233                parent_session_id: p.parent_session_id.clone(),
234                role: p.role,
235                isolation: p.isolation,
236                context_inheritance: p.context_inheritance,
237                result,
238            })
239        });
240
241        Some(Tcb {
242            id: snapshot.id.clone(),
243            parent: snapshot.parent.clone(),
244            state: snapshot.state,
245            budget: crate::scheduler::tcb::BudgetLedger {
246                limits,
247                turns: snapshot.turns,
248                total_tokens: snapshot.total_tokens,
249                started_at_ms: snapshot.started_at_ms,
250            },
251            wait,
252            caps: snapshot.caps.clone(),
253            proc,
254            deferred_until: snapshot.deferred_until,
255        })
256    }
257}
258
259impl From<&Tcb> for TcbSnapshot {
260    fn from(tcb: &Tcb) -> Self {
261        Self {
262            id: tcb.id.clone(),
263            parent: tcb.parent.clone(),
264            state: tcb.state.clone(),
265            turns: tcb.budget.turns,
266            total_tokens: tcb.budget.total_tokens,
267            started_at_ms: tcb.budget.started_at_ms,
268            max_tokens: tcb.budget.limits.max_tokens,
269            max_turns: tcb.budget.limits.max_turns,
270            max_total_tokens: tcb.budget.limits.max_total_tokens,
271            max_wall_ms: tcb.budget.limits.max_wall_ms,
272            wait_reason: tcb.wait.as_ref().map(|w| w.label().to_string()),
273            wait_children: match &tcb.wait {
274                Some(crate::scheduler::tcb::WaitReason::SubAgentJoin(children)) => {
275                    Some(children.clone())
276                }
277                _ => None,
278            },
279            deferred_until: tcb.deferred_until,
280            caps: tcb.caps.clone(),
281            proc: tcb.proc.as_ref().map(|p| ProcInfoSnapshot {
282                parent_session_id: p.parent_session_id.clone(),
283                role: p.role,
284                isolation: p.isolation,
285                context_inheritance: p.context_inheritance,
286                result: p.result.as_ref().map(|r| ResultSnapshot {
287                    termination: format!("{:?}", r.result.termination),
288                }),
289            }),
290        }
291    }
292}
293
294#[cfg(test)]
295mod tests {
296    use super::*;
297    use crate::scheduler::policy::SchedulerBudget;
298
299    #[test]
300    fn tcb_snapshot_roundtrip() {
301        let mut tcb = Tcb::root("test-task", SchedulerBudget {
302            max_tokens: 128_000,
303            max_turns: 10,
304            max_total_tokens: 1000,
305            max_wall_ms: Some(60000),
306        });
307        tcb.budget.turns = 5;
308        tcb.budget.total_tokens = 500;
309        tcb.deferred_until = Some(1000);
310
311        let snapshot = TcbSnapshot::from(&tcb);
312        assert_eq!(snapshot.id.as_str(), "test-task");
313        assert_eq!(snapshot.turns, 5);
314        assert_eq!(snapshot.total_tokens, 500);
315        assert_eq!(snapshot.deferred_until, Some(1000));
316    }
317
318    #[test]
319    fn kernel_snapshot_serializes() {
320        let snap = KernelSnapshot {
321            turn: 1,
322            total_tokens: 100,
323            tasks: vec![],
324            context: ContextSnapshot::default(),
325            run_spec: None,
326        };
327
328        let json = serde_json::to_string(&snap).expect("serialize");
329        let restored: KernelSnapshot = serde_json::from_str(&json).expect("deserialize");
330
331        assert_eq!(restored.turn, 1);
332        assert_eq!(restored.total_tokens, 100);
333    }
334
335    #[test]
336    fn context_snapshot_captures_fields() {
337        let ctx = ContextSnapshot {
338            system_messages: vec![Message::system("You are helpful")],
339            task_goal: Some("Build something".to_string()),
340            ..Default::default()
341        };
342
343        assert_eq!(ctx.system_messages.len(), 1);
344        assert_eq!(ctx.task_goal.as_deref(), Some("Build something"));
345    }
346
347    #[test]
348    fn snapshot_from_state_captures_tasks() {
349        let mut table = TaskTable::new();
350        table.insert(Tcb::root("root", SchedulerBudget::default()));
351        table.insert(Tcb::root("child", SchedulerBudget::default()));
352
353        let ctx = ContextSnapshot::default();
354        let snap = KernelSnapshot::from_state(5, 1000, &table, &ctx, None);
355
356        assert_eq!(snap.turn, 5);
357        assert_eq!(snap.total_tokens, 1000);
358        assert_eq!(snap.tasks.len(), 2);
359    }
360
361    #[test]
362    fn context_snapshot_from_manager() {
363        use crate::context::manager::ContextManager;
364        use crate::types::message::Message;
365
366        let mut ctx = ContextManager::new(1000);
367        ctx.partitions.system.push(Message::system("You are helpful"), 10);
368        ctx.partitions.task_state.goal = "Test goal".to_string();
369
370        let snap = ContextSnapshot::from_context(&ctx);
371        assert_eq!(snap.system_messages.len(), 1);
372        assert_eq!(snap.task_goal.as_deref(), Some("Test goal"));
373        assert_eq!(snap.max_tokens, 1000);
374        // Empty plan becomes None
375        assert!(snap.task_plan.is_none());
376    }
377
378    #[test]
379    fn tcb_restore_roundtrip() {
380        let original = Tcb::root("test-task", SchedulerBudget {
381            max_tokens: 128_000,
382            max_turns: 10,
383            max_total_tokens: 1000,
384            max_wall_ms: Some(60000),
385        });
386
387        let snap = TcbSnapshot::from(&original);
388        let kernel_snap = KernelSnapshot {
389            turn: 1,
390            total_tokens: 100,
391            tasks: vec![snap.clone()],
392            context: ContextSnapshot::default(),
393            run_spec: None,
394        };
395
396        let restored = kernel_snap.restore_tcb(&snap);
397        assert!(restored.is_some());
398
399        let tcb = restored.unwrap();
400        assert_eq!(tcb.id.as_str(), "test-task");
401        assert_eq!(tcb.state, original.state);
402        assert_eq!(tcb.budget.turns, original.budget.turns);
403    }
404
405    #[test]
406    fn tcb_restore_with_wait_reason() {
407        let mut tcb = Tcb::root("waiting-task", SchedulerBudget::default());
408        tcb.wait = Some(crate::scheduler::tcb::WaitReason::SubAgentJoin(
409            vec!["child-1".into(), "child-2".into()]
410        ));
411
412        let snap = TcbSnapshot::from(&tcb);
413        let kernel_snap = KernelSnapshot {
414            turn: 1,
415            total_tokens: 100,
416            tasks: vec![snap.clone()],
417            context: ContextSnapshot::default(),
418            run_spec: None,
419        };
420
421        let restored = kernel_snap.restore_tcb(&snap).expect("restore should succeed");
422        match restored.wait {
423            Some(crate::scheduler::tcb::WaitReason::SubAgentJoin(children)) => {
424                assert_eq!(children.len(), 2);
425                assert_eq!(children[0].as_str(), "child-1");
426                assert_eq!(children[1].as_str(), "child-2");
427            }
428            other => panic!("Expected SubAgentJoin, got {:?}", other),
429        }
430    }
431
432    #[test]
433    fn kernel_snapshot_to_os_process_records() {
434        // Create a snapshot with root + sub-agent tasks
435        let snap = KernelSnapshot {
436            turn: 5,
437            total_tokens: 1000,
438            tasks: vec![
439                // Root task (no ProcInfo - should be filtered out)
440                TcbSnapshot {
441                    id: "root".into(),
442                    parent: None,
443                    state: TaskState::Running,
444                    turns: 5,
445                    total_tokens: 500,
446                    started_at_ms: Some(0),
447                    max_tokens: 128_000,
448                    max_turns: 100,
449                    max_total_tokens: 1_000_000,
450                    max_wall_ms: None,
451                    wait_reason: None,
452                    wait_children: None,
453                    deferred_until: None,
454                    caps: vec![],
455                    proc: None,
456                },
457                // Sub-agent task (has ProcInfo - should be included)
458                TcbSnapshot {
459                    id: "child-1".into(),
460                    parent: Some("root".into()),
461                    state: TaskState::Done(crate::types::result::TerminationReason::Completed),
462                    turns: 3,
463                    total_tokens: 300,
464                    started_at_ms: Some(100),
465                    max_tokens: 64_000,
466                    max_turns: 50,
467                    max_total_tokens: 500_000,
468                    max_wall_ms: None,
469                    wait_reason: None,
470                    wait_children: None,
471                    deferred_until: None,
472                    caps: vec![],
473                    proc: Some(ProcInfoSnapshot {
474                        parent_session_id: "root".into(),
475                        role: AgentRole::Implement,
476                        isolation: AgentIsolation::Shared,
477                        context_inheritance: ContextInheritance::None,
478                        result: None,
479                    }),
480                },
481            ],
482            context: ContextSnapshot::default(),
483            run_spec: None,
484        };
485
486        let records = snap.to_os_process_records();
487        assert_eq!(records.len(), 1); // Root task filtered out
488        assert_eq!(records[0].agent_id, "child-1");
489        assert_eq!(records[0].parent_session_id, "root");
490        assert_eq!(records[0].turn, 3);
491        assert_eq!(records[0].state, "done"); // Done state -> "done" label
492    }
493
494    #[test]
495    fn kernel_snapshot_to_os_records_matches_state_machine() {
496        use crate::scheduler::state_machine::LoopStateMachine;
497        use crate::scheduler::policy::LoopPolicy;
498        use crate::types::agent::{AgentIdentity, AgentRole, AgentRunSpec};
499
500        // Create a state machine and spawn a sub-agent
501        let mut sm = LoopStateMachine::new(LoopPolicy {
502            max_tokens: 128_000,
503            ..Default::default()
504        });
505        sm.start(crate::types::task::RuntimeTask::new("parent task"));
506
507        // Spawn sub-agent
508        let _ = sm.spawn_sub_agent(
509            AgentRunSpec::new(
510                AgentIdentity::sub_agent("child", "child-session"),
511                AgentRole::Implement,
512                "child task",
513            ),
514            "parent-sess",
515        );
516
517        // Take snapshot and convert to OsSnapshot records
518        let snap = sm.snapshot();
519        let records = snap.to_os_process_records();
520
521        // Should have exactly one sub-agent record
522        assert_eq!(records.len(), 1);
523        assert_eq!(records[0].agent_id, "child");
524        assert_eq!(records[0].parent_session_id, "parent-sess");
525        // State should be "running" or "suspended" (depends on when snapshot was taken)
526        assert!(
527            records[0].state == "running" || records[0].state == "suspended",
528            "unexpected state: {}",
529            records[0].state
530        );
531    }
532
533    #[test]
534    fn pre_k1_snapshot_without_entries_meta_still_loads() {
535        // K1 back-compat: a snapshot serialized before `knowledge_entries_meta` existed must
536        // deserialize (serde default) and restore every knowledge entry unkeyed/unpinned.
537        let json = serde_json::json!({
538            "system_messages": [],
539            "knowledge_messages": [
540                { "role": "system", "content": "legacy knowledge" }
541            ],
542            "task_goal": "g",
543            "task_plan": null,
544            "task_progress": "p",
545            "task_open_steps": [],
546            "history_messages": [],
547            "signals": [],
548            "max_tokens": 1000,
549            "sprint": 0
550        });
551        let snap: ContextSnapshot = serde_json::from_value(json).expect("old snapshot loads");
552        assert!(snap.knowledge_entries_meta.is_empty());
553
554        let kernel_snap = KernelSnapshot {
555            turn: 0,
556            total_tokens: 0,
557            tasks: vec![],
558            context: snap,
559            run_spec: None,
560        };
561        let sm = crate::scheduler::state_machine::LoopStateMachine::restore(&kernel_snap);
562        assert_eq!(sm.ctx.partitions.knowledge.len(), 1);
563        let entry = &sm.ctx.partitions.knowledge.entries[0];
564        assert!(entry.key.is_none());
565        assert!(!entry.pinned);
566    }
567
568    #[test]
569    fn keyed_pinned_knowledge_round_trips_through_snapshot() {
570        // K1: key + pinned survive snapshot → restore (index-parallel meta vec).
571        let mut ctx = crate::context::manager::ContextManager::new(1000);
572        ctx.push_knowledge_entry(Some("ref".into()), Message::system("keyed doc"), 5, true);
573        ctx.push_knowledge(Message::system("legacy"), 3);
574
575        let context_snap = ContextSnapshot::from_context(&ctx);
576        assert_eq!(context_snap.knowledge_entries_meta.len(), 2);
577
578        let kernel_snap = KernelSnapshot {
579            turn: 0,
580            total_tokens: 0,
581            tasks: vec![],
582            context: context_snap,
583            run_spec: None,
584        };
585        let sm = crate::scheduler::state_machine::LoopStateMachine::restore(&kernel_snap);
586        let entries = &sm.ctx.partitions.knowledge.entries;
587        assert_eq!(entries.len(), 2);
588        assert_eq!(entries[0].key.as_deref(), Some("ref"));
589        assert!(entries[0].pinned);
590        assert!(entries[1].key.is_none());
591        assert!(!entries[1].pinned);
592    }
593}
deepstrike_core/runtime/snapshot.rs

deepstrike_core/runtime/
snapshot.rs
