use std::fs;
use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use anyhow::{Context, Result};
use rusqlite::{Connection, OptionalExtension, params};
use uuid::Uuid;
use crate::model::{ContentBlock, ConversationItem, Role, Usage};
use crate::pricing::Schedule;
use crate::tools::patch::{FileChange, FileSnapshot};
const SCHEMA: &str = r#"
CREATE TABLE IF NOT EXISTS sessions (
id TEXT PRIMARY KEY,
cwd TEXT NOT NULL,
client_session_key TEXT,
title TEXT,
head_item_id TEXT,
provider_type TEXT NOT NULL,
model TEXT NOT NULL,
model_profile TEXT,
effort TEXT,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_sessions_cwd_updated ON sessions(cwd, updated_at DESC);
CREATE TABLE IF NOT EXISTS conversation_items (
id TEXT PRIMARY KEY,
session_id TEXT NOT NULL,
parent_id TEXT,
role TEXT NOT NULL,
kind TEXT NOT NULL,
content_json TEXT NOT NULL,
usage_json TEXT,
created_at INTEGER NOT NULL,
FOREIGN KEY(session_id) REFERENCES sessions(id)
);
CREATE INDEX IF NOT EXISTS idx_items_session ON conversation_items(session_id);
CREATE INDEX IF NOT EXISTS idx_items_parent ON conversation_items(parent_id);
CREATE TABLE IF NOT EXISTS provider_state (
session_id TEXT NOT NULL,
key TEXT NOT NULL,
value_json TEXT NOT NULL,
PRIMARY KEY(session_id, key),
FOREIGN KEY(session_id) REFERENCES sessions(id)
);
CREATE TABLE IF NOT EXISTS shell_history (
id INTEGER PRIMARY KEY AUTOINCREMENT,
cwd TEXT NOT NULL,
client_session_key TEXT,
command TEXT NOT NULL,
exit_code INTEGER,
pipe_status TEXT,
started_at INTEGER NOT NULL,
duration_ms INTEGER
);
CREATE INDEX IF NOT EXISTS idx_shell_cwd_time ON shell_history(cwd, started_at DESC);
CREATE TABLE IF NOT EXISTS input_history (
id INTEGER PRIMARY KEY AUTOINCREMENT,
text TEXT NOT NULL,
created_at INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS file_snapshots (
id INTEGER PRIMARY KEY AUTOINCREMENT,
session_id TEXT NOT NULL,
turn_item_id TEXT NOT NULL,
path TEXT NOT NULL,
change TEXT NOT NULL,
before TEXT,
after_len INTEGER NOT NULL,
after_hash TEXT NOT NULL,
added INTEGER NOT NULL,
removed INTEGER NOT NULL,
restorable INTEGER NOT NULL,
created_at INTEGER NOT NULL,
FOREIGN KEY(session_id) REFERENCES sessions(id)
);
CREATE INDEX IF NOT EXISTS idx_snapshots_turn ON file_snapshots(session_id, turn_item_id);
CREATE TABLE IF NOT EXISTS model_prices (
key TEXT PRIMARY KEY,
source TEXT NOT NULL,
schedule_json TEXT NOT NULL,
fetched_at INTEGER NOT NULL
);
"#;
pub const TURN_INTERRUPTED_NOTICE: &str = "[The user interrupted the previous turn. Do not continue or retry its unfinished task unless the user explicitly asks you to.]";
pub const CONVERSATION_SUMMARY_PREFIX: &str = "[Compacted conversation summary]";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Session {
pub id: String,
pub cwd: String,
pub title: Option<String>,
pub head_item_id: Option<String>,
pub provider_type: String,
pub model: String,
pub model_profile: Option<String>,
pub effort: Option<String>,
pub created_at: i64,
pub updated_at: i64,
}
#[derive(Debug, Clone)]
pub struct NewSession {
pub cwd: String,
pub provider_type: String,
pub model: String,
pub client_session_key: Option<String>,
pub model_profile: Option<String>,
pub effort: Option<String>,
}
impl NewSession {
pub fn new(
cwd: impl Into<String>,
provider_type: impl Into<String>,
model: impl Into<String>,
) -> Self {
Self {
cwd: cwd.into(),
provider_type: provider_type.into(),
model: model.into(),
client_session_key: None,
model_profile: None,
effort: None,
}
}
pub fn with_client_session_key(mut self, key: impl Into<String>) -> Self {
self.client_session_key = Some(key.into());
self
}
pub fn with_model_selection(
mut self,
profile: impl Into<String>,
effort: Option<&str>,
) -> Self {
self.model_profile = Some(profile.into());
self.effort = effort.map(str::to_owned);
self
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ShellHistoryItem {
pub id: i64,
pub cwd: String,
pub command: String,
pub exit_code: Option<i32>,
pub pipe_status: Option<String>,
pub started_at: i64,
pub duration_ms: Option<i64>,
}
pub struct SessionStore {
connection: Connection,
}
impl SessionStore {
pub fn open(path: &Path) -> Result<Self> {
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)
.with_context(|| format!("create state directory {}", parent.display()))?;
}
let connection = Connection::open(path)
.with_context(|| format!("open session database {}", path.display()))?;
Self::initialize(connection)
}
pub fn open_in_memory() -> Result<Self> {
Self::initialize(Connection::open_in_memory()?)
}
fn initialize(connection: Connection) -> Result<Self> {
connection.busy_timeout(std::time::Duration::from_millis(1000))?;
connection.pragma_update(None, "foreign_keys", "ON")?;
connection.pragma_update(None, "synchronous", "NORMAL")?;
connection.pragma_update(None, "journal_mode", "WAL")?;
connection.execute_batch(SCHEMA)?;
migrate_schema(&connection)?;
Ok(Self { connection })
}
pub fn create_session(&mut self, new: NewSession) -> Result<Session> {
let id = format!("a_{}", Uuid::new_v4().simple());
let now = now_millis();
self.connection.execute(
"INSERT INTO sessions (id, cwd, client_session_key, provider_type, model, model_profile, effort, created_at, updated_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?8)",
params![id, new.cwd, new.client_session_key, new.provider_type, new.model, new.model_profile, new.effort, now],
)?;
self.get_session(&id)?.context("new session disappeared")
}
pub fn get_session(&self, id: &str) -> Result<Option<Session>> {
self.connection
.query_row(
"SELECT id, cwd, title, head_item_id, provider_type, model, model_profile, effort, created_at, updated_at FROM sessions WHERE id = ?1",
[id],
row_to_session,
)
.optional()
.map_err(Into::into)
}
pub fn find_latest_session(&self, cwd: &str) -> Result<Option<Session>> {
self.connection
.query_row(
"SELECT id, cwd, title, head_item_id, provider_type, model, model_profile, effort, created_at, updated_at FROM sessions WHERE cwd = ?1 ORDER BY updated_at DESC, rowid DESC LIMIT 1",
[cwd],
row_to_session,
)
.optional()
.map_err(Into::into)
}
pub fn find_client_session(&self, cwd: &str, key: &str) -> Result<Option<Session>> {
self.connection
.query_row(
"SELECT id, cwd, title, head_item_id, provider_type, model, model_profile, effort, created_at, updated_at FROM sessions WHERE cwd = ?1 AND client_session_key = ?2 LIMIT 1",
params![cwd, key],
row_to_session,
)
.optional()
.map_err(Into::into)
}
pub fn recent_sessions(&self, cwd: &str, limit: usize) -> Result<Vec<Session>> {
let mut statement = self.connection.prepare(
"SELECT id, cwd, title, head_item_id, provider_type, model, model_profile, effort, created_at, updated_at FROM sessions WHERE cwd = ?1 AND EXISTS (SELECT 1 FROM conversation_items WHERE conversation_items.session_id = sessions.id AND conversation_items.kind = 'user_message') ORDER BY updated_at DESC, rowid DESC LIMIT ?2",
)?;
let rows = statement.query_map(params![cwd, limit], row_to_session)?;
rows.collect::<rusqlite::Result<Vec<_>>>()
.map_err(Into::into)
}
pub fn first_user_prompt(&self, session_id: &str) -> Result<Option<String>> {
let content = self
.connection
.query_row(
"SELECT content_json FROM conversation_items WHERE session_id = ?1 AND kind = 'user_message' ORDER BY created_at ASC, rowid ASC LIMIT 1",
[session_id],
|row| row.get::<_, String>(0),
)
.optional()?;
content
.map(|content| {
let blocks: Vec<ContentBlock> = serde_json::from_str(&content)?;
Ok(blocks.into_iter().find_map(|block| match block {
ContentBlock::Text(text) => Some(text),
_ => None,
}))
})
.transpose()
.map(Option::flatten)
}
pub fn rebind_client_session_key(
&mut self,
cwd: &str,
key: &str,
target_session_id: &str,
) -> Result<()> {
let transaction = self.connection.transaction()?;
let (target_cwd, target_key): (String, Option<String>) = transaction
.query_row(
"SELECT cwd, client_session_key FROM sessions WHERE id = ?1",
[target_session_id],
|row| Ok((row.get(0)?, row.get(1)?)),
)
.optional()?
.with_context(|| format!("session not found: {target_session_id}"))?;
if target_cwd != cwd {
anyhow::bail!("cannot resume a session from a different cwd");
}
if target_key.as_deref().is_some_and(|target| target != key) {
anyhow::bail!("session is active in another Fish process");
}
transaction.execute(
"UPDATE sessions SET client_session_key = NULL WHERE cwd = ?1 AND client_session_key = ?2",
params![cwd, key],
)?;
transaction.execute(
"UPDATE sessions SET client_session_key = ?1, updated_at = ?2 WHERE id = ?3",
params![key, now_millis(), target_session_id],
)?;
transaction.commit()?;
Ok(())
}
pub fn append_item(
&mut self,
session_id: &str,
role: Role,
blocks: Vec<ContentBlock>,
) -> Result<ConversationItem> {
let kind = item_kind(role, &blocks);
self.append_item_with_kind(session_id, role, blocks, kind, None)
}
pub fn append_assistant_item(
&mut self,
session_id: &str,
blocks: Vec<ContentBlock>,
usage: Option<Usage>,
) -> Result<ConversationItem> {
let kind = item_kind(Role::Assistant, &blocks);
self.append_item_with_kind(session_id, Role::Assistant, blocks, kind, usage)
}
pub fn append_turn_interrupted(&mut self, session_id: &str) -> Result<ConversationItem> {
self.append_item_with_kind(
session_id,
Role::User,
vec![ContentBlock::Text(TURN_INTERRUPTED_NOTICE.into())],
"turn_interrupted",
None,
)
}
pub fn replace_branch_with_summary(
&mut self,
session_id: &str,
summary: &str,
) -> Result<ConversationItem> {
self.get_session(session_id)?
.with_context(|| format!("session not found: {session_id}"))?;
let id = format!("i_{}", Uuid::new_v4().simple());
let created_at = now_millis();
let blocks = vec![ContentBlock::Text(format!(
"{CONVERSATION_SUMMARY_PREFIX}\n{summary}"
))];
let content_json = serde_json::to_string(&blocks)?;
let transaction = self.connection.transaction()?;
transaction.execute(
"INSERT INTO conversation_items (id, session_id, parent_id, role, kind, content_json, usage_json, created_at) VALUES (?1, ?2, NULL, 'user', 'conversation_summary', ?3, NULL, ?4)",
params![id, session_id, content_json, created_at],
)?;
transaction.execute(
"UPDATE sessions SET head_item_id = ?1, updated_at = ?2 WHERE id = ?3",
params![id, created_at, session_id],
)?;
transaction.commit()?;
Ok(ConversationItem {
id,
session_id: session_id.into(),
parent_id: None,
role: Role::User,
blocks,
usage: None,
created_at,
})
}
fn append_item_with_kind(
&mut self,
session_id: &str,
role: Role,
blocks: Vec<ContentBlock>,
kind: &str,
usage: Option<Usage>,
) -> Result<ConversationItem> {
let parent_id = self
.get_session(session_id)?
.with_context(|| format!("session not found: {session_id}"))?
.head_item_id;
let id = format!("i_{}", Uuid::new_v4().simple());
let created_at = now_millis();
let content_json = serde_json::to_string(&blocks)?;
let usage_json = usage
.map(|usage| serde_json::to_string(&usage))
.transpose()?;
let role_text = role_to_str(role);
let transaction = self.connection.transaction()?;
transaction.execute(
"INSERT INTO conversation_items (id, session_id, parent_id, role, kind, content_json, usage_json, created_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
params![id, session_id, parent_id, role_text, kind, content_json, usage_json, created_at],
)?;
transaction.execute(
"UPDATE sessions SET head_item_id = ?1, updated_at = ?2 WHERE id = ?3",
params![id, created_at, session_id],
)?;
transaction.commit()?;
Ok(ConversationItem {
id,
session_id: session_id.into(),
parent_id,
role,
blocks,
usage,
created_at,
})
}
pub fn get_item(&self, id: &str) -> Result<Option<ConversationItem>> {
self.connection
.query_row(
"SELECT id, session_id, parent_id, role, content_json, usage_json, created_at FROM conversation_items WHERE id = ?1",
[id],
row_to_item,
)
.optional()
.map_err(Into::into)
}
pub fn active_branch(&self, session_id: &str) -> Result<Vec<ConversationItem>> {
let mut current = self
.get_session(session_id)?
.with_context(|| format!("session not found: {session_id}"))?
.head_item_id;
let mut branch = Vec::new();
while let Some(id) = current {
let item = self
.get_item(&id)?
.with_context(|| format!("conversation item not found: {id}"))?;
current = item.parent_id.clone();
branch.push(item);
}
branch.reverse();
Ok(branch)
}
pub fn rewind(&mut self, session_id: &str, item_id: &str) -> Result<()> {
let item = self
.get_item(item_id)?
.with_context(|| format!("rewind item not found: {item_id}"))?;
if item.session_id != session_id {
anyhow::bail!("rewind item does not belong to session {session_id}");
}
self.connection.execute(
"UPDATE sessions SET head_item_id = ?1, updated_at = ?2 WHERE id = ?3",
params![item_id, now_millis(), session_id],
)?;
Ok(())
}
pub fn user_checkpoints(&self, session_id: &str) -> Result<Vec<ConversationItem>> {
let mut statement = self.connection.prepare(
"SELECT id, session_id, parent_id, role, content_json, usage_json, created_at FROM conversation_items WHERE session_id = ?1 AND kind = 'user_message' ORDER BY created_at ASC, rowid ASC",
)?;
let rows = statement.query_map([session_id], row_to_item)?;
rows.collect::<rusqlite::Result<Vec<_>>>()
.map_err(Into::into)
}
pub fn count_items(&self, session_id: &str) -> Result<usize> {
Ok(self.connection.query_row(
"SELECT COUNT(*) FROM conversation_items WHERE session_id = ?1",
[session_id],
|row| row.get(0),
)?)
}
pub fn update_model_selection(
&self,
session_id: &str,
provider_type: &str,
model: &str,
model_profile: &str,
effort: Option<&str>,
) -> Result<()> {
self.connection.execute(
"UPDATE sessions SET provider_type = ?1, model = ?2, model_profile = ?3, effort = ?4, updated_at = ?5 WHERE id = ?6",
params![provider_type, model, model_profile, effort, now_millis(), session_id],
)?;
Ok(())
}
pub fn clear_session(&self, session_id: &str) -> Result<()> {
self.connection.execute(
"UPDATE sessions SET head_item_id = NULL, updated_at = ?1 WHERE id = ?2",
params![now_millis(), session_id],
)?;
Ok(())
}
pub fn record_input_history(&self, text: &str) -> Result<()> {
self.connection.execute(
"INSERT INTO input_history (text, created_at) VALUES (?1, ?2)",
params![text, now_millis()],
)?;
Ok(())
}
pub fn recent_input_history(&self, limit: usize) -> Result<Vec<String>> {
let mut statement = self
.connection
.prepare("SELECT text FROM input_history ORDER BY id DESC LIMIT ?1")?;
let rows = statement.query_map([limit], |row| row.get::<_, String>(0))?;
let mut entries = rows.collect::<rusqlite::Result<Vec<_>>>()?;
entries.reverse();
Ok(entries)
}
pub fn prune_input_history(&self, maximum: usize) -> Result<()> {
self.connection.execute(
"DELETE FROM input_history WHERE id NOT IN (SELECT id FROM input_history ORDER BY id DESC LIMIT ?1)",
[maximum],
)?;
Ok(())
}
pub fn record_file_snapshots(
&self,
session_id: &str,
turn_item_id: &str,
snapshots: &[FileSnapshot],
) -> Result<()> {
let now = now_millis();
for snapshot in snapshots {
self.connection.execute(
"INSERT INTO file_snapshots (session_id, turn_item_id, path, change, before, after_len, after_hash, added, removed, restorable, created_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)",
params![
session_id,
turn_item_id,
snapshot.path.to_string_lossy(),
snapshot.change.as_str(),
snapshot.before,
snapshot.after_len as i64,
snapshot.after_hash.to_string(),
snapshot.added as i64,
snapshot.removed as i64,
i64::from(snapshot.restorable),
now
],
)?;
}
Ok(())
}
pub fn snapshots_from_turn(
&self,
session_id: &str,
item_id: &str,
) -> Result<Vec<FileSnapshot>> {
let mut statement = self.connection.prepare(
"SELECT path, change, before, after_len, after_hash, added, removed, restorable FROM file_snapshots
WHERE session_id = ?1 AND created_at >= (
SELECT created_at FROM conversation_items WHERE id = ?2
)
ORDER BY id DESC",
)?;
let rows = statement.query_map(params![session_id, item_id], |row| {
let change: String = row.get(1)?;
let hash: String = row.get(4)?;
Ok(FileSnapshot {
path: PathBuf::from(row.get::<_, String>(0)?),
change: FileChange::parse(&change).unwrap_or(FileChange::Modified),
before: row.get(2)?,
after_len: row.get::<_, i64>(3)? as u64,
after_hash: hash.parse().unwrap_or_default(),
added: row.get::<_, i64>(5)? as usize,
removed: row.get::<_, i64>(6)? as usize,
restorable: row.get::<_, i64>(7)? != 0,
})
})?;
rows.collect::<rusqlite::Result<Vec<_>>>()
.map_err(Into::into)
}
pub fn delete_snapshots_from_turn(&self, session_id: &str, item_id: &str) -> Result<()> {
self.connection.execute(
"DELETE FROM file_snapshots WHERE session_id = ?1 AND created_at >= (
SELECT created_at FROM conversation_items WHERE id = ?2
)",
params![session_id, item_id],
)?;
Ok(())
}
pub fn session_request_usage(&self, session_id: &str) -> Result<Vec<Usage>> {
let mut statement = self.connection.prepare(
"SELECT usage_json FROM conversation_items WHERE session_id = ?1 AND usage_json IS NOT NULL ORDER BY created_at ASC, rowid ASC",
)?;
let rows = statement.query_map([session_id], |row| row.get::<_, String>(0))?;
let mut requests = Vec::new();
for row in rows {
if let Ok(usage) = serde_json::from_str::<Usage>(&row?) {
requests.push(usage);
}
}
Ok(requests)
}
pub fn cached_prices(&self, key: &str, ttl: Duration) -> Result<Option<(Schedule, String)>> {
let cutoff = now_millis() - ttl.as_millis() as i64;
let mut statement = self.connection.prepare(
"SELECT source, schedule_json FROM model_prices WHERE key = ?1 AND fetched_at >= ?2",
)?;
let mut rows = statement.query(params![key, cutoff])?;
let Some(row) = rows.next()? else {
return Ok(None);
};
let source: String = row.get(0)?;
let schedule: String = row.get(1)?;
let Ok(schedule) = serde_json::from_str::<Schedule>(&schedule) else {
return Ok(None);
};
Ok(Some((schedule, source)))
}
pub fn cache_prices(&self, key: &str, source: &str, schedule: &Schedule) -> Result<()> {
self.connection.execute(
"INSERT INTO model_prices (key, source, schedule_json, fetched_at) VALUES (?1, ?2, ?3, ?4) ON CONFLICT(key) DO UPDATE SET source = excluded.source, schedule_json = excluded.schedule_json, fetched_at = excluded.fetched_at",
params![
key,
source,
serde_json::to_string(schedule)?,
now_millis()
],
)?;
Ok(())
}
pub fn set_provider_state(
&self,
session_id: &str,
key: &str,
value: &serde_json::Value,
) -> Result<()> {
self.connection.execute(
"INSERT INTO provider_state (session_id, key, value_json) VALUES (?1, ?2, ?3) ON CONFLICT(session_id, key) DO UPDATE SET value_json = excluded.value_json",
params![session_id, key, serde_json::to_string(value)?],
)?;
Ok(())
}
pub fn provider_state(&self, session_id: &str, key: &str) -> Result<Option<serde_json::Value>> {
let value = self
.connection
.query_row(
"SELECT value_json FROM provider_state WHERE session_id = ?1 AND key = ?2",
params![session_id, key],
|row| row.get::<_, String>(0),
)
.optional()?;
value
.map(|value| serde_json::from_str(&value).map_err(Into::into))
.transpose()
}
#[allow(clippy::too_many_arguments)]
pub fn record_shell_history(
&self,
cwd: &str,
client_session_key: Option<&str>,
command: &str,
exit_code: Option<i32>,
started_at: i64,
duration_ms: Option<i64>,
pipe_status: Option<&str>,
) -> Result<()> {
self.connection.execute(
"INSERT INTO shell_history (cwd, client_session_key, command, exit_code, pipe_status, started_at, duration_ms) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
params![cwd, client_session_key, command, exit_code, pipe_status, started_at, duration_ms],
)?;
Ok(())
}
pub fn recent_shell_history(
&self,
cwd: &str,
client_session_key: Option<&str>,
limit: usize,
) -> Result<Vec<ShellHistoryItem>> {
let query = if client_session_key.is_some() {
"SELECT id, cwd, command, exit_code, pipe_status, started_at, duration_ms FROM shell_history WHERE cwd = ?1 AND client_session_key = ?2 ORDER BY started_at DESC, id DESC LIMIT ?3"
} else {
"SELECT id, cwd, command, exit_code, pipe_status, started_at, duration_ms FROM shell_history WHERE cwd = ?1 ORDER BY started_at DESC, id DESC LIMIT ?3"
};
let mut statement = self.connection.prepare(query)?;
let rows = statement.query_map(params![cwd, client_session_key, limit], |row| {
Ok(ShellHistoryItem {
id: row.get(0)?,
cwd: row.get(1)?,
command: row.get(2)?,
exit_code: row.get(3)?,
pipe_status: row.get(4)?,
started_at: row.get(5)?,
duration_ms: row.get(6)?,
})
})?;
rows.collect::<rusqlite::Result<Vec<_>>>()
.map_err(Into::into)
}
pub fn prune_shell_history(&self, maximum: usize) -> Result<()> {
self.connection.execute(
"DELETE FROM shell_history WHERE id NOT IN (SELECT id FROM shell_history ORDER BY started_at DESC, id DESC LIMIT ?1)",
[maximum],
)?;
Ok(())
}
pub fn shell_history_count(&self) -> Result<usize> {
Ok(self
.connection
.query_row("SELECT COUNT(*) FROM shell_history", [], |row| row.get(0))?)
}
}
pub fn default_database_path(home: &Path) -> PathBuf {
std::env::var_os("XDG_STATE_HOME")
.map(PathBuf::from)
.unwrap_or_else(|| home.join(".local/state"))
.join("a/sessions.db")
}
fn migrate_schema(connection: &Connection) -> Result<()> {
ensure_column(connection, "sessions", "client_session_key", "TEXT")?;
ensure_column(connection, "sessions", "model_profile", "TEXT")?;
ensure_column(connection, "sessions", "effort", "TEXT")?;
ensure_column(connection, "shell_history", "client_session_key", "TEXT")?;
ensure_column(connection, "conversation_items", "usage_json", "TEXT")?;
rebuild_cache_table(connection, "model_prices", "schedule_json")?;
connection.execute_batch(
"CREATE UNIQUE INDEX IF NOT EXISTS idx_sessions_client_key ON sessions(cwd, client_session_key) WHERE client_session_key IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_shell_client_time ON shell_history(cwd, client_session_key, started_at DESC);
PRAGMA user_version = 4;",
)?;
Ok(())
}
fn rebuild_cache_table(connection: &Connection, table: &str, expected_column: &str) -> Result<()> {
let mut statement = connection.prepare(&format!("PRAGMA table_info({table})"))?;
let columns = statement
.query_map([], |row| row.get::<_, String>(1))?
.collect::<rusqlite::Result<Vec<_>>>()?;
if columns.is_empty() || columns.iter().any(|column| column == expected_column) {
return Ok(());
}
drop(statement);
connection.execute_batch(&format!("DROP TABLE {table}"))?;
connection.execute_batch(SCHEMA)?;
Ok(())
}
fn ensure_column(connection: &Connection, table: &str, column: &str, kind: &str) -> Result<()> {
let mut statement = connection.prepare(&format!("PRAGMA table_info({table})"))?;
let columns = statement
.query_map([], |row| row.get::<_, String>(1))?
.collect::<rusqlite::Result<Vec<_>>>()?;
if !columns.iter().any(|existing| existing == column) {
connection.execute(
&format!("ALTER TABLE {table} ADD COLUMN {column} {kind}"),
[],
)?;
}
Ok(())
}
fn row_to_session(row: &rusqlite::Row<'_>) -> rusqlite::Result<Session> {
Ok(Session {
id: row.get(0)?,
cwd: row.get(1)?,
title: row.get(2)?,
head_item_id: row.get(3)?,
provider_type: row.get(4)?,
model: row.get(5)?,
model_profile: row.get(6)?,
effort: row.get(7)?,
created_at: row.get(8)?,
updated_at: row.get(9)?,
})
}
fn row_to_item(row: &rusqlite::Row<'_>) -> rusqlite::Result<ConversationItem> {
let role: String = row.get(3)?;
let content: String = row.get(4)?;
let usage: Option<String> = row.get(5)?;
Ok(ConversationItem {
id: row.get(0)?,
session_id: row.get(1)?,
parent_id: row.get(2)?,
role: str_to_role(&role).map_err(|error| {
rusqlite::Error::FromSqlConversionFailure(3, rusqlite::types::Type::Text, error.into())
})?,
blocks: serde_json::from_str(&content).map_err(|error| {
rusqlite::Error::FromSqlConversionFailure(4, rusqlite::types::Type::Text, error.into())
})?,
usage: usage
.map(|usage| serde_json::from_str(&usage))
.transpose()
.map_err(|error| {
rusqlite::Error::FromSqlConversionFailure(
5,
rusqlite::types::Type::Text,
error.into(),
)
})?,
created_at: row.get(6)?,
})
}
fn role_to_str(role: Role) -> &'static str {
match role {
Role::System => "system",
Role::User => "user",
Role::Assistant => "assistant",
Role::Tool => "tool",
}
}
fn str_to_role(role: &str) -> Result<Role, String> {
match role {
"system" => Ok(Role::System),
"user" => Ok(Role::User),
"assistant" => Ok(Role::Assistant),
"tool" => Ok(Role::Tool),
_ => Err(format!("unknown conversation role: {role}")),
}
}
fn item_kind(role: Role, blocks: &[ContentBlock]) -> &'static str {
match role {
Role::User => "user_message",
Role::Tool => "tool_result",
Role::System => "system_checkpoint",
Role::Assistant
if blocks
.iter()
.any(|block| matches!(block, ContentBlock::ToolCall(_))) =>
{
"tool_call"
}
Role::Assistant
if blocks
.iter()
.any(|block| matches!(block, ContentBlock::Reasoning(_))) =>
{
"assistant_reasoning"
}
Role::Assistant => "assistant_text",
}
}
fn now_millis() -> i64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as i64
}