414105d419
- topics 表迁移新增 provider/model 列，记录话题级显式模型选择 - storage 层新增 update_topic_model / list_topic_model_selections - 新增 POST /api/topic/select-model 与 GET /api/topic/selected-model - 双写 key 取 topic 行自带 session_id（不信任请求体，防污染其他会话） - 清除时同步清 session 级选择，避免'重置为默认'被物化逻辑重新覆盖
2505 lines
92 KiB
Rust
2505 lines
92 KiB
Rust
#[cfg(not(test))]
|
||
use std::path::{Path, PathBuf};
|
||
|
||
use std::collections::HashMap;
|
||
|
||
use crate::utils::current_timestamp;
|
||
|
||
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_MEMORY_NAMESPACES, GLOBAL_SCOPE_KEY, MemoryRecord, MemoryUpsert, PendingSubagentRecord,
|
||
SchedulerJobRecord, SchedulerJobState, SchedulerJobStatus, SchedulerJobUpsert, SessionRecord,
|
||
SessionTokenStats, SkillEventRecord, TodoRecord, TopicRecord, allowed_namespace_names,
|
||
get_namespace_description, is_valid_namespace,
|
||
};
|
||
|
||
#[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,
|
||
tool_duration_ms INTEGER,
|
||
prompt_tokens INTEGER,
|
||
completion_tokens INTEGER,
|
||
total_tokens INTEGER,
|
||
context_window_tokens INTEGER,
|
||
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_topics_schema(&conn)?;
|
||
ensure_scheduler_schema(&conn)?;
|
||
ensure_memory_scope_key_migration(&conn)?;
|
||
ensure_todos_schema(&conn)?;
|
||
ensure_pending_subagents_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, provider, model 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)?,
|
||
provider: row.get(8)?,
|
||
model: row.get(9)?,
|
||
})
|
||
})
|
||
.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, provider, model 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)?,
|
||
provider: row.get(8)?,
|
||
model: row.get(9)?,
|
||
})
|
||
})?;
|
||
|
||
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(())
|
||
}
|
||
|
||
/// 设置/清除话题级模型选择。provider 与 model 均为 None 时清除（恢复继承）。
|
||
pub fn update_topic_model(
|
||
&self,
|
||
topic_id: &str,
|
||
provider: Option<&str>,
|
||
model: Option<&str>,
|
||
) -> Result<(), StorageError> {
|
||
let now = current_timestamp();
|
||
let conn = self.pool.get()?;
|
||
conn.execute(
|
||
"UPDATE topics SET provider = ?2, model = ?3, updated_at = ?4 WHERE id = ?1",
|
||
params![topic_id, provider, model, now],
|
||
)?;
|
||
Ok(())
|
||
}
|
||
|
||
/// 全量读出话题级模型选择，供启动时预热内存缓存。
|
||
pub fn list_topic_model_selections(
|
||
&self,
|
||
) -> Result<Vec<(String, Option<String>, Option<String>)>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let mut stmt = conn.prepare(
|
||
"SELECT id, provider, model FROM topics WHERE provider IS NOT NULL OR model IS NOT NULL",
|
||
)?;
|
||
let rows = stmt.query_map([], |row| {
|
||
Ok((row.get(0)?, row.get(1)?, row.get(2)?))
|
||
})?;
|
||
let mut result = Vec::new();
|
||
for row in rows {
|
||
result.push(row?);
|
||
}
|
||
Ok(result)
|
||
}
|
||
|
||
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, prompt_tokens, completion_tokens, total_tokens, context_window_tokens, created_at
|
||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)
|
||
",
|
||
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.usage.as_ref().map(|u| u.prompt_tokens as i64),
|
||
message.usage.as_ref().map(|u| u.completion_tokens as i64),
|
||
message.usage.as_ref().map(|u| u.total_tokens as i64),
|
||
message.usage.as_ref().and_then(|u| u.context_window_tokens.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, prompt_tokens, completion_tokens, total_tokens, context_window_tokens, created_at
|
||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)
|
||
",
|
||
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.usage.as_ref().map(|u| u.prompt_tokens as i64),
|
||
message.usage.as_ref().map(|u| u.completion_tokens as i64),
|
||
message.usage.as_ref().map(|u| u.total_tokens as i64),
|
||
message.usage.as_ref().and_then(|u| u.context_window_tokens.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(())
|
||
}
|
||
|
||
/// 压缩该 topic 的历史：保留原始消息（标记 is_compacted=1，前端可见、LLM 不可见），
|
||
/// 并插入压缩摘要消息（is_compacted=0，LLM 可见，前端通过 system_context 过滤排除）。
|
||
///
|
||
/// 与 `replace_topic_history` 的区别：不删除原消息，仅打标记，从而让前端仍能展示
|
||
/// 完整原始对话，同时 LLM 只看到压缩后的精简历史。
|
||
///
|
||
/// `new_messages` 是 `compress_two_segment` 的输出，包含：
|
||
/// - 保留原样的消息（system_guards / 最新 user，保留原 ID）
|
||
/// - 压缩摘要消息（system_context = history_compaction_*，新 ID）
|
||
pub fn compact_topic_history(
|
||
&self,
|
||
session_id: &str,
|
||
topic_id: &str,
|
||
new_messages: &[ChatMessage],
|
||
) -> Result<(), StorageError> {
|
||
let mut conn = self.pool.get()?;
|
||
let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
|
||
let now = current_timestamp();
|
||
|
||
// 分离摘要消息与保留消息（保留消息携带原 ID，摘要消息是新构造的）
|
||
let (summaries, preserved): (Vec<&ChatMessage>, Vec<&ChatMessage>) =
|
||
new_messages.iter().partition(|m| {
|
||
m.system_context
|
||
.as_deref()
|
||
.map_or(false, |sc| sc.starts_with("history_compaction"))
|
||
});
|
||
|
||
// 先删除该 topic 下已有的旧压缩摘要（system_context LIKE 'history_compaction%'）。
|
||
// 旧摘要已被新摘要替代，保留它们只会累积垃圾行（前端和 LLM 都看不到，但占存储）。
|
||
tx.execute(
|
||
"DELETE FROM messages \
|
||
WHERE topic_id = ?1 AND session_id = ?2 \
|
||
AND system_context LIKE 'history_compaction%'",
|
||
params![topic_id, session_id],
|
||
)?;
|
||
|
||
// 将该 topic 中未被保留的原消息标记为 is_compacted=1（仅更新尚未标记的行，避免重复写）。
|
||
// 保留消息（system_guards / 最新 user）保持 is_compacted=0，不重复插入。
|
||
let preserved_ids: Vec<String> = preserved.iter().map(|m| m.id.clone()).collect();
|
||
if preserved_ids.is_empty() {
|
||
tx.execute(
|
||
"UPDATE messages SET is_compacted = 1 \
|
||
WHERE topic_id = ?1 AND session_id = ?2 AND is_compacted = 0",
|
||
params![topic_id, session_id],
|
||
)?;
|
||
} else {
|
||
let placeholders = (0..preserved_ids.len())
|
||
.map(|_| "?")
|
||
.collect::<Vec<_>>()
|
||
.join(",");
|
||
let sql = format!(
|
||
"UPDATE messages SET is_compacted = 1 \
|
||
WHERE topic_id = ? AND session_id = ? AND is_compacted = 0 \
|
||
AND id NOT IN ({})",
|
||
placeholders
|
||
);
|
||
let mut params_vec: Vec<String> = vec![topic_id.to_string(), session_id.to_string()];
|
||
params_vec.extend(preserved_ids.iter().cloned());
|
||
tx.execute(&sql, rusqlite::params_from_iter(params_vec))?;
|
||
}
|
||
|
||
// 插入压缩摘要消息（is_compacted=0，由列默认值保证）
|
||
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 summaries.iter().enumerate() {
|
||
let seq = start_seq + i as i64;
|
||
insert_message_with_topic_seq(&tx, session_id, topic_id, seq, message)?;
|
||
}
|
||
|
||
// 更新 topic / session 计数（基于该 topic 全部消息，含被压缩的原始消息）
|
||
let topic_count: i64 = tx.query_row(
|
||
"SELECT COUNT(*) FROM messages WHERE session_id = ?1 AND topic_id = ?2",
|
||
params![session_id, topic_id],
|
||
|row| row.get(0),
|
||
)?;
|
||
tx.execute(
|
||
"UPDATE topics SET message_count = ?2, last_active_at = ?3, updated_at = ?3 WHERE id = ?1",
|
||
params![topic_id, topic_count, now],
|
||
)?;
|
||
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)
|
||
}
|
||
|
||
/// LLM 视角：只返回 is_compacted = 0 的消息（压缩摘要 + 未被压缩的新消息）。
|
||
/// 被压缩消费掉的原始消息（is_compacted = 1）对 LLM 不可见，以节省 context。
|
||
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, prompt_tokens, completion_tokens, total_tokens, context_window_tokens
|
||
FROM messages
|
||
WHERE topic_id = ?1 AND session_id = ?2 AND is_compacted = 0
|
||
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, prompt_tokens, completion_tokens, total_tokens, context_window_tokens
|
||
FROM messages
|
||
WHERE topic_id = ?1 AND is_compacted = 0
|
||
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)
|
||
}
|
||
}
|
||
|
||
/// UI 视角：返回原始消息（含被压缩消费的 is_compacted=1 消息）+ 未压缩新消息，
|
||
/// 排除压缩摘要消息（system_context LIKE 'history_compaction%'）。
|
||
/// 用于前端历史展示、/current、/save-topic、topic 描述生成等场景。
|
||
pub fn load_messages_for_topic_full(
|
||
&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, prompt_tokens, completion_tokens, total_tokens, context_window_tokens
|
||
FROM messages
|
||
WHERE topic_id = ?1 AND session_id = ?2
|
||
AND (system_context IS NULL OR system_context NOT LIKE 'history_compaction%')
|
||
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, prompt_tokens, completion_tokens, total_tokens, context_window_tokens
|
||
FROM messages
|
||
WHERE topic_id = ?1
|
||
AND (system_context IS NULL OR system_context NOT LIKE 'history_compaction%')
|
||
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)
|
||
}
|
||
}
|
||
|
||
/// 获取指定话题的消息数量。
|
||
///
|
||
/// 使用 `SELECT COUNT(*)` 在数据库侧计数，避免将所有消息
|
||
/// （含 content、tool_calls_json 等大字段反序列化）加载到内存。
|
||
/// 查询命中 `idx_messages_topic_seq(topic_id, seq)` 索引。
|
||
pub fn get_topic_message_count(&self, topic_id: &str) -> Result<usize, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let count: i64 = conn.query_row(
|
||
"SELECT COUNT(*) FROM messages WHERE topic_id = ?1",
|
||
params![topic_id],
|
||
|row| row.get(0),
|
||
)?;
|
||
Ok(count as usize)
|
||
}
|
||
|
||
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)
|
||
}
|
||
|
||
/// 批量查询多个 topic 的 token 消耗统计（cost 累计 + context 瞬时）。
|
||
///
|
||
/// 按 `topic_id` 聚合而非 `session_id`：一个 session 可包含多个 topic，
|
||
/// 若按 session_id 聚合会导致同 session 下的所有 topic 显示相同的总和。
|
||
///
|
||
/// 子代理隔离：子代理消息持久化时 session_id='sub:...'，topic_id=父 topic_id
|
||
/// （见 task::runtime PersistingEmittedMessageHandler 构造），因此不能仅靠
|
||
/// topic_id 隔离。此处用 `session_id NOT LIKE 'sub:%'` 显式排除子代理消息，
|
||
/// 与项目约定一致（session 列表同样过滤 'sub:%'）。子代理 token 不计入父 topic，
|
||
/// 保持"子代理分别计算"语义。
|
||
pub fn batch_topic_token_stats(
|
||
&self,
|
||
topic_ids: &[&str],
|
||
) -> Result<HashMap<String, SessionTokenStats>, StorageError> {
|
||
if topic_ids.is_empty() {
|
||
return Ok(HashMap::new());
|
||
}
|
||
let conn = self.pool.get()?;
|
||
|
||
let placeholders = (0..topic_ids.len())
|
||
.map(|i| format!("?{}", i + 1))
|
||
.collect::<Vec<_>>()
|
||
.join(", ");
|
||
// topic_id IN (...) 自动排除 NULL topic_id 的旧消息；
|
||
// session_id NOT LIKE 'sub:%' 排除子代理消息（其 topic_id=父 topic_id）
|
||
let sum_sql = format!(
|
||
"SELECT topic_id, \
|
||
COALESCE(SUM(prompt_tokens), 0) AS sum_prompt, \
|
||
COALESCE(SUM(completion_tokens), 0) AS sum_completion, \
|
||
COALESCE(SUM(total_tokens), 0) AS sum_total \
|
||
FROM messages \
|
||
WHERE topic_id IN ({placeholders}) AND role = 'assistant' \
|
||
AND session_id NOT LIKE 'sub:%' \
|
||
GROUP BY topic_id"
|
||
);
|
||
|
||
let mut stmt = conn.prepare(&sum_sql)?;
|
||
let params: Vec<&dyn rusqlite::ToSql> = topic_ids
|
||
.iter()
|
||
.map(|s| s as &dyn rusqlite::ToSql)
|
||
.collect();
|
||
let sum_rows = stmt.query_map(params.as_slice(), |row| {
|
||
Ok((
|
||
row.get::<_, String>(0)?,
|
||
SessionTokenStats {
|
||
prompt_tokens: row.get::<_, i64>(1)? as u64,
|
||
completion_tokens: row.get::<_, i64>(2)? as u64,
|
||
total_tokens: row.get::<_, i64>(3)? as u64,
|
||
last_prompt_tokens: None,
|
||
context_window_tokens: None,
|
||
},
|
||
))
|
||
})?;
|
||
|
||
let mut stats: HashMap<String, SessionTokenStats> = HashMap::new();
|
||
for row in sum_rows {
|
||
let (tid, s) = row?;
|
||
stats.insert(tid, s);
|
||
}
|
||
|
||
// 查找每个 topic 中最新的**有 usage 数据的** assistant 消息，
|
||
// 读取其 prompt_tokens 和 context_window_tokens。
|
||
// 过滤 prompt_tokens IS NOT NULL 确保跳过 error/cancel 消息（usage 为 NULL）；
|
||
// session_id NOT LIKE 'sub:%' 排除子代理消息，避免取到子代理的 context_window。
|
||
//
|
||
// 注意：seq 是 session 级递增（见 append_message_with_topic），主 session 与
|
||
// 子代理 session 各自独立计数，可能存在相同 seq。外层 WHERE 必须再次过滤
|
||
// session_id NOT LIKE 'sub:%'，否则 JOIN 会同时匹配主消息和子代理消息，
|
||
// 导致重复行并使 stats.entry(tid) 被覆盖，结果不确定。
|
||
let last_sql = format!(
|
||
"SELECT m.topic_id, m.prompt_tokens, m.context_window_tokens \
|
||
FROM messages m \
|
||
INNER JOIN ( \
|
||
SELECT topic_id, MAX(seq) AS max_seq \
|
||
FROM messages \
|
||
WHERE topic_id IN ({placeholders}) AND role = 'assistant' \
|
||
AND prompt_tokens IS NOT NULL \
|
||
AND session_id NOT LIKE 'sub:%' \
|
||
GROUP BY topic_id \
|
||
) latest ON m.topic_id = latest.topic_id AND m.seq = latest.max_seq \
|
||
WHERE m.session_id NOT LIKE 'sub:%'"
|
||
);
|
||
let mut stmt2 = conn.prepare(&last_sql)?;
|
||
let params2: Vec<&dyn rusqlite::ToSql> = topic_ids
|
||
.iter()
|
||
.map(|s| s as &dyn rusqlite::ToSql)
|
||
.collect();
|
||
let last_rows = stmt2.query_map(params2.as_slice(), |row| {
|
||
Ok((
|
||
row.get::<_, String>(0)?,
|
||
row.get::<_, Option<i64>>(1)?,
|
||
row.get::<_, Option<i64>>(2)?,
|
||
))
|
||
})?;
|
||
for row in last_rows {
|
||
let (tid, last_prompt, last_ctx_window) = row?;
|
||
let entry = stats.entry(tid).or_insert(SessionTokenStats {
|
||
prompt_tokens: 0,
|
||
completion_tokens: 0,
|
||
total_tokens: 0,
|
||
last_prompt_tokens: None,
|
||
context_window_tokens: None,
|
||
});
|
||
entry.last_prompt_tokens = last_prompt.map(|v| v as u32);
|
||
entry.context_window_tokens = last_ctx_window.map(|v| v as u32);
|
||
}
|
||
|
||
Ok(stats)
|
||
}
|
||
|
||
/// 查询单个 session 的 token 消耗统计（cost 累计 + context 瞬时）。
|
||
///
|
||
/// 按 `session_id` 精确匹配查询，**不过滤** `sub:%`——专门用于子代理 session
|
||
/// （session_id = `sub:...`）的 token 统计。子代理没有 topic，一个 session 即
|
||
/// 一个完整执行单元，因此按 session 维度聚合而非 topic 维度。
|
||
///
|
||
/// 与 `batch_topic_token_stats` 共享同一套 SQL 模式（SUM + last），差异仅在
|
||
/// WHERE 条件：单 session 精确匹配，无 `NOT LIKE 'sub:%'` 过滤，无 topic 维度。
|
||
pub fn get_session_token_stats(
|
||
&self,
|
||
session_id: &str,
|
||
) -> Result<Option<SessionTokenStats>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
|
||
// 1. SUM 查询：累计 prompt/completion/total
|
||
let sum_sql = "SELECT \
|
||
COALESCE(SUM(prompt_tokens), 0), \
|
||
COALESCE(SUM(completion_tokens), 0), \
|
||
COALESCE(SUM(total_tokens), 0) \
|
||
FROM messages \
|
||
WHERE session_id = ?1 AND role = 'assistant'";
|
||
let mut stmt = conn.prepare(sum_sql)?;
|
||
let sum_row = stmt.query_row(params![session_id], |row| {
|
||
Ok((
|
||
row.get::<_, i64>(0)? as u64,
|
||
row.get::<_, i64>(1)? as u64,
|
||
row.get::<_, i64>(2)? as u64,
|
||
))
|
||
})?;
|
||
let (prompt_tokens, completion_tokens, total_tokens) = sum_row;
|
||
|
||
// 无 assistant 消息时直接返回 None
|
||
if total_tokens == 0 && prompt_tokens == 0 && completion_tokens == 0 {
|
||
// 需要二次确认是否真的没有 assistant 消息（usage 全 0 也可能是合法的）
|
||
let count_sql =
|
||
"SELECT COUNT(*) FROM messages WHERE session_id = ?1 AND role = 'assistant'";
|
||
let count: i64 = conn.query_row(count_sql, params![session_id], |row| row.get(0))?;
|
||
if count == 0 {
|
||
return Ok(None);
|
||
}
|
||
}
|
||
|
||
// 2. last 查询：最新有 usage 的 assistant 消息的 prompt_tokens + context_window_tokens
|
||
let last_sql = "SELECT prompt_tokens, context_window_tokens \
|
||
FROM messages \
|
||
WHERE session_id = ?1 AND role = 'assistant' AND prompt_tokens IS NOT NULL \
|
||
ORDER BY seq DESC LIMIT 1";
|
||
let mut stmt2 = conn.prepare(last_sql)?;
|
||
let last_row = stmt2
|
||
.query_row(params![session_id], |row| {
|
||
Ok((row.get::<_, Option<i64>>(0)?, row.get::<_, Option<i64>>(1)?))
|
||
})
|
||
.optional()?;
|
||
|
||
let (last_prompt_tokens, context_window_tokens) = match last_row {
|
||
Some((lp, lcw)) => (lp.map(|v| v as u32), lcw.map(|v| v as u32)),
|
||
None => (None, None),
|
||
};
|
||
|
||
Ok(Some(SessionTokenStats {
|
||
prompt_tokens,
|
||
completion_tokens,
|
||
total_tokens,
|
||
last_prompt_tokens,
|
||
context_window_tokens,
|
||
}))
|
||
}
|
||
|
||
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)
|
||
}
|
||
|
||
// ==================== pending_subagents ====================
|
||
|
||
/// 插入一条 pending_subagent 记录（task 工具 spawn 时调用）。
|
||
pub fn insert_pending_subagent(
|
||
&self,
|
||
record: &PendingSubagentRecord,
|
||
) -> Result<(), StorageError> {
|
||
let conn = self.pool.get()?;
|
||
conn.execute(
|
||
"INSERT OR REPLACE INTO pending_subagents
|
||
(task_id, parent_session_id, parent_topic_id, parent_chat_id, parent_channel, def_name, spawned_at, status)
|
||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
|
||
params![
|
||
record.task_id,
|
||
record.parent_session_id,
|
||
record.parent_topic_id,
|
||
record.parent_chat_id,
|
||
record.parent_channel,
|
||
record.def_name,
|
||
record.spawned_at,
|
||
record.status,
|
||
],
|
||
)?;
|
||
Ok(())
|
||
}
|
||
|
||
/// 查询指定 topic 下匹配状态的 pending_subagent 记录。
|
||
/// `status` 为 None 时查询所有状态。
|
||
pub fn list_pending_subagents(
|
||
&self,
|
||
topic_id: &str,
|
||
status: Option<&str>,
|
||
) -> Result<Vec<PendingSubagentRecord>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let sql = if status.is_some() {
|
||
"SELECT task_id, parent_session_id, parent_topic_id, parent_chat_id, parent_channel, def_name, spawned_at, status
|
||
FROM pending_subagents
|
||
WHERE parent_topic_id = ?1 AND status = ?2
|
||
ORDER BY spawned_at ASC"
|
||
} else {
|
||
"SELECT task_id, parent_session_id, parent_topic_id, parent_chat_id, parent_channel, def_name, spawned_at, status
|
||
FROM pending_subagents
|
||
WHERE parent_topic_id = ?1
|
||
ORDER BY spawned_at ASC"
|
||
};
|
||
let mut stmt = conn.prepare(sql)?;
|
||
let rows = if let Some(s) = status {
|
||
stmt.query_map(params![topic_id, s], map_pending_subagent_record)?
|
||
} else {
|
||
stmt.query_map(params![topic_id], map_pending_subagent_record)?
|
||
};
|
||
let mut result = Vec::new();
|
||
for row in rows {
|
||
result.push(row?);
|
||
}
|
||
Ok(result)
|
||
}
|
||
|
||
/// 获取指定 task_id 的 pending_subagent 记录。
|
||
pub fn get_pending_subagent(
|
||
&self,
|
||
task_id: &str,
|
||
) -> Result<Option<PendingSubagentRecord>, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let mut stmt = conn.prepare(
|
||
"SELECT task_id, parent_session_id, parent_topic_id, parent_chat_id, parent_channel, def_name, spawned_at, status
|
||
FROM pending_subagents
|
||
WHERE task_id = ?1",
|
||
)?;
|
||
let mut rows = stmt.query_map(params![task_id], map_pending_subagent_record)?;
|
||
match rows.next() {
|
||
Some(row) => Ok(Some(row?)),
|
||
None => Ok(None),
|
||
}
|
||
}
|
||
|
||
/// 更新指定 task_id 的状态（子代理完成或取消时调用）。
|
||
pub fn update_pending_subagent_status(
|
||
&self,
|
||
task_id: &str,
|
||
new_status: &str,
|
||
) -> Result<(), StorageError> {
|
||
let conn = self.pool.get()?;
|
||
conn.execute(
|
||
"UPDATE pending_subagents SET status = ?1 WHERE task_id = ?2",
|
||
params![new_status, task_id],
|
||
)?;
|
||
Ok(())
|
||
}
|
||
|
||
/// 条件更新状态：仅在当前状态为 `expected_current` 时才更新为 `new_status`。
|
||
///
|
||
/// 实现状态机不可逆性不变量：避免 cancel 路径覆盖 spawn 已写入的终态
|
||
/// （completed → cancelled 是非法转换）。
|
||
///
|
||
/// 返回是否实际更新（affected rows > 0）。false 表示状态已被其他路径更新，
|
||
/// 调用方应跳过后续基于该假设的操作。
|
||
pub fn try_update_pending_subagent_status(
|
||
&self,
|
||
task_id: &str,
|
||
expected_current: &str,
|
||
new_status: &str,
|
||
) -> Result<bool, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let affected = conn.execute(
|
||
"UPDATE pending_subagents SET status = ?1 WHERE task_id = ?2 AND status = ?3",
|
||
params![new_status, task_id, expected_current],
|
||
)?;
|
||
Ok(affected > 0)
|
||
}
|
||
|
||
/// 将所有 running 状态的 pending_subagent 标记为 interrupted（启动时崩溃恢复调用）。
|
||
pub fn mark_all_running_as_interrupted(&self) -> Result<usize, StorageError> {
|
||
let conn = self.pool.get()?;
|
||
let affected = conn.execute(
|
||
"UPDATE pending_subagents SET status = 'interrupted' WHERE status = 'running'",
|
||
[],
|
||
)?;
|
||
Ok(affected)
|
||
}
|
||
}
|
||
|
||
fn map_pending_subagent_record(row: &rusqlite::Row<'_>) -> rusqlite::Result<PendingSubagentRecord> {
|
||
Ok(PendingSubagentRecord {
|
||
task_id: row.get(0)?,
|
||
parent_session_id: row.get(1)?,
|
||
parent_topic_id: row.get(2)?,
|
||
parent_chat_id: row.get(3)?,
|
||
parent_channel: row.get(4)?,
|
||
def_name: row.get(5)?,
|
||
spawned_at: row.get(6)?,
|
||
status: row.get(7)?,
|
||
})
|
||
}
|
||
|
||
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 = crate::platform::picobot_home_dir();
|
||
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, prompt_tokens, completion_tokens, total_tokens, context_window_tokens, created_at
|
||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17)
|
||
",
|
||
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.usage.as_ref().map(|u| u.prompt_tokens as i64),
|
||
message.usage.as_ref().map(|u| u.completion_tokens as i64),
|
||
message.usage.as_ref().map(|u| u.total_tokens as i64),
|
||
message.usage.as_ref().and_then(|u| u.context_window_tokens.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, prompt_tokens, completion_tokens, total_tokens, context_window_tokens, created_at
|
||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)
|
||
",
|
||
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.usage.as_ref().map(|u| u.prompt_tokens as i64),
|
||
message.usage.as_ref().map(|u| u.completion_tokens as i64),
|
||
message.usage.as_ref().map(|u| u.total_tokens as i64),
|
||
message.usage.as_ref().and_then(|u| u.context_window_tokens.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(),
|
||
// 压缩克隆不保留 usage：压缩产生的是合成消息，不代表真实 LLM 调用
|
||
usage: None,
|
||
}
|
||
}
|
||
|
||
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, prompt_tokens, completion_tokens, total_tokens, context_window_tokens
|
||
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,
|
||
usage: map_usage_row(row, 11, 12, 13, 14)?,
|
||
})
|
||
},
|
||
)?;
|
||
|
||
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, prompt_tokens, completion_tokens, total_tokens, context_window_tokens
|
||
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,
|
||
usage: map_usage_row(row, 11, 12, 13, 14)?,
|
||
})
|
||
})?;
|
||
|
||
let mut messages = Vec::new();
|
||
for row in rows {
|
||
messages.push(row?);
|
||
}
|
||
Ok(messages)
|
||
}
|
||
|
||
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;
|