1use 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#[derive(Clone)]
18pub struct CliAgentSessionQuery {
19 storage: SqliteStorage,
20}
21
22impl CliAgentSessionQuery {
23 #[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}