cognee_database/traits/
session_lifecycle_db.rs

1//! `SessionLifecycleDb` trait — repository for the `/api/v1/sessions/*`
2//! HTTP endpoints (E-09, E-10, E-11, E-12).
3//!
4//! Wraps the seven dashboard operations that live in Python's
5//! `cognee/modules/session_lifecycle/metrics.py`:
6//!   * `ensure_and_touch_session` — upsert + bump activity (idempotent)
7//!   * `accumulate_usage` — atomic counter add to session row + per-model
8//!     `session_model_usage` row
9//!   * `get_session_row` — visibility-checked single-row read
10//!   * `list_session_rows` — paginated list with status / since filters
11//!   * `aggregate_stats` — totals / durations / status buckets
12//!   * `cost_by_model` — grouped per-model attribution
13//!
14//! The `effective_status` value for a session is computed at read time —
15//! `running` rows whose `last_activity_at` is older than
16//! `SESSION_ABANDON_AFTER_SECONDS` (default 1800s — Decision 12) report
17//! as `abandoned` *without* mutating the row. This mirrors Python's
18//! `get_effective_status_sql` and keeps abandonment cheap (no sweeper,
19//! no writes on read).
20//!
21//! Conventions match LIB-03's entity layer: UUIDs are persisted as
22//! 32-char hex strings (`uuid_hex.rs`), timestamps as `DateTimeUtc`.
23//! The trait's public Rust signatures take `Uuid` and convert at the
24//! boundary so callers don't see hex strings.
25
26use async_trait::async_trait;
27use chrono::{DateTime, Utc};
28use serde::{Deserialize, Serialize};
29use uuid::Uuid;
30
31use crate::entities::session_record;
32use crate::types::DatabaseError;
33
34/// Filters for `SessionLifecycleDb::list_session_rows`. Field-for-field
35/// parity with Python's `list_session_rows` keyword arguments at
36/// `cognee/modules/session_lifecycle/metrics.py:365-374`.
37#[derive(Debug, Clone)]
38pub struct SessionListFilters {
39    /// Visibility scope: caller's own sessions are always included.
40    pub user_id: Uuid,
41    /// Additional dataset scope — sessions whose `dataset_id` is in this
42    /// list are included via OR'd visibility predicate.
43    pub permitted_dataset_ids: Vec<Uuid>,
44    /// Optional `last_activity_at >= since` filter.
45    pub since: Option<DateTime<Utc>>,
46    /// Optional effective-status filter (`completed` / `failed` /
47    /// `abandoned` / `running`). The repository applies the
48    /// `effective_status` SQL expression so `abandoned` matches running
49    /// rows past the idle threshold.
50    pub status_filter: Option<String>,
51    /// Page size. Caller-validated upstream (E-09 enforces `1..=500`).
52    pub limit: u32,
53    /// Page offset.
54    pub offset: u32,
55    /// Column to sort by. Recognized: `last_activity_at`, `started_at`,
56    /// `ended_at`, `cost_usd`, `tokens_in`, `tokens_out`. Anything else
57    /// silently falls back to `last_activity_at` (mirrors Python's
58    /// `sortable.get(order_by, ...)` lookup at `metrics.py:415-423`).
59    pub order_by: String,
60    /// Direction. `true` → DESC.
61    pub descending: bool,
62}
63
64/// Wraps a stored `session_records` row plus the read-time effective
65/// status (`abandoned` for stale running rows). Mirrors Python's
66/// `SessionRowWithStatus` dataclass at
67/// `cognee/modules/session_lifecycle/metrics.py:336-348`.
68#[derive(Debug, Clone)]
69pub struct SessionRowWithStatus {
70    pub record: session_record::Model,
71    pub effective_status: String,
72}
73
74impl SessionRowWithStatus {
75    /// Render to a JSON object whose key order matches Python's
76    /// `to_dict()` — entity dict + `effective_status`.
77    pub fn to_dict(&self) -> serde_json::Value {
78        let mut value = self.record.to_dict();
79        if let Some(map) = value.as_object_mut() {
80            map.insert(
81                "effective_status".to_string(),
82                serde_json::Value::String(self.effective_status.clone()),
83            );
84        }
85        value
86    }
87}
88
89/// Paginated envelope returned by `list_session_rows`. Parity with
90/// Python's `SessionListPage` dataclass at `metrics.py:351-362`.
91#[derive(Debug, Clone)]
92pub struct SessionListPage {
93    pub sessions: Vec<SessionRowWithStatus>,
94    pub total: i64,
95    pub limit: u32,
96    pub offset: u32,
97}
98
99impl SessionListPage {
100    /// `true` when pagination has more rows beyond the current page.
101    /// Matches Python's `has_more` property at `metrics.py:360-362`.
102    pub fn has_more(&self) -> bool {
103        let returned = i64::try_from(self.sessions.len()).unwrap_or(i64::MAX);
104        let offset = i64::from(self.offset);
105        offset.saturating_add(returned) < self.total
106    }
107}
108
109/// Aggregate counters for `GET /api/v1/sessions/stats`. Field-for-field
110/// parity with the Python response body at
111/// `cognee/api/v1/sessions/routers/get_sessions_router.py:179-196`.
112#[derive(Debug, Clone, Serialize, Deserialize)]
113pub struct SessionStats {
114    pub sessions: i64,
115    pub total_spend_usd: f64,
116    pub avg_spend_per_session_usd: f64,
117    pub tokens_in: i64,
118    pub tokens_out: i64,
119    pub tokens_total: i64,
120    pub agent_time_s: f64,
121    pub avg_session_s: f64,
122    pub success_rate: f64,
123    pub completed: i64,
124    pub failed: i64,
125    pub abandoned: i64,
126    pub running: i64,
127}
128
129/// Single per-model row for `GET /api/v1/sessions/cost-by-model`.
130/// Parity with the Python response items at
131/// `cognee/api/v1/sessions/routers/get_sessions_router.py:241-251`.
132#[derive(Debug, Clone, Serialize, Deserialize)]
133pub struct CostByModelRow {
134    pub model: String,
135    pub session_count: i64,
136    pub cost_usd: f64,
137    pub tokens_in: i64,
138    pub tokens_out: i64,
139}
140
141/// Repository trait for `/api/v1/sessions/*`. See module docs.
142#[async_trait]
143#[allow(clippy::too_many_arguments)] // accumulate_usage mirrors Python kw-args at metrics.py:133-141
144pub trait SessionLifecycleDb: Send + Sync {
145    /// Upsert a session row, bumping `last_activity_at` if the row is
146    /// already running. Mirrors `metrics.py:62-130`.
147    async fn ensure_and_touch_session(
148        &self,
149        session_id: &str,
150        user_id: Uuid,
151        dataset_id: Option<Uuid>,
152    ) -> Result<(), DatabaseError>;
153
154    /// Atomically add usage counters to the session row + per-model
155    /// row. Mirrors `metrics.py:133-241`.
156    async fn accumulate_usage(
157        &self,
158        session_id: &str,
159        user_id: Uuid,
160        model: Option<&str>,
161        tokens_in: i64,
162        tokens_out: i64,
163        cost_usd: f64,
164        errored: bool,
165    ) -> Result<(), DatabaseError>;
166
167    /// Visibility-checked single-row read. Mirrors `metrics.py:295-333`.
168    async fn get_session_row(
169        &self,
170        session_id: &str,
171        user_id: Uuid,
172        permitted_dataset_ids: &[Uuid],
173        prefer_other_owner: bool,
174    ) -> Result<Option<SessionRowWithStatus>, DatabaseError>;
175
176    /// Paginated list with `effective_status` filter support. Mirrors
177    /// `metrics.py:365-438`.
178    async fn list_session_rows(
179        &self,
180        filters: SessionListFilters,
181    ) -> Result<SessionListPage, DatabaseError>;
182
183    /// Dashboard counters for `GET /sessions/stats`. Mirrors
184    /// `get_sessions_router.py:112-196`.
185    async fn aggregate_stats(
186        &self,
187        user_id: Uuid,
188        permitted_dataset_ids: &[Uuid],
189        since: Option<DateTime<Utc>>,
190    ) -> Result<SessionStats, DatabaseError>;
191
192    /// Per-model attribution for `GET /sessions/cost-by-model`.
193    /// Mirrors `get_sessions_router.py:198-252`.
194    async fn cost_by_model(
195        &self,
196        user_id: Uuid,
197        permitted_dataset_ids: &[Uuid],
198        since: Option<DateTime<Utc>>,
199    ) -> Result<Vec<CostByModelRow>, DatabaseError>;
200}
cognee_database/traits/session_lifecycle_db.rs

cognee_database/traits/
session_lifecycle_db.rs
