1use async_trait::async_trait;
7use chrono::{DateTime, Utc};
8use serde::{Deserialize, Serialize};
9use stakpak_shared::models::integrations::openai::ChatMessage;
10use uuid::Uuid;
11
12pub use crate::local::storage::LocalStorage;
14pub use crate::stakpak::storage::StakpakStorage;
15
16#[async_trait]
25pub trait SessionStorage: Send + Sync {
26 async fn list_sessions(
32 &self,
33 query: &ListSessionsQuery,
34 ) -> Result<ListSessionsResult, StorageError>;
35
36 async fn get_session(&self, session_id: Uuid) -> Result<Session, StorageError>;
38
39 async fn create_session(
41 &self,
42 request: &CreateSessionRequest,
43 ) -> Result<CreateSessionResult, StorageError>;
44
45 async fn update_session(
47 &self,
48 session_id: Uuid,
49 request: &UpdateSessionRequest,
50 ) -> Result<Session, StorageError>;
51
52 async fn delete_session(&self, session_id: Uuid) -> Result<(), StorageError>;
54
55 async fn list_checkpoints(
61 &self,
62 session_id: Uuid,
63 query: &ListCheckpointsQuery,
64 ) -> Result<ListCheckpointsResult, StorageError>;
65
66 async fn get_checkpoint(&self, checkpoint_id: Uuid) -> Result<Checkpoint, StorageError>;
68
69 async fn create_checkpoint(
71 &self,
72 session_id: Uuid,
73 request: &CreateCheckpointRequest,
74 ) -> Result<Checkpoint, StorageError>;
75
76 async fn get_active_checkpoint(&self, session_id: Uuid) -> Result<Checkpoint, StorageError> {
82 let session = self.get_session(session_id).await?;
83 session
84 .active_checkpoint
85 .ok_or(StorageError::NotFound("No active checkpoint".to_string()))
86 }
87
88 async fn get_session_stats(&self, _session_id: Uuid) -> Result<SessionStats, StorageError> {
90 Ok(SessionStats::default())
91 }
92}
93
94pub type BoxedSessionStorage = Box<dyn SessionStorage>;
96
97#[derive(Debug, Clone, PartialEq)]
103pub enum StorageError {
104 NotFound(String),
106 InvalidRequest(String),
108 Unauthorized(String),
110 RateLimited(String),
112 Internal(String),
114 Connection(String),
116}
117
118impl std::fmt::Display for StorageError {
119 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
120 match self {
121 StorageError::NotFound(msg) => write!(f, "Not found: {}", msg),
122 StorageError::InvalidRequest(msg) => write!(f, "Invalid request: {}", msg),
123 StorageError::Unauthorized(msg) => write!(f, "Unauthorized: {}", msg),
124 StorageError::RateLimited(msg) => write!(f, "Rate limited: {}", msg),
125 StorageError::Internal(msg) => write!(f, "Internal error: {}", msg),
126 StorageError::Connection(msg) => write!(f, "Connection error: {}", msg),
127 }
128 }
129}
130
131impl std::error::Error for StorageError {}
132
133impl From<String> for StorageError {
134 fn from(s: String) -> Self {
135 StorageError::Internal(s)
136 }
137}
138
139impl From<&str> for StorageError {
140 fn from(s: &str) -> Self {
141 StorageError::Internal(s.to_string())
142 }
143}
144
145#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
151#[serde(rename_all = "UPPERCASE")]
152pub enum SessionVisibility {
153 #[default]
154 Private,
155 Public,
156}
157
158impl std::fmt::Display for SessionVisibility {
159 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
160 match self {
161 SessionVisibility::Private => write!(f, "PRIVATE"),
162 SessionVisibility::Public => write!(f, "PUBLIC"),
163 }
164 }
165}
166
167#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
169#[serde(rename_all = "UPPERCASE")]
170pub enum SessionStatus {
171 #[default]
172 Active,
173 Deleted,
174}
175
176impl std::fmt::Display for SessionStatus {
177 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
178 match self {
179 SessionStatus::Active => write!(f, "ACTIVE"),
180 SessionStatus::Deleted => write!(f, "DELETED"),
181 }
182 }
183}
184
185#[derive(Debug, Clone, Serialize, Deserialize)]
187pub struct Session {
188 pub id: Uuid,
189 pub title: String,
190 pub visibility: SessionVisibility,
191 pub status: SessionStatus,
192 pub cwd: Option<String>,
193 pub created_at: DateTime<Utc>,
194 pub updated_at: DateTime<Utc>,
195 pub active_checkpoint: Option<Checkpoint>,
196}
197
198#[derive(Debug, Clone, Serialize, Deserialize)]
200pub struct SessionSummary {
201 pub id: Uuid,
202 pub title: String,
203 pub visibility: SessionVisibility,
204 pub status: SessionStatus,
205 pub cwd: Option<String>,
206 pub created_at: DateTime<Utc>,
207 pub updated_at: DateTime<Utc>,
208 pub message_count: u32,
209 pub active_checkpoint_id: Option<Uuid>,
210 pub last_message_at: Option<DateTime<Utc>>,
211}
212
213#[derive(Debug, Clone, Serialize, Deserialize)]
219pub struct Checkpoint {
220 pub id: Uuid,
221 pub session_id: Uuid,
222 pub parent_id: Option<Uuid>,
223 pub state: CheckpointState,
224 pub created_at: DateTime<Utc>,
225 pub updated_at: DateTime<Utc>,
226}
227
228#[derive(Debug, Clone, Serialize, Deserialize)]
230pub struct CheckpointSummary {
231 pub id: Uuid,
232 pub session_id: Uuid,
233 pub parent_id: Option<Uuid>,
234 pub message_count: u32,
235 pub created_at: DateTime<Utc>,
236 pub updated_at: DateTime<Utc>,
237}
238
239#[derive(Debug, Clone, Default, Serialize, Deserialize)]
241pub struct CheckpointState {
242 #[serde(default)]
243 pub messages: Vec<ChatMessage>,
244}
245
246#[derive(Debug, Clone, Serialize)]
252pub struct CreateSessionRequest {
253 pub title: String,
254 pub visibility: SessionVisibility,
255 pub cwd: Option<String>,
256 pub initial_state: CheckpointState,
257}
258
259impl CreateSessionRequest {
260 pub fn new(title: impl Into<String>, messages: Vec<ChatMessage>) -> Self {
261 Self {
262 title: title.into(),
263 visibility: SessionVisibility::Private,
264 cwd: None,
265 initial_state: CheckpointState { messages },
266 }
267 }
268
269 pub fn with_visibility(mut self, visibility: SessionVisibility) -> Self {
270 self.visibility = visibility;
271 self
272 }
273
274 pub fn with_cwd(mut self, cwd: impl Into<String>) -> Self {
275 self.cwd = Some(cwd.into());
276 self
277 }
278}
279
280#[derive(Debug, Clone, Default, Serialize)]
282pub struct UpdateSessionRequest {
283 pub title: Option<String>,
284 pub visibility: Option<SessionVisibility>,
285}
286
287impl UpdateSessionRequest {
288 pub fn new() -> Self {
289 Self::default()
290 }
291
292 pub fn with_title(mut self, title: impl Into<String>) -> Self {
293 self.title = Some(title.into());
294 self
295 }
296
297 pub fn with_visibility(mut self, visibility: SessionVisibility) -> Self {
298 self.visibility = Some(visibility);
299 self
300 }
301}
302
303#[derive(Debug, Clone, Serialize)]
305pub struct CreateCheckpointRequest {
306 pub state: CheckpointState,
307 pub parent_id: Option<Uuid>,
308}
309
310impl CreateCheckpointRequest {
311 pub fn new(messages: Vec<ChatMessage>) -> Self {
312 Self {
313 state: CheckpointState { messages },
314 parent_id: None,
315 }
316 }
317
318 pub fn with_parent(mut self, parent_id: Uuid) -> Self {
319 self.parent_id = Some(parent_id);
320 self
321 }
322}
323
324#[derive(Debug, Clone, Default, Serialize)]
326pub struct ListSessionsQuery {
327 pub limit: Option<u32>,
328 pub offset: Option<u32>,
329 pub search: Option<String>,
330 pub status: Option<SessionStatus>,
331 pub visibility: Option<SessionVisibility>,
332}
333
334impl ListSessionsQuery {
335 pub fn new() -> Self {
336 Self::default()
337 }
338
339 pub fn with_limit(mut self, limit: u32) -> Self {
340 self.limit = Some(limit);
341 self
342 }
343
344 pub fn with_offset(mut self, offset: u32) -> Self {
345 self.offset = Some(offset);
346 self
347 }
348
349 pub fn with_search(mut self, search: impl Into<String>) -> Self {
350 self.search = Some(search.into());
351 self
352 }
353}
354
355#[derive(Debug, Clone, Default, Serialize)]
357pub struct ListCheckpointsQuery {
358 pub limit: Option<u32>,
359 pub offset: Option<u32>,
360 pub include_state: Option<bool>,
361}
362
363impl ListCheckpointsQuery {
364 pub fn new() -> Self {
365 Self::default()
366 }
367
368 pub fn with_limit(mut self, limit: u32) -> Self {
369 self.limit = Some(limit);
370 self
371 }
372
373 pub fn with_state(mut self) -> Self {
374 self.include_state = Some(true);
375 self
376 }
377}
378
379#[derive(Debug, Clone)]
385pub struct CreateSessionResult {
386 pub session_id: Uuid,
387 pub checkpoint: Checkpoint,
388}
389
390#[derive(Debug, Clone)]
392pub struct ListSessionsResult {
393 pub sessions: Vec<SessionSummary>,
394 pub total: Option<u32>,
395}
396
397#[derive(Debug, Clone)]
399pub struct ListCheckpointsResult {
400 pub checkpoints: Vec<CheckpointSummary>,
401 pub total: Option<u32>,
402}
403
404#[derive(Debug, Clone, Default, Serialize, Deserialize)]
410pub struct SessionStats {
411 pub total_sessions: u32,
412 pub total_tool_calls: u32,
413 pub successful_tool_calls: u32,
414 pub failed_tool_calls: u32,
415 pub aborted_tool_calls: u32,
416 pub sessions_with_activity: u32,
417 pub total_time_saved_seconds: Option<u32>,
418 pub tools_usage: Vec<ToolUsageStats>,
419}
420
421#[derive(Debug, Clone, Serialize, Deserialize)]
423pub struct ToolUsageStats {
424 pub tool_name: String,
425 pub display_name: String,
426 pub usage_counts: ToolUsageCounts,
427 pub time_saved_per_call: Option<f64>,
428 pub time_saved_seconds: Option<u32>,
429}
430
431#[derive(Debug, Clone, Serialize, Deserialize)]
433pub struct ToolUsageCounts {
434 pub total: u32,
435 pub successful: u32,
436 pub failed: u32,
437 pub aborted: u32,
438}