d9ad58b84b
Replace the dual task/monitor model, NO_REPLY string protocol, and Agent self-delivery with a single Scheduled Run path: claim-time JobRun snapshots, isolated Root/named Agent execution, exactly-once complete_scheduled_run termination, and Scheduler-owned policy delivery through a persistent outbox. - SQLite v11: drop job_kind/model/delete_after_run, add job_runs with status/outcome joint constraints and delivery lease columns; one-shot BEGIN IMMEDIATE migration with atomic rollback. - Non-blocking JoinSet event loop with bounded run/delivery concurrency; terminal commit before any channel I/O; recover unfinished runs as unknown. - ExecutionOrigin::Scheduled propagates to descendants, completion sink is top-level only, background delegation downgrades to foreground. - Typed delivery receipts, fixed target_session_id, idempotent scheduled:<job_run_id> history insert. - New cron_runs read-only tool; cron_add/update drop kind/model; WebUI and Health consume the same JobRun projection. - Bump version to 1.22.0.
2783 lines
100 KiB
Rust
2783 lines
100 KiB
Rust
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<Sqlite>,
|
||
}
|
||
|
||
impl Storage {
|
||
/// 打开或创建数据库
|
||
pub async fn new(db_path: &Path) -> Result<Self, StorageError> {
|
||
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::<String, _>("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<Sqlite> {
|
||
&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<crate::storage::session::SessionMeta, 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 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<Vec<crate::storage::session::SessionMeta>, 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<Option<crate::storage::session::SessionMeta>, 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<i64, StorageError> {
|
||
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<bool, StorageError> {
|
||
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<crate::storage::SessionUsageTotals, StorageError> {
|
||
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::<i64, _>("prompt_tokens")).unwrap_or_default(),
|
||
completion_tokens: u64::try_from(row.get::<i64, _>("completion_tokens"))
|
||
.unwrap_or_default(),
|
||
total_tokens: u64::try_from(row.get::<i64, _>("total_tokens")).unwrap_or_default(),
|
||
cached_input_tokens: row
|
||
.get::<Option<i64>, _>("cached_input_tokens")
|
||
.and_then(|value| u64::try_from(value).ok()),
|
||
request_count: u64::try_from(row.get::<i64, _>("request_count")).unwrap_or_default(),
|
||
turn_count: u64::try_from(row.get::<i64, _>("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::<i64, _>("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<Vec<crate::storage::message::MessageMeta>, 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<Option<crate::storage::message::MessageMeta>, 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<i64, StorageError> {
|
||
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::<i64, _>("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<Vec<crate::storage::message::MessageMeta>, 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<Vec<crate::storage::message::MessageMeta>, 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<crate::storage::session::SessionMeta>, 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<Vec<crate::storage::message::MessageMeta>, 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<i64>,
|
||
after_time: Option<i64>,
|
||
offset: i64,
|
||
limit: i64,
|
||
) -> Result<(Vec<crate::storage::message::MessageMeta>, 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<i64, StorageError> {
|
||
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<Option<String>, StorageError> {
|
||
let context: Option<String> = 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<String>)> {
|
||
sqlx::query_as::<_, (String, String, Option<String>)>(
|
||
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<Sqlite>) {
|
||
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<_>>(),
|
||
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::<String, _>("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::<String, _>("name"))
|
||
.collect::<Vec<_>>();
|
||
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::<String, _>("name"))
|
||
.collect::<Vec<_>>();
|
||
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::<String, _>("name"))
|
||
.collect::<Vec<_>>();
|
||
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::<String, _>("name"))
|
||
.collect::<Vec<_>>();
|
||
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<String, String> =
|
||
serde_json::from_str(&loaded).unwrap();
|
||
assert_eq!(parsed.get("feishu.thread_id").unwrap(), "thread-9");
|
||
}
|
||
}
|