pub mod agent_inbox; pub mod agent_run; pub mod context_checkpoint; pub mod error; pub mod memory; pub mod message; pub mod scheduler; pub mod session; pub mod usage; pub use context_checkpoint::{ContextCheckpoint, ContextCheckpointState, NewContextCheckpoint}; pub use error::StorageError; pub use scheduler::{ ClaimedScheduledRun, DeliveryPolicy, JobRun, ScheduledDeliveryStatus, ScheduledJob, ScheduledJobUpdate, ScheduledOutcomeKind, ScheduledRunCompletion, ScheduledRunStatus, }; pub use usage::{SessionUsageTotals, TurnUsageRecord}; use sqlx::sqlite::{ SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteRow, SqliteSynchronous, }; use sqlx::{Pool, Row, Sqlite}; use std::path::Path; use tokio::time::{Duration, sleep}; const SCHEMA_VERSION: i64 = 11; const INSERT_MESSAGE_SQL: &str = r#" INSERT INTO messages ( id, session_id, seq, role, content, reasoning_content, provider_state, turn_id, iteration, completion_status, client_visibility, turn_origin, media_refs, tool_call_id, tool_name, tool_calls, source, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) "#; const INSERT_MESSAGE_IF_ABSENT_SQL: &str = r#" INSERT INTO messages ( id, session_id, seq, role, content, reasoning_content, provider_state, turn_id, iteration, completion_status, client_visibility, turn_origin, media_refs, tool_call_id, tool_name, tool_calls, source, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO NOTHING "#; pub(crate) fn insert_message_query<'a>( session_id: &'a str, msg: &'a crate::storage::message::MessageMeta, ) -> sqlx::query::Query<'a, Sqlite, sqlx::sqlite::SqliteArguments> { sqlx::query(INSERT_MESSAGE_SQL) .bind(&msg.id) .bind(session_id) .bind(msg.seq) .bind(&msg.role) .bind(&msg.content) .bind(&msg.reasoning_content) .bind(&msg.provider_state) .bind(&msg.turn_id) .bind(msg.iteration) .bind(msg.completion_status.as_str()) .bind(msg.client_visibility.as_str()) .bind(msg.turn_origin.as_str()) .bind(&msg.media_refs) .bind(&msg.tool_call_id) .bind(&msg.tool_name) .bind(&msg.tool_calls) .bind(&msg.source) .bind(msg.created_at) } pub(crate) fn insert_message_if_absent_query<'a>( session_id: &'a str, msg: &'a crate::storage::message::MessageMeta, ) -> sqlx::query::Query<'a, Sqlite, sqlx::sqlite::SqliteArguments> { sqlx::query(INSERT_MESSAGE_IF_ABSENT_SQL) .bind(&msg.id) .bind(session_id) .bind(msg.seq) .bind(&msg.role) .bind(&msg.content) .bind(&msg.reasoning_content) .bind(&msg.provider_state) .bind(&msg.turn_id) .bind(msg.iteration) .bind(msg.completion_status.as_str()) .bind(msg.client_visibility.as_str()) .bind(msg.turn_origin.as_str()) .bind(&msg.media_refs) .bind(&msg.tool_call_id) .bind(&msg.tool_name) .bind(&msg.tool_calls) .bind(&msg.source) .bind(msg.created_at) } fn message_meta_from_row(row: SqliteRow) -> crate::storage::message::MessageMeta { let completion_status: String = row.get("completion_status"); crate::storage::message::MessageMeta { id: row.get("id"), session_id: row.get("session_id"), seq: row.get("seq"), role: row.get("role"), content: row.get("content"), reasoning_content: row.get("reasoning_content"), provider_state: row.get("provider_state"), turn_id: row.get("turn_id"), iteration: row.get("iteration"), completion_status: crate::bus::CompletionStatus::from_storage(&completion_status), client_visibility: parse_visibility(row.get("client_visibility")), turn_origin: parse_turn_origin(row.get("turn_origin")), media_refs: row.get("media_refs"), tool_call_id: row.get("tool_call_id"), tool_name: row.get("tool_name"), tool_calls: row.get("tool_calls"), source: row.get("source"), created_at: row.get("created_at"), } } fn parse_visibility(value: String) -> crate::bus::ClientVisibility { if value == "hidden" { crate::bus::ClientVisibility::Hidden } else { crate::bus::ClientVisibility::Visible } } fn parse_turn_origin(value: String) -> crate::bus::TurnOrigin { match value.as_str() { "agent_continuation" => crate::bus::TurnOrigin::AgentContinuation, "scheduled" => crate::bus::TurnOrigin::Scheduled, _ => crate::bus::TurnOrigin::User, } } pub struct Storage { pub(crate) pool: Pool, } impl Storage { /// 打开或创建数据库 pub async fn new(db_path: &Path) -> Result { let options = SqliteConnectOptions::new() .filename(db_path) .create_if_missing(true) .journal_mode(SqliteJournalMode::Wal) .synchronous(SqliteSynchronous::Normal) .busy_timeout(Duration::from_secs(5)) .foreign_keys(true); let pool = SqlitePoolOptions::new() .max_connections(8) .connect_with(options) .await?; let storage = Self { pool }; storage.init_schema().await?; Ok(storage) } /// 初始化数据库 schema async fn init_schema(&self) -> Result<(), StorageError> { sqlx::query( r#" CREATE TABLE IF NOT EXISTS sessions ( id TEXT PRIMARY KEY, channel TEXT NOT NULL, chat_id TEXT NOT NULL, dialog_id TEXT NOT NULL, title TEXT NOT NULL DEFAULT '新对话', created_at INTEGER NOT NULL, last_active_at INTEGER NOT NULL, message_count INTEGER DEFAULT 0, routing_info TEXT, archived_at INTEGER, deleted_at INTEGER, last_consolidated_at INTEGER, last_compressed_message_at INTEGER, active_context_checkpoint_id TEXT, context_generation INTEGER NOT NULL DEFAULT 0, delivery_context TEXT, delivery_context_updated_at INTEGER, UNIQUE(channel, chat_id, dialog_id) ) "#, ) .execute(&self.pool) .await?; sqlx::query( r#" CREATE TABLE IF NOT EXISTS context_checkpoints ( id TEXT PRIMARY KEY, session_id TEXT NOT NULL, generation INTEGER NOT NULL, parent_checkpoint_id TEXT, summary TEXT NOT NULL, first_retained_seq INTEGER NOT NULL, source_max_seq INTEGER NOT NULL, trigger_reason TEXT NOT NULL, provider_kind TEXT NOT NULL, model TEXT NOT NULL, tokens_before INTEGER NOT NULL, tokens_after INTEGER NOT NULL, degraded INTEGER NOT NULL DEFAULT 0, created_at INTEGER NOT NULL, FOREIGN KEY(session_id) REFERENCES sessions(id) ON DELETE CASCADE, UNIQUE(session_id, generation) ) "#, ) .execute(&self.pool) .await?; sqlx::query( "CREATE INDEX IF NOT EXISTS idx_context_checkpoints_session_created ON context_checkpoints(session_id, created_at DESC)", ) .execute(&self.pool) .await?; sqlx::query( r#" CREATE INDEX IF NOT EXISTS idx_sessions_chat ON sessions(channel, chat_id, deleted_at) "#, ) .execute(&self.pool) .await?; sqlx::query( r#" CREATE TABLE IF NOT EXISTS messages ( id TEXT PRIMARY KEY, session_id TEXT NOT NULL, seq INTEGER NOT NULL, role TEXT NOT NULL, content TEXT NOT NULL, media_refs TEXT, tool_call_id TEXT, tool_name TEXT, tool_calls TEXT, source TEXT, reasoning_content TEXT, provider_state TEXT, turn_id TEXT, iteration INTEGER, completion_status TEXT NOT NULL DEFAULT 'completed', client_visibility TEXT NOT NULL DEFAULT 'visible', turn_origin TEXT NOT NULL DEFAULT 'user', created_at INTEGER NOT NULL, FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE ) "#, ) .execute(&self.pool) .await?; sqlx::query( r#" CREATE INDEX IF NOT EXISTS idx_messages_session_seq ON messages(session_id, seq) "#, ) .execute(&self.pool) .await?; // Session-scoped task plans. A session may have at most one active plan, // while independent items can be executed concurrently. sqlx::query( r#" CREATE TABLE IF NOT EXISTS task_plans ( id TEXT PRIMARY KEY, session_id TEXT NOT NULL, objective TEXT NOT NULL, status TEXT NOT NULL DEFAULT 'active', version INTEGER NOT NULL DEFAULT 1, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL, closed_at INTEGER ) "#, ) .execute(&self.pool) .await?; sqlx::query( r#" CREATE UNIQUE INDEX IF NOT EXISTS idx_task_plans_active_session ON task_plans(session_id) WHERE status = 'active' "#, ) .execute(&self.pool) .await?; sqlx::query( r#" CREATE TABLE IF NOT EXISTS task_items ( id TEXT NOT NULL, plan_id TEXT NOT NULL, ordinal INTEGER NOT NULL, title TEXT NOT NULL, status TEXT NOT NULL DEFAULT 'pending', executor_kind TEXT, execution_id TEXT, result_summary TEXT, error TEXT, version INTEGER NOT NULL DEFAULT 1, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL, PRIMARY KEY(plan_id, id), FOREIGN KEY(plan_id) REFERENCES task_plans(id) ON DELETE CASCADE, UNIQUE(plan_id, ordinal) ) "#, ) .execute(&self.pool) .await?; sqlx::query( "CREATE INDEX IF NOT EXISTS idx_task_items_plan ON task_items(plan_id, ordinal)", ) .execute(&self.pool) .await?; sqlx::query( r#" CREATE TABLE IF NOT EXISTS memories ( id TEXT PRIMARY KEY, key TEXT NOT NULL UNIQUE, content TEXT NOT NULL, category TEXT NOT NULL DEFAULT 'knowledge', importance REAL NOT NULL DEFAULT 0.5, session_id TEXT, created_at TEXT NOT NULL, updated_at TEXT NOT NULL ) "#, ) .execute(&self.pool) .await?; let memory_fts_exists: bool = sqlx::query_scalar( "SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'memory_fts')", ) .fetch_one(&self.pool) .await?; // FTS5 virtual table for full-text search on memories sqlx::query( r#" CREATE VIRTUAL TABLE IF NOT EXISTS memory_fts USING fts5( key, content, content=memories, content_rowid=rowid ) "#, ) .execute(&self.pool) .await?; // Triggers to keep FTS5 index in sync with memories table sqlx::query( r#" CREATE TRIGGER IF NOT EXISTS memories_ai AFTER INSERT ON memories BEGIN INSERT INTO memory_fts(rowid, key, content) VALUES (new.rowid, new.key, new.content); END "#, ) .execute(&self.pool) .await?; sqlx::query( r#" CREATE TRIGGER IF NOT EXISTS memories_ad AFTER DELETE ON memories BEGIN INSERT INTO memory_fts(memory_fts, rowid, key, content) VALUES ('delete', old.rowid, old.key, old.content); END "#, ) .execute(&self.pool) .await?; sqlx::query( r#" CREATE TRIGGER IF NOT EXISTS memories_au AFTER UPDATE ON memories BEGIN INSERT INTO memory_fts(memory_fts, rowid, key, content) VALUES ('delete', old.rowid, old.key, old.content); INSERT INTO memory_fts(rowid, key, content) VALUES (new.rowid, new.key, new.content); END "#, ) .execute(&self.pool) .await?; // Only a newly-created index needs a backfill. Triggers keep an // existing index current, so rebuilding it on every startup is wasted // work proportional to the total memory corpus. if !memory_fts_exists { sqlx::query("INSERT INTO memory_fts(memory_fts) VALUES ('rebuild')") .execute(&self.pool) .await?; } sqlx::query( r#" CREATE TABLE IF NOT EXISTS llm_calls ( id INTEGER PRIMARY KEY AUTOINCREMENT, created_at INTEGER NOT NULL, provider TEXT NOT NULL, model TEXT NOT NULL, request_body TEXT NOT NULL, response_body TEXT, error TEXT, duration_ms INTEGER ) "#, ) .execute(&self.pool) .await?; sqlx::query( r#" CREATE TABLE IF NOT EXISTS session_turn_usage ( turn_id TEXT PRIMARY KEY, session_id TEXT NOT NULL, provider TEXT NOT NULL, model TEXT NOT NULL, prompt_tokens INTEGER NOT NULL, completion_tokens INTEGER NOT NULL, total_tokens INTEGER NOT NULL, cached_input_tokens INTEGER, request_count INTEGER NOT NULL, last_prompt_tokens INTEGER NOT NULL, created_at INTEGER NOT NULL, FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE ) "#, ) .execute(&self.pool) .await?; sqlx::query( "CREATE INDEX IF NOT EXISTS idx_session_turn_usage_session_created ON session_turn_usage(session_id, created_at)", ) .execute(&self.pool) .await?; self.migrate_schema().await?; Ok(()) } /// Apply ordered, atomic migrations. Existing installations predate /// `user_version`, so each step also checks the actual table shape. async fn migrate_schema(&self) -> Result<(), StorageError> { let current: i64 = sqlx::query_scalar("PRAGMA user_version") .fetch_one(&self.pool) .await?; if current > SCHEMA_VERSION { return Err(StorageError::Migration(format!( "database schema version {current} is newer than supported version {SCHEMA_VERSION}" ))); } if current == SCHEMA_VERSION { return Ok(()); } // Acquire the migration write lock before inspecting or modifying any // legacy shape. This keeps the v11 rebuild on one connection and makes // a concurrently running older Gateway fail startup cleanly instead of // partially racing the schema migration. let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; // The legacy drops below are a pre-v8 rebuild concern: the batch // "group" concept was removed in v8 and the old `background_tasks` // table in v7. Gate them on `current < 8` so a v8 -> v9 upgrade only // adds the new transcript table and preserves existing run history. if current < 8 { // Legacy table removed in schema v7; drop it so old databases do // not keep dead rows around. sqlx::query("DROP TABLE IF EXISTS background_tasks") .execute(&mut *tx) .await?; // Schema v8 removes the batch "group" concept entirely: the // `agent_run_groups` table is gone, and the run/inbox tables are // rebuilt without their `group_id`/`scope_kind`/`scope_id` columns. // Drop the transcript table before runs and the remaining tables in // dependency order (messages -> inbox -> runs -> groups) so // foreign-key enforcement never blocks the implicit row delete. sqlx::query("DROP TABLE IF EXISTS agent_run_messages") .execute(&mut *tx) .await?; sqlx::query("DROP TABLE IF EXISTS agent_inbox_events") .execute(&mut *tx) .await?; sqlx::query("DROP TABLE IF EXISTS agent_runs") .execute(&mut *tx) .await?; sqlx::query("DROP TABLE IF EXISTS agent_run_groups") .execute(&mut *tx) .await?; } let legacy_scheduler_exists: i64 = sqlx::query_scalar( "SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = 'scheduled_jobs'", ) .fetch_one(&mut *tx) .await?; let legacy_scheduler_exists = legacy_scheduler_exists == 1; if legacy_scheduler_exists { sqlx::query( r#" CREATE TABLE IF NOT EXISTS job_runs ( id INTEGER PRIMARY KEY AUTOINCREMENT, job_id TEXT NOT NULL REFERENCES scheduled_jobs(id) ON DELETE CASCADE, started_at INTEGER NOT NULL, finished_at INTEGER NOT NULL, status TEXT NOT NULL, output TEXT, error TEXT, duration_ms INTEGER NOT NULL, result_kind TEXT, delivery_status TEXT, delivery_error TEXT ) "#, ) .execute(&mut *tx) .await?; } for (table, column, definition) in [ ("messages", "source", "source TEXT"), ("messages", "reasoning_content", "reasoning_content TEXT"), ("messages", "provider_state", "provider_state TEXT"), ("messages", "turn_id", "turn_id TEXT"), ("messages", "iteration", "iteration INTEGER"), ( "messages", "completion_status", "completion_status TEXT NOT NULL DEFAULT 'completed'", ), ( "messages", "client_visibility", "client_visibility TEXT NOT NULL DEFAULT 'visible'", ), ( "messages", "turn_origin", "turn_origin TEXT NOT NULL DEFAULT 'user'", ), ("sessions", "archived_at", "archived_at INTEGER"), ( "sessions", "last_consolidated_at", "last_consolidated_at INTEGER", ), ( "sessions", "last_compressed_message_at", "last_compressed_message_at INTEGER", ), ( "sessions", "active_context_checkpoint_id", "active_context_checkpoint_id TEXT", ), ( "sessions", "context_generation", "context_generation INTEGER NOT NULL DEFAULT 0", ), ("sessions", "delivery_context", "delivery_context TEXT"), ( "sessions", "delivery_context_updated_at", "delivery_context_updated_at INTEGER", ), ("scheduled_jobs", "locked_at", "locked_at INTEGER"), ("scheduled_jobs", "lock_owner", "lock_owner TEXT"), ("scheduled_jobs", "lease_until", "lease_until INTEGER"), ( "scheduled_jobs", "job_kind", "job_kind TEXT NOT NULL DEFAULT 'task'", ), ( "scheduled_jobs", "delivery_policy", "delivery_policy TEXT NOT NULL DEFAULT 'direct'", ), ("job_runs", "result_kind", "result_kind TEXT"), ("job_runs", "delivery_status", "delivery_status TEXT"), ("job_runs", "delivery_error", "delivery_error TEXT"), ] { let pragma = format!("PRAGMA table_info({table})"); // `table` comes exclusively from the fixed migration list above. let columns = sqlx::query(sqlx::AssertSqlSafe(pragma)) .fetch_all(&mut *tx) .await?; if !columns.is_empty() && !columns .iter() .any(|row| row.get::("name") == column) { let alter = format!("ALTER TABLE {table} ADD COLUMN {definition}"); // All identifiers and definitions come from the fixed migration list above. sqlx::query(sqlx::AssertSqlSafe(alter)) .execute(&mut *tx) .await?; } } let duplicate: Option<(String, i64, i64)> = sqlx::query_as( r#" SELECT session_id, seq, COUNT(*) FROM messages GROUP BY session_id, seq HAVING COUNT(*) > 1 LIMIT 1 "#, ) .fetch_optional(&mut *tx) .await?; if let Some((session_id, seq, count)) = duplicate { return Err(StorageError::Migration(format!( "cannot enforce unique message sequence: session {session_id} has {count} rows at seq {seq}" ))); } sqlx::query( "CREATE UNIQUE INDEX IF NOT EXISTS idx_messages_session_seq_unique ON messages(session_id, seq)", ) .execute(&mut *tx) .await?; sqlx::query( r#" CREATE TABLE IF NOT EXISTS session_turn_usage ( turn_id TEXT PRIMARY KEY, session_id TEXT NOT NULL, provider TEXT NOT NULL, model TEXT NOT NULL, prompt_tokens INTEGER NOT NULL, completion_tokens INTEGER NOT NULL, total_tokens INTEGER NOT NULL, cached_input_tokens INTEGER, request_count INTEGER NOT NULL, last_prompt_tokens INTEGER NOT NULL, created_at INTEGER NOT NULL, FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE ) "#, ) .execute(&mut *tx) .await?; sqlx::query( "CREATE INDEX IF NOT EXISTS idx_session_turn_usage_session_created ON session_turn_usage(session_id, created_at)", ) .execute(&mut *tx) .await?; sqlx::query( "CREATE INDEX IF NOT EXISTS idx_messages_session_visibility_seq ON messages(session_id, client_visibility, seq)", ) .execute(&mut *tx) .await?; sqlx::query( r#" CREATE TABLE IF NOT EXISTS context_checkpoints ( id TEXT PRIMARY KEY, session_id TEXT NOT NULL, generation INTEGER NOT NULL, parent_checkpoint_id TEXT, summary TEXT NOT NULL, first_retained_seq INTEGER NOT NULL, source_max_seq INTEGER NOT NULL, trigger_reason TEXT NOT NULL, provider_kind TEXT NOT NULL, model TEXT NOT NULL, tokens_before INTEGER NOT NULL, tokens_after INTEGER NOT NULL, degraded INTEGER NOT NULL DEFAULT 0, created_at INTEGER NOT NULL, FOREIGN KEY(session_id) REFERENCES sessions(id) ON DELETE CASCADE, UNIQUE(session_id, generation) ) "#, ) .execute(&mut *tx) .await?; sqlx::query( "CREATE INDEX IF NOT EXISTS idx_context_checkpoints_session_created ON context_checkpoints(session_id, created_at DESC)", ) .execute(&mut *tx) .await?; for statement in agent_run::AGENT_SCHEMA_STATEMENTS { sqlx::query(*statement).execute(&mut *tx).await?; } scheduler::migrate_scheduler_v11(&mut tx, legacy_scheduler_exists).await?; sqlx::query(sqlx::AssertSqlSafe(format!( "PRAGMA user_version = {SCHEMA_VERSION}" ))) .execute(&mut *tx) .await?; tx.commit().await?; Ok(()) } pub async fn append_llm_call( &self, provider: &str, model: &str, request_body: &str, response_body: Option<&str>, error: Option<&str>, duration_ms: u64, ) -> Result<(), StorageError> { let now = chrono::Utc::now().timestamp_millis(); sqlx::query( r#" INSERT INTO llm_calls (created_at, provider, model, request_body, response_body, error, duration_ms) VALUES (?, ?, ?, ?, ?, ?, ?) "#, ) .bind(now) .bind(provider) .bind(model) .bind(request_body) .bind(response_body) .bind(error) .bind(duration_ms as i64) .execute(self.pool()) .await?; // Prune to keep last 1000 records self.prune_llm_calls(1000).await?; Ok(()) } async fn prune_llm_calls(&self, max_records: i64) -> Result<(), StorageError> { sqlx::query( r#" DELETE FROM llm_calls WHERE id <= ( SELECT COALESCE(MAX(id), 0) - ? FROM llm_calls ) "#, ) .bind(max_records) .execute(self.pool()) .await?; Ok(()) } /// 获取连接池引用（供内部 CRUD 使用） pub(crate) fn pool(&self) -> &Pool { &self.pool } pub async fn upsert_session( &self, meta: &crate::storage::session::SessionMeta, ) -> Result<(), StorageError> { sqlx::query( r#" INSERT INTO sessions (id, channel, chat_id, dialog_id, title, created_at, last_active_at, message_count, routing_info, archived_at, deleted_at, last_consolidated_at, last_compressed_message_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET title = excluded.title, last_active_at = excluded.last_active_at, message_count = excluded.message_count, routing_info = excluded.routing_info, archived_at = excluded.archived_at, deleted_at = excluded.deleted_at, last_consolidated_at = excluded.last_consolidated_at, last_compressed_message_at = excluded.last_compressed_message_at "#, ) .bind(&meta.id) .bind(&meta.channel) .bind(&meta.chat_id) .bind(&meta.dialog_id) .bind(&meta.title) .bind(meta.created_at) .bind(meta.last_active_at) .bind(meta.message_count) .bind(&meta.routing_info) .bind(meta.archived_at) .bind(meta.deleted_at) .bind(meta.last_consolidated_at) .bind(meta.last_compressed_message_at) .execute(self.pool()) .await?; Ok(()) } pub async fn get_session( &self, id: &str, ) -> Result { let row = sqlx::query( r#" SELECT id, channel, chat_id, dialog_id, title, created_at, last_active_at, message_count, routing_info, archived_at, deleted_at, last_consolidated_at, last_compressed_message_at FROM sessions WHERE id = ? AND deleted_at IS NULL "#, ) .bind(id) .fetch_optional(self.pool()) .await? .ok_or_else(|| StorageError::NotFound(id.to_string()))?; Ok(crate::storage::session::SessionMeta { id: row.get("id"), channel: row.get("channel"), chat_id: row.get("chat_id"), dialog_id: row.get("dialog_id"), title: row.get("title"), created_at: row.get("created_at"), last_active_at: row.get("last_active_at"), message_count: row.get("message_count"), routing_info: row.get("routing_info"), archived_at: row.get("archived_at"), deleted_at: row.get("deleted_at"), last_consolidated_at: row.get("last_consolidated_at"), last_compressed_message_at: row.get("last_compressed_message_at"), }) } pub async fn list_sessions( &self, channel: &str, chat_id: &str, limit: i64, include_archived: bool, ) -> Result, StorageError> { let rows = sqlx::query( r#" SELECT id, channel, chat_id, dialog_id, title, created_at, last_active_at, message_count, routing_info, archived_at, deleted_at, last_consolidated_at, last_compressed_message_at FROM sessions WHERE channel = ? AND chat_id = ? AND deleted_at IS NULL AND (? OR archived_at IS NULL) ORDER BY last_active_at DESC LIMIT ? "#, ) .bind(channel) .bind(chat_id) .bind(include_archived) .bind(limit) .fetch_all(self.pool()) .await?; Ok(rows .into_iter() .map(|row| crate::storage::session::SessionMeta { id: row.get("id"), channel: row.get("channel"), chat_id: row.get("chat_id"), dialog_id: row.get("dialog_id"), title: row.get("title"), created_at: row.get("created_at"), last_active_at: row.get("last_active_at"), message_count: row.get("message_count"), routing_info: row.get("routing_info"), archived_at: row.get("archived_at"), deleted_at: row.get("deleted_at"), last_consolidated_at: row.get("last_consolidated_at"), last_compressed_message_at: row.get("last_compressed_message_at"), }) .collect()) } pub async fn touch_session( &self, id: &str, message_count: i64, last_active_at: i64, ) -> Result<(), StorageError> { sqlx::query( r#" UPDATE sessions SET message_count = ?, last_active_at = ? WHERE id = ? "#, ) .bind(message_count) .bind(last_active_at) .bind(id) .execute(self.pool()) .await?; Ok(()) } pub async fn soft_delete_session(&self, id: &str) -> Result<(), StorageError> { let now = chrono::Utc::now().timestamp_millis(); sqlx::query(r#"UPDATE sessions SET deleted_at = ? WHERE id = ?"#) .bind(now) .bind(id) .execute(self.pool()) .await?; Ok(()) } pub async fn archive_session(&self, id: &str) -> Result<(), StorageError> { let now = chrono::Utc::now().timestamp_millis(); sqlx::query(r#"UPDATE sessions SET archived_at = ? WHERE id = ? AND deleted_at IS NULL"#) .bind(now) .bind(id) .execute(self.pool()) .await?; Ok(()) } pub async fn find_most_recent_session( &self, channel: &str, chat_id: &str, ) -> Result, StorageError> { let row = sqlx::query( r#" SELECT id, channel, chat_id, dialog_id, title, created_at, last_active_at, message_count, routing_info, archived_at, deleted_at, last_consolidated_at, last_compressed_message_at FROM sessions WHERE channel = ? AND chat_id = ? AND deleted_at IS NULL AND archived_at IS NULL ORDER BY last_active_at DESC LIMIT 1 "#, ) .bind(channel) .bind(chat_id) .fetch_optional(self.pool()) .await?; match row { Some(row) => Ok(Some(crate::storage::session::SessionMeta { id: row.get("id"), channel: row.get("channel"), chat_id: row.get("chat_id"), dialog_id: row.get("dialog_id"), title: row.get("title"), created_at: row.get("created_at"), last_active_at: row.get("last_active_at"), message_count: row.get("message_count"), routing_info: row.get("routing_info"), archived_at: row.get("archived_at"), deleted_at: row.get("deleted_at"), last_consolidated_at: row.get("last_consolidated_at"), last_compressed_message_at: row.get("last_compressed_message_at"), })), None => Ok(None), } } pub async fn append_message( &self, session_id: &str, msg: &crate::storage::message::MessageMeta, ) -> Result { insert_message_query(session_id, msg) .execute(self.pool()) .await?; Ok(msg.seq) } pub async fn persist_message_if_absent( &self, session_id: &str, msg: &crate::storage::message::MessageMeta, meta: &crate::storage::session::SessionMeta, ) -> Result { let mut tx = self.pool.begin().await?; let inserted = insert_message_if_absent_query(session_id, msg) .execute(&mut *tx) .await? .rows_affected() == 1; if inserted { sqlx::query( r#" INSERT INTO sessions (id, channel, chat_id, dialog_id, title, created_at, last_active_at, message_count, routing_info, archived_at, deleted_at, last_consolidated_at, last_compressed_message_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET title = excluded.title, last_active_at = excluded.last_active_at, message_count = excluded.message_count, routing_info = excluded.routing_info, archived_at = excluded.archived_at, deleted_at = excluded.deleted_at, last_consolidated_at = excluded.last_consolidated_at, last_compressed_message_at = excluded.last_compressed_message_at "#, ) .bind(&meta.id) .bind(&meta.channel) .bind(&meta.chat_id) .bind(&meta.dialog_id) .bind(&meta.title) .bind(meta.created_at) .bind(meta.last_active_at) .bind(meta.message_count) .bind(&meta.routing_info) .bind(meta.archived_at) .bind(meta.deleted_at) .bind(meta.last_consolidated_at) .bind(meta.last_compressed_message_at) .execute(&mut *tx) .await?; } tx.commit().await?; Ok(inserted) } /// Atomically persist all messages produced by one logical turn together /// with the resulting session metadata. A turn is either fully visible /// after restart or not visible at all. async fn persist_message_batch_inner( &self, session_id: &str, msgs: &[crate::storage::message::MessageMeta], meta: &crate::storage::session::SessionMeta, usage: Option<&crate::storage::TurnUsageRecord>, steer: Option<&crate::storage::agent_inbox::SteerConsumption>, ) -> Result<(), StorageError> { let now = chrono::Utc::now().timestamp_millis(); let mut tx = self.pool.begin().await?; for msg in msgs { insert_message_query(session_id, msg) .execute(&mut *tx) .await?; } sqlx::query( r#" INSERT INTO sessions (id, channel, chat_id, dialog_id, title, created_at, last_active_at, message_count, routing_info, archived_at, deleted_at, last_consolidated_at, last_compressed_message_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET title = excluded.title, last_active_at = excluded.last_active_at, message_count = excluded.message_count, routing_info = excluded.routing_info, archived_at = excluded.archived_at, deleted_at = excluded.deleted_at, last_consolidated_at = excluded.last_consolidated_at, last_compressed_message_at = excluded.last_compressed_message_at "#, ) .bind(&meta.id) .bind(&meta.channel) .bind(&meta.chat_id) .bind(&meta.dialog_id) .bind(&meta.title) .bind(meta.created_at) .bind(meta.last_active_at) .bind(meta.message_count) .bind(&meta.routing_info) .bind(meta.archived_at) .bind(meta.deleted_at) .bind(meta.last_consolidated_at) .bind(meta.last_compressed_message_at) .execute(&mut *tx) .await?; if let Some(usage) = usage { debug_assert_eq!(session_id, usage.session_id); let request_count = msgs .iter() .filter_map(|message| message.iteration) .max() .map_or(1_i64, |iteration| iteration.saturating_add(1)); sqlx::query( r#" INSERT INTO session_turn_usage ( turn_id, session_id, provider, model, prompt_tokens, completion_tokens, total_tokens, cached_input_tokens, request_count, last_prompt_tokens, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(turn_id) DO NOTHING "#, ) .bind(&usage.turn_id) .bind(&usage.session_id) .bind(&usage.provider) .bind(&usage.model) .bind(i64::from(usage.usage.prompt_tokens)) .bind(i64::from(usage.usage.completion_tokens)) .bind(i64::from(usage.usage.total_tokens)) .bind(usage.usage.cached_tokens.map(i64::from)) .bind(request_count) .bind(i64::from(usage.last_prompt_tokens)) .bind(usage.created_at) .execute(&mut *tx) .await?; } if let Some(steer) = steer { let mut consumed = 0i64; for event_id in &steer.event_ids { let rows = sqlx::query( "UPDATE agent_inbox_events SET status = 'consumed', consumed_at = ?, \ updated_at = ? \ WHERE id = ? AND status = 'admitted' AND admitted_turn_id = ?", ) .bind(now) .bind(now) .bind(event_id) .bind(&steer.admitted_turn_id) .execute(&mut *tx) .await? .rows_affected(); consumed += rows as i64; } if consumed > 0 { sqlx::query( "UPDATE agent_session_state \ SET pending_event_count = MAX(pending_event_count - ?, 0), \ revision = revision + 1, updated_at = ? \ WHERE root_session_id = ?", ) .bind(consumed) .bind(now) .bind(session_id) .execute(&mut *tx) .await?; } } tx.commit().await?; Ok(()) } pub async fn persist_message_batch( &self, session_id: &str, msgs: &[crate::storage::message::MessageMeta], meta: &crate::storage::session::SessionMeta, ) -> Result<(), StorageError> { self.persist_message_batch_inner(session_id, msgs, meta, None, None) .await } pub async fn persist_turn_batch( &self, session_id: &str, msgs: &[crate::storage::message::MessageMeta], meta: &crate::storage::session::SessionMeta, usage: &crate::storage::TurnUsageRecord, ) -> Result<(), StorageError> { self.persist_message_batch_inner(session_id, msgs, meta, Some(usage), None) .await } /// Persist a Turn and consume the admitted steer events of the same Turn /// in one transaction: the steer inputs appear in history exactly when /// their durable events become `consumed`. pub async fn persist_turn_batch_with_steer_consumption( &self, session_id: &str, msgs: &[crate::storage::message::MessageMeta], meta: &crate::storage::session::SessionMeta, usage: &crate::storage::TurnUsageRecord, steer: &crate::storage::agent_inbox::SteerConsumption, ) -> Result<(), StorageError> { self.persist_message_batch_inner(session_id, msgs, meta, Some(usage), Some(steer)) .await } /// Persist a turn with bounded retry. Retrying the whole transaction keeps /// message rows and metadata consistent on transient SQLite failures. pub async fn persist_message_batch_with_retry( &self, session_id: &str, msgs: &[crate::storage::message::MessageMeta], meta: &crate::storage::session::SessionMeta, ) -> Result<(), StorageError> { let delays = [100, 200, 300]; for (attempt, delay) in delays.iter().enumerate() { match self.persist_message_batch(session_id, msgs, meta).await { Ok(()) => return Ok(()), Err(error) if attempt < delays.len() - 1 && error.is_transient() => { tracing::warn!(attempt = attempt + 1, error = %error, "Turn persistence failed; retrying"); sleep(Duration::from_millis(*delay)).await; } Err(error) => return Err(error), } } unreachable!() } pub async fn persist_turn_batch_with_retry( &self, session_id: &str, msgs: &[crate::storage::message::MessageMeta], meta: &crate::storage::session::SessionMeta, usage: &crate::storage::TurnUsageRecord, ) -> Result<(), StorageError> { let delays = [100, 200, 300]; for (attempt, delay) in delays.iter().enumerate() { match self.persist_turn_batch(session_id, msgs, meta, usage).await { Ok(()) => return Ok(()), Err(error) if attempt < delays.len() - 1 && error.is_transient() => { tracing::warn!(attempt = attempt + 1, error = %error, "Turn persistence failed; retrying"); sleep(Duration::from_millis(*delay)).await; } Err(error) => return Err(error), } } unreachable!() } /// Persist a Turn with steer consumption and bounded retry. pub async fn persist_turn_with_steer_with_retry( &self, session_id: &str, msgs: &[crate::storage::message::MessageMeta], meta: &crate::storage::session::SessionMeta, usage: &crate::storage::TurnUsageRecord, steer: &crate::storage::agent_inbox::SteerConsumption, ) -> Result<(), StorageError> { let delays = [100, 200, 300]; for (attempt, delay) in delays.iter().enumerate() { match self .persist_turn_batch_with_steer_consumption(session_id, msgs, meta, usage, steer) .await { Ok(()) => return Ok(()), Err(error) if attempt < delays.len() - 1 && error.is_transient() => { tracing::warn!(attempt = attempt + 1, error = %error, "Turn persistence failed; retrying"); sleep(Duration::from_millis(*delay)).await; } Err(error) => return Err(error), } } unreachable!() } pub async fn get_session_usage_totals( &self, session_id: &str, ) -> Result { let row = sqlx::query( r#" SELECT COALESCE(SUM(prompt_tokens), 0) AS prompt_tokens, COALESCE(SUM(completion_tokens), 0) AS completion_tokens, COALESCE(SUM(total_tokens), 0) AS total_tokens, SUM(cached_input_tokens) AS cached_input_tokens, COALESCE(SUM(request_count), 0) AS request_count, COUNT(*) AS turn_count, MIN(created_at) AS tracked_since FROM session_turn_usage WHERE session_id = ? "#, ) .bind(session_id) .fetch_one(self.pool()) .await?; let last = sqlx::query( r#" SELECT last_prompt_tokens, created_at FROM session_turn_usage WHERE session_id = ? ORDER BY created_at DESC, rowid DESC LIMIT 1 "#, ) .bind(session_id) .fetch_optional(self.pool()) .await?; Ok(crate::storage::SessionUsageTotals { prompt_tokens: u64::try_from(row.get::("prompt_tokens")).unwrap_or_default(), completion_tokens: u64::try_from(row.get::("completion_tokens")) .unwrap_or_default(), total_tokens: u64::try_from(row.get::("total_tokens")).unwrap_or_default(), cached_input_tokens: row .get::, _>("cached_input_tokens") .and_then(|value| u64::try_from(value).ok()), request_count: u64::try_from(row.get::("request_count")).unwrap_or_default(), turn_count: u64::try_from(row.get::("turn_count")).unwrap_or_default(), tracked_since: row.get("tracked_since"), last_prompt_tokens: last .as_ref() .and_then(|value| u64::try_from(value.get::("last_prompt_tokens")).ok()), last_observed_at: last.map(|value| value.get("created_at")), }) } pub async fn load_messages( &self, session_id: &str, from_seq: i64, ) -> Result, StorageError> { let rows = sqlx::query( r#" SELECT id, session_id, seq, role, content, reasoning_content, provider_state, turn_id, iteration, completion_status, client_visibility, turn_origin, media_refs, tool_call_id, tool_name, tool_calls, source, created_at FROM messages WHERE session_id = ? AND seq >= ? ORDER BY seq ASC "#, ) .bind(session_id) .bind(from_seq) .fetch_all(self.pool()) .await?; Ok(rows.into_iter().map(message_meta_from_row).collect()) } pub async fn get_message( &self, session_id: &str, message_id: &str, ) -> Result, StorageError> { let row = sqlx::query( r#" SELECT id, session_id, seq, role, content, reasoning_content, provider_state, turn_id, iteration, completion_status, client_visibility, turn_origin, media_refs, tool_call_id, tool_name, tool_calls, source, created_at FROM messages WHERE session_id = ? AND id = ? "#, ) .bind(session_id) .bind(message_id) .fetch_optional(self.pool()) .await?; Ok(row.map(message_meta_from_row)) } pub async fn get_max_message_seq(&self, session_id: &str) -> Result { let row = sqlx::query( "SELECT COALESCE(MAX(seq), 0) as max_seq FROM messages WHERE session_id = ?", ) .bind(session_id) .fetch_one(self.pool()) .await?; Ok(row.get::("max_seq")) } /// Load a bounded tail of one session while preserving chronological order. pub async fn load_recent_session_messages( &self, session_id: &str, limit: u32, ) -> Result, StorageError> { let limit = limit.clamp(1, 2_000); let rows = sqlx::query( r#" SELECT id, session_id, seq, role, content, reasoning_content, provider_state, turn_id, iteration, completion_status, client_visibility, turn_origin, media_refs, tool_call_id, tool_name, tool_calls, source, created_at FROM ( SELECT id, session_id, seq, role, content, reasoning_content, provider_state, turn_id, iteration, completion_status, client_visibility, turn_origin, media_refs, tool_call_id, tool_name, tool_calls, source, created_at FROM messages WHERE session_id = ? AND client_visibility = 'visible' ORDER BY seq DESC LIMIT ? ) ORDER BY seq ASC "#, ) .bind(session_id) .bind(i64::from(limit)) .fetch_all(self.pool()) .await?; Ok(rows.into_iter().map(message_meta_from_row).collect()) } pub async fn load_messages_after_timestamp( &self, session_id: &str, after_ts: i64, ) -> Result, StorageError> { let rows = sqlx::query( r#" SELECT id, session_id, seq, role, content, reasoning_content, provider_state, turn_id, iteration, completion_status, client_visibility, turn_origin, media_refs, tool_call_id, tool_name, tool_calls, source, created_at FROM messages WHERE session_id = ? AND created_at > ? ORDER BY seq ASC "#, ) .bind(session_id) .bind(after_ts) .fetch_all(self.pool()) .await?; Ok(rows.into_iter().map(message_meta_from_row).collect()) } pub async fn query_sessions_range( &self, offset: i64, limit: i64, ) -> Result<(Vec, i64), StorageError> { let count_row = sqlx::query("SELECT COUNT(*) as total FROM sessions WHERE deleted_at IS NULL") .fetch_one(self.pool()) .await?; let total: i64 = count_row.get("total"); let rows = sqlx::query( r#" SELECT id, channel, chat_id, dialog_id, title, created_at, last_active_at, message_count, routing_info, archived_at, deleted_at, last_consolidated_at, last_compressed_message_at FROM sessions WHERE deleted_at IS NULL ORDER BY last_active_at DESC LIMIT ? OFFSET ? "#, ) .bind(limit) .bind(offset) .fetch_all(self.pool()) .await?; let sessions: Vec<_> = rows .into_iter() .map(|row| crate::storage::session::SessionMeta { id: row.get("id"), channel: row.get("channel"), chat_id: row.get("chat_id"), dialog_id: row.get("dialog_id"), title: row.get("title"), created_at: row.get("created_at"), last_active_at: row.get("last_active_at"), message_count: row.get("message_count"), routing_info: row.get("routing_info"), archived_at: row.get("archived_at"), deleted_at: row.get("deleted_at"), last_consolidated_at: row.get("last_consolidated_at"), last_compressed_message_at: row.get("last_compressed_message_at"), }) .collect(); Ok((sessions, total)) } pub async fn list_recent_messages( &self, session_id: &str, count: i64, ) -> Result, StorageError> { let rows = sqlx::query( r#" SELECT id, session_id, seq, role, content, reasoning_content, provider_state, turn_id, iteration, completion_status, client_visibility, turn_origin, media_refs, tool_call_id, tool_name, tool_calls, source, created_at FROM messages WHERE session_id = ? ORDER BY seq DESC LIMIT ? "#, ) .bind(session_id) .bind(count) .fetch_all(self.pool()) .await?; let mut messages: Vec<_> = rows.into_iter().map(message_meta_from_row).collect(); messages.reverse(); Ok(messages) } pub async fn query_messages_range( &self, session_id: &str, before_time: Option, after_time: Option, offset: i64, limit: i64, ) -> Result<(Vec, i64), StorageError> { let mut where_extra = String::new(); if before_time.is_some() { where_extra.push_str(" AND created_at < ?"); } if after_time.is_some() { where_extra.push_str(" AND created_at > ?"); } let count_sql = format!( "SELECT COUNT(*) as total FROM messages WHERE session_id = ? AND client_visibility = 'visible'{}", where_extra ); let select_sql = format!( r#" SELECT id, session_id, seq, role, content, reasoning_content, provider_state, turn_id, iteration, completion_status, client_visibility, turn_origin, media_refs, tool_call_id, tool_name, tool_calls, source, created_at FROM messages WHERE session_id = ? AND client_visibility = 'visible'{} ORDER BY seq ASC LIMIT ? OFFSET ? "#, where_extra ); // `where_extra` contains only the fixed timestamp predicates constructed above. let mut count_query = sqlx::query(sqlx::AssertSqlSafe(count_sql)).bind(session_id); if let Some(bt) = before_time { count_query = count_query.bind(bt); } if let Some(at) = after_time { count_query = count_query.bind(at); } let count_row = count_query.fetch_one(self.pool()).await?; let total: i64 = count_row.get("total"); let mut select_query = sqlx::query(sqlx::AssertSqlSafe(select_sql)).bind(session_id); if let Some(bt) = before_time { select_query = select_query.bind(bt); } if let Some(at) = after_time { select_query = select_query.bind(at); } let rows = select_query .bind(limit) .bind(offset) .fetch_all(self.pool()) .await?; let messages: Vec<_> = rows.into_iter().map(message_meta_from_row).collect(); Ok((messages, total)) } pub async fn clear_messages(&self, session_id: &str) -> Result<(), StorageError> { let mut tx = self.pool.begin().await?; sqlx::query(r#"DELETE FROM messages WHERE session_id = ?"#) .bind(session_id) .execute(&mut *tx) .await?; let updated = sqlx::query( r#" UPDATE sessions SET active_context_checkpoint_id = NULL, context_generation = context_generation + 1, last_compressed_message_at = NULL WHERE id = ? "#, ) .bind(session_id) .execute(&mut *tx) .await?; if updated.rows_affected() != 1 { tx.rollback().await?; return Err(StorageError::NotFound(session_id.to_string())); } tx.commit().await?; Ok(()) } /// 追加消息，带重试逻辑 /// 重试 3 次（100/200/300ms 退避），仍失败返回错误 pub async fn append_message_with_retry( &self, session_id: &str, msg: &crate::storage::message::MessageMeta, ) -> Result { let delays = [100, 200, 300]; for (i, delay) in delays.iter().enumerate() { match self.append_message(session_id, msg).await { Ok(seq) => return Ok(seq), Err(e) if i < delays.len() - 1 => { sleep(Duration::from_millis(*delay)).await; tracing::warn!("Storage write failed, retrying: {}", e); } Err(e) => { tracing::error!("Storage write failed after retries: {}", e); return Err(e); } } } unreachable!() } /// Persist the channel's durable delivery context for a session. Only /// channel-declared reusable values (thread/root identity) ever reach /// this column; one-shot reply/reaction ids never do. pub async fn update_session_delivery_context( &self, session_id: &str, context_json: &str, now: i64, ) -> Result<(), StorageError> { sqlx::query( "UPDATE sessions SET delivery_context = ?, delivery_context_updated_at = ? \ WHERE id = ?", ) .bind(context_json) .bind(now) .bind(session_id) .execute(self.pool()) .await?; Ok(()) } /// Durable delivery context previously saved for the session. pub async fn get_session_delivery_context( &self, session_id: &str, ) -> Result, StorageError> { let context: Option = sqlx::query_scalar( "SELECT delivery_context FROM sessions \ WHERE id = ? AND delivery_context IS NOT NULL", ) .bind(session_id) .fetch_optional(self.pool()) .await?; Ok(context) } } #[cfg(test)] mod tests { use super::*; use tempfile::TempDir; async fn create_test_storage() -> (Storage, TempDir) { let dir = tempfile::tempdir().unwrap(); let db_path = dir.path().join("test.db"); let storage = Storage::new(&db_path).await.unwrap(); (storage, dir) } async fn scheduler_schema_sql(storage: &Storage) -> Vec<(String, String, Option)> { sqlx::query_as::<_, (String, String, Option)>( r#" SELECT type, name, sql FROM sqlite_master WHERE name IN ( 'scheduled_jobs', 'job_runs', 'idx_jobs_claimable', 'idx_job_runs_job_finished', 'idx_job_runs_recovery', 'idx_job_runs_delivery' ) ORDER BY type, name "#, ) .fetch_all(storage.pool()) .await .unwrap() } async fn create_v10_scheduler_schema(pool: &Pool) { sqlx::query( r#" CREATE TABLE scheduled_jobs ( id TEXT PRIMARY KEY, name TEXT NOT NULL, schedule TEXT NOT NULL, prompt TEXT NOT NULL, channel TEXT NOT NULL, chat_id TEXT NOT NULL, model TEXT, job_kind TEXT NOT NULL DEFAULT 'task', delivery_policy TEXT NOT NULL DEFAULT 'direct', enabled INTEGER NOT NULL DEFAULT 1, delete_after_run INTEGER NOT NULL DEFAULT 0, next_run_at INTEGER NOT NULL, last_run_at INTEGER, last_status TEXT, last_error TEXT, locked_at INTEGER, lock_owner TEXT, lease_until INTEGER, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL ) "#, ) .execute(pool) .await .unwrap(); sqlx::query( r#" CREATE TABLE job_runs ( id INTEGER PRIMARY KEY AUTOINCREMENT, job_id TEXT NOT NULL REFERENCES scheduled_jobs(id) ON DELETE CASCADE, started_at INTEGER NOT NULL, finished_at INTEGER NOT NULL, status TEXT NOT NULL, output TEXT, error TEXT, duration_ms INTEGER NOT NULL, result_kind TEXT, delivery_status TEXT, delivery_error TEXT ) "#, ) .execute(pool) .await .unwrap(); } #[tokio::test] async fn sqlite_runtime_guards_are_enabled() { let (storage, _dir) = create_test_storage().await; let journal_mode: String = sqlx::query_scalar("PRAGMA journal_mode") .fetch_one(storage.pool()) .await .unwrap(); let foreign_keys: i64 = sqlx::query_scalar("PRAGMA foreign_keys") .fetch_one(storage.pool()) .await .unwrap(); let busy_timeout: i64 = sqlx::query_scalar("PRAGMA busy_timeout") .fetch_one(storage.pool()) .await .unwrap(); let schema_version: i64 = sqlx::query_scalar("PRAGMA user_version") .fetch_one(storage.pool()) .await .unwrap(); assert_eq!(journal_mode, "wal"); assert_eq!(foreign_keys, 1); assert_eq!(busy_timeout, 5000); assert_eq!(schema_version, SCHEMA_VERSION); let orphan = sqlx::query( r#" INSERT INTO messages (id, session_id, seq, role, content, created_at) VALUES ('orphan', 'missing', 1, 'user', 'no parent', 1) "#, ) .execute(storage.pool()) .await; assert!(orphan.is_err()); } #[tokio::test] async fn committed_turn_usage_is_aggregated_and_idempotent() { let (storage, _dir) = create_test_storage().await; let meta = crate::storage::session::SessionMeta { id: "cli_chat:chat:dialog".to_string(), channel: "cli_chat".to_string(), chat_id: "chat".to_string(), dialog_id: "dialog".to_string(), title: "usage".to_string(), created_at: 1, last_active_at: 2, message_count: 1, routing_info: None, archived_at: None, deleted_at: None, last_consolidated_at: None, last_compressed_message_at: None, }; let first = crate::storage::TurnUsageRecord { session_id: meta.id.clone(), turn_id: "turn-1".to_string(), provider: "test".to_string(), model: "model".to_string(), usage: crate::providers::Usage { prompt_tokens: 100, completion_tokens: 20, total_tokens: 120, cached_tokens: Some(40), cache_read_input_tokens: None, cache_creation_input_tokens: None, }, last_prompt_tokens: 75, created_at: 10, }; storage .persist_turn_batch(&meta.id, &[], &meta, &first) .await .unwrap(); storage .persist_turn_batch(&meta.id, &[], &meta, &first) .await .unwrap(); let second = crate::storage::TurnUsageRecord { turn_id: "turn-2".to_string(), usage: crate::providers::Usage { prompt_tokens: 50, completion_tokens: 10, total_tokens: 60, cached_tokens: None, cache_read_input_tokens: None, cache_creation_input_tokens: None, }, last_prompt_tokens: 45, created_at: 20, ..first }; storage .persist_turn_batch(&meta.id, &[], &meta, &second) .await .unwrap(); let totals = storage.get_session_usage_totals(&meta.id).await.unwrap(); assert_eq!(totals.prompt_tokens, 150); assert_eq!(totals.completion_tokens, 30); assert_eq!(totals.total_tokens, 180); assert_eq!(totals.cached_input_tokens, Some(40)); assert_eq!(totals.request_count, 2); assert_eq!(totals.turn_count, 2); assert_eq!(totals.tracked_since, Some(10)); assert_eq!(totals.last_prompt_tokens, Some(45)); assert_eq!(totals.last_observed_at, Some(20)); } #[tokio::test] async fn reopening_database_does_not_rebuild_existing_fts_index() { let dir = tempfile::tempdir().unwrap(); let db_path = dir.path().join("fts.db"); let storage = Storage::new(&db_path).await.unwrap(); sqlx::query( "INSERT INTO memory_fts(rowid, key, content) VALUES (999999, 'startup_sentinel', 'startup_sentinel')", ) .execute(storage.pool()) .await .unwrap(); drop(storage); let reopened = Storage::new(&db_path).await.unwrap(); let sentinel_count: i64 = sqlx::query_scalar( "SELECT count(*) FROM memory_fts WHERE memory_fts MATCH 'startup_sentinel'", ) .fetch_one(reopened.pool()) .await .unwrap(); assert_eq!(sentinel_count, 1); } #[tokio::test] async fn webui_lists_and_filters_memories_without_search_text() { let (storage, _dir) = create_test_storage().await; for (key, category, updated_at) in [ ( "fact", crate::memory::MemoryCategory::Knowledge, "2026-01-01T00:00:00Z", ), ( "summary", crate::memory::MemoryCategory::Timeline, "2026-01-02T00:00:00Z", ), ] { storage .upsert_memory(&crate::memory::MemoryEntry { id: key.into(), key: key.into(), content: format!("content {key}"), category, importance: 0.5, session_id: Some("cli:test:dialog".into()), created_at: updated_at.into(), updated_at: updated_at.into(), }) .await .unwrap(); } let all = storage.list_memories(None, None, 10).await.unwrap(); assert_eq!( all.iter() .map(|entry| entry.key.as_str()) .collect::>(), vec!["summary", "fact"] ); let knowledge = storage .list_memories(Some(&crate::memory::MemoryCategory::Knowledge), None, 10) .await .unwrap(); assert_eq!(knowledge.len(), 1); assert_eq!(knowledge[0].key, "fact"); } #[tokio::test] async fn legacy_schema_is_migrated_to_canonical_v11() { let dir = tempfile::tempdir().unwrap(); let db_path = dir.path().join("legacy.db"); let pool = SqlitePoolOptions::new() .connect_with( SqliteConnectOptions::new() .filename(&db_path) .create_if_missing(true), ) .await .unwrap(); sqlx::query( r#" CREATE TABLE sessions ( id TEXT PRIMARY KEY, channel TEXT NOT NULL, chat_id TEXT NOT NULL, dialog_id TEXT NOT NULL, title TEXT NOT NULL DEFAULT 'new', created_at INTEGER NOT NULL, last_active_at INTEGER NOT NULL, message_count INTEGER DEFAULT 0, routing_info TEXT, deleted_at INTEGER, UNIQUE(channel, chat_id, dialog_id) ) "#, ) .execute(&pool) .await .unwrap(); sqlx::query( r#" CREATE TABLE messages ( id TEXT PRIMARY KEY, session_id TEXT NOT NULL, seq INTEGER NOT NULL, role TEXT NOT NULL, content TEXT NOT NULL, media_refs TEXT, tool_call_id TEXT, tool_name TEXT, tool_calls TEXT, created_at INTEGER NOT NULL ) "#, ) .execute(&pool) .await .unwrap(); sqlx::query( r#" CREATE TABLE scheduled_jobs ( id TEXT PRIMARY KEY, name TEXT NOT NULL, schedule TEXT NOT NULL, prompt TEXT NOT NULL, channel TEXT NOT NULL, chat_id TEXT NOT NULL, model TEXT, enabled INTEGER NOT NULL DEFAULT 1, delete_after_run INTEGER NOT NULL DEFAULT 0, next_run_at INTEGER NOT NULL, last_run_at INTEGER, last_status TEXT, last_error TEXT, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL ) "#, ) .execute(&pool) .await .unwrap(); sqlx::query( "INSERT INTO sessions (id, channel, chat_id, dialog_id, created_at, last_active_at) VALUES ('cli:c:d', 'cli', 'c', 'd', 1, 1)", ) .execute(&pool) .await .unwrap(); sqlx::query( "INSERT INTO messages (id, session_id, seq, role, content, created_at) VALUES ('m1', 'cli:c:d', 1, 'user', 'legacy message', 1)", ) .execute(&pool) .await .unwrap(); drop(pool); let storage = Storage::new(&db_path).await.unwrap(); for (table, expected) in [ ( "messages", vec![ "source", "reasoning_content", "provider_state", "turn_id", "iteration", "completion_status", "client_visibility", "turn_origin", ], ), ( "sessions", vec![ "archived_at", "last_consolidated_at", "last_compressed_message_at", "active_context_checkpoint_id", "context_generation", "delivery_context", "delivery_context_updated_at", ], ), ( "scheduled_jobs", vec![ "agent_id", "delivery_policy", "last_outcome", "locked_at", "lock_owner", "lease_until", ], ), ] { let columns = sqlx::query(sqlx::AssertSqlSafe(format!("PRAGMA table_info({table})"))) .fetch_all(storage.pool()) .await .unwrap(); for column in expected { assert!( columns .iter() .any(|row| row.get::("name") == column), "missing migrated column {table}.{column}" ); } } let scheduler_columns = sqlx::query("PRAGMA table_info(scheduled_jobs)") .fetch_all(storage.pool()) .await .unwrap() .into_iter() .map(|row| row.get::("name")) .collect::>(); for removed in [ "model", "job_kind", "delete_after_run", "last_status", "last_error", ] { assert!( !scheduler_columns.iter().any(|column| column == removed), "legacy scheduler column survived v11: {removed}" ); } let schema_version: i64 = sqlx::query_scalar("PRAGMA user_version") .fetch_one(storage.pool()) .await .unwrap(); assert_eq!(schema_version, SCHEMA_VERSION); for table in [ "task_plans", "task_items", "session_turn_usage", "agent_runs", "agent_session_state", "agent_inbox_events", "context_checkpoints", ] { let exists: i64 = sqlx::query_scalar( "SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = ?", ) .bind(table) .fetch_one(storage.pool()) .await .unwrap(); assert_eq!(exists, 1, "missing migrated table {table}"); } let visibility: String = sqlx::query_scalar("SELECT client_visibility FROM messages ORDER BY seq LIMIT 1") .fetch_one(storage.pool()) .await .unwrap_or_default(); assert_eq!(visibility, "visible"); let origin: String = sqlx::query_scalar("SELECT turn_origin FROM messages ORDER BY seq LIMIT 1") .fetch_one(storage.pool()) .await .unwrap_or_default(); assert_eq!(origin, "user"); } #[tokio::test] async fn v10_scheduler_history_and_active_lease_migrate_to_v11() { let dir = tempfile::tempdir().unwrap(); let db_path = dir.path().join("v10-scheduler.db"); let pool = SqlitePoolOptions::new() .connect_with( SqliteConnectOptions::new() .filename(&db_path) .create_if_missing(true), ) .await .unwrap(); create_v10_scheduler_schema(&pool).await; sqlx::query( r#" INSERT INTO scheduled_jobs (id, name, schedule, prompt, channel, chat_id, model, job_kind, delivery_policy, enabled, next_run_at, last_run_at, last_status, last_error, created_at, updated_at) VALUES ('legacy-direct', 'Legacy direct', '{"type":"every","every_ms":60000}', 'check it', 'cli_chat', 'chat', 'obsolete-model', 'monitor', 'direct', 1, 2000, 1000, 'delivery_error', 'legacy diagnostic', 1, 2) "#, ) .execute(&pool) .await .unwrap(); sqlx::query( r#" INSERT INTO job_runs (job_id, started_at, finished_at, status, output, error, duration_ms, result_kind, delivery_status, delivery_error) VALUES ('legacy-direct', 1000, 1100, 'delivery_error', 'legacy result', NULL, 100, NULL, 'failed', 'channel rejected target') "#, ) .execute(&pool) .await .unwrap(); sqlx::query( r#" INSERT INTO scheduled_jobs (id, name, schedule, prompt, channel, chat_id, job_kind, delivery_policy, enabled, next_run_at, last_run_at, last_status, locked_at, lock_owner, lease_until, created_at, updated_at) VALUES ('legacy-locked', 'Legacy locked', '{"type":"every","every_ms":60000}', 'check lock', 'cli_chat', 'chat', 'task', 'never', 1, 2000, 1500, NULL, 1500, 'old-owner', 900000, 1, 2) "#, ) .execute(&pool) .await .unwrap(); sqlx::query("PRAGMA user_version = 10") .execute(&pool) .await .unwrap(); drop(pool); let storage = Storage::new(&db_path).await.unwrap(); let migrated_columns = sqlx::query("PRAGMA table_info(scheduled_jobs)") .fetch_all(storage.pool()) .await .unwrap() .into_iter() .map(|row| row.get::("name")) .collect::>(); assert!( migrated_columns .iter() .any(|column| column == "last_outcome"), "unexpected post-migration scheduler schema: {migrated_columns:?}" ); let direct = storage.get_scheduled_job("legacy-direct").await.unwrap(); assert_eq!( direct.delivery_policy, crate::storage::DeliveryPolicy::Always ); assert_eq!( direct.last_outcome, Some(crate::storage::ScheduledOutcomeKind::Ok) ); let historical = storage .list_scheduled_job_runs("legacy-direct", 10) .await .unwrap(); assert_eq!(historical.len(), 1); assert_eq!( historical[0].status, crate::storage::ScheduledRunStatus::Completed ); assert_eq!( historical[0].outcome, Some(crate::storage::ScheduledOutcomeKind::Ok) ); assert_eq!( historical[0].delivery_status, crate::storage::ScheduledDeliveryStatus::Failed ); assert_eq!( historical[0].diagnostic.as_deref(), Some("legacy diagnostic") ); let recovered = storage .list_scheduled_job_runs("legacy-locked", 10) .await .unwrap(); assert_eq!(recovered.len(), 1); assert_eq!( recovered[0].status, crate::storage::ScheduledRunStatus::Unknown ); assert_eq!( recovered[0].delivery_status, crate::storage::ScheduledDeliveryStatus::NotRequested ); let locked = storage.get_scheduled_job("legacy-locked").await.unwrap(); assert!(locked.lock_owner.is_none()); let columns = sqlx::query("PRAGMA table_info(scheduled_jobs)") .fetch_all(storage.pool()) .await .unwrap() .into_iter() .map(|row| row.get::("name")) .collect::>(); for removed in [ "model", "job_kind", "delete_after_run", "last_status", "last_error", ] { assert!(!columns.iter().any(|column| column == removed)); } let violations = sqlx::query("PRAGMA foreign_key_check") .fetch_all(storage.pool()) .await .unwrap(); assert!(violations.is_empty()); let fresh_dir = tempfile::tempdir().unwrap(); let fresh = Storage::new(&fresh_dir.path().join("fresh.db")) .await .unwrap(); assert_eq!( scheduler_schema_sql(&storage).await, scheduler_schema_sql(&fresh).await ); } #[tokio::test] async fn invalid_v10_scheduler_data_rolls_back_v11_migration() { let dir = tempfile::tempdir().unwrap(); let db_path = dir.path().join("invalid-v10-scheduler.db"); let pool = SqlitePoolOptions::new() .connect_with( SqliteConnectOptions::new() .filename(&db_path) .create_if_missing(true), ) .await .unwrap(); create_v10_scheduler_schema(&pool).await; sqlx::query( r#" INSERT INTO scheduled_jobs (id, name, schedule, prompt, channel, chat_id, delivery_policy, enabled, next_run_at, created_at, updated_at) VALUES ('broken', 'Broken', '{not-json}', 'check', 'cli_chat', 'chat', 'always', 1, 1000, 1, 1) "#, ) .execute(&pool) .await .unwrap(); sqlx::query("PRAGMA user_version = 10") .execute(&pool) .await .unwrap(); drop(pool); let error = match Storage::new(&db_path).await { Ok(_) => panic!("invalid v10 scheduler data unexpectedly migrated"), Err(error) => error, }; assert!(matches!(error, StorageError::Migration(_))); let pool = SqlitePoolOptions::new() .connect_with(SqliteConnectOptions::new().filename(&db_path)) .await .unwrap(); let version: i64 = sqlx::query_scalar("PRAGMA user_version") .fetch_one(&pool) .await .unwrap(); assert_eq!(version, 10); let columns = sqlx::query("PRAGMA table_info(scheduled_jobs)") .fetch_all(&pool) .await .unwrap() .into_iter() .map(|row| row.get::("name")) .collect::>(); assert!(columns.iter().any(|column| column == "job_kind")); assert!(!columns.iter().any(|column| column == "agent_id")); for table in ["scheduled_jobs_v10_legacy", "job_runs_v10_legacy"] { let exists: i64 = sqlx::query_scalar( "SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = ?", ) .bind(table) .fetch_one(&pool) .await .unwrap(); assert_eq!(exists, 0); } } #[tokio::test] async fn newer_schema_version_is_rejected() { let dir = tempfile::tempdir().unwrap(); let db_path = dir.path().join("future.db"); let pool = SqlitePoolOptions::new() .connect_with( SqliteConnectOptions::new() .filename(&db_path) .create_if_missing(true), ) .await .unwrap(); sqlx::query("PRAGMA user_version = 12") .execute(&pool) .await .unwrap(); drop(pool); let error = match Storage::new(&db_path).await { Ok(_) => panic!("newer schema version unexpectedly opened"), Err(error) => error, }; assert!(matches!(error, StorageError::Migration(_))); } #[tokio::test] async fn v3_migration_preserves_existing_reasoning_and_defaults_completion() { let dir = tempfile::tempdir().unwrap(); let db_path = dir.path().join("v3.db"); let pool = SqlitePoolOptions::new() .connect_with( SqliteConnectOptions::new() .filename(&db_path) .create_if_missing(true), ) .await .unwrap(); sqlx::query( r#" CREATE TABLE sessions ( id TEXT PRIMARY KEY, channel TEXT NOT NULL, chat_id TEXT NOT NULL, dialog_id TEXT NOT NULL, title TEXT NOT NULL DEFAULT 'new', created_at INTEGER NOT NULL, last_active_at INTEGER NOT NULL, message_count INTEGER DEFAULT 0, routing_info TEXT, archived_at INTEGER, deleted_at INTEGER, last_consolidated_at INTEGER, last_compressed_message_at INTEGER, UNIQUE(channel, chat_id, dialog_id) ) "#, ) .execute(&pool) .await .unwrap(); sqlx::query( r#" CREATE TABLE messages ( id TEXT PRIMARY KEY, session_id TEXT NOT NULL, seq INTEGER NOT NULL, role TEXT NOT NULL, content TEXT NOT NULL, reasoning_content TEXT, media_refs TEXT, tool_call_id TEXT, tool_name TEXT, tool_calls TEXT, source TEXT, created_at INTEGER NOT NULL ) "#, ) .execute(&pool) .await .unwrap(); sqlx::query( "INSERT INTO sessions (id, channel, chat_id, dialog_id, created_at, last_active_at) VALUES ('cli:c:d', 'cli', 'c', 'd', 1, 1)", ) .execute(&pool) .await .unwrap(); sqlx::query( "INSERT INTO messages (id, session_id, seq, role, content, reasoning_content, created_at) VALUES ('m1', 'cli:c:d', 1, 'assistant', 'answer', 'existing reasoning', 1)", ) .execute(&pool) .await .unwrap(); sqlx::query("PRAGMA user_version = 3") .execute(&pool) .await .unwrap(); drop(pool); let storage = Storage::new(&db_path).await.unwrap(); let messages = storage.load_messages("cli:c:d", 0).await.unwrap(); assert_eq!(messages.len(), 1); assert_eq!( messages[0].reasoning_content.as_deref(), Some("existing reasoning") ); assert_eq!(messages[0].provider_state, None); assert_eq!(messages[0].turn_id, None); assert_eq!(messages[0].iteration, None); assert_eq!( messages[0].completion_status, crate::bus::CompletionStatus::Completed ); } #[tokio::test] async fn v8_migration_preserves_agent_runs_and_adds_transcript_table() { let dir = tempfile::tempdir().unwrap(); let db_path = dir.path().join("v8.db"); let pool = SqlitePoolOptions::new() .connect_with( SqliteConnectOptions::new() .filename(&db_path) .create_if_missing(true), ) .await .unwrap(); // The v8 `agent_runs` shape is unchanged in v9: v9 only adds the // transcript table. Build a v8 database holding a durable run so the // upgrade must preserve it rather than dropping the table. sqlx::query(agent_run::AGENT_SCHEMA_STATEMENTS[0]) .execute(&pool) .await .unwrap(); sqlx::query( "INSERT INTO agent_runs (id, root_session_id, caller_agent_id, caller_scope_id, \ agent_id, definition_hash, provider_profile, provider_name, model_id, mode, \ depth, execution_id, task, budget_json, status, runtime_generation, attempt, \ completion_slot_reserved, deadline_at, revision, created_at, updated_at) \ VALUES ('run-1', 'cli:c:d', 'ROOT', 'turn-1', 'researcher', 'hash', 'profile', \ 'test', 'model', 'foreground', 1, 'exec-1', 'task', '{}', 'completed', 1, 1, \ 0, 1000, 0, 1, 1)", ) .execute(&pool) .await .unwrap(); sqlx::query("PRAGMA user_version = 8") .execute(&pool) .await .unwrap(); drop(pool); let storage = Storage::new(&db_path).await.unwrap(); let run = storage.get_agent_run("run-1").await.unwrap(); assert!( run.is_some(), "v8 agent run must survive the current upgrade without a rebuild" ); assert_eq!(run.unwrap().status.as_str(), "completed"); let version: i64 = sqlx::query_scalar("PRAGMA user_version") .fetch_one(storage.pool()) .await .unwrap(); assert_eq!(version, SCHEMA_VERSION); let exists: i64 = sqlx::query_scalar( "SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = 'agent_run_messages'", ) .fetch_one(storage.pool()) .await .unwrap(); assert_eq!(exists, 1, "agent_run_messages table must be created"); } #[tokio::test] async fn test_upsert_and_get_session() { let (storage, _dir) = create_test_storage().await; let meta = crate::storage::session::SessionMeta { id: "cli_chat:sid123:dialog1".to_string(), channel: "cli_chat".to_string(), chat_id: "sid123".to_string(), dialog_id: "dialog1".to_string(), title: "测试会话".to_string(), created_at: 1000, last_active_at: 1000, message_count: 0, routing_info: Some(r#"{"type":"cli"}"#.to_string()), archived_at: None, deleted_at: None, last_consolidated_at: None, last_compressed_message_at: None, }; storage.upsert_session(&meta).await.unwrap(); let loaded = storage.get_session(&meta.id).await.unwrap(); assert_eq!(loaded.title, "测试会话"); assert_eq!(loaded.channel, "cli_chat"); } #[tokio::test] async fn test_get_nonexistent_session() { let (storage, _dir) = create_test_storage().await; let result = storage.get_session("nonexistent").await; assert!(result.is_err()); matches!(result.unwrap_err(), StorageError::NotFound(_)); } #[tokio::test] async fn test_list_sessions() { let (storage, _dir) = create_test_storage().await; for i in 0..5 { let meta = crate::storage::session::SessionMeta { id: format!("cli_chat:sid123:dialog{}", i), channel: "cli_chat".to_string(), chat_id: "sid123".to_string(), dialog_id: format!("dialog{}", i), title: format!("会话{}", i), created_at: i * 1000, last_active_at: i * 1000, message_count: i, routing_info: None, archived_at: None, deleted_at: None, last_consolidated_at: None, last_compressed_message_at: None, }; storage.upsert_session(&meta).await.unwrap(); } let sessions = storage .list_sessions("cli_chat", "sid123", 10, false) .await .unwrap(); assert_eq!(sessions.len(), 5); // 按 last_active_at DESC 排序 assert_eq!(sessions[0].dialog_id, "dialog4"); } #[tokio::test] async fn test_soft_delete() { let (storage, _dir) = create_test_storage().await; let meta = crate::storage::session::SessionMeta { id: "cli_chat:sid123:dialog1".to_string(), channel: "cli_chat".to_string(), chat_id: "sid123".to_string(), dialog_id: "dialog1".to_string(), title: "测试".to_string(), created_at: 1000, last_active_at: 1000, message_count: 0, routing_info: None, archived_at: None, deleted_at: None, last_consolidated_at: None, last_compressed_message_at: None, }; storage.upsert_session(&meta).await.unwrap(); storage.soft_delete_session(&meta.id).await.unwrap(); let result = storage.get_session(&meta.id).await; assert!(result.is_err()); matches!(result.unwrap_err(), StorageError::NotFound(_)); } #[tokio::test] async fn test_append_and_load_messages() { let (storage, _dir) = create_test_storage().await; let session_meta = crate::storage::session::SessionMeta { id: "cli_chat:sid123:dialog1".to_string(), channel: "cli_chat".to_string(), chat_id: "sid123".to_string(), dialog_id: "dialog1".to_string(), title: "测试".to_string(), created_at: 1000, last_active_at: 1000, message_count: 0, routing_info: None, archived_at: None, deleted_at: None, last_consolidated_at: None, last_compressed_message_at: None, }; storage.upsert_session(&session_meta).await.unwrap(); let msg = crate::storage::message::MessageMeta { id: "msg1".to_string(), session_id: session_meta.id.clone(), seq: 1, role: "user".to_string(), content: "你好".to_string(), reasoning_content: None, provider_state: None, turn_id: None, iteration: None, completion_status: crate::bus::CompletionStatus::Completed, client_visibility: crate::bus::ClientVisibility::Visible, turn_origin: crate::bus::TurnOrigin::User, media_refs: None, tool_call_id: None, tool_name: None, tool_calls: None, source: None, created_at: 1000, }; let seq = storage .append_message(&session_meta.id, &msg) .await .unwrap(); assert_eq!(seq, 1); let loaded = storage.load_messages(&session_meta.id, 0).await.unwrap(); assert_eq!(loaded.len(), 1); assert_eq!(loaded[0].content, "你好"); for seq in 2..=5 { let mut message = msg.clone(); message.id = format!("msg{seq}"); message.seq = seq; message.content = format!("message {seq}"); storage .append_message(&session_meta.id, &message) .await .unwrap(); } let recent = storage .load_recent_session_messages(&session_meta.id, 2) .await .unwrap(); assert_eq!(recent.len(), 2); assert_eq!(recent[0].seq, 4); assert_eq!(recent[1].seq, 5); } #[tokio::test] async fn streaming_message_metadata_round_trips() { let (storage, _dir) = create_test_storage().await; let session_meta = crate::storage::session::SessionMeta { id: "cli_chat:stream:dialog1".to_string(), channel: "cli_chat".to_string(), chat_id: "stream".to_string(), dialog_id: "dialog1".to_string(), title: "Stream metadata".to_string(), created_at: 1000, last_active_at: 1000, message_count: 0, routing_info: None, archived_at: None, deleted_at: None, last_consolidated_at: None, last_compressed_message_at: None, }; storage.upsert_session(&session_meta).await.unwrap(); let message = crate::storage::message::MessageMeta { id: "assistant-1".to_string(), session_id: session_meta.id.clone(), seq: 1, role: "assistant".to_string(), content: "partial".to_string(), reasoning_content: Some("visible reasoning".to_string()), provider_state: Some( serde_json::json!({"provider":"anthropic","payload":{"signature":"opaque"}}) .to_string(), ), turn_id: Some("turn-1".to_string()), iteration: Some(2), completion_status: crate::bus::CompletionStatus::Interrupted, client_visibility: crate::bus::ClientVisibility::Visible, turn_origin: crate::bus::TurnOrigin::User, media_refs: None, tool_call_id: None, tool_name: None, tool_calls: None, source: None, created_at: 1001, }; storage .append_message(&session_meta.id, &message) .await .unwrap(); let loaded = storage.load_messages(&session_meta.id, 0).await.unwrap(); assert_eq!(loaded.len(), 1); assert_eq!( loaded[0].reasoning_content.as_deref(), Some("visible reasoning") ); assert_eq!(loaded[0].provider_state, message.provider_state); assert_eq!(loaded[0].turn_id.as_deref(), Some("turn-1")); assert_eq!(loaded[0].iteration, Some(2)); assert_eq!( loaded[0].completion_status, crate::bus::CompletionStatus::Interrupted ); } #[tokio::test] async fn test_persist_message_batch_is_atomic() { let (storage, _dir) = create_test_storage().await; let session_meta = crate::storage::session::SessionMeta { id: "cli_chat:atomic:dialog1".to_string(), channel: "cli_chat".to_string(), chat_id: "atomic".to_string(), dialog_id: "dialog1".to_string(), title: "Atomic turn".to_string(), created_at: 1000, last_active_at: 2000, message_count: 1, routing_info: None, archived_at: None, deleted_at: None, last_consolidated_at: None, last_compressed_message_at: None, }; storage.upsert_session(&session_meta).await.unwrap(); let message = crate::storage::message::MessageMeta { id: "duplicate-id".to_string(), session_id: session_meta.id.clone(), seq: 1, role: "assistant".to_string(), content: "must roll back".to_string(), reasoning_content: None, provider_state: None, turn_id: None, iteration: None, completion_status: crate::bus::CompletionStatus::Completed, client_visibility: crate::bus::ClientVisibility::Visible, turn_origin: crate::bus::TurnOrigin::User, media_refs: None, tool_call_id: None, tool_name: None, tool_calls: None, source: None, created_at: 2000, }; let result = storage .persist_message_batch(&session_meta.id, &[message.clone(), message], &session_meta) .await; assert!(result.is_err()); assert!( storage .load_messages(&session_meta.id, 0) .await .unwrap() .is_empty() ); } #[tokio::test] async fn test_touch_session() { let (storage, _dir) = create_test_storage().await; let meta = crate::storage::session::SessionMeta { id: "cli_chat:sid123:dialog1".to_string(), channel: "cli_chat".to_string(), chat_id: "sid123".to_string(), dialog_id: "dialog1".to_string(), title: "测试".to_string(), created_at: 1000, last_active_at: 1000, message_count: 0, routing_info: None, archived_at: None, deleted_at: None, last_consolidated_at: None, last_compressed_message_at: None, }; storage.upsert_session(&meta).await.unwrap(); storage.touch_session(&meta.id, 5, 2000).await.unwrap(); let loaded = storage.get_session(&meta.id).await.unwrap(); assert_eq!(loaded.message_count, 5); assert_eq!(loaded.last_active_at, 2000); } #[tokio::test] async fn durable_delivery_context_round_trips() { let (storage, _dir) = create_test_storage().await; let meta = crate::storage::session::SessionMeta { id: "feishu:chat-1:dialog".to_string(), channel: "feishu".to_string(), chat_id: "chat-1".to_string(), dialog_id: "dialog".to_string(), title: "t".to_string(), created_at: 1, last_active_at: 2, message_count: 0, routing_info: None, archived_at: None, deleted_at: None, last_consolidated_at: None, last_compressed_message_at: None, }; storage.upsert_session(&meta).await.unwrap(); assert!( storage .get_session_delivery_context("feishu:chat-1:dialog") .await .unwrap() .is_none() ); let context = serde_json::json!({ "feishu.thread_id": "thread-9", "feishu.chat_type": "group", }) .to_string(); storage .update_session_delivery_context("feishu:chat-1:dialog", &context, 50) .await .unwrap(); let loaded = storage .get_session_delivery_context("feishu:chat-1:dialog") .await .unwrap() .unwrap(); let parsed: std::collections::HashMap = serde_json::from_str(&loaded).unwrap(); assert_eq!(parsed.get("feishu.thread_id").unwrap(), "thread-9"); } }