starweaver_cli/local_store/
archive.rs

1//! CLI adapter over the shared `SQLite` stream archive.
2
3use async_trait::async_trait;
4use starweaver_core::{RunId, SessionId};
5use starweaver_runtime::AgentStreamRecord;
6use starweaver_storage::{SqliteStorage, SqliteStreamArchive};
7use starweaver_stream::{
8    DisplayMessage, ReplayCursor, ReplayCursorFamily, ReplayError, ReplayResult, ReplayScope,
9    ReplaySnapshot, StreamArchive,
10};
11
12use super::{DisplayReplayWindow, LocalStore};
13use crate::{CliResult, config::CliConfig};
14
15async fn run_blocking<T, F>(operation: F) -> ReplayResult<T>
16where
17    T: Send + 'static,
18    F: FnOnce() -> ReplayResult<T> + Send + 'static,
19{
20    tokio::task::spawn_blocking(operation)
21        .await
22        .map_err(|error| ReplayError::Failed(format!("blocking archive task failed: {error}")))?
23}
24
25/// Stream archive bound to a resolved CLI database path.
26#[derive(Clone, Debug)]
27pub struct LocalStreamArchive {
28    config: CliConfig,
29    storage: SqliteStorage,
30    archive: SqliteStreamArchive,
31}
32
33impl LocalStreamArchive {
34    /// Open a CLI archive adapter before it is used by async runtime code.
35    pub fn new(config: CliConfig) -> ReplayResult<Self> {
36        crate::config::ensure_config_dirs(&config)
37            .map_err(|error| ReplayError::Failed(error.to_string()))?;
38        let storage = SqliteStorage::open(&config.database_path)
39            .map_err(|error| ReplayError::Failed(error.to_string()))?;
40        let archive = storage.stream_archive();
41        Ok(Self {
42            config,
43            storage,
44            archive,
45        })
46    }
47
48    fn validate_display_messages(
49        storage: &SqliteStorage,
50        scope: &ReplayScope,
51        messages: &[DisplayMessage],
52    ) -> ReplayResult<()> {
53        if let Some(run_id) = scope.as_str().strip_prefix("run:") {
54            let sessions = storage
55                .list_sessions()
56                .map_err(|error| ReplayError::Failed(error.to_string()))?;
57            let mut owners = Vec::new();
58            for session in sessions {
59                let owns_run = storage
60                    .list_runs(&session.session_id)
61                    .map_err(|error| ReplayError::Failed(error.to_string()))?
62                    .iter()
63                    .any(|run| run.run_id.as_str() == run_id);
64                if owns_run {
65                    owners.push(session.session_id);
66                }
67            }
68            let owner = match owners.as_slice() {
69                [owner] => owner,
70                [] => return Err(ReplayError::NotFound(scope.as_str().to_string())),
71                _ => {
72                    return Err(ReplayError::Failed(format!(
73                        "run scope {} is ambiguous across sessions",
74                        scope.as_str()
75                    )));
76                }
77            };
78            for (index, message) in messages.iter().enumerate() {
79                if message.session_id != *owner {
80                    return Err(ReplayError::Failed(format!(
81                        "display message at index {index} has session_id {}, but run scope belongs to session_id {}",
82                        message.session_id.as_str(),
83                        owner.as_str()
84                    )));
85                }
86            }
87            return Ok(());
88        }
89        if let Some(session_id) = scope.as_str().strip_prefix("session:") {
90            let session_id = SessionId::from_string(session_id);
91            let run_ids = storage
92                .list_runs(&session_id)
93                .map_err(|error| ReplayError::Failed(error.to_string()))?
94                .into_iter()
95                .map(|run| run.run_id)
96                .collect::<Vec<_>>();
97            for (index, message) in messages.iter().enumerate() {
98                if message.session_id != session_id {
99                    return Err(ReplayError::Failed(format!(
100                        "display message at index {index} has session_id {}, but session scope is session:{}",
101                        message.session_id.as_str(),
102                        session_id.as_str()
103                    )));
104                }
105                if !run_ids.contains(&message.run_id) {
106                    return Err(ReplayError::Failed(format!(
107                        "display message at index {index} has run_id {}, which is not a run in session scope session:{}",
108                        message.run_id.as_str(),
109                        session_id.as_str()
110                    )));
111                }
112            }
113            return Ok(());
114        }
115        Err(ReplayError::InvalidCursor(format!(
116            "unsupported replay scope {}",
117            scope.as_str()
118        )))
119    }
120
121    /// Replay display messages as a CLI projection window.
122    pub fn replay_display_window(
123        &self,
124        session_id: &str,
125        run_id: Option<&str>,
126        cursor: Option<&ReplayCursor>,
127    ) -> CliResult<DisplayReplayWindow> {
128        LocalStore::open(&self.config)?.replay_display_window(session_id, run_id, cursor)
129    }
130}
131
132#[async_trait]
133impl StreamArchive for LocalStreamArchive {
134    async fn append_raw_records(
135        &self,
136        session_id: &SessionId,
137        run_id: &RunId,
138        records: Vec<AgentStreamRecord>,
139    ) -> ReplayResult<()> {
140        if records.is_empty() {
141            return Ok(());
142        }
143        self.archive
144            .append_raw_records(session_id, run_id, records)
145            .await
146    }
147
148    async fn replay_raw_after(
149        &self,
150        session_id: &SessionId,
151        run_id: &RunId,
152        cursor: Option<ReplayCursor>,
153    ) -> ReplayResult<Vec<AgentStreamRecord>> {
154        self.archive
155            .replay_raw_after(session_id, run_id, cursor)
156            .await
157    }
158
159    async fn append_display_messages(
160        &self,
161        scope: ReplayScope,
162        messages: Vec<DisplayMessage>,
163    ) -> ReplayResult<()> {
164        if messages.is_empty() {
165            return Ok(());
166        }
167        let storage = self.storage.clone();
168        let validation_scope = scope.clone();
169        let validation_messages = messages.clone();
170        run_blocking(move || {
171            Self::validate_display_messages(&storage, &validation_scope, &validation_messages)
172        })
173        .await?;
174        self.archive.append_display_messages(scope, messages).await
175    }
176
177    async fn replay_display_after(
178        &self,
179        scope: &ReplayScope,
180        cursor: Option<ReplayCursor>,
181    ) -> ReplayResult<Vec<DisplayMessage>> {
182        if let Some(session_id) = scope.as_str().strip_prefix("session:") {
183            if let Some(cursor) = cursor.as_ref() {
184                cursor.validate(ReplayCursorFamily::Display, scope)?;
185            }
186            let session_id = SessionId::from_string(session_id);
187            let after = cursor.as_ref().map(|cursor| cursor.sequence);
188            let storage = self.storage.clone();
189            return run_blocking(move || {
190                storage
191                    .load_display_messages(&session_id, None, after)
192                    .map_err(|error| ReplayError::Failed(error.to_string()))
193            })
194            .await;
195        }
196        self.archive.replay_display_after(scope, cursor).await
197    }
198
199    async fn append_snapshot(
200        &self,
201        scope: ReplayScope,
202        snapshot: ReplaySnapshot,
203    ) -> ReplayResult<()> {
204        self.archive.append_snapshot(scope, snapshot).await
205    }
206
207    async fn latest_snapshot(&self, scope: &ReplayScope) -> ReplayResult<Option<ReplaySnapshot>> {
208        self.archive.latest_snapshot(scope).await
209    }
210
211    async fn cursor_range(
212        &self,
213        scope: &ReplayScope,
214    ) -> ReplayResult<Option<(ReplayCursor, ReplayCursor)>> {
215        if let Some(session_id) = scope.as_str().strip_prefix("session:") {
216            let session_id = SessionId::from_string(session_id);
217            let storage = self.storage.clone();
218            let message_count = run_blocking(move || {
219                storage
220                    .load_display_messages(&session_id, None, None)
221                    .map(|messages| messages.len())
222                    .map_err(|error| ReplayError::Failed(error.to_string()))
223            })
224            .await?;
225            let Some(last_sequence) = message_count.checked_sub(1) else {
226                return Ok(None);
227            };
228            return Ok(Some((
229                ReplayCursor::display(scope.clone(), 0),
230                ReplayCursor::display(scope.clone(), last_sequence),
231            )));
232        }
233        self.archive.cursor_range(scope).await
234    }
235}
starweaver_cli/local_store/archive.rs

starweaver_cli/local_store/
archive.rs
