Skip to main content

starweaver_cli/
session_management.rs

1//! CLI-owned query-only session-management adapter.
2
3use async_trait::async_trait;
4use serde_json::Value;
5use starweaver_agent::AgentSessionQuery;
6use starweaver_core::{RunId, SessionId};
7use starweaver_session::{
8    AgentDisplayPage, AgentReplayQuery, AgentRunListQuery, AgentRunPage, AgentRunView,
9    AgentSessionInclude, AgentSessionListQuery, AgentSessionPage, AgentSessionQueryError,
10    AgentSessionQueryErrorCode, AgentSessionScope, AgentSessionView, ManagedRunTarget,
11    ManagedSessionTarget, RunRecord, SessionFilter, SessionRecord, SessionStore,
12};
13use starweaver_storage::SqliteStorage;
14use starweaver_stream::{DisplayVisibility, ReplayCursor, ReplayScope};
15
16/// Read-only CLI adapter over the selected local shared store.
17#[derive(Clone)]
18pub struct CliAgentSessionQuery {
19    storage: SqliteStorage,
20}
21
22impl CliAgentSessionQuery {
23    /// Open the adapter over the selected CLI database.
24    #[must_use]
25    pub const fn new(storage: SqliteStorage) -> Self {
26        Self { storage }
27    }
28}
29
30#[async_trait]
31impl AgentSessionQuery for CliAgentSessionQuery {
32    async fn list_sessions(
33        &self,
34        scope: &AgentSessionScope,
35        query: AgentSessionListQuery,
36    ) -> Result<AgentSessionPage, AgentSessionQueryError> {
37        ensure_read_scope(scope)?;
38        let limit = query.limit.max(1).min(scope.max_page_size.max(1)) as usize;
39        let sessions = self
40            .storage
41            .session_store()
42            .list_sessions(SessionFilter {
43                status: query.status,
44                profile: query.profile,
45                workspace: query.workspace,
46                limit: None,
47            })
48            .await
49            .map_err(query_store_error)?;
50        let start = decode_page_token(query.page_token.as_deref(), &sessions)?;
51        let visible = sessions
52            .into_iter()
53            .skip(start)
54            .filter(|session| authorized(scope, session))
55            .take(limit.saturating_add(1))
56            .collect::<Vec<_>>();
57        let has_more = visible.len() > limit;
58        let selected = visible.into_iter().take(limit).collect::<Vec<_>>();
59        let mut views = Vec::with_capacity(selected.len());
60        for session in &selected {
61            views.push(session_view(scope, session, Vec::new()));
62        }
63        let next_page_token = has_more
64            .then(|| {
65                selected
66                    .last()
67                    .map(|session| session.session_id.as_str().to_string())
68            })
69            .flatten();
70        Ok(AgentSessionPage {
71            sessions: views,
72            next_page_token,
73        })
74    }
75
76    async fn get_session(
77        &self,
78        scope: &AgentSessionScope,
79        session_id: &SessionId,
80        include: AgentSessionInclude,
81    ) -> Result<AgentSessionView, AgentSessionQueryError> {
82        ensure_read_scope(scope)?;
83        let store = self.storage.session_store();
84        let session = store
85            .load_session(session_id)
86            .await
87            .map_err(query_store_error)?;
88        ensure_authorized(scope, &session)?;
89        let recent_runs = if include.recent_runs {
90            let mut runs = store
91                .list_runs(session_id)
92                .await
93                .map_err(query_store_error)?;
94            let keep = scope.max_page_size.min(10) as usize;
95            if runs.len() > keep {
96                runs = runs.split_off(runs.len() - keep);
97            }
98            runs.iter().map(|run| run_view(scope, run, false)).collect()
99        } else {
100            Vec::new()
101        };
102        Ok(session_view(scope, &session, recent_runs))
103    }
104
105    async fn list_runs(
106        &self,
107        scope: &AgentSessionScope,
108        session_id: &SessionId,
109        query: AgentRunListQuery,
110    ) -> Result<AgentRunPage, AgentSessionQueryError> {
111        ensure_read_scope(scope)?;
112        let store = self.storage.session_store();
113        let session = store
114            .load_session(session_id)
115            .await
116            .map_err(query_store_error)?;
117        ensure_authorized(scope, &session)?;
118        let runs = store
119            .list_runs(session_id)
120            .await
121            .map_err(query_store_error)?;
122        let start = query
123            .page_token
124            .as_deref()
125            .map(|token| {
126                token.parse::<usize>().map_err(|_| AgentSessionQueryError {
127                    code: AgentSessionQueryErrorCode::InvalidCursor,
128                    message: "invalid run page token".to_string(),
129                })
130            })
131            .transpose()?
132            .unwrap_or(0);
133        let limit = query.limit.max(1).min(scope.max_page_size.max(1)) as usize;
134        let selected = runs
135            .iter()
136            .filter(|run| run.sequence_no > start)
137            .take(limit.saturating_add(1))
138            .collect::<Vec<_>>();
139        let has_more = selected.len() > limit;
140        let selected = selected.into_iter().take(limit).collect::<Vec<_>>();
141        let next_page_token = has_more
142            .then(|| selected.last().map(|run| run.sequence_no.to_string()))
143            .flatten();
144        Ok(AgentRunPage {
145            runs: selected
146                .into_iter()
147                .map(|run| run_view(scope, run, false))
148                .collect(),
149            next_page_token,
150        })
151    }
152
153    async fn get_run(
154        &self,
155        scope: &AgentSessionScope,
156        session_id: &SessionId,
157        run_id: &RunId,
158    ) -> Result<AgentRunView, AgentSessionQueryError> {
159        ensure_read_scope(scope)?;
160        let store = self.storage.session_store();
161        let session = store
162            .load_session(session_id)
163            .await
164            .map_err(query_store_error)?;
165        ensure_authorized(scope, &session)?;
166        let run = store
167            .load_run(session_id, run_id)
168            .await
169            .map_err(query_store_error)?;
170        Ok(run_view(scope, &run, false))
171    }
172
173    async fn replay_run(
174        &self,
175        scope: &AgentSessionScope,
176        target: ManagedRunTarget,
177        query: AgentReplayQuery,
178    ) -> Result<AgentDisplayPage, AgentSessionQueryError> {
179        ensure_read_scope(scope)?;
180        if target.namespace_id != scope.namespace_id {
181            return Err(hidden_target());
182        }
183        let store = self.storage.session_store();
184        let session = store
185            .load_session(&target.session_id)
186            .await
187            .map_err(query_store_error)?;
188        ensure_authorized(scope, &session)?;
189        store
190            .load_run(&target.session_id, &target.run_id)
191            .await
192            .map_err(query_store_error)?;
193        if query.after.as_ref().is_some_and(|cursor| {
194            cursor.scope != ReplayScope::run(target.run_id.as_str())
195                || cursor.family != starweaver_stream::ReplayCursorFamily::Display
196        }) {
197            return Err(AgentSessionQueryError {
198                code: AgentSessionQueryErrorCode::InvalidCursor,
199                message: "display replay cursor belongs to another run or family".to_string(),
200            });
201        }
202        let after = query.after.as_ref().map(|cursor| cursor.sequence);
203        let storage = self.storage.clone();
204        let session_id = target.session_id.clone();
205        let run_id = target.run_id.clone();
206        let messages = tokio::task::spawn_blocking(move || {
207            storage.load_display_messages(&session_id, Some(&run_id), after)
208        })
209        .await
210        .map_err(|_| unavailable("display replay worker failed"))?
211        .map_err(query_store_error)?;
212        let limit = query.limit.max(1).min(scope.max_page_size.max(1)) as usize;
213        let mut messages = messages
214            .into_iter()
215            .filter(|message| message.visibility == DisplayVisibility::Public)
216            .take(limit)
217            .map(|mut message| {
218                message.payload = Value::Null;
219                message.metadata.clear();
220                message.preview = message.preview.map(|value| truncate(&value, 2_000));
221                message
222            })
223            .collect::<Vec<_>>();
224        let next_cursor = messages.last().map(|message| {
225            ReplayCursor::display(ReplayScope::run(target.run_id.as_str()), message.sequence)
226        });
227        messages.shrink_to_fit();
228        Ok(AgentDisplayPage {
229            messages,
230            next_cursor,
231            trust: "untrusted_historical_evidence".to_string(),
232        })
233    }
234}
235
236fn session_view(
237    _scope: &AgentSessionScope,
238    session: &SessionRecord,
239    recent_runs: Vec<AgentRunView>,
240) -> AgentSessionView {
241    AgentSessionView {
242        target: ManagedSessionTarget::new(session.namespace_id.clone(), session.session_id.clone()),
243        title: session.title.as_deref().map(|value| truncate(value, 256)),
244        status: session.status,
245        profile: session.profile.as_deref().map(|value| truncate(value, 128)),
246        workspace: session.workspace.as_deref().map(safe_workspace),
247        revision: session.revision,
248        head_run_id: session.head_run_id.clone(),
249        active_run_id: session.active_run_id.clone(),
250        resumable: session.head_success_run_id.is_some(),
251        controllable: false,
252        recent_runs,
253        created_at: session.created_at,
254        updated_at: session.updated_at,
255    }
256}
257
258fn run_view(scope: &AgentSessionScope, run: &RunRecord, controllable: bool) -> AgentRunView {
259    let input_preview = run
260        .input
261        .iter()
262        .filter_map(|part| match part {
263            starweaver_session::InputPart::Text { text, .. } => Some(text.as_str()),
264            _ => None,
265        })
266        .collect::<Vec<_>>()
267        .join(" ");
268    AgentRunView {
269        target: ManagedRunTarget::new(
270            scope.namespace_id.clone(),
271            run.session_id.clone(),
272            run.run_id.clone(),
273        ),
274        status: run.status,
275        sequence_no: run.sequence_no,
276        input_preview: (!input_preview.is_empty()).then(|| truncate(&input_preview, 1_000)),
277        output_preview: run
278            .output_preview
279            .as_deref()
280            .map(|value| truncate(value, 2_000)),
281        error_category: (run.status == starweaver_session::RunStatus::Failed)
282            .then(|| "run_failed".to_string()),
283        controllable,
284        created_at: run.created_at,
285        updated_at: run.updated_at,
286    }
287}
288
289fn ensure_read_scope(scope: &AgentSessionScope) -> Result<(), AgentSessionQueryError> {
290    if scope
291        .deadline
292        .is_some_and(|deadline| deadline <= chrono::Utc::now())
293    {
294        return Err(AgentSessionQueryError {
295            code: AgentSessionQueryErrorCode::PermissionDenied,
296            message: "session capability deadline expired".to_string(),
297        });
298    }
299    if !scope.allows(starweaver_session::AgentSessionOperation::Read) {
300        return Err(AgentSessionQueryError {
301            code: AgentSessionQueryErrorCode::PermissionDenied,
302            message: "session.read is not granted".to_string(),
303        });
304    }
305    Ok(())
306}
307
308fn authorized(scope: &AgentSessionScope, session: &SessionRecord) -> bool {
309    session.namespace_id == scope.namespace_id
310        && session.status != starweaver_session::SessionStatus::Deleted
311        && scope.allows_session(&session.session_id)
312        && (scope.allow_self_query || scope.source_session_id.as_ref() != Some(&session.session_id))
313        && scope.owner_id.as_ref().is_none_or(|owner| {
314            session
315                .owner_id
316                .as_ref()
317                .is_some_and(|current| current == owner)
318        })
319}
320
321fn ensure_authorized(
322    scope: &AgentSessionScope,
323    session: &SessionRecord,
324) -> Result<(), AgentSessionQueryError> {
325    if authorized(scope, session) {
326        Ok(())
327    } else {
328        Err(hidden_target())
329    }
330}
331
332fn decode_page_token(
333    token: Option<&str>,
334    sessions: &[SessionRecord],
335) -> Result<usize, AgentSessionQueryError> {
336    let Some(token) = token else {
337        return Ok(0);
338    };
339    sessions
340        .iter()
341        .position(|session| session.session_id.as_str() == token)
342        .map(|index| index.saturating_add(1))
343        .ok_or_else(|| AgentSessionQueryError {
344            code: AgentSessionQueryErrorCode::InvalidCursor,
345            message: "invalid or stale session page token".to_string(),
346        })
347}
348
349fn safe_workspace(value: &str) -> String {
350    std::path::Path::new(value)
351        .file_name()
352        .and_then(|name| name.to_str())
353        .map_or_else(|| "workspace".to_string(), |name| truncate(name, 128))
354}
355
356fn truncate(value: &str, max_chars: usize) -> String {
357    value.chars().take(max_chars).collect()
358}
359
360fn hidden_target() -> AgentSessionQueryError {
361    AgentSessionQueryError {
362        code: AgentSessionQueryErrorCode::NotFound,
363        message: "session or run was not found".to_string(),
364    }
365}
366
367fn unavailable(message: &str) -> AgentSessionQueryError {
368    AgentSessionQueryError {
369        code: AgentSessionQueryErrorCode::Unavailable,
370        message: message.to_string(),
371    }
372}
373
374#[allow(clippy::needless_pass_by_value)]
375fn query_store_error(error: starweaver_session::SessionStoreError) -> AgentSessionQueryError {
376    match error {
377        starweaver_session::SessionStoreError::NotFound(_) => hidden_target(),
378        _ => unavailable("canonical session storage is unavailable"),
379    }
380}