Session manager (#4648)
Co-authored-by: Douwe Osinga <douwe@squareup.com>
This commit is contained in:
Douwe Osinga
@@ -1,99 +0,0 @@
|
||||
use crate::session::{self, SessionMetadata};
|
||||
use anyhow::Result;
|
||||
use serde::Serialize;
|
||||
use std::cmp::Ordering;
|
||||
use utoipa::ToSchema;
|
||||
|
||||
#[derive(Clone, Serialize, ToSchema)]
|
||||
pub struct SessionInfo {
|
||||
pub id: String,
|
||||
pub path: String,
|
||||
pub modified: String,
|
||||
pub metadata: SessionMetadata,
|
||||
}
|
||||
|
||||
/// Sort order for listing sessions
|
||||
pub enum SortOrder {
|
||||
Ascending,
|
||||
Descending,
|
||||
}
|
||||
|
||||
pub fn get_valid_sorted_sessions(sort_order: SortOrder) -> Result<Vec<SessionInfo>> {
|
||||
let sessions = match session::list_sessions() {
|
||||
Ok(sessions) => sessions,
|
||||
Err(e) => {
|
||||
tracing::error!("Failed to list sessions: {:?}", e);
|
||||
return Err(anyhow::anyhow!("Failed to list sessions"));
|
||||
}
|
||||
};
|
||||
|
||||
let mut session_infos: Vec<SessionInfo> = Vec::new();
|
||||
let mut corrupted_count = 0;
|
||||
|
||||
for (id, path) in sessions {
|
||||
// Get file modification time with fallback
|
||||
let modified = path
|
||||
.metadata()
|
||||
.and_then(|m| m.modified())
|
||||
.map(|time| {
|
||||
chrono::DateTime::<chrono::Utc>::from(time)
|
||||
.format("%Y-%m-%d %H:%M:%S UTC")
|
||||
.to_string()
|
||||
})
|
||||
.unwrap_or_else(|_| {
|
||||
tracing::warn!("Failed to get modification time for session: {}", id);
|
||||
"Unknown".to_string()
|
||||
});
|
||||
|
||||
// Try to read metadata with error handling
|
||||
match session::read_metadata(&path) {
|
||||
Ok(metadata) => {
|
||||
session_infos.push(SessionInfo {
|
||||
id,
|
||||
path: path.to_string_lossy().to_string(),
|
||||
modified,
|
||||
metadata,
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
corrupted_count += 1;
|
||||
tracing::warn!(
|
||||
"Failed to read metadata for session '{}': {}. Skipping corrupted session.",
|
||||
id,
|
||||
e
|
||||
);
|
||||
|
||||
// Optionally, we could create a placeholder entry for corrupted sessions
|
||||
// to show them in the UI with an error indicator, but for now we skip them
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if corrupted_count > 0 {
|
||||
tracing::warn!(
|
||||
"Skipped {} corrupted sessions during listing",
|
||||
corrupted_count
|
||||
);
|
||||
}
|
||||
|
||||
// Sort sessions by modified date
|
||||
// Since all dates are in ISO format (YYYY-MM-DD HH:MM:SS UTC), we can just use string comparison
|
||||
// This works because the ISO format ensures lexicographical ordering matches chronological ordering
|
||||
session_infos.sort_by(|a, b| {
|
||||
if a.modified == "Unknown" && b.modified == "Unknown" {
|
||||
return Ordering::Equal;
|
||||
} else if a.modified == "Unknown" {
|
||||
return Ordering::Greater; // Unknown dates go last
|
||||
} else if b.modified == "Unknown" {
|
||||
return Ordering::Less;
|
||||
}
|
||||
|
||||
match sort_order {
|
||||
SortOrder::Ascending => a.modified.cmp(&b.modified),
|
||||
SortOrder::Descending => b.modified.cmp(&a.modified),
|
||||
}
|
||||
});
|
||||
|
||||
Ok(session_infos)
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
use crate::conversation::Conversation;
|
||||
use crate::session::Session;
|
||||
use anyhow::Result;
|
||||
use chrono::NaiveDateTime;
|
||||
use std::fs;
|
||||
use std::io::{self, BufRead};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::SystemTime;
|
||||
|
||||
const MAX_FILE_SIZE: u64 = 50 * 1024 * 1024;
|
||||
|
||||
pub fn list_sessions(session_dir: &PathBuf) -> Result<Vec<(String, PathBuf)>> {
|
||||
let entries = fs::read_dir(session_dir)?
|
||||
.filter_map(|entry| {
|
||||
let entry = entry.ok()?;
|
||||
let path = entry.path();
|
||||
|
||||
if path.extension().is_some_and(|ext| ext == "jsonl") {
|
||||
let name = path.file_stem()?.to_string_lossy().to_string();
|
||||
Some((name, path))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
Ok(entries)
|
||||
}
|
||||
|
||||
pub fn load_session(session_name: &str, session_path: &Path) -> Result<Session> {
|
||||
let file = fs::File::open(session_path).map_err(|e| {
|
||||
anyhow::anyhow!(
|
||||
"Failed to open session file {}: {}",
|
||||
session_path.display(),
|
||||
e
|
||||
)
|
||||
})?;
|
||||
|
||||
let file_metadata = file.metadata()?;
|
||||
|
||||
if file_metadata.len() > MAX_FILE_SIZE {
|
||||
return Err(anyhow::anyhow!("Session file too large"));
|
||||
}
|
||||
if file_metadata.len() == 0 {
|
||||
return Err(anyhow::anyhow!("Empty session file"));
|
||||
}
|
||||
|
||||
let modified_time = file_metadata.modified().unwrap_or(SystemTime::now());
|
||||
let created_time = file_metadata
|
||||
.created()
|
||||
.unwrap_or_else(|_| parse_session_timestamp(session_name).unwrap_or(modified_time));
|
||||
|
||||
let reader = io::BufReader::new(file);
|
||||
let mut lines = reader.lines();
|
||||
let mut messages = Vec::new();
|
||||
let mut session = Session {
|
||||
id: session_name.to_string(),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
if let Some(Ok(line)) = lines.next() {
|
||||
let mut metadata_json: serde_json::Value = serde_json::from_str(&line)
|
||||
.map_err(|_| anyhow::anyhow!("Invalid session metadata JSON"))?;
|
||||
|
||||
if let Some(obj) = metadata_json.as_object_mut() {
|
||||
obj.entry("id").or_insert(serde_json::json!(session_name));
|
||||
obj.entry("created_at")
|
||||
.or_insert(serde_json::json!(format_timestamp(created_time)?));
|
||||
obj.entry("updated_at")
|
||||
.or_insert(serde_json::json!(format_timestamp(modified_time)?));
|
||||
obj.entry("extension_data").or_insert(serde_json::json!({}));
|
||||
obj.entry("message_count").or_insert(serde_json::json!(0));
|
||||
|
||||
if let Some(desc) = obj.get_mut("description") {
|
||||
if let Some(desc_str) = desc.as_str() {
|
||||
*desc = serde_json::json!(desc_str
|
||||
.split_whitespace()
|
||||
.collect::<Vec<_>>()
|
||||
.join(" "));
|
||||
}
|
||||
}
|
||||
}
|
||||
session = serde_json::from_value(metadata_json)?;
|
||||
session.id = session_name.to_string();
|
||||
}
|
||||
|
||||
for line in lines.map_while(Result::ok) {
|
||||
if let Ok(message) = serde_json::from_str(&line) {
|
||||
messages.push(message);
|
||||
}
|
||||
}
|
||||
|
||||
if !messages.is_empty() {
|
||||
session.conversation = Some(Conversation::new_unvalidated(messages));
|
||||
}
|
||||
|
||||
Ok(session)
|
||||
}
|
||||
|
||||
fn format_timestamp(time: SystemTime) -> Result<String> {
|
||||
let duration = time.duration_since(std::time::UNIX_EPOCH)?;
|
||||
let timestamp = chrono::DateTime::from_timestamp(duration.as_secs() as i64, 0)
|
||||
.unwrap_or_default()
|
||||
.format("%Y-%m-%d %H:%M:%S")
|
||||
.to_string();
|
||||
Ok(timestamp)
|
||||
}
|
||||
|
||||
fn parse_session_timestamp(session_name: &str) -> Option<SystemTime> {
|
||||
NaiveDateTime::parse_from_str(session_name, "%Y%m%d_%H%M%S")
|
||||
.ok()
|
||||
.map(|dt| SystemTime::from(dt.and_utc()))
|
||||
}
|
||||
@@ -1,14 +1,5 @@
|
||||
pub mod extension_data;
|
||||
pub mod info;
|
||||
pub mod storage;
|
||||
mod legacy;
|
||||
pub mod session_manager;
|
||||
|
||||
// Re-export common session types and functions
|
||||
pub use storage::{
|
||||
ensure_session_dir, generate_description, generate_description_with_schedule_id,
|
||||
generate_session_id, get_most_recent_session, get_path, list_sessions, persist_messages,
|
||||
persist_messages_with_schedule_id, read_messages, read_metadata, update_metadata, Identifier,
|
||||
SessionMetadata,
|
||||
};
|
||||
|
||||
pub use extension_data::{ExtensionData, ExtensionState, TodoState};
|
||||
pub use info::{get_valid_sorted_sessions, SessionInfo};
|
||||
pub use session_manager::{Session, SessionInsights, SessionManager};
|
||||
|
||||
@@ -0,0 +1,858 @@
|
||||
use crate::config::APP_STRATEGY;
|
||||
use crate::conversation::message::Message;
|
||||
use crate::conversation::Conversation;
|
||||
use crate::providers::base::{Provider, MSG_COUNT_FOR_SESSION_NAME_GENERATION};
|
||||
use crate::recipe::Recipe;
|
||||
use crate::session::extension_data::ExtensionData;
|
||||
use anyhow::Result;
|
||||
use etcetera::{choose_app_strategy, AppStrategy};
|
||||
use rmcp::model::Role;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::sqlite::SqliteConnectOptions;
|
||||
use sqlx::{Pool, Sqlite};
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::OnceCell;
|
||||
use tracing::{info, warn};
|
||||
use utoipa::ToSchema;
|
||||
|
||||
const CURRENT_SCHEMA_VERSION: i32 = 1;
|
||||
|
||||
static SESSION_STORAGE: OnceCell<Arc<SessionStorage>> = OnceCell::const_new();
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
|
||||
pub struct Session {
|
||||
pub id: String,
|
||||
#[schema(value_type = String)]
|
||||
pub working_dir: PathBuf,
|
||||
pub description: String,
|
||||
pub created_at: String,
|
||||
pub updated_at: String,
|
||||
pub extension_data: ExtensionData,
|
||||
pub total_tokens: Option<i32>,
|
||||
pub input_tokens: Option<i32>,
|
||||
pub output_tokens: Option<i32>,
|
||||
pub accumulated_total_tokens: Option<i32>,
|
||||
pub accumulated_input_tokens: Option<i32>,
|
||||
pub accumulated_output_tokens: Option<i32>,
|
||||
pub schedule_id: Option<String>,
|
||||
pub recipe: Option<Recipe>,
|
||||
pub conversation: Option<Conversation>,
|
||||
pub message_count: usize,
|
||||
}
|
||||
|
||||
pub struct SessionUpdateBuilder {
|
||||
session_id: String,
|
||||
description: Option<String>,
|
||||
working_dir: Option<PathBuf>,
|
||||
extension_data: Option<ExtensionData>,
|
||||
total_tokens: Option<Option<i32>>,
|
||||
input_tokens: Option<Option<i32>>,
|
||||
output_tokens: Option<Option<i32>>,
|
||||
accumulated_total_tokens: Option<Option<i32>>,
|
||||
accumulated_input_tokens: Option<Option<i32>>,
|
||||
accumulated_output_tokens: Option<Option<i32>>,
|
||||
schedule_id: Option<Option<String>>,
|
||||
recipe: Option<Option<Recipe>>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, ToSchema, Debug)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SessionInsights {
|
||||
/// Total number of sessions
|
||||
total_sessions: usize,
|
||||
/// Total tokens used across all sessions
|
||||
total_tokens: i64,
|
||||
}
|
||||
|
||||
impl SessionUpdateBuilder {
|
||||
fn new(session_id: String) -> Self {
|
||||
Self {
|
||||
session_id,
|
||||
description: None,
|
||||
working_dir: None,
|
||||
extension_data: None,
|
||||
total_tokens: None,
|
||||
input_tokens: None,
|
||||
output_tokens: None,
|
||||
accumulated_total_tokens: None,
|
||||
accumulated_input_tokens: None,
|
||||
accumulated_output_tokens: None,
|
||||
schedule_id: None,
|
||||
recipe: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn description(mut self, description: impl Into<String>) -> Self {
|
||||
self.description = Some(description.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn working_dir(mut self, working_dir: PathBuf) -> Self {
|
||||
self.working_dir = Some(working_dir);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn extension_data(mut self, data: ExtensionData) -> Self {
|
||||
self.extension_data = Some(data);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn total_tokens(mut self, tokens: Option<i32>) -> Self {
|
||||
self.total_tokens = Some(tokens);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn input_tokens(mut self, tokens: Option<i32>) -> Self {
|
||||
self.input_tokens = Some(tokens);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn output_tokens(mut self, tokens: Option<i32>) -> Self {
|
||||
self.output_tokens = Some(tokens);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn accumulated_total_tokens(mut self, tokens: Option<i32>) -> Self {
|
||||
self.accumulated_total_tokens = Some(tokens);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn accumulated_input_tokens(mut self, tokens: Option<i32>) -> Self {
|
||||
self.accumulated_input_tokens = Some(tokens);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn accumulated_output_tokens(mut self, tokens: Option<i32>) -> Self {
|
||||
self.accumulated_output_tokens = Some(tokens);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn schedule_id(mut self, schedule_id: Option<String>) -> Self {
|
||||
self.schedule_id = Some(schedule_id);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn recipe(mut self, recipe: Option<Recipe>) -> Self {
|
||||
self.recipe = Some(recipe);
|
||||
self
|
||||
}
|
||||
|
||||
pub async fn apply(self) -> Result<()> {
|
||||
SessionManager::apply_update(self).await
|
||||
}
|
||||
}
|
||||
|
||||
pub struct SessionManager;
|
||||
|
||||
impl SessionManager {
|
||||
pub async fn instance() -> Result<Arc<SessionStorage>> {
|
||||
SESSION_STORAGE
|
||||
.get_or_try_init(|| async { SessionStorage::new().await.map(Arc::new) })
|
||||
.await
|
||||
.map(Arc::clone)
|
||||
}
|
||||
|
||||
pub async fn create_session(working_dir: PathBuf, description: String) -> Result<Session> {
|
||||
let today = chrono::Utc::now().format("%Y%m%d").to_string();
|
||||
let storage = Self::instance().await?;
|
||||
|
||||
let mut tx = storage.pool.begin().await?;
|
||||
|
||||
let max_idx = sqlx::query_scalar::<_, Option<i32>>(
|
||||
"SELECT MAX(CAST(SUBSTR(id, 10) AS INTEGER)) FROM sessions WHERE id LIKE ?",
|
||||
)
|
||||
.bind(format!("{}_%", today))
|
||||
.fetch_one(&mut *tx)
|
||||
.await?
|
||||
.unwrap_or(0);
|
||||
|
||||
let session_id = format!("{}_{}", today, max_idx + 1);
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO sessions (id, description, working_dir, extension_data)
|
||||
VALUES (?, ?, ?, '{}')
|
||||
"#,
|
||||
)
|
||||
.bind(&session_id)
|
||||
.bind(&description)
|
||||
.bind(working_dir.to_string_lossy().as_ref())
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
tx.commit().await?;
|
||||
|
||||
Self::get_session(&session_id, false).await
|
||||
}
|
||||
|
||||
pub async fn get_session(id: &str, include_messages: bool) -> Result<Session> {
|
||||
Self::instance()
|
||||
.await?
|
||||
.get_session(id, include_messages)
|
||||
.await
|
||||
}
|
||||
|
||||
pub fn update_session(id: &str) -> SessionUpdateBuilder {
|
||||
SessionUpdateBuilder::new(id.to_string())
|
||||
}
|
||||
|
||||
async fn apply_update(builder: SessionUpdateBuilder) -> Result<()> {
|
||||
Self::instance().await?.apply_update(builder).await
|
||||
}
|
||||
|
||||
pub async fn add_message(id: &str, message: &Message) -> Result<()> {
|
||||
Self::instance().await?.add_message(id, message).await
|
||||
}
|
||||
|
||||
pub async fn replace_conversation(id: &str, conversation: &Conversation) -> Result<()> {
|
||||
Self::instance()
|
||||
.await?
|
||||
.replace_conversation(id, conversation)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn list_sessions() -> Result<Vec<Session>> {
|
||||
Self::instance().await?.list_sessions().await
|
||||
}
|
||||
|
||||
pub async fn delete_session(id: &str) -> Result<()> {
|
||||
Self::instance().await?.delete_session(id).await
|
||||
}
|
||||
|
||||
pub async fn get_insights() -> Result<SessionInsights> {
|
||||
Self::instance().await?.get_insights().await
|
||||
}
|
||||
|
||||
pub async fn maybe_update_description(id: &str, provider: Arc<dyn Provider>) -> Result<()> {
|
||||
let session = Self::get_session(id, true).await?;
|
||||
let conversation = session
|
||||
.conversation
|
||||
.ok_or_else(|| anyhow::anyhow!("No messages found"))?;
|
||||
|
||||
let user_message_count = conversation
|
||||
.messages()
|
||||
.iter()
|
||||
.filter(|m| matches!(m.role, Role::User))
|
||||
.count();
|
||||
|
||||
if user_message_count <= MSG_COUNT_FOR_SESSION_NAME_GENERATION {
|
||||
let description = provider.generate_session_name(&conversation).await?;
|
||||
Self::update_session(id)
|
||||
.description(description)
|
||||
.apply()
|
||||
.await
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct SessionStorage {
|
||||
pool: Pool<Sqlite>,
|
||||
}
|
||||
|
||||
pub fn ensure_session_dir() -> Result<PathBuf> {
|
||||
let data_dir = choose_app_strategy(APP_STRATEGY.clone())
|
||||
.expect("goose requires a home dir")
|
||||
.data_dir()
|
||||
.join("sessions");
|
||||
|
||||
if !data_dir.exists() {
|
||||
fs::create_dir_all(&data_dir)?;
|
||||
}
|
||||
|
||||
Ok(data_dir)
|
||||
}
|
||||
|
||||
fn role_to_string(role: &Role) -> &'static str {
|
||||
match role {
|
||||
Role::User => "user",
|
||||
Role::Assistant => "assistant",
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for Session {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
id: String::new(),
|
||||
working_dir: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
|
||||
description: String::new(),
|
||||
created_at: String::new(),
|
||||
updated_at: String::new(),
|
||||
extension_data: ExtensionData::default(),
|
||||
total_tokens: None,
|
||||
input_tokens: None,
|
||||
output_tokens: None,
|
||||
accumulated_total_tokens: None,
|
||||
accumulated_input_tokens: None,
|
||||
accumulated_output_tokens: None,
|
||||
schedule_id: None,
|
||||
recipe: None,
|
||||
conversation: None,
|
||||
message_count: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Session {
|
||||
pub fn without_messages(mut self) -> Self {
|
||||
self.conversation = None;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl sqlx::FromRow<'_, sqlx::sqlite::SqliteRow> for Session {
|
||||
fn from_row(row: &sqlx::sqlite::SqliteRow) -> Result<Self, sqlx::Error> {
|
||||
use sqlx::Row;
|
||||
|
||||
let recipe_json: Option<String> = row.try_get("recipe_json")?;
|
||||
let recipe = recipe_json.and_then(|json| serde_json::from_str(&json).ok());
|
||||
|
||||
Ok(Session {
|
||||
id: row.try_get("id")?,
|
||||
working_dir: PathBuf::from(row.try_get::<String, _>("working_dir")?),
|
||||
description: row.try_get("description")?,
|
||||
created_at: row.try_get("created_at")?,
|
||||
updated_at: row.try_get("updated_at")?,
|
||||
extension_data: serde_json::from_str(&row.try_get::<String, _>("extension_data")?)
|
||||
.unwrap_or_default(),
|
||||
total_tokens: row.try_get("total_tokens")?,
|
||||
input_tokens: row.try_get("input_tokens")?,
|
||||
output_tokens: row.try_get("output_tokens")?,
|
||||
accumulated_total_tokens: row.try_get("accumulated_total_tokens")?,
|
||||
accumulated_input_tokens: row.try_get("accumulated_input_tokens")?,
|
||||
accumulated_output_tokens: row.try_get("accumulated_output_tokens")?,
|
||||
schedule_id: row.try_get("schedule_id")?,
|
||||
recipe,
|
||||
conversation: None,
|
||||
message_count: row.try_get("message_count").unwrap_or(0) as usize,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl SessionStorage {
|
||||
async fn new() -> Result<Self> {
|
||||
let session_dir = ensure_session_dir()?;
|
||||
let db_path = session_dir.join("sessions.db");
|
||||
|
||||
let storage = if db_path.exists() {
|
||||
Self::open(&db_path).await?
|
||||
} else {
|
||||
let storage = Self::create(&db_path).await?;
|
||||
|
||||
if let Err(e) = storage.import_legacy(&session_dir).await {
|
||||
warn!("Failed to import some legacy sessions: {}", e);
|
||||
}
|
||||
|
||||
storage
|
||||
};
|
||||
|
||||
Ok(storage)
|
||||
}
|
||||
|
||||
async fn get_pool(db_path: &Path, create_if_missing: bool) -> Result<Pool<Sqlite>> {
|
||||
let options = SqliteConnectOptions::new()
|
||||
.filename(db_path)
|
||||
.create_if_missing(create_if_missing);
|
||||
|
||||
sqlx::SqlitePool::connect_with(options).await.map_err(|e| {
|
||||
anyhow::anyhow!(
|
||||
"Failed to open SQLite database at '{}': {}",
|
||||
db_path.display(),
|
||||
e
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
async fn open(db_path: &Path) -> Result<Self> {
|
||||
let pool = Self::get_pool(db_path, false).await?;
|
||||
|
||||
let storage = Self { pool };
|
||||
storage.run_migrations().await?;
|
||||
Ok(storage)
|
||||
}
|
||||
|
||||
async fn create(db_path: &Path) -> Result<Self> {
|
||||
let pool = Self::get_pool(db_path, true).await?;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
CREATE TABLE schema_version (
|
||||
version INTEGER PRIMARY KEY,
|
||||
applied_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.execute(&pool)
|
||||
.await?;
|
||||
|
||||
sqlx::query("INSERT INTO schema_version (version) VALUES (?)")
|
||||
.bind(CURRENT_SCHEMA_VERSION)
|
||||
.execute(&pool)
|
||||
.await?;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
CREATE TABLE sessions (
|
||||
id TEXT PRIMARY KEY,
|
||||
description TEXT NOT NULL DEFAULT '',
|
||||
working_dir TEXT NOT NULL,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
extension_data TEXT DEFAULT '{}',
|
||||
total_tokens INTEGER,
|
||||
input_tokens INTEGER,
|
||||
output_tokens INTEGER,
|
||||
accumulated_total_tokens INTEGER,
|
||||
accumulated_input_tokens INTEGER,
|
||||
accumulated_output_tokens INTEGER,
|
||||
schedule_id TEXT,
|
||||
recipe_json TEXT
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.execute(&pool)
|
||||
.await?;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
CREATE TABLE messages (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
session_id TEXT NOT NULL REFERENCES sessions(id),
|
||||
role TEXT NOT NULL,
|
||||
content_json TEXT NOT NULL,
|
||||
created_timestamp INTEGER NOT NULL,
|
||||
timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
tokens INTEGER
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.execute(&pool)
|
||||
.await?;
|
||||
|
||||
sqlx::query("CREATE INDEX idx_messages_session ON messages(session_id)")
|
||||
.execute(&pool)
|
||||
.await?;
|
||||
sqlx::query("CREATE INDEX idx_messages_timestamp ON messages(timestamp)")
|
||||
.execute(&pool)
|
||||
.await?;
|
||||
sqlx::query("CREATE INDEX idx_sessions_updated ON sessions(updated_at DESC)")
|
||||
.execute(&pool)
|
||||
.await?;
|
||||
|
||||
Ok(Self { pool })
|
||||
}
|
||||
|
||||
async fn import_legacy(&self, session_dir: &PathBuf) -> Result<()> {
|
||||
use crate::session::legacy;
|
||||
|
||||
let sessions = match legacy::list_sessions(session_dir) {
|
||||
Ok(sessions) => sessions,
|
||||
Err(_) => {
|
||||
warn!("No legacy sessions found to import");
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
|
||||
if sessions.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let mut imported_count = 0;
|
||||
let mut failed_count = 0;
|
||||
|
||||
for (session_name, session_path) in sessions {
|
||||
match legacy::load_session(&session_name, &session_path) {
|
||||
Ok(session) => match self.import_legacy_session(&session).await {
|
||||
Ok(_) => {
|
||||
imported_count += 1;
|
||||
info!(" ✓ Imported: {}", session_name);
|
||||
}
|
||||
Err(e) => {
|
||||
failed_count += 1;
|
||||
info!(" ✗ Failed to import {}: {}", session_name, e);
|
||||
}
|
||||
},
|
||||
Err(e) => {
|
||||
failed_count += 1;
|
||||
info!(" ✗ Failed to load {}: {}", session_name, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
info!(
|
||||
"Import complete: {} successful, {} failed",
|
||||
imported_count, failed_count
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn import_legacy_session(&self, session: &Session) -> Result<()> {
|
||||
let recipe_json = match &session.recipe {
|
||||
Some(recipe) => Some(serde_json::to_string(recipe)?),
|
||||
None => None,
|
||||
};
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO sessions (
|
||||
id, description, working_dir, created_at, updated_at, extension_data,
|
||||
total_tokens, input_tokens, output_tokens,
|
||||
accumulated_total_tokens, accumulated_input_tokens, accumulated_output_tokens,
|
||||
schedule_id, recipe_json
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
"#,
|
||||
)
|
||||
.bind(&session.id)
|
||||
.bind(&session.description)
|
||||
.bind(session.working_dir.to_string_lossy().as_ref())
|
||||
.bind(&session.created_at)
|
||||
.bind(&session.updated_at)
|
||||
.bind(serde_json::to_string(&session.extension_data)?)
|
||||
.bind(session.total_tokens)
|
||||
.bind(session.input_tokens)
|
||||
.bind(session.output_tokens)
|
||||
.bind(session.accumulated_total_tokens)
|
||||
.bind(session.accumulated_input_tokens)
|
||||
.bind(session.accumulated_output_tokens)
|
||||
.bind(&session.schedule_id)
|
||||
.bind(recipe_json)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
if let Some(conversation) = &session.conversation {
|
||||
self.replace_conversation(&session.id, conversation).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_migrations(&self) -> Result<()> {
|
||||
let current_version = self.get_schema_version().await?;
|
||||
|
||||
if current_version < CURRENT_SCHEMA_VERSION {
|
||||
info!(
|
||||
"Running database migrations from v{} to v{}...",
|
||||
current_version, CURRENT_SCHEMA_VERSION
|
||||
);
|
||||
|
||||
for version in (current_version + 1)..=CURRENT_SCHEMA_VERSION {
|
||||
info!(" Applying migration v{}...", version);
|
||||
self.apply_migration(version).await?;
|
||||
self.update_schema_version(version).await?;
|
||||
info!(" ✓ Migration v{} complete", version);
|
||||
}
|
||||
|
||||
info!("All migrations complete");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_schema_version(&self) -> Result<i32> {
|
||||
let table_exists = sqlx::query_scalar::<_, bool>(
|
||||
r#"
|
||||
SELECT EXISTS (
|
||||
SELECT name FROM sqlite_master
|
||||
WHERE type='table' AND name='schema_version'
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
|
||||
if !table_exists {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let version = sqlx::query_scalar::<_, i32>("SELECT MAX(version) FROM schema_version")
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(version)
|
||||
}
|
||||
|
||||
async fn update_schema_version(&self, version: i32) -> Result<()> {
|
||||
sqlx::query("INSERT INTO schema_version (version) VALUES (?)")
|
||||
.bind(version)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn apply_migration(&self, version: i32) -> Result<()> {
|
||||
match version {
|
||||
1 => {
|
||||
sqlx::query(
|
||||
r#"
|
||||
CREATE TABLE IF NOT EXISTS schema_version (
|
||||
version INTEGER PRIMARY KEY,
|
||||
applied_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
}
|
||||
_ => {
|
||||
anyhow::bail!("Unknown migration version: {}", version);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_session(&self, id: &str, include_messages: bool) -> Result<Session> {
|
||||
let mut session = sqlx::query_as::<_, Session>(
|
||||
r#"
|
||||
SELECT id, working_dir, description, created_at, updated_at, extension_data,
|
||||
total_tokens, input_tokens, output_tokens,
|
||||
accumulated_total_tokens, accumulated_input_tokens, accumulated_output_tokens,
|
||||
schedule_id, recipe_json
|
||||
FROM sessions
|
||||
WHERE id = ?
|
||||
"#,
|
||||
)
|
||||
.bind(id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?
|
||||
.ok_or_else(|| anyhow::anyhow!("Session not found"))?;
|
||||
|
||||
if include_messages {
|
||||
let conv = self.get_conversation(&session.id).await?;
|
||||
session.message_count = conv.messages().len();
|
||||
session.conversation = Some(conv);
|
||||
} else {
|
||||
let count =
|
||||
sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM messages WHERE session_id = ?")
|
||||
.bind(&session.id)
|
||||
.fetch_one(&self.pool)
|
||||
.await? as usize;
|
||||
session.message_count = count;
|
||||
}
|
||||
|
||||
Ok(session)
|
||||
}
|
||||
|
||||
async fn apply_update(&self, builder: SessionUpdateBuilder) -> Result<()> {
|
||||
let mut updates = Vec::new();
|
||||
let mut query = String::from("UPDATE sessions SET ");
|
||||
|
||||
macro_rules! add_update {
|
||||
($field:expr, $name:expr) => {
|
||||
if $field.is_some() {
|
||||
if !updates.is_empty() {
|
||||
query.push_str(", ");
|
||||
}
|
||||
updates.push($name);
|
||||
query.push_str($name);
|
||||
query.push_str(" = ?");
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
add_update!(builder.description, "description");
|
||||
add_update!(builder.working_dir, "working_dir");
|
||||
add_update!(builder.extension_data, "extension_data");
|
||||
add_update!(builder.total_tokens, "total_tokens");
|
||||
add_update!(builder.input_tokens, "input_tokens");
|
||||
add_update!(builder.output_tokens, "output_tokens");
|
||||
add_update!(builder.accumulated_total_tokens, "accumulated_total_tokens");
|
||||
add_update!(builder.accumulated_input_tokens, "accumulated_input_tokens");
|
||||
add_update!(
|
||||
builder.accumulated_output_tokens,
|
||||
"accumulated_output_tokens"
|
||||
);
|
||||
add_update!(builder.schedule_id, "schedule_id");
|
||||
add_update!(builder.recipe, "recipe_json");
|
||||
|
||||
if updates.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if !updates.is_empty() {
|
||||
query.push_str(", ");
|
||||
}
|
||||
query.push_str("updated_at = datetime('now') WHERE id = ?");
|
||||
|
||||
let mut q = sqlx::query(&query);
|
||||
|
||||
if let Some(desc) = builder.description {
|
||||
q = q.bind(desc);
|
||||
}
|
||||
if let Some(wd) = builder.working_dir {
|
||||
q = q.bind(wd.to_string_lossy().to_string());
|
||||
}
|
||||
if let Some(ed) = builder.extension_data {
|
||||
q = q.bind(serde_json::to_string(&ed)?);
|
||||
}
|
||||
if let Some(tt) = builder.total_tokens {
|
||||
q = q.bind(tt);
|
||||
}
|
||||
if let Some(it) = builder.input_tokens {
|
||||
q = q.bind(it);
|
||||
}
|
||||
if let Some(ot) = builder.output_tokens {
|
||||
q = q.bind(ot);
|
||||
}
|
||||
if let Some(att) = builder.accumulated_total_tokens {
|
||||
q = q.bind(att);
|
||||
}
|
||||
if let Some(ait) = builder.accumulated_input_tokens {
|
||||
q = q.bind(ait);
|
||||
}
|
||||
if let Some(aot) = builder.accumulated_output_tokens {
|
||||
q = q.bind(aot);
|
||||
}
|
||||
if let Some(sid) = builder.schedule_id {
|
||||
q = q.bind(sid);
|
||||
}
|
||||
if let Some(recipe) = builder.recipe {
|
||||
let recipe_json = recipe.map(|r| serde_json::to_string(&r)).transpose()?;
|
||||
q = q.bind(recipe_json);
|
||||
}
|
||||
|
||||
q = q.bind(&builder.session_id);
|
||||
q.execute(&self.pool).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_conversation(&self, session_id: &str) -> Result<Conversation> {
|
||||
let rows = sqlx::query_as::<_, (String, String, i64)>(
|
||||
"SELECT role, content_json, created_timestamp FROM messages WHERE session_id = ? ORDER BY timestamp",
|
||||
)
|
||||
.bind(session_id)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
let mut messages = Vec::new();
|
||||
for (role_str, content_json, created_timestamp) in rows {
|
||||
let role = match role_str.as_str() {
|
||||
"user" => Role::User,
|
||||
"assistant" => Role::Assistant,
|
||||
_ => continue,
|
||||
};
|
||||
|
||||
let content = serde_json::from_str(&content_json)?;
|
||||
let message = Message::new(role, created_timestamp, content);
|
||||
messages.push(message);
|
||||
}
|
||||
|
||||
Ok(Conversation::new_unvalidated(messages))
|
||||
}
|
||||
|
||||
async fn add_message(&self, session_id: &str, message: &Message) -> Result<()> {
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO messages (session_id, role, content_json, created_timestamp)
|
||||
VALUES (?, ?, ?, ?)
|
||||
"#,
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(role_to_string(&message.role))
|
||||
.bind(serde_json::to_string(&message.content)?)
|
||||
.bind(message.created)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
sqlx::query("UPDATE sessions SET updated_at = datetime('now') WHERE id = ?")
|
||||
.bind(session_id)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn replace_conversation(
|
||||
&self,
|
||||
session_id: &str,
|
||||
conversation: &Conversation,
|
||||
) -> Result<()> {
|
||||
let mut tx = self.pool.begin().await?;
|
||||
|
||||
sqlx::query("DELETE FROM messages WHERE session_id = ?")
|
||||
.bind(session_id)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
for message in conversation.messages() {
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO messages (session_id, role, content_json, created_timestamp)
|
||||
VALUES (?, ?, ?, ?)
|
||||
"#,
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(role_to_string(&message.role))
|
||||
.bind(serde_json::to_string(&message.content)?)
|
||||
.bind(message.created)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
}
|
||||
|
||||
tx.commit().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn list_sessions(&self) -> Result<Vec<Session>> {
|
||||
sqlx::query_as::<_, Session>(
|
||||
r#"
|
||||
SELECT s.id, s.working_dir, s.description, s.created_at, s.updated_at, s.extension_data,
|
||||
s.total_tokens, s.input_tokens, s.output_tokens,
|
||||
s.accumulated_total_tokens, s.accumulated_input_tokens, s.accumulated_output_tokens,
|
||||
s.schedule_id, s.recipe_json,
|
||||
COUNT(m.id) as message_count
|
||||
FROM sessions s
|
||||
INNER JOIN messages m ON s.id = m.session_id
|
||||
GROUP BY s.id
|
||||
ORDER BY s.updated_at DESC
|
||||
"#,
|
||||
)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_err(Into::into)
|
||||
}
|
||||
|
||||
async fn delete_session(&self, session_id: &str) -> Result<()> {
|
||||
let exists =
|
||||
sqlx::query_scalar::<_, bool>("SELECT EXISTS(SELECT 1 FROM sessions WHERE id = ?)")
|
||||
.bind(session_id)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
|
||||
if !exists {
|
||||
return Err(anyhow::anyhow!("Session not found"));
|
||||
}
|
||||
|
||||
sqlx::query("DELETE FROM messages WHERE session_id = ?")
|
||||
.bind(session_id)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
sqlx::query("DELETE FROM sessions WHERE id = ?")
|
||||
.bind(session_id)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_insights(&self) -> Result<SessionInsights> {
|
||||
let row = sqlx::query_as::<_, (i64, Option<i64>)>(
|
||||
r#"
|
||||
SELECT COUNT(*) as total_sessions,
|
||||
COALESCE(SUM(COALESCE(accumulated_total_tokens, total_tokens, 0)), 0) as total_tokens
|
||||
FROM sessions
|
||||
"#,
|
||||
)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(SessionInsights {
|
||||
total_sessions: row.0 as usize,
|
||||
total_tokens: row.1.unwrap_or(0),
|
||||
})
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user