deepstrike_sdk/
memory.rs

1use async_trait::async_trait;
2use deepstrike_core::memory::durable::SessionData;
3use deepstrike_core::mm::memory::{
4    MemoryAuthor, MemoryKind, MemoryProvenance, MemoryQuery, MemoryRecall, MemoryRecord,
5    MemoryScope, MemoryTrustLevel,
6};
7
8/// Durable-memory host storage. `upsert` is the only mutation and is called only after the
9/// kernel's `WriteMemory` gate accepts the record.
10#[async_trait]
11pub trait DreamStore: Send + Sync {
12    async fn upsert(&self, agent_id: &str, record: MemoryRecord) -> crate::Result<()>;
13
14    /// Semantic search over the agent's long-term memories.
15    /// Called on demand during a session when the LLM invokes the `memory` meta-tool.
16    async fn search(
17        &self,
18        agent_id: &str,
19        query: &MemoryQuery,
20    ) -> crate::Result<Vec<MemoryRecall>>;
21
22    /// Persist a completed session before the runner's one extraction pass.
23    async fn save_session(
24        &self,
25        data: deepstrike_core::memory::durable::SessionData,
26    ) -> crate::Result<()>;
27}
28
29pub(crate) fn parse_extracted_memories(
30    output: &str,
31    session: &SessionData,
32    scope: &MemoryScope,
33) -> Vec<MemoryRecord> {
34    let cleaned = output
35        .trim()
36        .strip_prefix("```json")
37        .or_else(|| output.trim().strip_prefix("```"))
38        .unwrap_or(output.trim())
39        .strip_suffix("```")
40        .unwrap_or(output.trim())
41        .trim();
42    let Ok(value) = serde_json::from_str::<serde_json::Value>(cleaned) else {
43        return Vec::new();
44    };
45    let Some(drafts) = value.get("memories").and_then(serde_json::Value::as_array) else {
46        return Vec::new();
47    };
48    drafts
49        .iter()
50        .take(10)
51        .filter_map(|draft| {
52            let name = draft.get("name")?.as_str()?.trim();
53            let content = draft.get("content")?.as_str()?.trim();
54            if name.is_empty() || content.is_empty() {
55                return None;
56            }
57            let kind = match draft.get("kind")?.as_str()? {
58                "user" => MemoryKind::User,
59                "feedback" => MemoryKind::Feedback,
60                "project" => MemoryKind::Project,
61                "reference" => MemoryKind::Reference,
62                _ => return None,
63            };
64            let confidence = draft
65                .get("confidence")
66                .and_then(serde_json::Value::as_f64)
67                .unwrap_or(0.5)
68                .clamp(0.0, 1.0);
69            let strings = |field: &str| {
70                draft
71                    .get(field)
72                    .and_then(serde_json::Value::as_array)
73                    .map(|values| {
74                        values
75                            .iter()
76                            .filter_map(serde_json::Value::as_str)
77                            .map(str::to_string)
78                            .collect()
79                    })
80                    .unwrap_or_default()
81            };
82            Some(MemoryRecord {
83                record_id: format!("{}:{}:{}:{name}", scope.tenant_id, scope.namespace, kind.label()),
84                scope: scope.clone(),
85                name: name.to_string(),
86                kind,
87                content: content.to_string(),
88                description: draft
89                    .get("description")
90                    .and_then(serde_json::Value::as_str)
91                    .unwrap_or_default()
92                    .trim()
93                    .to_string(),
94                provenance: MemoryProvenance {
95                    session_id: Some(session.session_id.clone()),
96                    author: MemoryAuthor::Extraction,
97                    trust: MemoryTrustLevel::Untrusted,
98                    evidence_refs: strings("evidence_refs"),
99                },
100                created_at: session.updated_at_ms,
101                updated_at: session.updated_at_ms,
102                last_recalled_at: None,
103                recall_count: 0,
104                confidence,
105                links: strings("links"),
106                pinned: draft.get("pinned").and_then(serde_json::Value::as_bool).unwrap_or(false),
107                ttl_days: draft
108                    .get("ttl_days")
109                    .and_then(serde_json::Value::as_u64)
110                    .and_then(|days| u32::try_from(days).ok())
111                    .filter(|days| *days > 0),
112            })
113        })
114        .collect()
115}
116
117/// In-process scratch pad for within-run state.
118#[derive(Default)]
119pub struct WorkingMemory {
120    store: std::collections::HashMap<String, serde_json::Value>,
121}
122
123impl WorkingMemory {
124    pub fn set(&mut self, key: impl Into<String>, value: impl Into<serde_json::Value>) {
125        self.store.insert(key.into(), value.into());
126    }
127    pub fn get(&self, key: &str) -> Option<&serde_json::Value> {
128        self.store.get(key)
129    }
130    pub fn clear(&mut self) {
131        self.store.clear();
132    }
133}
134
135/// `InMemoryDreamStore` — a lightweight `DreamStore` backed by per-agent in-memory maps.
136///
137/// Rust port of node/src/memory/in-memory-store.ts. Use for benchmarks, unit tests, and local
138/// development where persistent memory isn't needed. `search()` is a deterministic reference
139/// ranker: distinct lexical overlap first, metadata recency second, insertion order last.
140pub struct InMemoryDreamStore {
141    memories: std::sync::Mutex<std::collections::HashMap<String, Vec<MemoryRecord>>>,
142    initial_memories: Vec<MemoryRecord>,
143    saved_sessions: std::sync::Mutex<Vec<SessionData>>,
144}
145
146impl InMemoryDreamStore {
147    pub fn new() -> Self {
148        Self::with_initial_memories(Vec::new())
149    }
150
151    pub fn with_initial_memories(initial: Vec<MemoryRecord>) -> Self {
152        Self {
153            memories: std::sync::Mutex::new(std::collections::HashMap::new()),
154            initial_memories: initial,
155            saved_sessions: std::sync::Mutex::new(Vec::new()),
156        }
157    }
158
159    pub fn saved_sessions(&self) -> Vec<SessionData> {
160        self.saved_sessions.lock().unwrap().clone()
161    }
162}
163
164impl Default for InMemoryDreamStore {
165    fn default() -> Self {
166        Self::new()
167    }
168}
169
170#[async_trait]
171impl DreamStore for InMemoryDreamStore {
172    async fn upsert(&self, agent_id: &str, incoming: MemoryRecord) -> crate::Result<()> {
173        let mut memories = self.memories.lock().unwrap();
174        let kept = memories
175            .entry(agent_id.to_string())
176            .or_insert_with(|| self.initial_memories.clone());
177        if let Some(index) = kept.iter().position(|record| {
178                record.scope == incoming.scope
179                    && record.kind == incoming.kind
180                    && record.name == incoming.name
181            }) {
182            kept[index] = incoming;
183        } else {
184            kept.push(incoming);
185        }
186        Ok(())
187    }
188
189    async fn search(
190        &self,
191        agent_id: &str,
192        query: &MemoryQuery,
193    ) -> crate::Result<Vec<MemoryRecall>> {
194        let all = {
195            let mut memories = self.memories.lock().unwrap();
196            memories
197                .entry(agent_id.to_string())
198                .or_insert_with(|| self.initial_memories.clone())
199                .clone()
200        };
201        let query_terms = memory_terms(&query.query);
202        let mut ranked = all
203            .into_iter()
204            .enumerate()
205            .filter(|(_, record)| {
206                record.scope == query.scope
207                    && (query.kinds.is_empty() || query.kinds.contains(&record.kind))
208                    && query.min_score.is_none_or(|minimum| record.confidence >= minimum)
209            })
210            .filter_map(|(insertion_index, record)| {
211                let searchable = format!("{} {} {}", record.name, record.description, record.content);
212                let candidate_terms = memory_terms(&searchable);
213                let lexical_matches = query_terms
214                    .iter()
215                    .filter(|term| candidate_terms.contains(*term))
216                    .count();
217                if !query_terms.is_empty() && lexical_matches == 0 {
218                    return None;
219                }
220                Some((record, lexical_matches, insertion_index))
221            })
222            .collect::<Vec<_>>();
223        ranked.sort_by(|left, right| {
224            right
225                .1
226                .cmp(&left.1)
227                .then_with(|| right.0.updated_at.cmp(&left.0.updated_at))
228                .then_with(|| left.2.cmp(&right.2))
229        });
230        Ok(ranked
231            .into_iter()
232            .take(query.top_k)
233            .map(|(record, _, _)| MemoryRecall {
234                score: record.confidence.clamp(0.0, 1.0),
235                record,
236                why: "deterministic lexical relevance with recency tie-breaking".into(),
237            })
238            .collect())
239    }
240
241    async fn save_session(&self, data: SessionData) -> crate::Result<()> {
242        self.saved_sessions.lock().unwrap().push(data);
243        Ok(())
244    }
245}
246
247fn memory_terms(text: &str) -> std::collections::HashSet<String> {
248    let mut terms = std::collections::HashSet::new();
249    let mut segment = String::new();
250    let flush = |segment: &mut String, terms: &mut std::collections::HashSet<String>| {
251        if segment.is_empty() {
252            return;
253        }
254        let lowered = segment.to_lowercase();
255        terms.insert(lowered.clone());
256        let characters = lowered.chars().collect::<Vec<_>>();
257        if characters.iter().any(|character| is_han(*character)) {
258            for pair in characters.windows(2) {
259                terms.insert(pair.iter().collect());
260            }
261        }
262        segment.clear();
263    };
264    for character in text.chars() {
265        if character.is_alphanumeric() {
266            segment.push(character);
267        } else {
268            flush(&mut segment, &mut terms);
269        }
270    }
271    flush(&mut segment, &mut terms);
272    terms
273}
274
275fn is_han(character: char) -> bool {
276    matches!(character as u32,
277        0x3400..=0x4DBF | 0x4E00..=0x9FFF | 0xF900..=0xFAFF | 0x20000..=0x3134F)
278}
279
280#[cfg(test)]
281mod ranking_tests {
282    use super::{DreamStore, InMemoryDreamStore};
283    use deepstrike_core::mm::memory::{
284        MemoryAuthor, MemoryKind, MemoryProvenance, MemoryQuery, MemoryRecord, MemoryScope,
285        MemoryTrustLevel,
286    };
287
288    fn entry(text: &str, updated_at: u64) -> MemoryRecord {
289        MemoryRecord {
290            record_id: format!("record-{updated_at}"),
291            scope: MemoryScope::new("tenant-test", "ranking"),
292            name: text.into(),
293            kind: MemoryKind::Project,
294            content: text.into(),
295            description: text.into(),
296            provenance: MemoryProvenance {
297                session_id: None,
298                author: MemoryAuthor::Host,
299                trust: MemoryTrustLevel::HostVerified,
300                evidence_refs: Vec::new(),
301            },
302            created_at: 1,
303            updated_at,
304            last_recalled_at: None,
305            recall_count: 0,
306            confidence: 1.0,
307            links: Vec::new(),
308            pinned: false,
309            ttl_days: None,
310        }
311    }
312
313    #[tokio::test]
314    async fn search_uses_query_and_never_falls_back_to_unrelated_entries() {
315        let store = InMemoryDreamStore::with_initial_memories(vec![
316            entry("database migration checklist", 1),
317            entry("rust scheduler fairness", 2),
318            entry("newer unrelated note", 3),
319        ]);
320
321        let query = |text: &str| MemoryQuery {
322            scope: MemoryScope::new("tenant-test", "ranking"),
323            query: text.into(),
324            top_k: 5,
325            kinds: Vec::new(),
326            min_score: None,
327        };
328        let hits = store.search("agent", &query("scheduler rust")).await.unwrap();
329        assert_eq!(hits.len(), 1);
330        assert_eq!(hits[0].record.content, "rust scheduler fairness");
331        assert!(store
332            .search("agent", &query("nonexistent"))
333            .await
334            .unwrap()
335            .is_empty());
336    }
337}
deepstrike_sdk/memory.rs

deepstrike_sdk/
memory.rs
