atman_runtime/
history_store.rs

1//! Abstract history store — tools call this trait, never touching storage directly.
2//!
3//! `HistoryStoreImpl` routes internally: SQLite (priority) → messages_full
4//! (fallback) → events.jsonl replay (last resort).
5
6use std::path::PathBuf;
7use std::sync::Arc;
8
9use crate::error::RuntimeError;
10use crate::event::EventEnvelope;
11use crate::index::AnchorIndex;
12use crate::message::Message;
13use crate::projection::message_window::replay_all_messages_with_seq;
14use crate::session::SessionOpenError;
15
16pub enum SearchScope {
17    Session,
18    Project,
19}
20
21pub struct HistoryQuery {
22    pub session_id: String,
23    pub offset: usize,
24    pub limit: usize,
25    pub role_filter: Option<Vec<String>>,
26}
27
28#[derive(Debug)]
29pub struct HistoryPage {
30    pub total: u64,
31    pub offset: usize,
32    pub limit: usize,
33    pub items: Vec<Message>,
34}
35
36#[derive(Debug)]
37pub struct SearchHit {
38    pub session_id: String,
39    pub seq: u64,
40    pub ts: String,
41    pub kind: String,
42    pub snippet: String,
43}
44
45#[derive(Debug)]
46pub struct SearchResult {
47    pub total: u64,
48    pub hits: Vec<SearchHit>,
49}
50
51pub trait HistoryStore: Send + Sync {
52    fn count(&self, session_id: &str, role_filter: Option<&[&str]>) -> Result<u64, RuntimeError>;
53    fn read(&self, query: HistoryQuery) -> Result<HistoryPage, RuntimeError>;
54    fn search(
55        &self,
56        query: &str,
57        scope: SearchScope,
58        limit: usize,
59    ) -> Result<SearchResult, RuntimeError>;
60    fn recent(&self, n: usize) -> Result<(u64, Vec<Message>), RuntimeError>;
61}
62
63pub struct HistoryStoreImpl {
64    project_index: Option<Arc<AnchorIndex>>,
65    session: Option<Arc<crate::session::Session>>,
66    current_session_id: Option<String>,
67    sessions_root: Option<PathBuf>,
68}
69
70impl HistoryStoreImpl {
71    pub fn new(
72        project_index: Option<Arc<AnchorIndex>>,
73        session: Option<Arc<crate::session::Session>>,
74        current_session_id: Option<String>,
75        sessions_root: Option<PathBuf>,
76    ) -> Self {
77        Self {
78            project_index,
79            session,
80            current_session_id,
81            sessions_root,
82        }
83    }
84
85    fn is_current_session(&self, session_id: &str) -> bool {
86        self.current_session_id
87            .as_deref()
88            .is_some_and(|sid| sid == session_id)
89    }
90
91    fn sqlite_available(&self) -> bool {
92        self.project_index.is_some()
93    }
94
95    fn session_dir(&self, session_id: &str) -> Result<PathBuf, RuntimeError> {
96        self.sessions_root
97            .as_ref()
98            .map(|root| root.join(session_id))
99            .ok_or_else(|| {
100                RuntimeError::ToolFailed(format!(
101                    "history: no sessions_root to resolve session `{session_id}`"
102                ))
103            })
104    }
105
106    fn replay_from_jsonl(
107        &self,
108        session_id: &str,
109        role_filter: Option<&[&str]>,
110    ) -> Result<Vec<Message>, RuntimeError> {
111        let dir = self.session_dir(session_id)?;
112        let path = dir.join("events.jsonl");
113        let msgs = replay_all_messages_with_seq(&path).map_err(|e| {
114            RuntimeError::ToolFailed(format!("history: replay {}: {e}", path.display()))
115        })?;
116        let msgs: Vec<Message> = msgs.into_iter().map(|(_, m)| m).collect();
117        Ok(filter_messages_by_role(msgs, role_filter))
118    }
119}
120
121fn role_to_kind(role: &str) -> Option<&'static str> {
122    match role {
123        "user" => Some("user_msg"),
124        "assistant" => Some("assistant_msg"),
125        "tool" => Some("tool_result_msg"),
126        "system" => Some("system_msg"),
127        _ => None,
128    }
129}
130
131const MESSAGE_KINDS: &[&str] = &["user_msg", "assistant_msg", "tool_result_msg", "system_msg"];
132
133fn roles_to_kinds(roles: Option<&[&str]>) -> Vec<&'static str> {
134    match roles {
135        Some(rs) if !rs.is_empty() => rs
136            .iter()
137            .filter_map(|r| role_to_kind(r))
138            .collect::<Vec<_>>(),
139        _ => MESSAGE_KINDS.to_vec(),
140    }
141}
142
143fn filter_messages_by_role(msgs: Vec<Message>, roles: Option<&[&str]>) -> Vec<Message> {
144    match roles {
145        Some(rs) if !rs.is_empty() => msgs
146            .into_iter()
147            .filter(|m| rs.iter().any(|r| *r == m.role.as_str()))
148            .collect(),
149        _ => msgs,
150    }
151}
152
153fn extract_message_from_payload(payload: &str) -> Option<Message> {
154    let env: EventEnvelope = serde_json::from_str(payload).ok()?;
155    match env.event {
156        crate::event::Event::UserMsg { message, .. }
157        | crate::event::Event::AssistantMsg { message, .. }
158        | crate::event::Event::ToolResultMsg { message, .. }
159        | crate::event::Event::SystemMsg { message, .. } => Some(message),
160        _ => None,
161    }
162}
163
164fn rows_to_messages(rows: Vec<crate::index::ProjectEventRow>) -> Vec<Message> {
165    rows.into_iter()
166        .filter_map(|r| extract_message_from_payload(&r.payload))
167        .collect()
168}
169
170impl HistoryStore for HistoryStoreImpl {
171    fn count(&self, session_id: &str, role_filter: Option<&[&str]>) -> Result<u64, RuntimeError> {
172        let kinds = roles_to_kinds(role_filter);
173
174        if let Some(idx) = &self.project_index
175            && self.sqlite_available()
176        {
177            return idx
178                .count_events(session_id, Some(&kinds))
179                .map_err(|e| RuntimeError::ToolFailed(format!("history.count: {e}")));
180        }
181
182        if self.is_current_session(session_id)
183            && let Some(session) = &self.session
184        {
185            let msgs = session.messages_full();
186            let filtered = filter_messages_by_role(msgs.to_vec(), role_filter);
187            return Ok(filtered.len() as u64);
188        }
189
190        let msgs = self.replay_from_jsonl(session_id, role_filter)?;
191        Ok(msgs.len() as u64)
192    }
193
194    fn read(&self, query: HistoryQuery) -> Result<HistoryPage, RuntimeError> {
195        let HistoryQuery {
196            session_id,
197            offset,
198            limit,
199            role_filter,
200        } = query;
201        let role_strs: Option<Vec<&str>> = role_filter
202            .as_ref()
203            .map(|rs| rs.iter().map(|s| s.as_str()).collect());
204        let kinds = roles_to_kinds(role_strs.as_deref());
205        let offset0 = offset.saturating_sub(1);
206
207        if let Some(idx) = &self.project_index
208            && self.sqlite_available()
209        {
210            let total = idx
211                .count_events(&session_id, Some(&kinds))
212                .map_err(|e| RuntimeError::ToolFailed(format!("history.read count: {e}")))?;
213            let rows = idx
214                .read_events_paginated(&session_id, offset0, limit, Some(&kinds))
215                .map_err(|e| RuntimeError::ToolFailed(format!("history.read: {e}")))?;
216            let items = rows_to_messages(rows);
217            return Ok(HistoryPage {
218                total,
219                offset,
220                limit,
221                items,
222            });
223        }
224
225        let msgs = if self.is_current_session(&session_id)
226            && let Some(session) = &self.session
227        {
228            session.messages_full().to_vec()
229        } else {
230            self.replay_from_jsonl(&session_id, None)?
231        };
232        let filtered = filter_messages_by_role(msgs, role_strs.as_deref());
233        let total = filtered.len() as u64;
234        let end = (offset0 + limit).min(filtered.len());
235        let items = if offset0 >= filtered.len() {
236            Vec::new()
237        } else {
238            filtered[offset0..end].to_vec()
239        };
240        Ok(HistoryPage {
241            total,
242            offset,
243            limit,
244            items,
245        })
246    }
247
248    fn search(
249        &self,
250        query: &str,
251        scope: SearchScope,
252        limit: usize,
253    ) -> Result<SearchResult, RuntimeError> {
254        if query.trim().is_empty() {
255            return Err(RuntimeError::ToolFailed(
256                "history.search: empty query".into(),
257            ));
258        }
259
260        let session_filter = match &scope {
261            SearchScope::Project => None,
262            SearchScope::Session => self.current_session_id.clone(),
263        };
264
265        if let Some(idx) = &self.project_index {
266            let total = idx
267                .count_search_hits(query, session_filter.as_deref())
268                .map_err(|e| RuntimeError::ToolFailed(format!("history.search count: {e}")))?;
269            let rows = idx
270                .fts_search_project_events(query, session_filter.as_deref(), limit)
271                .map_err(|e| RuntimeError::ToolFailed(format!("history.search: {e}")))?;
272            let hits = rows
273                .into_iter()
274                .map(|r| {
275                    let snippet: String = r
276                        .payload
277                        .chars()
278                        .take(200)
279                        .collect::<String>()
280                        .replace('\n', " ");
281                    SearchHit {
282                        session_id: r.session_id,
283                        seq: r.seq,
284                        ts: r.ts,
285                        kind: r.kind,
286                        snippet,
287                    }
288                })
289                .collect();
290            return Ok(SearchResult { total, hits });
291        }
292
293        if matches!(scope, SearchScope::Project) {
294            return Err(RuntimeError::ToolFailed(
295                "history.search: project scope requires project index".into(),
296            ));
297        }
298
299        if let Some(session) = &self.session {
300            let msgs = session.messages_full();
301            let query_lower = query.to_lowercase();
302            let mut hits = Vec::new();
303            let sid = self.current_session_id.clone().unwrap_or_default();
304            for (i, msg) in msgs.iter().enumerate() {
305                let text = msg.text_concat();
306                if text.to_lowercase().contains(&query_lower) {
307                    let snippet: String = text.chars().take(200).collect();
308                    hits.push(SearchHit {
309                        session_id: sid.clone(),
310                        seq: i as u64,
311                        ts: String::new(),
312                        kind: msg.role.as_str().to_string(),
313                        snippet,
314                    });
315                    if hits.len() >= limit {
316                        break;
317                    }
318                }
319            }
320            let total = hits.len() as u64;
321            return Ok(SearchResult { total, hits });
322        }
323
324        Err(RuntimeError::ToolFailed(
325            "history.search: no project index or session on context".into(),
326        ))
327    }
328
329    fn recent(&self, n: usize) -> Result<(u64, Vec<Message>), RuntimeError> {
330        if let Some(session) = &self.session {
331            let msgs = session.messages_full();
332            let total = msgs.len() as u64;
333            if n == 0 {
334                return Ok((total, Vec::new()));
335            }
336            let start = msgs.len().saturating_sub(n);
337            let items = msgs[start..].to_vec();
338            return Ok((total, items));
339        }
340
341        if let Some(session_id) = &self.current_session_id {
342            let msgs = self.replay_from_jsonl(session_id, None)?;
343            let total = msgs.len() as u64;
344            if n == 0 {
345                return Ok((total, Vec::new()));
346            }
347            let start = msgs.len().saturating_sub(n);
348            let items = msgs[start..].to_vec();
349            return Ok((total, items));
350        }
351
352        Err(RuntimeError::ToolFailed(
353            "memory.recent_turns: no session available".into(),
354        ))
355    }
356}
357
358impl From<SessionOpenError> for RuntimeError {
359    fn from(e: SessionOpenError) -> Self {
360        RuntimeError::ToolFailed(format!("{e}"))
361    }
362}
363
364#[cfg(test)]
365mod tests {
366    use super::*;
367    use crate::index::{AnchorIndex, ProjectEventInsert};
368    use crate::message::{Message, MessageOrigin, MessagePart, MessageRole};
369    use tempfile::TempDir;
370
371    fn user_msg(text: &str) -> Message {
372        Message {
373            role: MessageRole::User,
374            parts: vec![MessagePart::Text {
375                text: text.to_string(),
376            }],
377            turn_id: crate::event::TurnId::now(),
378            origin: MessageOrigin::User,
379        }
380    }
381
382    fn assistant_msg(text: &str) -> Message {
383        Message {
384            role: MessageRole::Assistant,
385            parts: vec![MessagePart::Text {
386                text: text.to_string(),
387            }],
388            turn_id: crate::event::TurnId::now(),
389            origin: MessageOrigin::User,
390        }
391    }
392
393    fn seed(idx: &AnchorIndex, sid: &str, seq: i64, kind: &str, text: &str) {
394        let payload = serde_json::json!({
395            "type": kind,
396            "seq": seq,
397            "turn_id": "019f0000-0000-7000-0000-000000000001",
398            "message": {
399                "role": if kind == "user_msg" { "user" } else { "assistant" },
400                "parts": [{"type": "text", "text": text}],
401                "turn_id": "019f0000-0000-7000-0000-000000000001"
402            },
403            "ts": "2026-07-08T00:00:00Z"
404        });
405        idx.insert_project_event_raw(ProjectEventInsert {
406            session_id: sid,
407            seq,
408            ts: "2026-07-08T00:00:00Z",
409            kind,
410            turn_id: Some("019f0000-0000-7000-0000-000000000001"),
411            flow_run_id: None,
412            text_content: text,
413            payload_json: &payload.to_string(),
414        })
415        .unwrap();
416    }
417
418    #[test]
419    fn count_via_sqlite() {
420        let dir = TempDir::new().unwrap();
421        let idx = Arc::new(AnchorIndex::open_project(dir.path()).unwrap());
422        let store = HistoryStoreImpl::new(Some(idx), None, Some("s1".into()), None);
423        seed(
424            store.project_index.as_ref().unwrap(),
425            "s1",
426            1,
427            "user_msg",
428            "hello",
429        );
430        seed(
431            store.project_index.as_ref().unwrap(),
432            "s1",
433            2,
434            "assistant_msg",
435            "hi",
436        );
437        seed(
438            store.project_index.as_ref().unwrap(),
439            "s1",
440            3,
441            "user_msg",
442            "bye",
443        );
444        assert_eq!(store.count("s1", None).unwrap(), 3);
445        assert_eq!(store.count("s1", Some(&["user"])).unwrap(), 2);
446        assert_eq!(store.count("s1", Some(&["assistant"])).unwrap(), 1);
447    }
448
449    #[test]
450    fn read_via_sqlite_paginated() {
451        let dir = TempDir::new().unwrap();
452        let idx = Arc::new(AnchorIndex::open_project(dir.path()).unwrap());
453        let store = HistoryStoreImpl::new(Some(idx), None, Some("s1".into()), None);
454        for i in 1..=5 {
455            seed(
456                store.project_index.as_ref().unwrap(),
457                "s1",
458                i,
459                "user_msg",
460                &format!("msg {i}"),
461            );
462        }
463        let page = store
464            .read(HistoryQuery {
465                session_id: "s1".into(),
466                offset: 2,
467                limit: 2,
468                role_filter: None,
469            })
470            .unwrap();
471        assert_eq!(page.total, 5);
472        assert_eq!(page.offset, 2);
473        assert_eq!(page.limit, 2);
474        assert_eq!(page.items.len(), 2);
475        assert_eq!(page.items[0].text_concat(), "msg 2");
476        assert_eq!(page.items[1].text_concat(), "msg 3");
477    }
478
479    #[test]
480    fn read_via_sqlite_role_filter() {
481        let dir = TempDir::new().unwrap();
482        let idx = Arc::new(AnchorIndex::open_project(dir.path()).unwrap());
483        let store = HistoryStoreImpl::new(Some(idx), None, Some("s1".into()), None);
484        seed(
485            store.project_index.as_ref().unwrap(),
486            "s1",
487            1,
488            "user_msg",
489            "u1",
490        );
491        seed(
492            store.project_index.as_ref().unwrap(),
493            "s1",
494            2,
495            "assistant_msg",
496            "a1",
497        );
498        seed(
499            store.project_index.as_ref().unwrap(),
500            "s1",
501            3,
502            "user_msg",
503            "u2",
504        );
505        let page = store
506            .read(HistoryQuery {
507                session_id: "s1".into(),
508                offset: 1,
509                limit: 100,
510                role_filter: Some(vec!["user".into()]),
511            })
512            .unwrap();
513        assert_eq!(page.total, 2);
514        assert_eq!(page.items.len(), 2);
515        assert_eq!(page.items[0].text_concat(), "u1");
516        assert_eq!(page.items[1].text_concat(), "u2");
517    }
518
519    #[test]
520    fn search_returns_total_and_hits() {
521        let dir = TempDir::new().unwrap();
522        let idx = Arc::new(AnchorIndex::open_project(dir.path()).unwrap());
523        let store = HistoryStoreImpl::new(Some(idx), None, Some("s1".into()), None);
524        seed(
525            store.project_index.as_ref().unwrap(),
526            "s1",
527            1,
528            "user_msg",
529            "hello world",
530        );
531        seed(
532            store.project_index.as_ref().unwrap(),
533            "s1",
534            2,
535            "user_msg",
536            "hello again",
537        );
538        seed(
539            store.project_index.as_ref().unwrap(),
540            "s1",
541            3,
542            "user_msg",
543            "goodbye",
544        );
545        let result = store.search("hello", SearchScope::Session, 10).unwrap();
546        assert_eq!(result.total, 2);
547        assert_eq!(result.hits.len(), 2);
548    }
549
550    #[test]
551    fn search_cjk() {
552        let dir = TempDir::new().unwrap();
553        let idx = Arc::new(AnchorIndex::open_project(dir.path()).unwrap());
554        let store = HistoryStoreImpl::new(Some(idx), None, Some("s1".into()), None);
555        seed(
556            store.project_index.as_ref().unwrap(),
557            "s1",
558            1,
559            "user_msg",
560            "浮动面板设计",
561        );
562        let result = store.search("浮动", SearchScope::Session, 10).unwrap();
563        assert_eq!(result.total, 1);
564        assert_eq!(result.hits.len(), 1);
565    }
566
567    #[test]
568    fn count_empty_session_via_sqlite() {
569        let dir = TempDir::new().unwrap();
570        let idx = Arc::new(AnchorIndex::open_project(dir.path()).unwrap());
571        let store = HistoryStoreImpl::new(Some(idx), None, Some("s1".into()), None);
572        assert_eq!(store.count("no-such", None).unwrap(), 0);
573    }
574
575    #[test]
576    fn read_empty_session_via_sqlite() {
577        let dir = TempDir::new().unwrap();
578        let idx = Arc::new(AnchorIndex::open_project(dir.path()).unwrap());
579        let store = HistoryStoreImpl::new(Some(idx), None, Some("s1".into()), None);
580        let page = store
581            .read(HistoryQuery {
582                session_id: "no-such".into(),
583                offset: 1,
584                limit: 10,
585                role_filter: None,
586            })
587            .unwrap();
588        assert_eq!(page.total, 0);
589        assert!(page.items.is_empty());
590    }
591
592    // Fallback: no project_index, but session has messages_full
593    // We can't easily construct a Session in unit tests without a tempdir + events.jsonl,
594    // so the messages_full fallback path is covered by integration tests.
595    // Here we test the jsonl replay fallback (no session, no index, but sessions_root points to jsonl).
596
597    #[test]
598    fn fallback_replay_from_jsonl() {
599        let dir = TempDir::new().unwrap();
600        let sessions_root = dir.path().join("sessions");
601        let sid = "test-sid";
602        let session_dir = sessions_root.join(sid);
603        std::fs::create_dir_all(&session_dir).unwrap();
604        let events_path = session_dir.join("events.jsonl");
605        let u1 = user_msg("first");
606        let a1 = assistant_msg("second");
607        let env1 = crate::event::EventEnvelope::new(
608            1,
609            crate::event::Event::UserMsg {
610                turn_id: u1.turn_id.clone(),
611                flow_run_id: None,
612                message: u1,
613            },
614        );
615        let env2 = crate::event::EventEnvelope::new(
616            2,
617            crate::event::Event::AssistantMsg {
618                turn_id: a1.turn_id.clone(),
619                flow_run_id: None,
620                message: a1,
621            },
622        );
623        let line1 = serde_json::to_string(&env1).unwrap();
624        let line2 = serde_json::to_string(&env2).unwrap();
625        std::fs::write(&events_path, format!("{line1}\n{line2}\n")).unwrap();
626
627        let store = HistoryStoreImpl::new(None, None, Some(sid.into()), Some(sessions_root));
628        assert_eq!(store.count(sid, None).unwrap(), 2);
629        assert_eq!(store.count(sid, Some(&["user"])).unwrap(), 1);
630        assert_eq!(store.count(sid, Some(&["assistant"])).unwrap(), 1);
631
632        let page = store
633            .read(HistoryQuery {
634                session_id: sid.into(),
635                offset: 1,
636                limit: 10,
637                role_filter: None,
638            })
639            .unwrap();
640        assert_eq!(page.total, 2);
641        assert_eq!(page.items.len(), 2);
642        assert_eq!(page.items[0].text_concat(), "first");
643        assert_eq!(page.items[1].text_concat(), "second");
644    }
645
646    #[test]
647    fn fallback_replay_paginated() {
648        let dir = TempDir::new().unwrap();
649        let sessions_root = dir.path().join("sessions");
650        let sid = "test-sid";
651        let session_dir = sessions_root.join(sid);
652        std::fs::create_dir_all(&session_dir).unwrap();
653        let events_path = session_dir.join("events.jsonl");
654        let mut lines = Vec::new();
655        for i in 1..=5 {
656            let m = user_msg(&format!("msg {i}"));
657            let env = crate::event::EventEnvelope::new(
658                i,
659                crate::event::Event::UserMsg {
660                    turn_id: m.turn_id.clone(),
661                    flow_run_id: None,
662                    message: m,
663                },
664            );
665            lines.push(serde_json::to_string(&env).unwrap());
666        }
667        std::fs::write(&events_path, lines.join("\n") + "\n").unwrap();
668
669        let store = HistoryStoreImpl::new(None, None, Some(sid.into()), Some(sessions_root));
670        let page = store
671            .read(HistoryQuery {
672                session_id: sid.into(),
673                offset: 2,
674                limit: 2,
675                role_filter: None,
676            })
677            .unwrap();
678        assert_eq!(page.total, 5);
679        assert_eq!(page.items.len(), 2);
680        assert_eq!(page.items[0].text_concat(), "msg 2");
681        assert_eq!(page.items[1].text_concat(), "msg 3");
682    }
683
684    #[test]
685    fn search_empty_query_errors() {
686        let dir = TempDir::new().unwrap();
687        let idx = Arc::new(AnchorIndex::open_project(dir.path()).unwrap());
688        let store = HistoryStoreImpl::new(Some(idx), None, None, None);
689        let err = store.search("", SearchScope::Session, 10).unwrap_err();
690        assert!(matches!(err, RuntimeError::ToolFailed(_)));
691    }
692
693    #[test]
694    fn search_project_scope_without_index_errors() {
695        let store = HistoryStoreImpl::new(None, None, None, None);
696        let err = store.search("hello", SearchScope::Project, 10).unwrap_err();
697        assert!(matches!(err, RuntimeError::ToolFailed(_)));
698    }
699}
atman_runtime/history_store.rs

atman_runtime/
history_store.rs
