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, 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        }
379    }
380
381    fn assistant_msg(text: &str) -> Message {
382        Message {
383            role: MessageRole::Assistant,
384            parts: vec![MessagePart::Text {
385                text: text.to_string(),
386            }],
387            turn_id: crate::event::TurnId::now(),
388        }
389    }
390
391    fn seed(idx: &AnchorIndex, sid: &str, seq: i64, kind: &str, text: &str) {
392        let payload = serde_json::json!({
393            "type": kind,
394            "seq": seq,
395            "turn_id": "019f0000-0000-7000-0000-000000000001",
396            "message": {
397                "role": if kind == "user_msg" { "user" } else { "assistant" },
398                "parts": [{"type": "text", "text": text}],
399                "turn_id": "019f0000-0000-7000-0000-000000000001"
400            },
401            "ts": "2026-07-08T00:00:00Z"
402        });
403        idx.insert_project_event_raw(ProjectEventInsert {
404            session_id: sid,
405            seq,
406            ts: "2026-07-08T00:00:00Z",
407            kind,
408            turn_id: Some("019f0000-0000-7000-0000-000000000001"),
409            flow_run_id: None,
410            text_content: text,
411            payload_json: &payload.to_string(),
412        })
413        .unwrap();
414    }
415
416    #[test]
417    fn count_via_sqlite() {
418        let dir = TempDir::new().unwrap();
419        let idx = Arc::new(AnchorIndex::open_project(dir.path()).unwrap());
420        let store = HistoryStoreImpl::new(Some(idx), None, Some("s1".into()), None);
421        seed(
422            store.project_index.as_ref().unwrap(),
423            "s1",
424            1,
425            "user_msg",
426            "hello",
427        );
428        seed(
429            store.project_index.as_ref().unwrap(),
430            "s1",
431            2,
432            "assistant_msg",
433            "hi",
434        );
435        seed(
436            store.project_index.as_ref().unwrap(),
437            "s1",
438            3,
439            "user_msg",
440            "bye",
441        );
442        assert_eq!(store.count("s1", None).unwrap(), 3);
443        assert_eq!(store.count("s1", Some(&["user"])).unwrap(), 2);
444        assert_eq!(store.count("s1", Some(&["assistant"])).unwrap(), 1);
445    }
446
447    #[test]
448    fn read_via_sqlite_paginated() {
449        let dir = TempDir::new().unwrap();
450        let idx = Arc::new(AnchorIndex::open_project(dir.path()).unwrap());
451        let store = HistoryStoreImpl::new(Some(idx), None, Some("s1".into()), None);
452        for i in 1..=5 {
453            seed(
454                store.project_index.as_ref().unwrap(),
455                "s1",
456                i,
457                "user_msg",
458                &format!("msg {i}"),
459            );
460        }
461        let page = store
462            .read(HistoryQuery {
463                session_id: "s1".into(),
464                offset: 2,
465                limit: 2,
466                role_filter: None,
467            })
468            .unwrap();
469        assert_eq!(page.total, 5);
470        assert_eq!(page.offset, 2);
471        assert_eq!(page.limit, 2);
472        assert_eq!(page.items.len(), 2);
473        assert_eq!(page.items[0].text_concat(), "msg 2");
474        assert_eq!(page.items[1].text_concat(), "msg 3");
475    }
476
477    #[test]
478    fn read_via_sqlite_role_filter() {
479        let dir = TempDir::new().unwrap();
480        let idx = Arc::new(AnchorIndex::open_project(dir.path()).unwrap());
481        let store = HistoryStoreImpl::new(Some(idx), None, Some("s1".into()), None);
482        seed(
483            store.project_index.as_ref().unwrap(),
484            "s1",
485            1,
486            "user_msg",
487            "u1",
488        );
489        seed(
490            store.project_index.as_ref().unwrap(),
491            "s1",
492            2,
493            "assistant_msg",
494            "a1",
495        );
496        seed(
497            store.project_index.as_ref().unwrap(),
498            "s1",
499            3,
500            "user_msg",
501            "u2",
502        );
503        let page = store
504            .read(HistoryQuery {
505                session_id: "s1".into(),
506                offset: 1,
507                limit: 100,
508                role_filter: Some(vec!["user".into()]),
509            })
510            .unwrap();
511        assert_eq!(page.total, 2);
512        assert_eq!(page.items.len(), 2);
513        assert_eq!(page.items[0].text_concat(), "u1");
514        assert_eq!(page.items[1].text_concat(), "u2");
515    }
516
517    #[test]
518    fn search_returns_total_and_hits() {
519        let dir = TempDir::new().unwrap();
520        let idx = Arc::new(AnchorIndex::open_project(dir.path()).unwrap());
521        let store = HistoryStoreImpl::new(Some(idx), None, Some("s1".into()), None);
522        seed(
523            store.project_index.as_ref().unwrap(),
524            "s1",
525            1,
526            "user_msg",
527            "hello world",
528        );
529        seed(
530            store.project_index.as_ref().unwrap(),
531            "s1",
532            2,
533            "user_msg",
534            "hello again",
535        );
536        seed(
537            store.project_index.as_ref().unwrap(),
538            "s1",
539            3,
540            "user_msg",
541            "goodbye",
542        );
543        let result = store.search("hello", SearchScope::Session, 10).unwrap();
544        assert_eq!(result.total, 2);
545        assert_eq!(result.hits.len(), 2);
546    }
547
548    #[test]
549    fn search_cjk() {
550        let dir = TempDir::new().unwrap();
551        let idx = Arc::new(AnchorIndex::open_project(dir.path()).unwrap());
552        let store = HistoryStoreImpl::new(Some(idx), None, Some("s1".into()), None);
553        seed(
554            store.project_index.as_ref().unwrap(),
555            "s1",
556            1,
557            "user_msg",
558            "浮动面板设计",
559        );
560        let result = store.search("浮动", SearchScope::Session, 10).unwrap();
561        assert_eq!(result.total, 1);
562        assert_eq!(result.hits.len(), 1);
563    }
564
565    #[test]
566    fn count_empty_session_via_sqlite() {
567        let dir = TempDir::new().unwrap();
568        let idx = Arc::new(AnchorIndex::open_project(dir.path()).unwrap());
569        let store = HistoryStoreImpl::new(Some(idx), None, Some("s1".into()), None);
570        assert_eq!(store.count("no-such", None).unwrap(), 0);
571    }
572
573    #[test]
574    fn read_empty_session_via_sqlite() {
575        let dir = TempDir::new().unwrap();
576        let idx = Arc::new(AnchorIndex::open_project(dir.path()).unwrap());
577        let store = HistoryStoreImpl::new(Some(idx), None, Some("s1".into()), None);
578        let page = store
579            .read(HistoryQuery {
580                session_id: "no-such".into(),
581                offset: 1,
582                limit: 10,
583                role_filter: None,
584            })
585            .unwrap();
586        assert_eq!(page.total, 0);
587        assert!(page.items.is_empty());
588    }
589
590    // Fallback: no project_index, but session has messages_full
591    // We can't easily construct a Session in unit tests without a tempdir + events.jsonl,
592    // so the messages_full fallback path is covered by integration tests.
593    // Here we test the jsonl replay fallback (no session, no index, but sessions_root points to jsonl).
594
595    #[test]
596    fn fallback_replay_from_jsonl() {
597        let dir = TempDir::new().unwrap();
598        let sessions_root = dir.path().join("sessions");
599        let sid = "test-sid";
600        let session_dir = sessions_root.join(sid);
601        std::fs::create_dir_all(&session_dir).unwrap();
602        let events_path = session_dir.join("events.jsonl");
603        let u1 = user_msg("first");
604        let a1 = assistant_msg("second");
605        let env1 = crate::event::EventEnvelope::new(
606            1,
607            crate::event::Event::UserMsg {
608                turn_id: u1.turn_id.clone(),
609                message: u1,
610            },
611        );
612        let env2 = crate::event::EventEnvelope::new(
613            2,
614            crate::event::Event::AssistantMsg {
615                turn_id: a1.turn_id.clone(),
616                flow_run_id: None,
617                message: a1,
618            },
619        );
620        let line1 = serde_json::to_string(&env1).unwrap();
621        let line2 = serde_json::to_string(&env2).unwrap();
622        std::fs::write(&events_path, format!("{line1}\n{line2}\n")).unwrap();
623
624        let store = HistoryStoreImpl::new(None, None, Some(sid.into()), Some(sessions_root));
625        assert_eq!(store.count(sid, None).unwrap(), 2);
626        assert_eq!(store.count(sid, Some(&["user"])).unwrap(), 1);
627        assert_eq!(store.count(sid, Some(&["assistant"])).unwrap(), 1);
628
629        let page = store
630            .read(HistoryQuery {
631                session_id: sid.into(),
632                offset: 1,
633                limit: 10,
634                role_filter: None,
635            })
636            .unwrap();
637        assert_eq!(page.total, 2);
638        assert_eq!(page.items.len(), 2);
639        assert_eq!(page.items[0].text_concat(), "first");
640        assert_eq!(page.items[1].text_concat(), "second");
641    }
642
643    #[test]
644    fn fallback_replay_paginated() {
645        let dir = TempDir::new().unwrap();
646        let sessions_root = dir.path().join("sessions");
647        let sid = "test-sid";
648        let session_dir = sessions_root.join(sid);
649        std::fs::create_dir_all(&session_dir).unwrap();
650        let events_path = session_dir.join("events.jsonl");
651        let mut lines = Vec::new();
652        for i in 1..=5 {
653            let m = user_msg(&format!("msg {i}"));
654            let env = crate::event::EventEnvelope::new(
655                i,
656                crate::event::Event::UserMsg {
657                    turn_id: m.turn_id.clone(),
658                    message: m,
659                },
660            );
661            lines.push(serde_json::to_string(&env).unwrap());
662        }
663        std::fs::write(&events_path, lines.join("\n") + "\n").unwrap();
664
665        let store = HistoryStoreImpl::new(None, None, Some(sid.into()), Some(sessions_root));
666        let page = store
667            .read(HistoryQuery {
668                session_id: sid.into(),
669                offset: 2,
670                limit: 2,
671                role_filter: None,
672            })
673            .unwrap();
674        assert_eq!(page.total, 5);
675        assert_eq!(page.items.len(), 2);
676        assert_eq!(page.items[0].text_concat(), "msg 2");
677        assert_eq!(page.items[1].text_concat(), "msg 3");
678    }
679
680    #[test]
681    fn search_empty_query_errors() {
682        let dir = TempDir::new().unwrap();
683        let idx = Arc::new(AnchorIndex::open_project(dir.path()).unwrap());
684        let store = HistoryStoreImpl::new(Some(idx), None, None, None);
685        let err = store.search("", SearchScope::Session, 10).unwrap_err();
686        assert!(matches!(err, RuntimeError::ToolFailed(_)));
687    }
688
689    #[test]
690    fn search_project_scope_without_index_errors() {
691        let store = HistoryStoreImpl::new(None, None, None, None);
692        let err = store.search("hello", SearchScope::Project, 10).unwrap_err();
693        assert!(matches!(err, RuntimeError::ToolFailed(_)));
694    }
695}
atman_runtime/history_store.rs

atman_runtime/
history_store.rs
