c7ee6bb519
将 7 个写事务从 BEGIN DEFERRED 改为 BEGIN IMMEDIATE,在事务开始即获取写锁,消除多 sub-agent 并发写入时的死锁路径。同时将 busy_timeout 从 5s 提升至 30s,为并发写者排队提供 100 倍余量。 根因:BEGIN DEFERRED 下多个事务可同时读 MAX(seq) 不持写锁,提交时互相阻塞,5s timeout 耗尽后返回 SQLITE_BUSY。BEGIN IMMEDIATE 强制写者串行排队,顺带消除 MAX(seq)+1 竞态导致的 UNIQUE 约束冲突。
1978 lines
68 KiB
Rust
1978 lines
68 KiB
Rust
#[cfg(not(test))]
|
||
use std::path::{Path, PathBuf};
|
||
|
||
use r2d2::Pool;
|
||
use r2d2_sqlite::SqliteConnectionManager;
|
||
use rusqlite::{Connection, OptionalExtension, TransactionBehavior, params};
|
||
|
||
use crate::bus::ChatMessage;
|
||
|
||
pub mod error;
|
||
pub mod ports;
|
||
pub mod records;
|
||
|
||
mod migrations;
|
||
mod row_mapping;
|
||
|
||
// Bring extracted helpers into scope for use by SessionStore methods below.
|
||
use migrations::*;
|
||
use row_mapping::*;
|
||
|
||
pub use error::StorageError;
|
||
pub use ports::{
|
||
ConversationRepository, MemoryRepository, PromptInjectionRepository, SchedulerJobRepository,
|
||
SkillEventRepository, TodoRepository,
|
||
};
|
||
pub use records::{
|
||
allowed_namespace_names, get_namespace_description, is_valid_namespace,
|
||
ALLOWED_MEMORY_NAMESPACES, GLOBAL_SCOPE_KEY, MemoryRecord, MemoryUpsert, SchedulerJobRecord,
|
||
SchedulerJobState, SchedulerJobStatus, SchedulerJobUpsert, SessionRecord, SkillEventRecord,
|
||
TodoRecord, TopicRecord,
|
||
};
|
||
|
||
#[derive(Clone)]
|
||
pub struct SessionStore {
|
||
pool: Pool<SqliteConnectionManager>,
|
||
}
|
||
|
||
impl SessionStore {
|
||
#[cfg(test)]
|
||
pub fn new() -> Result<Self, StorageError> {
|
||
Self::in_memory()
|
||
}
|
||
|
||
#[cfg(not(test))]
|
||
pub fn new() -> Result<Self, StorageError> {
|
||
let db_path = default_session_db_path()?;
|
||
Self::open_at_path(&db_path)
|
||
}
|
||
|
||
#[cfg(not(test))]
|
||
fn open_at_path(path: &Path) -> Result<Self, StorageError> {
|
||
if let Some(parent) = path.parent() {
|
||
std::fs::create_dir_all(parent)?;
|
||
}
|
||
|
||
let path_str = path.to_string_lossy().to_string();
|
||
let conn = Connection::open(&path_str)?;
|
||
Self::from_connection(conn, &path_str)
|
||
}
|
||
|
||
/// Initialize a SessionStore from a connection and its file path.
|
||
/// The connection is used for schema initialization only; the pool
|
||
/// manages subsequent connections using the same file path.
|
||
fn from_connection(conn: Connection, db_uri: &str) -> Result<Self, StorageError> {
|
||
conn.busy_timeout(std::time::Duration::from_secs(30))?;
|
||
conn.execute_batch(
|
||
"
|
||
PRAGMA journal_mode = WAL;
|
||
PRAGMA foreign_keys = ON;
|
||
|
||
CREATE TABLE IF NOT EXISTS sessions (
|
||
id TEXT PRIMARY KEY,
|
||
title TEXT NOT NULL,
|
||
channel_name TEXT NOT NULL,
|
||
chat_id TEXT NOT NULL,
|
||
summary TEXT,
|
||
created_at INTEGER NOT NULL,
|
||
updated_at INTEGER NOT NULL,
|
||
last_active_at INTEGER NOT NULL,
|
||
archived_at INTEGER,
|
||
deleted_at INTEGER,
|
||
message_count INTEGER NOT NULL DEFAULT 0,
|
||
user_turn_count INTEGER NOT NULL DEFAULT 0,
|
||
agent_prompt_reinjection_count INTEGER NOT NULL DEFAULT 0
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_sessions_channel_archived
|
||
ON sessions(channel_name, archived_at, last_active_at DESC);
|
||
CREATE INDEX IF NOT EXISTS idx_sessions_updated_at
|
||
ON sessions(updated_at DESC);
|
||
|
||
CREATE TABLE IF NOT EXISTS messages (
|
||
id TEXT PRIMARY KEY,
|
||
session_id TEXT NOT NULL,
|
||
topic_id TEXT,
|
||
seq INTEGER NOT NULL,
|
||
role TEXT NOT NULL,
|
||
content TEXT NOT NULL,
|
||
system_context TEXT,
|
||
reasoning_content TEXT,
|
||
media_refs_json TEXT NOT NULL,
|
||
tool_call_id TEXT,
|
||
tool_name TEXT,
|
||
tool_calls_json TEXT,
|
||
created_at INTEGER NOT NULL,
|
||
FOREIGN KEY(session_id) REFERENCES sessions(id) ON DELETE CASCADE,
|
||
FOREIGN KEY(topic_id) REFERENCES topics(id) ON DELETE SET NULL,
|
||
UNIQUE(session_id, seq)
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_messages_session_seq
|
||
ON messages(session_id, seq);
|
||
CREATE INDEX IF NOT EXISTS idx_messages_session_created
|
||
ON messages(session_id, created_at);
|
||
|
||
CREATE TABLE IF NOT EXISTS topics (
|
||
id TEXT PRIMARY KEY,
|
||
session_id TEXT NOT NULL,
|
||
title TEXT NOT NULL,
|
||
description TEXT,
|
||
created_at INTEGER NOT NULL,
|
||
updated_at INTEGER NOT NULL,
|
||
last_active_at INTEGER NOT NULL,
|
||
message_count INTEGER NOT NULL DEFAULT 0,
|
||
FOREIGN KEY(session_id) REFERENCES sessions(id) ON DELETE CASCADE
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_topics_session
|
||
ON topics(session_id, last_active_at DESC);
|
||
|
||
CREATE TABLE IF NOT EXISTS skill_events (
|
||
id TEXT PRIMARY KEY,
|
||
session_id TEXT,
|
||
event_type TEXT NOT NULL,
|
||
skill_name TEXT,
|
||
payload_json TEXT NOT NULL,
|
||
created_at INTEGER NOT NULL,
|
||
FOREIGN KEY(session_id) REFERENCES sessions(id) ON DELETE CASCADE
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_skill_events_session_created
|
||
ON skill_events(session_id, created_at DESC);
|
||
CREATE INDEX IF NOT EXISTS idx_skill_events_type_created
|
||
ON skill_events(event_type, created_at DESC);
|
||
|
||
CREATE TABLE IF NOT EXISTS memories (
|
||
id TEXT PRIMARY KEY,
|
||
scope_kind TEXT NOT NULL,
|
||
scope_key TEXT NOT NULL,
|
||
namespace TEXT NOT NULL,
|
||
memory_key TEXT NOT NULL,
|
||
content TEXT NOT NULL,
|
||
source_type TEXT NOT NULL,
|
||
source_session_id TEXT,
|
||
source_message_id TEXT,
|
||
source_message_seq INTEGER,
|
||
source_channel_name TEXT,
|
||
source_chat_id TEXT,
|
||
created_at INTEGER NOT NULL,
|
||
updated_at INTEGER NOT NULL,
|
||
UNIQUE(scope_kind, scope_key, namespace, memory_key)
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_memories_scope_updated
|
||
ON memories(scope_kind, scope_key, updated_at DESC);
|
||
CREATE INDEX IF NOT EXISTS idx_memories_scope_namespace_updated
|
||
ON memories(scope_kind, scope_key, namespace, updated_at DESC);
|
||
CREATE INDEX IF NOT EXISTS idx_memories_source_session
|
||
ON memories(source_session_id, updated_at DESC);
|
||
|
||
CREATE TABLE IF NOT EXISTS scheduler_jobs (
|
||
id TEXT PRIMARY KEY,
|
||
kind TEXT NOT NULL,
|
||
schedule_json TEXT NOT NULL DEFAULT '{}',
|
||
interval_secs INTEGER NOT NULL DEFAULT 0,
|
||
startup_delay_secs INTEGER NOT NULL DEFAULT 0,
|
||
target_json TEXT NOT NULL,
|
||
payload_json TEXT NOT NULL,
|
||
enabled INTEGER NOT NULL DEFAULT 1,
|
||
state TEXT NOT NULL DEFAULT 'scheduled',
|
||
last_status TEXT,
|
||
last_error TEXT,
|
||
run_count INTEGER NOT NULL DEFAULT 0,
|
||
max_runs INTEGER,
|
||
last_fired_at INTEGER,
|
||
next_fire_at INTEGER,
|
||
paused_at INTEGER,
|
||
completed_at INTEGER,
|
||
created_at INTEGER NOT NULL,
|
||
updated_at INTEGER NOT NULL
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_scheduler_jobs_enabled_next_fire
|
||
ON scheduler_jobs(enabled, state, next_fire_at ASC);
|
||
|
||
CREATE VIRTUAL TABLE IF NOT EXISTS memories_fts USING fts5(
|
||
namespace,
|
||
memory_key,
|
||
content,
|
||
content='memories',
|
||
content_rowid='rowid'
|
||
);
|
||
|
||
CREATE TRIGGER IF NOT EXISTS memories_ai AFTER INSERT ON memories BEGIN
|
||
INSERT INTO memories_fts(rowid, namespace, memory_key, content)
|
||
VALUES (new.rowid, new.namespace, new.memory_key, new.content);
|
||
END;
|
||
|
||
CREATE TRIGGER IF NOT EXISTS memories_ad AFTER DELETE ON memories BEGIN
|
||
INSERT INTO memories_fts(memories_fts, rowid, namespace, memory_key, content)
|
||
VALUES ('delete', old.rowid, old.namespace, old.memory_key, old.content);
|
||
END;
|
||
|
||
CREATE TRIGGER IF NOT EXISTS memories_au AFTER UPDATE ON memories BEGIN
|
||
INSERT INTO memories_fts(memories_fts, rowid, namespace, memory_key, content)
|
||
VALUES ('delete', old.rowid, old.namespace, old.memory_key, old.content);
|
||
INSERT INTO memories_fts(rowid, namespace, memory_key, content)
|
||
VALUES (new.rowid, new.namespace, new.memory_key, new.content);
|
||
END;
|
||
",
|
||
)?;
|
||
|
||
ensure_sessions_schema(&conn)?;
|
||
ensure_messages_schema(&conn)?;
|
||
ensure_scheduler_schema(&conn)?;
|
||
ensure_memory_scope_key_migration(&conn)?;
|
||
ensure_todos_schema(&conn)?;
|
||
|
||
drop(conn);
|
||
|
||
let manager = SqliteConnectionManager::file(db_uri)
|
||
.with_init(|c| {
|
||
c.busy_timeout(std::time::Duration::from_secs(30))?;
|
||
Ok(())
|
||
});
|
||
let pool = Pool::builder()
|
||
.max_size(8)
|
||
.build(manager)?;
|
||
|
||
Ok(Self { pool })
|
||
}
|
||
|
||
#[cfg(test)]
|
||
pub(crate) fn in_memory() -> Result<Self, StorageError> {
|
||
// Use a temp file so the database survives across pool connections.
|
||
// Temp dir is cleaned by the OS eventually; tests that need cleanup
|
||
// can call std::fs::remove_file on the path.
|
||
let path = std::env::temp_dir()
|
||
.join(format!("picobot_test_{}.db", uuid::Uuid::new_v4()));
|
||
let conn = Connection::open(&path)?;
|
||
let path_str = path.to_string_lossy().to_string();
|
||
// ignore unused mut warning for manager in tests
|
||
#[allow(unused_mut)]
|
||
let store = Self::from_connection(conn, &path_str)?;
|
||
// Clean up temp file when the store is dropped
|
||
// We can't easily do this automatically, but the files are small.
|
||
Ok(store)
|
||
}
|
||
|
||
pub fn create_session(
|
||
&self,
|
||
channel_name: &str,
|
||
title: Option<&str>,
|
||
) -> Result<SessionRecord, StorageError> {
|
||
let now = current_timestamp();
|
||
let id = uuid::Uuid::new_v4().to_string();
|
||
// 统一使用 persistent_session_id 格式
|
||
let session_id = persistent_session_id(channel_name, &id);
|
||
let title = title
|
||
.map(str::trim)
|
||
.filter(|value| !value.is_empty())
|
||
.map(ToOwned::to_owned)
|
||
.unwrap_or_else(|| {
|
||
if channel_name == "cli" {
|
||
format!("CLI Session {}", &id[..8])
|
||
} else {
|
||
format!("Session {}", &id[..8])
|
||
}
|
||
});
|
||
|
||
let conn = self.pool.get()?;
|
||
conn.execute(
|
||
"
|
||
INSERT INTO sessions (
|
||
id, title, channel_name, chat_id, summary,
|
||
created_at, updated_at, last_active_at, archived_at, deleted_at, message_count,
|
||
user_turn_count, agent_prompt_reinjection_count
|
||
) VALUES (?1, ?2, ?3, ?4, NULL, ?5, ?5, ?5, NULL, NULL, 0, 0, 0)
|
||
",
|
||
params![&session_id, title, channel_name, id, now],
|
||
)?;
|
||
|
||
get_session_with_conn(&conn, &session_id)?
|
||
.ok_or_else(|| rusqlite::Error::QueryReturnedNoRows.into())
|
||
}
|
||
|
||
pub fn create_cli_session(&self, title: Option<&str>) -> Result<SessionRecord, StorageError> {
|
||
self.create_session("cli", title)
|
||
}
|
||
|
||
pub fn ensure_channel_session(
|
||
&self,
|
||
channel_name: &str,
|
||
chat_id: &str,
|
||
) -> Result<SessionRecord, StorageError> {
|
||
let session_id = persistent_session_id(channel_name, chat_id);
|
||
self.ensure_session(&session_id, channel_name, chat_id, &format!("{}:{}", channel_name, chat_id))
|
||
}
|
||
|
||
/// 确保指定 session_id 的会话存在（如果不存在则创建）
|
||
pub fn ensure_session(
|
||
&self,
|
||
session_id: &str,
|
||
channel_name: &str,
|
||
chat_id: &str,
|
||
title: &str,
|
||
) -> Result<SessionRecord, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
if let Some(record) = get_session_with_conn(&conn, session_id)? {
|
||
return Ok(record);
|
||
}
|
||
|
||
let now = current_timestamp();
|
||
conn.execute(
|
||
"
|
||
INSERT INTO sessions (
|
||
id, title, channel_name, chat_id, summary,
|
||
created_at, updated_at, last_active_at, archived_at, deleted_at, message_count,
|
||
user_turn_count, agent_prompt_reinjection_count
|
||
) VALUES (?1, ?2, ?3, ?4, NULL, ?5, ?5, ?5, NULL, NULL, 0, 0, 0)
|
||
",
|
||
params![session_id, title, channel_name, chat_id, now],
|
||
)?;
|
||
|
||
get_session_with_conn(&conn, session_id)?
|
||
.ok_or_else(|| rusqlite::Error::QueryReturnedNoRows.into())
|
||
}
|
||
|
||
pub fn get_session(&self, session_id: &str) -> Result<Option<SessionRecord>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
get_session_with_conn(&conn, session_id)
|
||
}
|
||
|
||
/// Find sessions whose id ends with the given suffix (used for task session lookup)
|
||
pub fn find_sessions_by_id_suffix(
|
||
&self,
|
||
suffix: &str,
|
||
) -> Result<Vec<SessionRecord>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let pattern = format!("%{}", suffix);
|
||
let mut stmt = conn.prepare(
|
||
"
|
||
SELECT id, title, channel_name, chat_id, summary,
|
||
created_at, updated_at, last_active_at,
|
||
archived_at, deleted_at, message_count,
|
||
user_turn_count, agent_prompt_reinjection_count
|
||
FROM sessions
|
||
WHERE id LIKE ?1 AND deleted_at IS NULL
|
||
ORDER BY last_active_at DESC
|
||
",
|
||
)?;
|
||
|
||
let rows = stmt.query_map(params![pattern], map_session_record)?;
|
||
let mut sessions = Vec::new();
|
||
for row in rows {
|
||
sessions.push(row?);
|
||
}
|
||
Ok(sessions)
|
||
}
|
||
|
||
pub fn list_sessions(
|
||
&self,
|
||
channel_name: &str,
|
||
include_archived: bool,
|
||
) -> Result<Vec<SessionRecord>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let mut sql = String::from(
|
||
"
|
||
SELECT id, title, channel_name, chat_id, summary,
|
||
created_at, updated_at, last_active_at,
|
||
archived_at, deleted_at, message_count,
|
||
user_turn_count, agent_prompt_reinjection_count
|
||
FROM sessions
|
||
WHERE channel_name = ?1
|
||
AND deleted_at IS NULL
|
||
AND id NOT LIKE 'sub:%'
|
||
",
|
||
);
|
||
|
||
if !include_archived {
|
||
sql.push_str(" AND archived_at IS NULL");
|
||
}
|
||
|
||
sql.push_str(" ORDER BY last_active_at DESC, created_at DESC");
|
||
|
||
let mut stmt = conn.prepare(&sql)?;
|
||
let rows = stmt.query_map(params![channel_name], map_session_record)?;
|
||
let mut sessions = Vec::new();
|
||
for row in rows {
|
||
sessions.push(row?);
|
||
}
|
||
Ok(sessions)
|
||
}
|
||
|
||
pub fn rename_session(&self, session_id: &str, title: &str) -> Result<(), StorageError> {
|
||
let now = current_timestamp();
|
||
let conn = self.pool.get()?;
|
||
conn.execute(
|
||
"UPDATE sessions SET title = ?2, updated_at = ?3 WHERE id = ?1 AND deleted_at IS NULL",
|
||
params![session_id, title.trim(), now],
|
||
)?;
|
||
Ok(())
|
||
}
|
||
|
||
pub fn archive_session(&self, session_id: &str) -> Result<(), StorageError> {
|
||
let now = current_timestamp();
|
||
let conn = self.pool.get()?;
|
||
conn.execute(
|
||
"UPDATE sessions SET archived_at = ?2, updated_at = ?2 WHERE id = ?1 AND deleted_at IS NULL",
|
||
params![session_id, now],
|
||
)?;
|
||
Ok(())
|
||
}
|
||
|
||
pub fn delete_session(&self, session_id: &str) -> Result<(), StorageError> {
|
||
let conn = self.pool.get()?;
|
||
conn.execute(
|
||
"DELETE FROM messages WHERE session_id = ?1",
|
||
params![session_id],
|
||
)?;
|
||
conn.execute("DELETE FROM sessions WHERE id = ?1", params![session_id])?;
|
||
Ok(())
|
||
}
|
||
|
||
// ==================== Topic Methods ====================
|
||
|
||
pub fn create_topic(
|
||
&self,
|
||
session_id: &str,
|
||
title: &str,
|
||
description: Option<&str>,
|
||
) -> Result<TopicRecord, StorageError> {
|
||
let now = current_timestamp();
|
||
let id = format!("topic:{}", uuid::Uuid::new_v4());
|
||
|
||
let conn = self.pool.get()?;
|
||
conn.execute(
|
||
"INSERT INTO topics (id, session_id, title, description, created_at, updated_at, last_active_at, message_count) VALUES (?1, ?2, ?3, ?4, ?5, ?5, ?5, 0)",
|
||
params![&id, session_id, title, description.unwrap_or(""), now],
|
||
)?;
|
||
drop(conn);
|
||
|
||
self.get_topic(&id)?
|
||
.ok_or_else(|| rusqlite::Error::QueryReturnedNoRows.into())
|
||
}
|
||
|
||
pub fn get_topic(&self, topic_id: &str) -> Result<Option<TopicRecord>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let mut stmt = conn.prepare(
|
||
"SELECT id, session_id, title, description, created_at, updated_at, last_active_at, message_count FROM topics WHERE id = ?1",
|
||
)?;
|
||
|
||
stmt.query_row(params![topic_id], |row| {
|
||
Ok(TopicRecord {
|
||
id: row.get(0)?,
|
||
session_id: row.get(1)?,
|
||
title: row.get(2)?,
|
||
description: row.get(3)?,
|
||
created_at: row.get(4)?,
|
||
updated_at: row.get(5)?,
|
||
last_active_at: row.get(6)?,
|
||
message_count: row.get(7)?,
|
||
})
|
||
})
|
||
.optional()
|
||
.map_err(StorageError::from)
|
||
}
|
||
|
||
pub fn list_topics(&self, session_id: &str) -> Result<Vec<TopicRecord>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let mut stmt = conn.prepare(
|
||
"SELECT id, session_id, title, description, created_at, updated_at, last_active_at, message_count FROM topics WHERE session_id = ?1 ORDER BY last_active_at DESC"
|
||
)?;
|
||
|
||
let rows = stmt.query_map(params![session_id], |row| {
|
||
Ok(TopicRecord {
|
||
id: row.get(0)?,
|
||
session_id: row.get(1)?,
|
||
title: row.get(2)?,
|
||
description: row.get(3)?,
|
||
created_at: row.get(4)?,
|
||
updated_at: row.get(5)?,
|
||
last_active_at: row.get(6)?,
|
||
message_count: row.get(7)?,
|
||
})
|
||
})?;
|
||
|
||
let mut topics = Vec::new();
|
||
for row in rows {
|
||
topics.push(row?);
|
||
}
|
||
Ok(topics)
|
||
}
|
||
|
||
pub fn update_topic_title(&self, topic_id: &str, title: &str) -> Result<(), StorageError> {
|
||
let now = current_timestamp();
|
||
let conn = self.pool.get()?;
|
||
conn.execute(
|
||
"UPDATE topics SET title = ?2, updated_at = ?3 WHERE id = ?1",
|
||
params![topic_id, title.trim(), now],
|
||
)?;
|
||
Ok(())
|
||
}
|
||
|
||
pub fn update_topic_description(&self, topic_id: &str, description: &str) -> Result<(), StorageError> {
|
||
let now = current_timestamp();
|
||
let conn = self.pool.get()?;
|
||
conn.execute(
|
||
"UPDATE topics SET description = ?2, updated_at = ?3 WHERE id = ?1",
|
||
params![topic_id, description, now],
|
||
)?;
|
||
Ok(())
|
||
}
|
||
|
||
pub fn delete_topic(&self, topic_id: &str) -> Result<(), StorageError> {
|
||
let conn = self.pool.get()?;
|
||
// Messages 的 topic_id 会被设为 NULL（ON DELETE SET NULL）
|
||
conn.execute("DELETE FROM topics WHERE id = ?1", params![topic_id])?;
|
||
Ok(())
|
||
}
|
||
|
||
pub fn touch_topic(&self, topic_id: &str) -> Result<(), StorageError> {
|
||
let now = current_timestamp();
|
||
let conn = self.pool.get()?;
|
||
conn.execute(
|
||
"UPDATE topics SET last_active_at = ?2 WHERE id = ?1",
|
||
params![topic_id, now],
|
||
)?;
|
||
Ok(())
|
||
}
|
||
|
||
pub fn clear_messages(&self, session_id: &str) -> Result<(), StorageError> {
|
||
let now = current_timestamp();
|
||
let conn = self.pool.get()?;
|
||
conn.execute(
|
||
"DELETE FROM messages WHERE session_id = ?1",
|
||
params![session_id],
|
||
)?;
|
||
conn.execute(
|
||
"
|
||
UPDATE sessions
|
||
SET message_count = 0,
|
||
updated_at = ?2,
|
||
last_active_at = ?2,
|
||
user_turn_count = 0,
|
||
agent_prompt_reinjection_count = 0
|
||
WHERE id = ?1 AND deleted_at IS NULL
|
||
",
|
||
params![session_id, now],
|
||
)?;
|
||
Ok(())
|
||
}
|
||
|
||
pub fn append_message(
|
||
&self,
|
||
session_id: &str,
|
||
message: &ChatMessage,
|
||
) -> Result<(), StorageError> {
|
||
self.append_message_with_topic(session_id, None, message)
|
||
}
|
||
|
||
pub fn append_message_with_topic(
|
||
&self,
|
||
session_id: &str,
|
||
topic_id: Option<&str>,
|
||
message: &ChatMessage,
|
||
) -> Result<(), StorageError> {
|
||
let mut conn = self.pool.get()?;
|
||
let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
|
||
|
||
let seq: i64 = tx.query_row(
|
||
"SELECT COALESCE(MAX(seq), 0) + 1 FROM messages WHERE session_id = ?1",
|
||
params![session_id],
|
||
|row| row.get(0),
|
||
)?;
|
||
|
||
let media_refs_json = serde_json::to_string(&message.media_refs)?;
|
||
let tool_calls_json = message
|
||
.tool_calls
|
||
.as_ref()
|
||
.map(serde_json::to_string)
|
||
.transpose()?;
|
||
tx.execute(
|
||
"
|
||
INSERT INTO messages (
|
||
id, session_id, topic_id, seq, role, content,
|
||
system_context, reasoning_content, media_refs_json, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, created_at
|
||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)
|
||
",
|
||
params![
|
||
message.id,
|
||
session_id,
|
||
topic_id,
|
||
seq,
|
||
message.role,
|
||
message.content,
|
||
message.system_context,
|
||
message.reasoning_content,
|
||
media_refs_json,
|
||
message.tool_call_id,
|
||
message.tool_name,
|
||
tool_calls_json,
|
||
message.tool_duration_ms.map(|v| v as i64),
|
||
message.timestamp,
|
||
],
|
||
)?;
|
||
|
||
let now = current_timestamp();
|
||
let is_user_message = message.role == "user";
|
||
tx.execute(
|
||
"
|
||
UPDATE sessions
|
||
SET message_count = message_count + 1,
|
||
user_turn_count = user_turn_count + ?3,
|
||
updated_at = ?2,
|
||
last_active_at = ?2,
|
||
archived_at = NULL
|
||
WHERE id = ?1 AND deleted_at IS NULL
|
||
",
|
||
params![session_id, now, if is_user_message { 1 } else { 0 }],
|
||
)?;
|
||
|
||
if let Some(tid) = topic_id {
|
||
tx.execute(
|
||
"UPDATE topics SET message_count = message_count + 1, last_active_at = ?2 WHERE id = ?1",
|
||
params![tid, now],
|
||
)?;
|
||
}
|
||
|
||
tx.commit()?;
|
||
Ok(())
|
||
}
|
||
|
||
pub fn append_messages_batch(
|
||
&self,
|
||
session_id: &str,
|
||
topic_id: Option<&str>,
|
||
messages: &[ChatMessage],
|
||
) -> Result<(), StorageError> {
|
||
if messages.is_empty() {
|
||
return Ok(());
|
||
}
|
||
|
||
let mut conn = self.pool.get()?;
|
||
let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
|
||
|
||
let mut seq: i64 = tx.query_row(
|
||
"SELECT COALESCE(MAX(seq), 0) + 1 FROM messages WHERE session_id = ?1",
|
||
params![session_id],
|
||
|row| row.get(0),
|
||
)?;
|
||
|
||
for message in messages {
|
||
let media_refs_json = serde_json::to_string(&message.media_refs)?;
|
||
let tool_calls_json = message
|
||
.tool_calls
|
||
.as_ref()
|
||
.map(serde_json::to_string)
|
||
.transpose()?;
|
||
tx.execute(
|
||
"
|
||
INSERT INTO messages (
|
||
id, session_id, topic_id, seq, role, content,
|
||
system_context, reasoning_content, media_refs_json,
|
||
tool_call_id, tool_name, tool_calls_json, tool_duration_ms, created_at
|
||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)
|
||
",
|
||
params![
|
||
message.id,
|
||
session_id,
|
||
topic_id,
|
||
seq,
|
||
message.role,
|
||
message.content,
|
||
message.system_context,
|
||
message.reasoning_content,
|
||
media_refs_json,
|
||
message.tool_call_id,
|
||
message.tool_name,
|
||
tool_calls_json,
|
||
message.tool_duration_ms.map(|v| v as i64),
|
||
message.timestamp,
|
||
],
|
||
)?;
|
||
seq += 1;
|
||
}
|
||
|
||
let now = current_timestamp();
|
||
let user_msg_count: i64 = messages
|
||
.iter()
|
||
.filter(|m| m.role == "user")
|
||
.count()
|
||
.try_into()
|
||
.unwrap_or(0);
|
||
let msg_count: i64 = messages.len() as i64;
|
||
|
||
tx.execute(
|
||
"
|
||
UPDATE sessions
|
||
SET message_count = message_count + ?2,
|
||
user_turn_count = user_turn_count + ?3,
|
||
updated_at = ?4,
|
||
last_active_at = ?4,
|
||
archived_at = NULL
|
||
WHERE id = ?1 AND deleted_at IS NULL
|
||
",
|
||
params![session_id, msg_count, user_msg_count, now],
|
||
)?;
|
||
|
||
if let Some(tid) = topic_id {
|
||
tx.execute(
|
||
"UPDATE topics SET message_count = message_count + ?2, last_active_at = ?3 WHERE id = ?1",
|
||
params![tid, msg_count, now],
|
||
)?;
|
||
}
|
||
|
||
tx.commit()?;
|
||
Ok(())
|
||
}
|
||
|
||
pub fn compact_active_history(
|
||
&self,
|
||
session_id: &str,
|
||
snapshot_end_seq: i64,
|
||
preserved_system_messages: &[ChatMessage],
|
||
summary_message: &ChatMessage,
|
||
preserved_messages: &[ChatMessage],
|
||
) -> Result<bool, StorageError> {
|
||
let mut conn = self.pool.get()?;
|
||
let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
|
||
|
||
let current_max_seq: i64 = tx.query_row(
|
||
"SELECT COALESCE(MAX(seq), 0) FROM messages WHERE session_id = ?1",
|
||
params![session_id],
|
||
|row| row.get(0),
|
||
)?;
|
||
|
||
if snapshot_end_seq > current_max_seq {
|
||
return Ok(false);
|
||
}
|
||
|
||
let delta_messages =
|
||
load_messages_between(&tx, session_id, snapshot_end_seq, current_max_seq)?;
|
||
let now = current_timestamp();
|
||
|
||
// Collect all new messages first, then sanitize incomplete tool call
|
||
// sequences before writing to DB. This prevents orphaned tool_calls
|
||
// (without corresponding tool results) from being persisted permanently
|
||
// when compaction preserves an incomplete sequence from the snapshot or
|
||
// captures a partial sequence from delta messages.
|
||
let mut new_messages: Vec<ChatMessage> = Vec::new();
|
||
|
||
for message in preserved_system_messages {
|
||
new_messages.push(clone_message_for_compaction(message, message.timestamp));
|
||
}
|
||
|
||
new_messages.push(clone_message_for_compaction(summary_message, now));
|
||
|
||
for message in preserved_messages.iter().chain(delta_messages.iter()) {
|
||
new_messages.push(clone_message_for_compaction(message, message.timestamp));
|
||
}
|
||
|
||
let removed =
|
||
crate::bus::message::sanitize_incomplete_tool_call_sequences(&mut new_messages);
|
||
if removed > 0 {
|
||
tracing::warn!(
|
||
removed_count = removed,
|
||
session_id = %session_id,
|
||
"Compaction removed incomplete tool call sequences from new history"
|
||
);
|
||
}
|
||
|
||
// Write sanitized messages to DB
|
||
let mut next_seq = current_max_seq + 1;
|
||
let mut inserted_count = 0_i64;
|
||
let mut active_user_turn_count = 0_i64;
|
||
|
||
for message in &new_messages {
|
||
if message.role == "user" {
|
||
active_user_turn_count += 1;
|
||
}
|
||
insert_message_with_seq(&tx, session_id, next_seq, message)?;
|
||
next_seq += 1;
|
||
inserted_count += 1;
|
||
}
|
||
|
||
// Delete all old messages (including delta messages that were just re-inserted)
|
||
tx.execute(
|
||
"DELETE FROM messages WHERE session_id = ?1 AND seq <= ?2",
|
||
params![session_id, current_max_seq],
|
||
)?;
|
||
|
||
tx.execute(
|
||
"
|
||
UPDATE sessions
|
||
SET message_count = ?2,
|
||
user_turn_count = ?3,
|
||
updated_at = ?4,
|
||
last_active_at = ?4,
|
||
archived_at = NULL
|
||
WHERE id = ?1 AND deleted_at IS NULL
|
||
",
|
||
params![
|
||
session_id,
|
||
inserted_count,
|
||
active_user_turn_count,
|
||
now,
|
||
],
|
||
)?;
|
||
|
||
tx.commit()?;
|
||
Ok(true)
|
||
}
|
||
|
||
/// Replace the entire active history for a session.
|
||
///
|
||
/// This is a simpler alternative to `compact_active_history` for when the
|
||
/// compressor has already produced a complete, validated message list
|
||
/// (e.g. two-segment compression). It replaces all existing messages
|
||
/// with the new list in a single transaction.
|
||
pub fn replace_active_history(
|
||
&self,
|
||
session_id: &str,
|
||
messages: &[ChatMessage],
|
||
) -> Result<(), StorageError> {
|
||
let mut conn = self.pool.get()?;
|
||
let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
|
||
let now = current_timestamp();
|
||
|
||
// Delete all existing messages for this session
|
||
tx.execute(
|
||
"DELETE FROM messages WHERE session_id = ?1",
|
||
params![session_id],
|
||
)?;
|
||
|
||
// Insert new messages with sequential seq numbers
|
||
let mut active_user_turn_count = 0_i64;
|
||
for (i, message) in messages.iter().enumerate() {
|
||
let seq = (i + 1) as i64;
|
||
if message.role == "user" {
|
||
active_user_turn_count += 1;
|
||
}
|
||
insert_message_with_seq(&tx, session_id, seq, message)?;
|
||
}
|
||
|
||
tx.execute(
|
||
"
|
||
UPDATE sessions
|
||
SET message_count = ?2,
|
||
user_turn_count = ?3,
|
||
updated_at = ?4,
|
||
last_active_at = ?4,
|
||
archived_at = NULL
|
||
WHERE id = ?1 AND deleted_at IS NULL
|
||
",
|
||
params![
|
||
session_id,
|
||
messages.len() as i64,
|
||
active_user_turn_count,
|
||
now,
|
||
],
|
||
)?;
|
||
|
||
tx.commit()?;
|
||
Ok(())
|
||
}
|
||
|
||
/// Replace the entire history for a specific topic.
|
||
///
|
||
/// Deletes only messages belonging to the given topic_id (preserving
|
||
/// other topics' messages), then inserts the new messages with topic_id
|
||
/// set correctly. Used by the compressor when it has produced a
|
||
/// complete, validated message list for a single topic.
|
||
///
|
||
/// Seq numbers continue from the current session-wide max (not reset to
|
||
/// 1) so we don't collide with other topics' messages. Gaps in seq
|
||
/// (from the deleted old messages) are harmless — per-topic loading
|
||
/// orders by seq and gaps don't affect ordering.
|
||
pub fn replace_topic_history(
|
||
&self,
|
||
session_id: &str,
|
||
topic_id: &str,
|
||
messages: &[ChatMessage],
|
||
) -> Result<(), StorageError> {
|
||
let mut conn = self.pool.get()?;
|
||
let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
|
||
let now = current_timestamp();
|
||
|
||
// Delete only messages belonging to this topic — other topics'
|
||
// messages are preserved (the pre-existing `replace_active_history`
|
||
// clobbered the entire session, which broke multi-topic isolation).
|
||
tx.execute(
|
||
"DELETE FROM messages WHERE session_id = ?1 AND topic_id = ?2",
|
||
params![session_id, topic_id],
|
||
)?;
|
||
|
||
// Continue seq from the session-wide max so we don't violate
|
||
// UNIQUE(session_id, seq). Other topics' messages keep their seqs.
|
||
let start_seq: i64 = tx.query_row(
|
||
"SELECT COALESCE(MAX(seq), 0) + 1 FROM messages WHERE session_id = ?1",
|
||
params![session_id],
|
||
|row| row.get(0),
|
||
)?;
|
||
|
||
for (i, message) in messages.iter().enumerate() {
|
||
let seq = start_seq + i as i64;
|
||
insert_message_with_topic_seq(&tx, session_id, topic_id, seq, message)?;
|
||
}
|
||
|
||
// Update this topic's message_count and timestamps.
|
||
tx.execute(
|
||
"UPDATE topics SET message_count = ?2, last_active_at = ?3, updated_at = ?3 WHERE id = ?1",
|
||
params![topic_id, messages.len() as i64, now],
|
||
)?;
|
||
|
||
// Recompute session-wide counts from the messages table so they stay
|
||
// consistent after a partial replacement (we only touched one topic,
|
||
// so we can't just set the session count to `messages.len()`).
|
||
let (total_count, user_turn_count): (i64, i64) = tx.query_row(
|
||
"SELECT COUNT(*), COALESCE(SUM(CASE WHEN role = 'user' THEN 1 ELSE 0 END), 0) \
|
||
FROM messages WHERE session_id = ?1",
|
||
params![session_id],
|
||
|row| Ok((row.get(0)?, row.get(1)?)),
|
||
)?;
|
||
tx.execute(
|
||
"UPDATE sessions SET message_count = ?2, user_turn_count = ?3, \
|
||
updated_at = ?4, last_active_at = ?4, archived_at = NULL \
|
||
WHERE id = ?1 AND deleted_at IS NULL",
|
||
params![session_id, total_count, user_turn_count, now],
|
||
)?;
|
||
|
||
tx.commit()?;
|
||
Ok(())
|
||
}
|
||
|
||
pub fn mark_agent_prompt_reinjected(&self, session_id: &str) -> Result<(), StorageError> {
|
||
let now = current_timestamp();
|
||
let conn = self.pool.get()?;
|
||
conn.execute(
|
||
"
|
||
UPDATE sessions
|
||
SET agent_prompt_reinjection_count = agent_prompt_reinjection_count + 1,
|
||
updated_at = ?2,
|
||
last_active_at = ?2,
|
||
archived_at = NULL
|
||
WHERE id = ?1 AND deleted_at IS NULL
|
||
",
|
||
params![session_id, now],
|
||
)?;
|
||
Ok(())
|
||
}
|
||
|
||
pub fn append_skill_event(
|
||
&self,
|
||
session_id: Option<&str>,
|
||
event_type: &str,
|
||
skill_name: Option<&str>,
|
||
payload: &serde_json::Value,
|
||
) -> Result<(), StorageError> {
|
||
let conn = self.pool.get()?;
|
||
conn.execute(
|
||
"
|
||
INSERT INTO skill_events (
|
||
id, session_id, event_type, skill_name, payload_json, created_at
|
||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6)
|
||
",
|
||
params![
|
||
uuid::Uuid::new_v4().to_string(),
|
||
session_id,
|
||
event_type,
|
||
skill_name,
|
||
serde_json::to_string(payload)?,
|
||
current_timestamp(),
|
||
],
|
||
)?;
|
||
Ok(())
|
||
}
|
||
|
||
pub fn list_skill_events(
|
||
&self,
|
||
session_id: Option<&str>,
|
||
) -> Result<Vec<SkillEventRecord>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let sql = if session_id.is_some() {
|
||
"
|
||
SELECT id, session_id, event_type, skill_name, payload_json, created_at
|
||
FROM skill_events
|
||
WHERE session_id = ?1
|
||
ORDER BY created_at ASC
|
||
"
|
||
} else {
|
||
"
|
||
SELECT id, session_id, event_type, skill_name, payload_json, created_at
|
||
FROM skill_events
|
||
WHERE session_id IS NULL
|
||
ORDER BY created_at ASC
|
||
"
|
||
};
|
||
|
||
let mut stmt = conn.prepare(sql)?;
|
||
let rows = if let Some(session_id) = session_id {
|
||
stmt.query_map(params![session_id], map_skill_event_record)?
|
||
} else {
|
||
stmt.query_map([], map_skill_event_record)?
|
||
};
|
||
|
||
let mut events = Vec::new();
|
||
for row in rows {
|
||
events.push(row?);
|
||
}
|
||
Ok(events)
|
||
}
|
||
|
||
pub fn put_memory(&self, input: &MemoryUpsert) -> Result<MemoryRecord, StorageError> {
|
||
let now = current_timestamp();
|
||
let mut conn = self.pool.get()?;
|
||
let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
|
||
|
||
let existing: Option<(String, i64)> = tx
|
||
.query_row(
|
||
"
|
||
SELECT id, created_at
|
||
FROM memories
|
||
WHERE scope_kind = ?1 AND scope_key = ?2 AND namespace = ?3 AND memory_key = ?4
|
||
",
|
||
params![
|
||
input.scope_kind,
|
||
input.scope_key,
|
||
input.namespace,
|
||
input.memory_key,
|
||
],
|
||
|row| Ok((row.get(0)?, row.get(1)?)),
|
||
)
|
||
.optional()?;
|
||
|
||
let (id, created_at) = existing.unwrap_or_else(|| (uuid::Uuid::new_v4().to_string(), now));
|
||
|
||
tx.execute(
|
||
"
|
||
INSERT INTO memories (
|
||
id, scope_kind, scope_key, namespace, memory_key, content,
|
||
source_type, source_session_id, source_message_id, source_message_seq,
|
||
source_channel_name, source_chat_id, created_at, updated_at
|
||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)
|
||
ON CONFLICT(scope_kind, scope_key, namespace, memory_key) DO UPDATE SET
|
||
content = excluded.content,
|
||
source_type = excluded.source_type,
|
||
source_session_id = excluded.source_session_id,
|
||
source_message_id = excluded.source_message_id,
|
||
source_message_seq = excluded.source_message_seq,
|
||
source_channel_name = excluded.source_channel_name,
|
||
source_chat_id = excluded.source_chat_id,
|
||
updated_at = excluded.updated_at
|
||
",
|
||
params![
|
||
id,
|
||
input.scope_kind,
|
||
input.scope_key,
|
||
input.namespace,
|
||
input.memory_key,
|
||
input.content,
|
||
input.source_type,
|
||
input.source_session_id,
|
||
input.source_message_id,
|
||
input.source_message_seq,
|
||
input.source_channel_name,
|
||
input.source_chat_id,
|
||
created_at,
|
||
now,
|
||
],
|
||
)?;
|
||
|
||
tx.commit()?;
|
||
|
||
get_memory_with_conn(
|
||
&conn,
|
||
&input.scope_kind,
|
||
&input.scope_key,
|
||
&input.namespace,
|
||
&input.memory_key,
|
||
)?
|
||
.ok_or_else(|| rusqlite::Error::QueryReturnedNoRows.into())
|
||
}
|
||
|
||
pub fn get_memory(
|
||
&self,
|
||
scope_kind: &str,
|
||
scope_key: &str,
|
||
namespace: &str,
|
||
memory_key: &str,
|
||
) -> Result<Option<MemoryRecord>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
get_memory_with_conn(&conn, scope_kind, scope_key, namespace, memory_key)
|
||
}
|
||
|
||
pub fn list_memories(
|
||
&self,
|
||
scope_kind: &str,
|
||
scope_key: &str,
|
||
namespace: Option<&str>,
|
||
limit: usize,
|
||
) -> Result<Vec<MemoryRecord>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let limit = limit.max(1) as i64;
|
||
let mut memories = Vec::new();
|
||
|
||
if let Some(namespace) = namespace {
|
||
let mut stmt = conn.prepare(
|
||
"
|
||
SELECT id, scope_kind, scope_key, namespace, memory_key, content,
|
||
source_type, source_session_id, source_message_id, source_message_seq,
|
||
source_channel_name, source_chat_id, created_at, updated_at
|
||
FROM memories
|
||
WHERE scope_kind = ?1 AND scope_key = ?2 AND namespace = ?3
|
||
ORDER BY updated_at DESC, created_at DESC
|
||
LIMIT ?4
|
||
",
|
||
)?;
|
||
let rows = stmt.query_map(
|
||
params![scope_kind, scope_key, namespace, limit],
|
||
map_memory_record,
|
||
)?;
|
||
for row in rows {
|
||
memories.push(row?);
|
||
}
|
||
} else {
|
||
let mut stmt = conn.prepare(
|
||
"
|
||
SELECT id, scope_kind, scope_key, namespace, memory_key, content,
|
||
source_type, source_session_id, source_message_id, source_message_seq,
|
||
source_channel_name, source_chat_id, created_at, updated_at
|
||
FROM memories
|
||
WHERE scope_kind = ?1 AND scope_key = ?2
|
||
ORDER BY updated_at DESC, created_at DESC
|
||
LIMIT ?3
|
||
",
|
||
)?;
|
||
let rows = stmt.query_map(params![scope_kind, scope_key, limit], map_memory_record)?;
|
||
for row in rows {
|
||
memories.push(row?);
|
||
}
|
||
}
|
||
|
||
Ok(memories)
|
||
}
|
||
|
||
pub fn list_memory_scope_keys(&self, scope_kind: &str) -> Result<Vec<String>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let mut stmt = conn.prepare(
|
||
"
|
||
SELECT DISTINCT scope_key
|
||
FROM memories
|
||
WHERE scope_kind = ?1
|
||
ORDER BY scope_key ASC
|
||
",
|
||
)?;
|
||
|
||
let rows = stmt.query_map(params![scope_kind], |row| row.get::<_, String>(0))?;
|
||
let mut scope_keys = Vec::new();
|
||
for row in rows {
|
||
scope_keys.push(row?);
|
||
}
|
||
Ok(scope_keys)
|
||
}
|
||
|
||
pub fn list_memories_for_scope(
|
||
&self,
|
||
scope_kind: &str,
|
||
scope_key: &str,
|
||
) -> Result<Vec<MemoryRecord>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let mut stmt = conn.prepare(
|
||
"
|
||
SELECT id, scope_kind, scope_key, namespace, memory_key, content,
|
||
source_type, source_session_id, source_message_id, source_message_seq,
|
||
source_channel_name, source_chat_id, created_at, updated_at
|
||
FROM memories
|
||
WHERE scope_kind = ?1 AND scope_key = ?2
|
||
ORDER BY updated_at DESC, namespace ASC, memory_key ASC
|
||
",
|
||
)?;
|
||
|
||
let rows = stmt.query_map(params![scope_kind, scope_key], map_memory_record)?;
|
||
let mut memories = Vec::new();
|
||
for row in rows {
|
||
memories.push(row?);
|
||
}
|
||
Ok(memories)
|
||
}
|
||
|
||
pub fn update_memory(
|
||
&self,
|
||
input: &MemoryUpsert,
|
||
) -> Result<Option<MemoryRecord>, StorageError> {
|
||
if self
|
||
.get_memory(
|
||
&input.scope_kind,
|
||
&input.scope_key,
|
||
&input.namespace,
|
||
&input.memory_key,
|
||
)?
|
||
.is_none()
|
||
{
|
||
return Ok(None);
|
||
}
|
||
|
||
self.put_memory(input).map(Some)
|
||
}
|
||
|
||
pub fn delete_memory(
|
||
&self,
|
||
scope_kind: &str,
|
||
scope_key: &str,
|
||
namespace: &str,
|
||
memory_key: &str,
|
||
) -> Result<bool, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let changed = conn.execute(
|
||
"
|
||
DELETE FROM memories
|
||
WHERE scope_kind = ?1 AND scope_key = ?2 AND namespace = ?3 AND memory_key = ?4
|
||
",
|
||
params![scope_kind, scope_key, namespace, memory_key],
|
||
)?;
|
||
Ok(changed > 0)
|
||
}
|
||
|
||
pub fn upsert_scheduler_job(
|
||
&self,
|
||
input: &SchedulerJobUpsert,
|
||
) -> Result<SchedulerJobRecord, StorageError> {
|
||
let now = current_timestamp();
|
||
let conn = self.pool.get()?;
|
||
conn.execute(
|
||
"
|
||
INSERT INTO scheduler_jobs (
|
||
id, kind, schedule_json, interval_secs, startup_delay_secs,
|
||
target_json, payload_json, enabled, state, last_status, last_error,
|
||
run_count, max_runs, last_fired_at, next_fire_at, paused_at, completed_at,
|
||
created_at, updated_at
|
||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?18)
|
||
ON CONFLICT(id) DO UPDATE SET
|
||
kind = excluded.kind,
|
||
schedule_json = excluded.schedule_json,
|
||
interval_secs = excluded.interval_secs,
|
||
startup_delay_secs = excluded.startup_delay_secs,
|
||
target_json = excluded.target_json,
|
||
payload_json = excluded.payload_json,
|
||
enabled = excluded.enabled,
|
||
state = excluded.state,
|
||
last_status = excluded.last_status,
|
||
last_error = excluded.last_error,
|
||
run_count = excluded.run_count,
|
||
max_runs = excluded.max_runs,
|
||
last_fired_at = excluded.last_fired_at,
|
||
next_fire_at = excluded.next_fire_at,
|
||
paused_at = excluded.paused_at,
|
||
completed_at = excluded.completed_at,
|
||
updated_at = excluded.updated_at
|
||
",
|
||
params![
|
||
input.id,
|
||
input.kind,
|
||
serde_json::to_string(&input.schedule)?,
|
||
input.interval_secs,
|
||
input.startup_delay_secs,
|
||
serde_json::to_string(&input.target)?,
|
||
serde_json::to_string(&input.payload)?,
|
||
if input.enabled { 1 } else { 0 },
|
||
input.state.as_str(),
|
||
input.last_status.as_ref().map(SchedulerJobStatus::as_str),
|
||
input.last_error,
|
||
input.run_count,
|
||
input.max_runs,
|
||
input.last_fired_at,
|
||
input.next_fire_at,
|
||
input.paused_at,
|
||
input.completed_at,
|
||
now,
|
||
],
|
||
)?;
|
||
get_scheduler_job_with_conn(&conn, &input.id)?
|
||
.ok_or_else(|| rusqlite::Error::QueryReturnedNoRows.into())
|
||
}
|
||
|
||
pub fn get_scheduler_job(
|
||
&self,
|
||
job_id: &str,
|
||
) -> Result<Option<SchedulerJobRecord>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
get_scheduler_job_with_conn(&conn, job_id)
|
||
}
|
||
|
||
pub fn list_scheduler_jobs(
|
||
&self,
|
||
enabled_only: bool,
|
||
) -> Result<Vec<SchedulerJobRecord>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let sql = if enabled_only {
|
||
"
|
||
SELECT id, kind, schedule_json, interval_secs, startup_delay_secs,
|
||
target_json, payload_json, enabled, state, last_status, last_error,
|
||
run_count, max_runs, last_fired_at, next_fire_at, paused_at, completed_at,
|
||
created_at, updated_at
|
||
FROM scheduler_jobs
|
||
WHERE enabled = 1
|
||
ORDER BY COALESCE(next_fire_at, created_at) ASC, id ASC
|
||
"
|
||
} else {
|
||
"
|
||
SELECT id, kind, schedule_json, interval_secs, startup_delay_secs,
|
||
target_json, payload_json, enabled, state, last_status, last_error,
|
||
run_count, max_runs, last_fired_at, next_fire_at, paused_at, completed_at,
|
||
created_at, updated_at
|
||
FROM scheduler_jobs
|
||
ORDER BY COALESCE(next_fire_at, created_at) ASC, id ASC
|
||
"
|
||
};
|
||
|
||
let mut stmt = conn.prepare(sql)?;
|
||
let rows = stmt.query_map([], map_scheduler_job_record)?;
|
||
let mut jobs = Vec::new();
|
||
for row in rows {
|
||
jobs.push(row?);
|
||
}
|
||
Ok(jobs)
|
||
}
|
||
|
||
pub fn list_running_scheduler_jobs(&self) -> Result<Vec<SchedulerJobRecord>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let sql = "
|
||
SELECT id, kind, schedule_json, interval_secs, startup_delay_secs,
|
||
target_json, payload_json, enabled, state, last_status, last_error,
|
||
run_count, max_runs, last_fired_at, next_fire_at, paused_at, completed_at,
|
||
created_at, updated_at
|
||
FROM scheduler_jobs
|
||
WHERE state = 'running'
|
||
ORDER BY COALESCE(next_fire_at, created_at) ASC, id ASC
|
||
";
|
||
|
||
let mut stmt = conn.prepare(sql)?;
|
||
let rows = stmt.query_map([], map_scheduler_job_record)?;
|
||
let mut jobs = Vec::new();
|
||
for row in rows {
|
||
jobs.push(row?);
|
||
}
|
||
Ok(jobs)
|
||
}
|
||
|
||
pub fn delete_scheduler_job(&self, job_id: &str) -> Result<(), StorageError> {
|
||
let conn = self.pool.get()?;
|
||
conn.execute("DELETE FROM scheduler_jobs WHERE id = ?1", params![job_id])?;
|
||
Ok(())
|
||
}
|
||
|
||
pub fn update_scheduler_job_runtime(
|
||
&self,
|
||
job_id: &str,
|
||
state: SchedulerJobState,
|
||
last_status: Option<SchedulerJobStatus>,
|
||
last_error: Option<&str>,
|
||
run_count: i64,
|
||
last_fired_at: Option<i64>,
|
||
next_fire_at: Option<i64>,
|
||
paused_at: Option<i64>,
|
||
completed_at: Option<i64>,
|
||
) -> Result<(), StorageError> {
|
||
let conn = self.pool.get()?;
|
||
conn.execute(
|
||
"
|
||
UPDATE scheduler_jobs
|
||
SET state = ?2,
|
||
last_status = ?3,
|
||
last_error = ?4,
|
||
run_count = ?5,
|
||
last_fired_at = ?6,
|
||
next_fire_at = ?7,
|
||
paused_at = ?8,
|
||
completed_at = ?9,
|
||
updated_at = ?10
|
||
WHERE id = ?1
|
||
",
|
||
params![
|
||
job_id,
|
||
state.as_str(),
|
||
last_status.as_ref().map(SchedulerJobStatus::as_str),
|
||
last_error,
|
||
run_count,
|
||
last_fired_at,
|
||
next_fire_at,
|
||
paused_at,
|
||
completed_at,
|
||
current_timestamp(),
|
||
],
|
||
)?;
|
||
Ok(())
|
||
}
|
||
|
||
pub fn search_memories(
|
||
&self,
|
||
scope_kind: &str,
|
||
scope_key: &str,
|
||
query: &str,
|
||
namespace: Option<&str>,
|
||
limit: usize,
|
||
) -> Result<Vec<MemoryRecord>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let limit = limit.max(1) as i64;
|
||
let query = quote_fts_query(query);
|
||
let mut memories = Vec::new();
|
||
|
||
if let Some(namespace) = namespace {
|
||
let mut stmt = conn.prepare(
|
||
"
|
||
SELECT m.id, m.scope_kind, m.scope_key, m.namespace, m.memory_key, m.content,
|
||
m.source_type, m.source_session_id, m.source_message_id, m.source_message_seq,
|
||
m.source_channel_name, m.source_chat_id, m.created_at, m.updated_at
|
||
FROM memories_fts f
|
||
JOIN memories m ON m.rowid = f.rowid
|
||
WHERE memories_fts MATCH ?1
|
||
AND m.scope_kind = ?2
|
||
AND m.scope_key = ?3
|
||
AND m.namespace = ?4
|
||
ORDER BY bm25(memories_fts), m.updated_at DESC
|
||
LIMIT ?5
|
||
",
|
||
)?;
|
||
let rows = stmt.query_map(
|
||
params![query, scope_kind, scope_key, namespace, limit],
|
||
map_memory_record,
|
||
)?;
|
||
for row in rows {
|
||
memories.push(row?);
|
||
}
|
||
} else {
|
||
let mut stmt = conn.prepare(
|
||
"
|
||
SELECT m.id, m.scope_kind, m.scope_key, m.namespace, m.memory_key, m.content,
|
||
m.source_type, m.source_session_id, m.source_message_id, m.source_message_seq,
|
||
m.source_channel_name, m.source_chat_id, m.created_at, m.updated_at
|
||
FROM memories_fts f
|
||
JOIN memories m ON m.rowid = f.rowid
|
||
WHERE memories_fts MATCH ?1
|
||
AND m.scope_kind = ?2
|
||
AND m.scope_key = ?3
|
||
ORDER BY bm25(memories_fts), m.updated_at DESC
|
||
LIMIT ?4
|
||
",
|
||
)?;
|
||
let rows = stmt.query_map(
|
||
params![query, scope_kind, scope_key, limit],
|
||
map_memory_record,
|
||
)?;
|
||
for row in rows {
|
||
memories.push(row?);
|
||
}
|
||
}
|
||
|
||
Ok(memories)
|
||
}
|
||
|
||
pub fn search_memories_any(
|
||
&self,
|
||
scope_kind: &str,
|
||
scope_key: &str,
|
||
queries: &[String],
|
||
namespace: Option<&str>,
|
||
limit: usize,
|
||
) -> Result<Vec<MemoryRecord>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let limit = limit.max(1) as i64;
|
||
let query = quote_fts_or_query(queries);
|
||
if query.is_empty() {
|
||
return Ok(Vec::new());
|
||
}
|
||
|
||
let mut memories = Vec::new();
|
||
|
||
if let Some(namespace) = namespace {
|
||
let mut stmt = conn.prepare(
|
||
"
|
||
SELECT m.id, m.scope_kind, m.scope_key, m.namespace, m.memory_key, m.content,
|
||
m.source_type, m.source_session_id, m.source_message_id, m.source_message_seq,
|
||
m.source_channel_name, m.source_chat_id, m.created_at, m.updated_at
|
||
FROM memories_fts f
|
||
JOIN memories m ON m.rowid = f.rowid
|
||
WHERE memories_fts MATCH ?1
|
||
AND m.scope_kind = ?2
|
||
AND m.scope_key = ?3
|
||
AND m.namespace = ?4
|
||
ORDER BY bm25(memories_fts), m.updated_at DESC
|
||
LIMIT ?5
|
||
",
|
||
)?;
|
||
let rows = stmt.query_map(
|
||
params![query, scope_kind, scope_key, namespace, limit],
|
||
map_memory_record,
|
||
)?;
|
||
for row in rows {
|
||
memories.push(row?);
|
||
}
|
||
} else {
|
||
let mut stmt = conn.prepare(
|
||
"
|
||
SELECT m.id, m.scope_kind, m.scope_key, m.namespace, m.memory_key, m.content,
|
||
m.source_type, m.source_session_id, m.source_message_id, m.source_message_seq,
|
||
m.source_channel_name, m.source_chat_id, m.created_at, m.updated_at
|
||
FROM memories_fts f
|
||
JOIN memories m ON m.rowid = f.rowid
|
||
WHERE memories_fts MATCH ?1
|
||
AND m.scope_kind = ?2
|
||
AND m.scope_key = ?3
|
||
ORDER BY bm25(memories_fts), m.updated_at DESC
|
||
LIMIT ?4
|
||
",
|
||
)?;
|
||
let rows = stmt.query_map(
|
||
params![query, scope_kind, scope_key, limit],
|
||
map_memory_record,
|
||
)?;
|
||
for row in rows {
|
||
memories.push(row?);
|
||
}
|
||
}
|
||
|
||
Ok(memories)
|
||
}
|
||
|
||
pub fn load_messages(&self, session_id: &str) -> Result<Vec<ChatMessage>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
load_messages_after(&conn, session_id, 0)
|
||
}
|
||
|
||
pub fn load_messages_for_topic(
|
||
&self,
|
||
topic_id: &str,
|
||
session_id: Option<&str>,
|
||
) -> Result<Vec<ChatMessage>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
|
||
if let Some(sid) = session_id {
|
||
let mut stmt = conn.prepare(
|
||
"
|
||
SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms
|
||
FROM messages
|
||
WHERE topic_id = ?1 AND session_id = ?2
|
||
ORDER BY seq ASC
|
||
",
|
||
)?;
|
||
let rows = stmt.query_map(params![topic_id, sid], map_chat_message_row)?;
|
||
let mut messages = Vec::new();
|
||
for row in rows {
|
||
messages.push(row?);
|
||
}
|
||
Ok(messages)
|
||
} else {
|
||
let mut stmt = conn.prepare(
|
||
"
|
||
SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms
|
||
FROM messages
|
||
WHERE topic_id = ?1
|
||
ORDER BY seq ASC
|
||
",
|
||
)?;
|
||
let rows = stmt.query_map(params![topic_id], map_chat_message_row)?;
|
||
let mut messages = Vec::new();
|
||
for row in rows {
|
||
messages.push(row?);
|
||
}
|
||
Ok(messages)
|
||
}
|
||
}
|
||
|
||
/// 获取指定话题的消息数量（动态计算，确保准确）
|
||
pub fn get_topic_message_count(&self, topic_id: &str) -> Result<usize, StorageError> {
|
||
self.load_messages_for_topic(topic_id, None).map(|msgs| msgs.len())
|
||
}
|
||
|
||
pub fn load_all_messages(&self, session_id: &str) -> Result<Vec<ChatMessage>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
load_messages_after(&conn, session_id, 0)
|
||
}
|
||
|
||
pub fn count_active_user_messages(&self, session_id: &str) -> Result<i64, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
conn.query_row(
|
||
"
|
||
SELECT COUNT(*)
|
||
FROM messages
|
||
WHERE session_id = ?1 AND role = 'user'
|
||
",
|
||
params![session_id],
|
||
|row| row.get(0),
|
||
)
|
||
.map_err(StorageError::from)
|
||
}
|
||
|
||
pub fn replace_todos(
|
||
&self,
|
||
scope_key: &str,
|
||
items: &[TodoRecord],
|
||
) -> Result<Vec<TodoRecord>, StorageError> {
|
||
let mut conn = self.pool.get()?;
|
||
// 用 BEGIN IMMEDIATE 事务保证严格语义：写锁在事务开始时获取，
|
||
// 避免并发写事务在提交时死锁导致 "database is locked"。
|
||
// 用户数据替换需保证原子性——中途失败必须回滚，避免 DELETE 后 INSERT
|
||
// 异常导致 todos 列表丢失且无法恢复。
|
||
let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
|
||
let now = current_timestamp();
|
||
|
||
// Delete existing todos for this scope_key
|
||
tx.execute(
|
||
"DELETE FROM todos WHERE scope_key = ?1",
|
||
params![scope_key],
|
||
)?;
|
||
|
||
// Insert new todos
|
||
for item in items {
|
||
tx.execute(
|
||
"INSERT OR REPLACE INTO todos (id, scope_key, session_id, topic_id, content, status, priority, created_at, updated_at, created_by_message_id)
|
||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
|
||
params![
|
||
item.id,
|
||
scope_key,
|
||
item.session_id,
|
||
item.topic_id,
|
||
item.content,
|
||
item.status,
|
||
item.priority,
|
||
item.created_at,
|
||
now,
|
||
item.created_by_message_id,
|
||
],
|
||
)?;
|
||
}
|
||
|
||
// 事务内复用同一连接查询返回值，避免 drop(conn) 后重新 pool.get()。
|
||
let mut stmt = tx.prepare(
|
||
"SELECT id, scope_key, session_id, topic_id, content, status, priority, created_at, updated_at, created_by_message_id
|
||
FROM todos
|
||
WHERE scope_key = ?1
|
||
ORDER BY created_at ASC",
|
||
)?;
|
||
let rows = stmt.query_map(params![scope_key], |row| {
|
||
Ok(TodoRecord {
|
||
id: row.get(0)?,
|
||
scope_key: row.get(1)?,
|
||
session_id: row.get(2)?,
|
||
topic_id: row.get(3)?,
|
||
content: row.get(4)?,
|
||
status: row.get(5)?,
|
||
priority: row.get(6)?,
|
||
created_at: row.get(7)?,
|
||
updated_at: row.get(8)?,
|
||
created_by_message_id: row.get(9)?,
|
||
})
|
||
})?;
|
||
let mut result = Vec::new();
|
||
for row in rows {
|
||
result.push(row?);
|
||
}
|
||
drop(stmt); // 释放 stmt 借用，才能 commit
|
||
tx.commit()?;
|
||
Ok(result)
|
||
}
|
||
|
||
pub fn list_todos(&self, scope_key: &str) -> Result<Vec<TodoRecord>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let mut stmt = conn.prepare(
|
||
"SELECT id, scope_key, session_id, topic_id, content, status, priority, created_at, updated_at, created_by_message_id
|
||
FROM todos
|
||
WHERE scope_key = ?1
|
||
ORDER BY created_at ASC",
|
||
)?;
|
||
|
||
let rows = stmt.query_map(params![scope_key], |row| {
|
||
Ok(TodoRecord {
|
||
id: row.get(0)?,
|
||
scope_key: row.get(1)?,
|
||
session_id: row.get(2)?,
|
||
topic_id: row.get(3)?,
|
||
content: row.get(4)?,
|
||
status: row.get(5)?,
|
||
priority: row.get(6)?,
|
||
created_at: row.get(7)?,
|
||
updated_at: row.get(8)?,
|
||
created_by_message_id: row.get(9)?,
|
||
})
|
||
})?;
|
||
|
||
let mut todos = Vec::new();
|
||
for row in rows {
|
||
todos.push(row?);
|
||
}
|
||
Ok(todos)
|
||
}
|
||
}
|
||
|
||
pub fn persistent_session_id(channel_name: &str, chat_id: &str) -> String {
|
||
// 幂等：循环去除已存在的 "{channel_name}:" 前缀，防止前缀累积
|
||
let prefix = format!("{}:", channel_name);
|
||
let mut chat_id = chat_id;
|
||
while chat_id.starts_with(&prefix) {
|
||
chat_id = &chat_id[prefix.len()..];
|
||
}
|
||
|
||
if channel_name == "cli" || channel_name == "websocket" {
|
||
chat_id.to_string()
|
||
} else {
|
||
format!("{}:{}", channel_name, chat_id)
|
||
}
|
||
}
|
||
|
||
#[cfg(not(test))]
|
||
fn default_session_db_path() -> Result<PathBuf, std::io::Error> {
|
||
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
|
||
Ok(home.join(".picobot").join("storage").join("sessions.db"))
|
||
}
|
||
|
||
fn insert_message_with_seq(
|
||
conn: &rusqlite::Transaction<'_>,
|
||
session_id: &str,
|
||
seq: i64,
|
||
message: &ChatMessage,
|
||
) -> Result<(), StorageError> {
|
||
let media_refs_json = serde_json::to_string(&message.media_refs)?;
|
||
let tool_calls_json = message
|
||
.tool_calls
|
||
.as_ref()
|
||
.map(serde_json::to_string)
|
||
.transpose()?;
|
||
conn.execute(
|
||
"
|
||
INSERT INTO messages (
|
||
id, session_id, seq, role, content,
|
||
system_context, reasoning_content, media_refs_json, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, created_at
|
||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13)
|
||
",
|
||
params![
|
||
message.id,
|
||
session_id,
|
||
seq,
|
||
message.role,
|
||
message.content,
|
||
message.system_context,
|
||
message.reasoning_content,
|
||
media_refs_json,
|
||
message.tool_call_id,
|
||
message.tool_name,
|
||
tool_calls_json,
|
||
message.tool_duration_ms.map(|v| v as i64),
|
||
message.timestamp,
|
||
],
|
||
)?;
|
||
Ok(())
|
||
}
|
||
|
||
/// Insert a message with an explicit `topic_id` and `seq`.
|
||
///
|
||
/// Used by `replace_topic_history` to insert compressed messages while
|
||
/// preserving topic association (the plain `insert_message_with_seq` would
|
||
/// set topic_id to NULL).
|
||
fn insert_message_with_topic_seq(
|
||
conn: &rusqlite::Transaction<'_>,
|
||
session_id: &str,
|
||
topic_id: &str,
|
||
seq: i64,
|
||
message: &ChatMessage,
|
||
) -> Result<(), StorageError> {
|
||
let media_refs_json = serde_json::to_string(&message.media_refs)?;
|
||
let tool_calls_json = message
|
||
.tool_calls
|
||
.as_ref()
|
||
.map(serde_json::to_string)
|
||
.transpose()?;
|
||
conn.execute(
|
||
"
|
||
INSERT INTO messages (
|
||
id, session_id, topic_id, seq, role, content,
|
||
system_context, reasoning_content, media_refs_json, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, created_at
|
||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)
|
||
",
|
||
params![
|
||
message.id,
|
||
session_id,
|
||
topic_id,
|
||
seq,
|
||
message.role,
|
||
message.content,
|
||
message.system_context,
|
||
message.reasoning_content,
|
||
media_refs_json,
|
||
message.tool_call_id,
|
||
message.tool_name,
|
||
tool_calls_json,
|
||
message.tool_duration_ms.map(|v| v as i64),
|
||
message.timestamp,
|
||
],
|
||
)?;
|
||
Ok(())
|
||
}
|
||
|
||
fn clone_message_for_compaction(message: &ChatMessage, timestamp: i64) -> ChatMessage {
|
||
ChatMessage {
|
||
id: uuid::Uuid::new_v4().to_string(),
|
||
role: message.role.clone(),
|
||
content: message.content.clone(),
|
||
media_refs: message.media_refs.clone(),
|
||
timestamp,
|
||
system_context: message.system_context.clone(),
|
||
reasoning_content: message.reasoning_content.clone(),
|
||
tool_call_id: message.tool_call_id.clone(),
|
||
tool_name: message.tool_name.clone(),
|
||
tool_state: message.tool_state.clone(),
|
||
tool_duration_ms: message.tool_duration_ms,
|
||
tool_calls: message.tool_calls.clone(),
|
||
}
|
||
}
|
||
|
||
fn load_messages_between(
|
||
conn: &rusqlite::Transaction<'_>,
|
||
session_id: &str,
|
||
start_seq_exclusive: i64,
|
||
end_seq_inclusive: i64,
|
||
) -> Result<Vec<ChatMessage>, StorageError> {
|
||
let mut stmt = conn.prepare(
|
||
"
|
||
SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms
|
||
FROM messages
|
||
WHERE session_id = ?1 AND seq > ?2 AND seq <= ?3
|
||
ORDER BY seq ASC
|
||
",
|
||
)?;
|
||
|
||
let rows = stmt.query_map(
|
||
params![session_id, start_seq_exclusive, end_seq_inclusive],
|
||
|row| {
|
||
let media_refs_json: String = row.get(5)?;
|
||
let media_refs: Vec<String> =
|
||
serde_json::from_str(&media_refs_json).map_err(|err| {
|
||
rusqlite::Error::FromSqlConversionFailure(
|
||
media_refs_json.len(),
|
||
rusqlite::types::Type::Text,
|
||
Box::new(err),
|
||
)
|
||
})?;
|
||
|
||
let tool_calls_json: Option<String> = row.get(9)?;
|
||
let tool_calls = tool_calls_json
|
||
.as_deref()
|
||
.map(serde_json::from_str)
|
||
.transpose()
|
||
.map_err(|err| {
|
||
rusqlite::Error::FromSqlConversionFailure(
|
||
9,
|
||
rusqlite::types::Type::Text,
|
||
Box::new(err),
|
||
)
|
||
})?;
|
||
|
||
Ok(ChatMessage {
|
||
id: row.get(0)?,
|
||
role: row.get(1)?,
|
||
content: row.get(2)?,
|
||
system_context: row.get(3)?,
|
||
reasoning_content: row.get(4)?,
|
||
media_refs,
|
||
timestamp: row.get(6)?,
|
||
tool_call_id: row.get(7)?,
|
||
tool_name: row.get(8)?,
|
||
tool_state: None,
|
||
tool_duration_ms: row.get::<_, Option<i64>>(10)?.map(|v| v as u64),
|
||
tool_calls,
|
||
})
|
||
},
|
||
)?;
|
||
|
||
let mut messages = Vec::new();
|
||
for row in rows {
|
||
messages.push(row?);
|
||
}
|
||
Ok(messages)
|
||
}
|
||
|
||
fn load_messages_after(
|
||
conn: &Connection,
|
||
session_id: &str,
|
||
cutoff_seq: i64,
|
||
) -> Result<Vec<ChatMessage>, StorageError> {
|
||
let mut stmt = conn.prepare(
|
||
"
|
||
SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms
|
||
FROM messages
|
||
WHERE session_id = ?1 AND seq > ?2
|
||
ORDER BY seq ASC
|
||
",
|
||
)?;
|
||
|
||
let rows = stmt.query_map(params![session_id, cutoff_seq], |row| {
|
||
let media_refs_json: String = row.get(5)?;
|
||
let media_refs: Vec<String> = serde_json::from_str(&media_refs_json).map_err(|err| {
|
||
rusqlite::Error::FromSqlConversionFailure(
|
||
media_refs_json.len(),
|
||
rusqlite::types::Type::Text,
|
||
Box::new(err),
|
||
)
|
||
})?;
|
||
|
||
let tool_calls_json: Option<String> = row.get(9)?;
|
||
let tool_calls = tool_calls_json
|
||
.as_deref()
|
||
.map(serde_json::from_str)
|
||
.transpose()
|
||
.map_err(|err| {
|
||
rusqlite::Error::FromSqlConversionFailure(
|
||
9,
|
||
rusqlite::types::Type::Text,
|
||
Box::new(err),
|
||
)
|
||
})?;
|
||
|
||
Ok(ChatMessage {
|
||
id: row.get(0)?,
|
||
role: row.get(1)?,
|
||
content: row.get(2)?,
|
||
system_context: row.get(3)?,
|
||
reasoning_content: row.get(4)?,
|
||
media_refs,
|
||
timestamp: row.get(6)?,
|
||
tool_call_id: row.get(7)?,
|
||
tool_name: row.get(8)?,
|
||
tool_state: None,
|
||
tool_duration_ms: row.get::<_, Option<i64>>(10)?.map(|v| v as u64),
|
||
tool_calls,
|
||
})
|
||
})?;
|
||
|
||
let mut messages = Vec::new();
|
||
for row in rows {
|
||
messages.push(row?);
|
||
}
|
||
Ok(messages)
|
||
}
|
||
|
||
fn current_timestamp() -> i64 {
|
||
std::time::SystemTime::now()
|
||
.duration_since(std::time::UNIX_EPOCH)
|
||
.expect("system clock before unix epoch")
|
||
.as_millis() as i64
|
||
}
|
||
|
||
fn quote_fts_query(query: &str) -> String {
|
||
format!("\"{}\"", query.replace('"', "\"\""))
|
||
}
|
||
|
||
fn quote_fts_or_query(queries: &[String]) -> String {
|
||
queries
|
||
.iter()
|
||
.map(|query| query.trim())
|
||
.filter(|query| !query.is_empty())
|
||
.map(quote_fts_query)
|
||
.collect::<Vec<_>>()
|
||
.join(" OR ")
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests;
|