1use std::collections::{BTreeMap, BTreeSet};
2
3use anyhow::{Context, Result};
4use rusqlite::{Connection, OptionalExtension};
5use sha2::{Digest, Sha256};
6
7use super::{ROLE_ASSISTANT, ROLE_USER};
8
9const SAMPLE_PREVIEW_CHARS: usize = 200;
10
11#[derive(Debug, Clone, Default)]
12pub struct RawSessionQuery {
13 pub since_epoch: Option<i64>,
14 pub until_epoch: Option<i64>,
15 pub project: Option<String>,
16 pub sample_user_messages: i64,
17 pub latest: Option<i64>,
18}
19
20#[derive(Debug, Clone, serde::Serialize, PartialEq, Eq)]
21pub struct RawSessionSummary {
22 pub session_ref: String,
23 pub host: String,
24 pub session_mode: String,
25 pub source_root: String,
26 pub project: String,
27 pub session_id: String,
28 pub first_epoch: i64,
29 pub last_epoch: i64,
30 pub message_count: i64,
31 pub user_message_count: i64,
32 pub assistant_message_count: i64,
33 pub content_hash: String,
34 pub user_message_samples: Vec<String>,
35 pub mmdd: Option<String>,
36 pub session_intent: Option<String>,
37 pub session_topic: Option<String>,
38 pub display_label: Option<String>,
39 pub session_intent_source: Option<String>,
40}
41
42#[derive(Debug, Clone, serde::Serialize, PartialEq, Eq, PartialOrd, Ord)]
43pub(crate) struct ExcludedSessionIdentity {
44 pub source_root: String,
45 pub project: String,
46 pub session_id: String,
47 pub host: Option<String>,
48}
49
50#[derive(Debug, Clone, PartialEq, Eq)]
51pub(crate) struct RawSessionListing {
52 pub(crate) sessions: Vec<RawSessionSummary>,
53 pub(crate) excluded_legacy_rows: usize,
54 pub(crate) excluded_legacy_sessions: usize,
55 pub(crate) excluded_legacy_identities: Vec<ExcludedSessionIdentity>,
56}
57
58impl std::ops::Deref for RawSessionListing {
59 type Target = [RawSessionSummary];
60
61 fn deref(&self) -> &Self::Target {
62 &self.sessions
63 }
64}
65
66#[derive(Debug, Clone, serde::Serialize)]
67pub struct RawSessionsJson {
68 pub since_epoch: Option<i64>,
69 pub until_epoch: Option<i64>,
70 pub project: Option<String>,
71 pub sample: i64,
72 pub latest: Option<i64>,
73 pub count: usize,
74 pub sessions: Vec<RawSessionSummary>,
75}
76
77pub fn build_sessions_json(
78 query: &RawSessionQuery,
79 sessions: Vec<RawSessionSummary>,
80) -> RawSessionsJson {
81 RawSessionsJson {
82 since_epoch: query.since_epoch,
83 until_epoch: query.until_epoch,
84 project: query.project.clone(),
85 sample: query.sample_user_messages,
86 latest: query.latest,
87 count: sessions.len(),
88 sessions,
89 }
90}
91
92#[derive(Debug, Clone, serde::Serialize)]
93pub(crate) struct RawSessionListingJson {
94 since_epoch: Option<i64>,
95 until_epoch: Option<i64>,
96 project: Option<String>,
97 sample: i64,
98 latest: Option<i64>,
99 count: usize,
100 excluded_legacy_rows: usize,
101 excluded_legacy_sessions: usize,
102 excluded_legacy_identities: Vec<ExcludedSessionIdentity>,
103 sessions: Vec<RawSessionSummary>,
104}
105
106pub(crate) fn build_session_listing_json(
107 query: &RawSessionQuery,
108 listing: RawSessionListing,
109) -> RawSessionListingJson {
110 RawSessionListingJson {
111 since_epoch: query.since_epoch,
112 until_epoch: query.until_epoch,
113 project: query.project.clone(),
114 sample: query.sample_user_messages,
115 latest: query.latest,
116 count: listing.sessions.len(),
117 excluded_legacy_rows: listing.excluded_legacy_rows,
118 excluded_legacy_sessions: listing.excluded_legacy_sessions,
119 excluded_legacy_identities: listing.excluded_legacy_identities,
120 sessions: listing.sessions,
121 }
122}
123
124pub fn list_sessions(conn: &Connection, query: &RawSessionQuery) -> Result<Vec<RawSessionSummary>> {
125 Ok(list_sessions_with_exclusions(conn, query)?.sessions)
126}
127
128pub(crate) fn list_sessions_with_exclusions(
129 conn: &Connection,
130 query: &RawSessionQuery,
131) -> Result<RawSessionListing> {
132 if query.latest.is_some_and(|latest| latest <= 0) {
133 anyhow::bail!("raw sessions latest must be positive");
134 }
135 let mut sql = String::from(
136 "SELECT r.transcript_identity_id, r.transcript_record_ordinal, \
137 r.source_root, r.project, r.session_id, r.role, \
138 r.content_hash, r.created_at_epoch, r.id, r.source, \
139 r.event_time_source, \
140 i.host, i.session_mode, i.status \
141 FROM raw_messages r \
142 LEFT JOIN raw_session_identities i ON i.id = r.transcript_identity_id \
143 WHERE NOT (r.source = 'hook' AND r.transcript_identity_id IS NULL)",
144 );
145 let mut binds: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
146 if let Some(project) = query.project.as_deref() {
147 sql.push_str(&format!(" AND r.project = ?{}", binds.len() + 1));
148 binds.push(Box::new(project.to_string()));
149 }
150 push_selector_window(&mut sql, &mut binds, query);
151 sql.push_str(" ORDER BY r.created_at_epoch ASC, r.id ASC");
152
153 let mut statement = conn.prepare(&sql)?;
154 let rows = statement.query_map(
155 rusqlite::params_from_iter(crate::db::to_sql_refs(&binds)),
156 |row| {
157 Ok((
158 row.get::<_, Option<i64>>(0)?,
159 row.get::<_, Option<i64>>(1)?,
160 row.get::<_, String>(2)?,
161 row.get::<_, String>(3)?,
162 row.get::<_, String>(4)?,
163 row.get::<_, String>(5)?,
164 row.get::<_, String>(6)?,
165 row.get::<_, i64>(7)?,
166 row.get::<_, i64>(8)?,
167 row.get::<_, String>(9)?,
168 row.get::<_, String>(10)?,
169 row.get::<_, Option<String>>(11)?,
170 row.get::<_, Option<String>>(12)?,
171 row.get::<_, Option<String>>(13)?,
172 ))
173 },
174 )?;
175
176 let mut grouped: BTreeMap<(String, String, String, String), Accumulator> = BTreeMap::new();
177 let mut exclusions = ExclusionState::default();
178 for row in rows {
179 let (
180 identity_id,
181 ordinal,
182 root,
183 project,
184 session_id,
185 role,
186 hash,
187 epoch,
188 row_id,
189 source,
190 event_time_source,
191 host,
192 session_mode,
193 status,
194 ) = row?;
195 let active = status.as_deref() == Some("active");
196 if identity_id.is_none() && source == "transcript" && event_time_source == "legacy_unknown"
197 {
198 exclusions.exclude(&root, &project, &session_id, None, ExclusionTaint::Tuple);
199 continue;
200 }
201 if !active {
202 let taint = if identity_id.is_none() || host.is_none() {
203 ExclusionTaint::Tuple
204 } else {
205 ExclusionTaint::Key
206 };
207 exclusions.exclude(&root, &project, &session_id, host.as_deref(), taint);
208 continue;
209 }
210 let Some(host) = host else {
211 exclusions.exclude(&root, &project, &session_id, None, ExclusionTaint::Tuple);
212 continue;
213 };
214 if crate::identity::InstallHost::parse(&host).is_err() {
215 exclusions.exclude(
216 &root,
217 &project,
218 &session_id,
219 Some(&host),
220 ExclusionTaint::Key,
221 );
222 continue;
223 }
224 let Some(session_mode) = session_mode else {
225 exclusions.exclude(
226 &root,
227 &project,
228 &session_id,
229 Some(&host),
230 ExclusionTaint::Key,
231 );
232 continue;
233 };
234 if !is_closed_session_mode(&session_mode) {
235 exclusions.exclude(
236 &root,
237 &project,
238 &session_id,
239 Some(&host),
240 ExclusionTaint::Key,
241 );
242 continue;
243 }
244 if exclusions.is_tainted(&root, &host, &project, &session_id) {
245 exclusions.exclude(
246 &root,
247 &project,
248 &session_id,
249 Some(&host),
250 ExclusionTaint::None,
251 );
252 continue;
253 }
254 let Some(identity_id) = identity_id else {
255 exclusions.exclude(
256 &root,
257 &project,
258 &session_id,
259 Some(&host),
260 ExclusionTaint::Key,
261 );
262 continue;
263 };
264 let Some(ordinal) = ordinal else {
265 exclusions.exclude(
266 &root,
267 &project,
268 &session_id,
269 Some(&host),
270 ExclusionTaint::Key,
271 );
272 continue;
273 };
274 let key = (
275 root.clone(),
276 host.clone(),
277 project.clone(),
278 session_id.clone(),
279 );
280 if let Some(accumulator) = grouped.get(&key) {
281 if accumulator.session_mode != session_mode {
282 exclusions.exclude(
283 &root,
284 &project,
285 &session_id,
286 Some(&host),
287 ExclusionTaint::Key,
288 );
289 continue;
290 }
291 }
292 let accumulator = grouped.entry(key).or_insert_with(|| {
293 Accumulator::new(root, host, session_mode.clone(), project, session_id, epoch)
294 });
295 accumulator.push(
296 identity_id,
297 ordinal,
298 &role,
299 &hash,
300 epoch,
301 row_id,
302 query.sample_user_messages.max(0),
303 );
304 }
305 grouped.retain(|key, accumulator| {
306 let (root, host, project, session_id) = key;
307 if !exclusions.is_tainted(root, host, project, session_id) {
308 return true;
309 }
310 exclusions.drop_accumulated(accumulator);
311 false
312 });
313
314 let mut accumulators = grouped.into_values().collect::<Vec<_>>();
315 if let Some(latest) = query.latest {
316 accumulators.sort_by(|left, right| {
317 right
318 .last_epoch
319 .cmp(&left.last_epoch)
320 .then_with(|| accumulator_selector_cmp(left, right))
321 });
322 accumulators.truncate(latest as usize);
323 } else {
324 accumulators.sort_by(|left, right| {
325 left.first_epoch
326 .cmp(&right.first_epoch)
327 .then_with(|| accumulator_selector_cmp(left, right))
328 });
329 }
330 let mut sample_statement = conn.prepare(
331 "SELECT substr(content, 1, ?2)
332 FROM raw_messages
333 WHERE id = ?1 AND role = 'user'",
334 )?;
335 let mut sessions = accumulators
336 .into_iter()
337 .map(|accumulator| {
338 let samples = accumulator
339 .sample_ids
340 .iter()
341 .map(|row_id| {
342 sample_statement
343 .query_row(
344 rusqlite::params![row_id, SAMPLE_PREVIEW_CHARS as i64],
345 |row| row.get::<_, String>(0),
346 )
347 .optional()?
348 .with_context(|| format!("raw session sample row {row_id} is missing"))
349 })
350 .collect::<Result<Vec<_>>>()?;
351 Ok(accumulator.finish(samples))
352 })
353 .collect::<Result<Vec<_>>>()?;
354 super::session_labels::attach_session_labels(conn, &mut sessions)?;
355 Ok(exclusions.into_listing(sessions))
356}
357
358fn is_closed_session_mode(session_mode: &str) -> bool {
359 matches!(
360 session_mode,
361 "interactive" | "unattended" | "subagent" | "unknown"
362 )
363}
364
365#[derive(Clone, Copy)]
366enum ExclusionTaint {
367 Tuple,
368 Key,
369 None,
370}
371
372#[derive(Default)]
373struct ExclusionState {
374 rows: usize,
375 identities: BTreeSet<ExcludedSessionIdentity>,
376 tainted_tuples: BTreeSet<(String, String, String)>,
377 tainted_keys: BTreeSet<(String, String, String, String)>,
378}
379
380impl ExclusionState {
381 fn exclude(
382 &mut self,
383 source_root: &str,
384 project: &str,
385 session_id: &str,
386 host: Option<&str>,
387 taint: ExclusionTaint,
388 ) {
389 self.rows += 1;
390 self.identities.insert(ExcludedSessionIdentity {
391 source_root: source_root.to_string(),
392 project: project.to_string(),
393 session_id: session_id.to_string(),
394 host: host.map(str::to_string),
395 });
396 match taint {
397 ExclusionTaint::Tuple => {
398 self.tainted_tuples.insert((
399 source_root.to_string(),
400 project.to_string(),
401 session_id.to_string(),
402 ));
403 }
404 ExclusionTaint::Key => {
405 if let Some(host) = host {
406 self.tainted_keys.insert((
407 source_root.to_string(),
408 host.to_string(),
409 project.to_string(),
410 session_id.to_string(),
411 ));
412 } else {
413 self.tainted_tuples.insert((
414 source_root.to_string(),
415 project.to_string(),
416 session_id.to_string(),
417 ));
418 }
419 }
420 ExclusionTaint::None => {}
421 }
422 }
423
424 fn is_tainted(&self, root: &str, host: &str, project: &str, session_id: &str) -> bool {
425 self.tainted_tuples
426 .iter()
427 .any(|(tainted_root, tainted_project, tainted_session)| {
428 tainted_root == root && tainted_project == project && tainted_session == session_id
429 })
430 || self.tainted_keys.contains(&(
431 root.to_string(),
432 host.to_string(),
433 project.to_string(),
434 session_id.to_string(),
435 ))
436 }
437
438 fn drop_accumulated(&mut self, accumulator: &Accumulator) {
439 self.rows += accumulator.message_count as usize;
440 self.identities.insert(ExcludedSessionIdentity {
441 source_root: accumulator.source_root.clone(),
442 project: accumulator.project.clone(),
443 session_id: accumulator.session_id.clone(),
444 host: Some(accumulator.host.clone()),
445 });
446 }
447
448 fn into_listing(self, sessions: Vec<RawSessionSummary>) -> RawSessionListing {
449 let excluded_legacy_sessions = self
450 .identities
451 .iter()
452 .map(|identity| {
453 (
454 identity.source_root.clone(),
455 identity.project.clone(),
456 identity.session_id.clone(),
457 )
458 })
459 .collect::<BTreeSet<_>>()
460 .len();
461 RawSessionListing {
462 sessions,
463 excluded_legacy_rows: self.rows,
464 excluded_legacy_sessions,
465 excluded_legacy_identities: self.identities.into_iter().collect(),
466 }
467 }
468}
469
470fn accumulator_selector_cmp(left: &Accumulator, right: &Accumulator) -> std::cmp::Ordering {
471 (
472 &left.source_root,
473 &left.host,
474 &left.project,
475 &left.session_id,
476 )
477 .cmp(&(
478 &right.source_root,
479 &right.host,
480 &right.project,
481 &right.session_id,
482 ))
483}
484
485fn push_selector_window(
486 sql: &mut String,
487 binds: &mut Vec<Box<dyn rusqlite::types::ToSql>>,
488 query: &RawSessionQuery,
489) {
490 sql.push_str(
491 " AND EXISTS (SELECT 1 FROM raw_messages w \
492 LEFT JOIN raw_session_identities wi ON wi.id = w.transcript_identity_id \
493 WHERE w.source_root = r.source_root AND w.project = r.project \
494 AND w.session_id = r.session_id \
495 AND NOT (w.source = 'hook' AND w.transcript_identity_id IS NULL) \
496 AND ((i.status = 'active' AND wi.status = 'active' AND wi.host = i.host) \
497 OR i.id IS NULL OR i.status != 'active' OR i.host IS NULL)",
498 );
499 if let Some(since) = query.since_epoch {
500 sql.push_str(&format!(" AND w.created_at_epoch >= ?{}", binds.len() + 1));
501 binds.push(Box::new(since));
502 }
503 if let Some(until) = query.until_epoch {
504 sql.push_str(&format!(" AND w.created_at_epoch <= ?{}", binds.len() + 1));
505 binds.push(Box::new(until));
506 }
507 sql.push(')');
508}
509
510struct Accumulator {
511 source_root: String,
512 host: String,
513 session_mode: String,
514 project: String,
515 session_id: String,
516 first_epoch: i64,
517 last_epoch: i64,
518 message_count: i64,
519 user_message_count: i64,
520 assistant_message_count: i64,
521 sample_ids: Vec<i64>,
522 fingerprint: SessionFingerprint,
523}
524
525impl Accumulator {
526 fn new(
527 root: String,
528 host: String,
529 session_mode: String,
530 project: String,
531 session: String,
532 epoch: i64,
533 ) -> Self {
534 let fingerprint = SessionFingerprint::new(&host, &root, &project, &session);
535 Self {
536 source_root: root,
537 host,
538 session_mode,
539 project,
540 session_id: session,
541 first_epoch: epoch,
542 last_epoch: epoch,
543 message_count: 0,
544 user_message_count: 0,
545 assistant_message_count: 0,
546 sample_ids: Vec::new(),
547 fingerprint,
548 }
549 }
550
551 fn push(
552 &mut self,
553 identity_id: i64,
554 ordinal: i64,
555 role: &str,
556 hash: &str,
557 epoch: i64,
558 row_id: i64,
559 limit: i64,
560 ) {
561 self.last_epoch = epoch;
562 self.message_count += 1;
563 if role == ROLE_USER {
564 self.user_message_count += 1;
565 if self.sample_ids.len() < limit as usize {
566 self.sample_ids.push(row_id);
567 }
568 } else if role == ROLE_ASSISTANT {
569 self.assistant_message_count += 1;
570 }
571 self.fingerprint
572 .push(identity_id, ordinal, role, hash, epoch);
573 }
574
575 fn finish(self, samples: Vec<String>) -> RawSessionSummary {
576 RawSessionSummary {
577 session_ref: session_ref(
578 &self.host,
579 &self.source_root,
580 &self.project,
581 &self.session_id,
582 ),
583 host: self.host,
584 session_mode: self.session_mode,
585 source_root: self.source_root,
586 project: self.project,
587 session_id: self.session_id,
588 first_epoch: self.first_epoch,
589 last_epoch: self.last_epoch,
590 message_count: self.message_count,
591 user_message_count: self.user_message_count,
592 assistant_message_count: self.assistant_message_count,
593 content_hash: self.fingerprint.finish(),
594 user_message_samples: samples,
595 mmdd: None,
596 session_intent: None,
597 session_topic: None,
598 display_label: None,
599 session_intent_source: None,
600 }
601 }
602}
603
604pub(crate) struct SessionFingerprint {
605 hasher: Sha256,
606}
607
608impl SessionFingerprint {
609 pub(crate) fn new(host: &str, root: &str, project: &str, session: &str) -> Self {
610 let mut hasher = Sha256::new();
611 for field in [
612 b"remem-raw-session-content-v1".as_slice(),
613 root.as_bytes(),
614 host.as_bytes(),
615 project.as_bytes(),
616 session.as_bytes(),
617 ] {
618 hash_field(&mut hasher, field);
619 }
620 Self { hasher }
621 }
622
623 pub(crate) fn push(
624 &mut self,
625 identity_id: i64,
626 ordinal: i64,
627 role: &str,
628 content_hash: &str,
629 epoch: i64,
630 ) {
631 for field in [
632 &identity_id.to_le_bytes()[..],
633 &ordinal.to_le_bytes(),
634 role.as_bytes(),
635 content_hash.as_bytes(),
636 &epoch.to_le_bytes(),
637 ] {
638 hash_field(&mut self.hasher, field);
639 }
640 }
641
642 pub(crate) fn finish(self) -> String {
643 format!("sha256:{:x}", self.hasher.finalize())
644 }
645}
646
647fn hash_field(hasher: &mut Sha256, value: &[u8]) {
648 hasher.update((value.len() as u64).to_le_bytes());
649 hasher.update(value);
650}
651
652fn session_ref(host: &str, root: &str, project: &str, session: &str) -> String {
653 format!(
654 "remem://raw-session/v2/{}/{}/{}/{}",
655 hex(host),
656 hex(root),
657 hex(project),
658 hex(session)
659 )
660}
661
662fn hex(value: &str) -> String {
663 const DIGITS: &[u8; 16] = b"0123456789abcdef";
664 let mut output = String::with_capacity(value.len() * 2);
665 for byte in value.bytes() {
666 output.push(char::from(DIGITS[(byte >> 4) as usize]));
667 output.push(char::from(DIGITS[(byte & 15) as usize]));
668 }
669 output
670}