Skip to main content

nexo_driver_loop/
compact_store.rs

1//! `CompactSummaryStore` implementations.
2//!
3//! `SqliteCompactSummaryStore` persists summaries via
4//! `nexo_memory::LongTermMemory::remember()` so resumed sessions can
5//! inject the last compact summary into the prompt without
6//! re-executing elided turns.
7
8use async_trait::async_trait;
9use nexo_driver_types::{CompactSummary, CompactSummaryStore, GoalId};
10use std::sync::Arc;
11
12use nexo_memory::LongTermMemory;
13
14/// Persists compact summaries via the long-term memory store.
15pub struct SqliteCompactSummaryStore {
16    ltm: Arc<LongTermMemory>,
17}
18
19impl SqliteCompactSummaryStore {
20    pub fn new(ltm: Arc<LongTermMemory>) -> Self {
21        Self { ltm }
22    }
23}
24
25#[async_trait]
26impl CompactSummaryStore for SqliteCompactSummaryStore {
27    async fn store(&self, summary: CompactSummary) -> Result<(), String> {
28        let json = serde_json::to_string(&summary).map_err(|e| e.to_string())?;
29        let goal_str = summary
30            .agent_id
31            .split("::")
32            .last()
33            .unwrap_or(&summary.agent_id);
34        let content = format!(
35            "compact_summary goal:{} turn:{} {}",
36            goal_str, summary.turn_index, json
37        );
38        self.ltm
39            .remember(&summary.agent_id, &content, &["compact_summary"])
40            .await
41            .map(|_| ())
42            .map_err(|e| e.to_string())
43    }
44
45    async fn load(
46        &self,
47        agent_id: &str,
48        goal_id: &GoalId,
49    ) -> Result<Option<CompactSummary>, String> {
50        let query = format!("compact_summary goal:{}", goal_id.0);
51        let entries = self
52            .ltm
53            .recall(agent_id, &query, 5)
54            .await
55            .map_err(|e| e.to_string())?;
56        // Find the most recent entry that deserializes correctly.
57        for entry in entries {
58            if let Ok(s) = serde_json::from_str::<CompactSummary>(&entry.content) {
59                return Ok(Some(s));
60            }
61            // The content has the prefix + JSON. Try stripping prefix.
62            if let Some(json_start) = entry.content.find("{\"agent_id\"") {
63                if let Ok(s) = serde_json::from_str::<CompactSummary>(&entry.content[json_start..])
64                {
65                    return Ok(Some(s));
66                }
67            }
68        }
69        Ok(None)
70    }
71
72    async fn forget(&self, goal_id: &GoalId) -> Result<(), String> {
73        // LongTermMemory::forget(id) needs a UUID. We don't track the
74        // memory entry UUID here, so clean up by relying on FTS recall
75        // \+ forget. For now this is best-effort.
76        let _ = goal_id;
77        Ok(())
78    }
79}
80
81/// No-op store for tests and when `store_in_long_term_memory` is false.
82pub struct NoopCompactSummaryStore;
83
84#[async_trait]
85impl CompactSummaryStore for NoopCompactSummaryStore {
86    async fn store(&self, _summary: CompactSummary) -> Result<(), String> {
87        Ok(())
88    }
89
90    async fn load(
91        &self,
92        _agent_id: &str,
93        _goal_id: &GoalId,
94    ) -> Result<Option<CompactSummary>, String> {
95        Ok(None)
96    }
97
98    async fn forget(&self, _goal_id: &GoalId) -> Result<(), String> {
99        Ok(())
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    #[tokio::test]
108    async fn noop_store_always_ok() {
109        let s = NoopCompactSummaryStore;
110        let summary = CompactSummary {
111            agent_id: "test".into(),
112            summary: "test summary".into(),
113            turn_index: 5,
114            before_tokens: 100_000,
115            after_tokens: 20_000,
116            stored_at: chrono::Utc::now(),
117            cache_pin_keys: Vec::new(),
118            truncated_tool_results: Vec::new(),
119        };
120        s.store(summary).await.unwrap();
121    }
122
123    #[tokio::test]
124    async fn noop_load_returns_none() {
125        let s = NoopCompactSummaryStore;
126        assert!(s.load("test", &GoalId::new()).await.unwrap().is_none());
127    }
128
129    #[tokio::test]
130    async fn noop_forget_always_ok() {
131        let s = NoopCompactSummaryStore;
132        s.forget(&GoalId::new()).await.unwrap();
133    }
134}