deepstrike_core/memory/
runtime.rs

1use crate::memory::durable::{SessionData, SessionStore};
2use crate::memory::extractor::MemoryExtractor;
3use crate::memory::semantic::SemanticMemory;
4use crate::memory::session::{RestoreConfig, RestorePolicy, restore};
5use crate::types::message::Message;
6
7pub struct MemoryRuntime {
8    pub session_store: Box<dyn SessionStore>,
9    pub semantic_store: Option<Box<dyn SemanticMemory>>,
10    pub extractor: MemoryExtractor,
11    pub restore_policy: RestorePolicy,
12    pub restore_config: RestoreConfig,
13}
14
15impl MemoryRuntime {
16    pub fn new(
17        session_store: Box<dyn SessionStore>,
18        semantic_store: Option<Box<dyn SemanticMemory>>,
19        extractor: MemoryExtractor,
20        restore_policy: RestorePolicy,
21    ) -> Self {
22        Self {
23            session_store,
24            semantic_store,
25            extractor,
26            restore_policy,
27            restore_config: RestoreConfig::default(),
28        }
29    }
30
31    /// Load session history and apply restore policy to build initial context.
32    pub fn on_run_start(&self, session_id: &str, _goal: &str) -> Vec<Message> {
33        let history = self
34            .session_store
35            .load(session_id)
36            .ok()
37            .flatten()
38            .map(|d| d.messages)
39            .unwrap_or_default();
40        restore(&self.restore_policy, &self.restore_config, &history)
41    }
42
43    /// Extract memories from the current turn and store in semantic memory.
44    pub fn on_turn_end(&mut self, user_msg: &Message, assistant_msg: &Message) {
45        if let Some(sem) = &self.semantic_store {
46            for entry in self
47                .extractor
48                .extract(&[user_msg.clone(), assistant_msg.clone()])
49            {
50                let _ = sem.store(entry);
51            }
52        }
53    }
54
55    /// Optionally store a tool result in semantic memory.
56    pub fn on_tool_result(&mut self, _tool_name: &str, result: &str) {
57        if let Some(sem) = &self.semantic_store {
58            let _ = sem.store(crate::memory::semantic::MemoryEntry {
59                text: result.to_string(),
60                score: 0.0,
61                metadata: serde_json::Value::Null,
62            });
63        }
64    }
65
66    /// Extract memories from the run and persist the session.
67    /// `now_ms` is injected by the SDK layer — the kernel never reads system time.
68    pub fn on_run_end(
69        &mut self,
70        session_id: &str,
71        agent_id: &str,
72        messages: &[Message],
73        now_ms: u64,
74    ) {
75        let extracted = self.extractor.extract(messages);
76        if let Some(sem) = &self.semantic_store {
77            for entry in extracted {
78                let _ = sem.store(entry);
79            }
80        }
81        let data = SessionData {
82            session_id: session_id.to_string(),
83            agent_id: agent_id.to_string(),
84            messages: messages.to_vec(),
85            metadata: serde_json::Value::Null,
86            created_at_ms: now_ms,
87            updated_at_ms: now_ms,
88        };
89        let _ = self.session_store.save(session_id, &data);
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    use crate::memory::durable::InMemoryStore;
97    use crate::memory::extractor::ExtractionPolicy;
98
99    fn make_runtime() -> MemoryRuntime {
100        MemoryRuntime::new(
101            Box::new(InMemoryStore::new()),
102            None,
103            MemoryExtractor::new(ExtractionPolicy::default()),
104            RestorePolicy::Window,
105        )
106    }
107
108    #[test]
109    fn on_run_start_returns_empty_for_new_session() {
110        let rt = make_runtime();
111        let msgs = rt.on_run_start("new-session", "do something");
112        assert!(msgs.is_empty());
113    }
114
115    #[test]
116    fn on_run_end_persists_and_run_start_restores() {
117        let mut rt = make_runtime();
118        let messages = vec![Message::user("hello"), Message::assistant("a".repeat(101))];
119        rt.on_run_end("s1", "agent1", &messages, 1_000_000);
120        let restored = rt.on_run_start("s1", "continue");
121        assert!(!restored.is_empty());
122    }
123}
deepstrike_core/memory/runtime.rs

deepstrike_core/memory/
runtime.rs
