use sqlx::SqlitePool;
use std::collections::hash_map::Entry;
use std::collections::{HashMap, HashSet, VecDeque};
use uuid::Uuid;
pub struct SyncMessageInput {
pub id: Option<String>,
pub role: String,
pub text: String,
pub blocks: Option<String>,
}
pub async fn insert_message(
db: &SqlitePool,
session_id: &str,
role: &str,
text: &str,
blocks: Option<&str>,
) -> Result<(), sqlx::Error> {
let id = Uuid::new_v4().to_string();
let now = chrono::Utc::now().to_rfc3339();
sqlx::query("INSERT INTO chat_messages (id, session_id, role, text, created_at, blocks) VALUES (?, ?, ?, ?, ?, ?)")
.bind(&id)
.bind(session_id)
.bind(role)
.bind(text)
.bind(&now)
.bind(blocks)
.execute(db)
.await?;
Ok(())
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MessageCursor {
pub created_at: String,
pub id: String,
}
impl MessageCursor {
pub fn encode(&self) -> String {
format!("{}|{}", self.created_at, self.id)
}
pub fn parse(raw: &str) -> Option<Self> {
let (created_at, id) = raw.split_once('|')?;
if created_at.is_empty() || id.is_empty() {
return None;
}
Some(Self { created_at: created_at.to_string(), id: id.to_string() })
}
}
pub struct MessagePage {
pub rows: Vec<ChatMessageRow>,
pub next_cursor: Option<MessageCursor>,
}
pub async fn list_messages_page(
db: &SqlitePool,
session_id: &str,
before: Option<&MessageCursor>,
limit: usize,
max_bytes: usize,
) -> Result<MessagePage, sqlx::Error> {
let limit = limit.max(1);
let (cursor_at, cursor_id) = match before {
Some(c) => (Some(c.created_at.as_str()), Some(c.id.as_str())),
None => (None, None),
};
let mut rows: Vec<ChatMessageRow> = sqlx::query_as(
"SELECT role, text, created_at, id, blocks, status, last_seq FROM chat_messages \
WHERE session_id = ? \
AND (? IS NULL OR created_at < ? OR (created_at = ? AND id < ?)) \
ORDER BY created_at DESC, id DESC LIMIT ?",
)
.bind(session_id)
.bind(cursor_at)
.bind(cursor_at)
.bind(cursor_at)
.bind(cursor_id)
.bind((limit + 1) as i64)
.fetch_all(db)
.await?;
let more_by_count = rows.len() > limit;
rows.truncate(limit);
let mut bytes = 0usize;
let mut kept = 0usize;
for row in &rows {
let (_, text, _, _, blocks, _, _) = row;
let row_bytes = text.len() + blocks.as_deref().map_or(0, str::len);
if kept > 0 && bytes + row_bytes > max_bytes {
break;
}
bytes += row_bytes;
kept += 1;
}
let more_by_bytes = kept < rows.len();
rows.truncate(kept);
rows.reverse();
let next_cursor = if more_by_count || more_by_bytes {
rows.first().map(|(_, _, created_at, id, _, _, _)| MessageCursor {
created_at: created_at.clone(),
id: id.clone(),
})
} else {
None
};
Ok(MessagePage { rows, next_cursor })
}
pub type ChatMessageRow = (String, String, String, String, Option<String>, String, Option<i64>);
pub async fn upsert_streaming_message(
db: &SqlitePool,
row_id: &str,
session_id: &str,
text: &str,
blocks: Option<&str>,
last_seq: i64,
) -> Result<(), sqlx::Error> {
let now = chrono::Utc::now().to_rfc3339();
sqlx::query(
"INSERT INTO chat_messages (id, session_id, role, text, created_at, blocks, status, last_seq) \
VALUES (?, ?, 'assistant', ?, ?, ?, 'streaming', ?) \
ON CONFLICT(id) DO UPDATE SET text = excluded.text, blocks = excluded.blocks, last_seq = excluded.last_seq",
)
.bind(row_id)
.bind(session_id)
.bind(text)
.bind(&now)
.bind(blocks)
.bind(last_seq)
.execute(db)
.await?;
Ok(())
}
pub async fn finalize_message(db: &SqlitePool, row_id: &str) -> Result<(), sqlx::Error> {
sqlx::query("UPDATE chat_messages SET status = 'complete' WHERE id = ?")
.bind(row_id)
.execute(db)
.await?;
Ok(())
}
pub async fn sync_messages(
db: &SqlitePool,
session_id: &str,
messages: &[SyncMessageInput],
) -> Result<(), sqlx::Error> {
let mut consumed: HashSet<String> = HashSet::new();
let mut candidates: HashMap<(&str, &str), VecDeque<String>> = HashMap::new();
for m in messages {
if let Some(row_id) = &m.id {
let Some(blocks) = &m.blocks else { continue };
let affected =
sqlx::query("UPDATE chat_messages SET blocks = ? WHERE id = ? AND session_id = ?")
.bind(blocks)
.bind(row_id)
.bind(session_id)
.execute(db)
.await?
.rows_affected();
if affected == 0 {
tracing::debug!(row_id = %row_id, "sync_messages: no such row, skipped");
} else {
consumed.insert(row_id.clone());
}
continue;
}
if m.text.is_empty() {
continue;
}
let key = (m.role.as_str(), m.text.as_str());
let queue = match candidates.entry(key) {
Entry::Occupied(e) => e.into_mut(),
Entry::Vacant(e) => {
let ids: Vec<String> = sqlx::query_scalar(
"SELECT id FROM chat_messages WHERE session_id = ? AND role = ? AND text = ? \
ORDER BY created_at, id",
)
.bind(session_id)
.bind(&m.role)
.bind(&m.text)
.fetch_all(db)
.await?;
e.insert(ids.into())
}
};
let matched = loop {
match queue.pop_front() {
Some(id) if consumed.contains(&id) => continue,
other => break other,
}
};
if let Some(id) = matched {
consumed.insert(id.clone());
if let Some(blocks) = &m.blocks {
sqlx::query("UPDATE chat_messages SET blocks = ? WHERE id = ?")
.bind(blocks)
.bind(&id)
.execute(db)
.await?;
}
continue;
}
let id = Uuid::new_v4().to_string();
let now = chrono::Utc::now().to_rfc3339();
sqlx::query("INSERT INTO chat_messages (id, session_id, role, text, created_at, blocks) VALUES (?, ?, ?, ?, ?, ?)")
.bind(&id)
.bind(session_id)
.bind(&m.role)
.bind(&m.text)
.bind(&now)
.bind(&m.blocks)
.execute(db)
.await?;
consumed.insert(id);
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use sqlx::sqlite::SqlitePoolOptions;
const SESSION: &str = "s1";
async fn fresh_db() -> SqlitePool {
let db = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("in-memory sqlite pool");
sqlx::migrate!("./migrations").run(&db).await.expect("run migrations");
sqlx::query(
"INSERT INTO projects (id, name, path, created_at) VALUES ('p1', 'p', '/tmp', '2026-08-10')",
)
.execute(&db)
.await
.expect("seed project");
sqlx::query(
"INSERT INTO sessions (id, project_id, name, created_at) VALUES (?, 'p1', 's', '2026-08-10')",
)
.bind(SESSION)
.execute(&db)
.await
.expect("seed session");
db
}
async fn seed(db: &SqlitePool, n: usize, blocks_bytes: usize) {
let blocks = "b".repeat(blocks_bytes);
for i in 0..n {
sqlx::query(
"INSERT INTO chat_messages (id, session_id, role, text, created_at, blocks) \
VALUES (?, ?, 'assistant', ?, ?, ?)",
)
.bind(format!("id-{i:04}"))
.bind(SESSION)
.bind(format!("msg-{i}"))
.bind(format!("2026-08-10T00:00:{:02}Z", i / 10))
.bind(&blocks)
.execute(db)
.await
.expect("seed message");
}
}
fn texts(page: &MessagePage) -> Vec<&str> {
page.rows.iter().map(|(_, text, _, _, _, _, _)| text.as_str()).collect()
}
#[tokio::test]
async fn newest_page_is_oldest_first_and_reports_more() {
let db = fresh_db().await;
seed(&db, 25, 0).await;
let page = list_messages_page(&db, SESSION, None, 10, usize::MAX).await.expect("page");
assert_eq!(
texts(&page),
[
"msg-15", "msg-16", "msg-17", "msg-18", "msg-19", "msg-20", "msg-21", "msg-22",
"msg-23", "msg-24"
],
"取最新 10 条且按旧→新排序"
);
let cursor = page.next_cursor.expect("还有更早的消息");
assert_eq!(cursor.id, "id-0015", "游标指向本页最旧一条");
}
#[tokio::test]
async fn cursor_walks_backwards_without_gap_or_overlap_across_same_second() {
let db = fresh_db().await;
seed(&db, 25, 0).await;
let mut seen: Vec<String> = Vec::new();
let mut cursor = None;
loop {
let page = list_messages_page(&db, SESSION, cursor.as_ref(), 7, usize::MAX)
.await
.expect("page");
let mut batch: Vec<String> = texts(&page).iter().map(|s| s.to_string()).collect();
batch.extend(seen);
seen = batch;
match page.next_cursor {
Some(c) => cursor = Some(c),
None => break,
}
}
let expected: Vec<String> = (0..25).map(|i| format!("msg-{i}")).collect();
assert_eq!(seen, expected, "逐页回溯应拼回完整历史(无重复、无缺失)");
}
#[tokio::test]
async fn byte_budget_cuts_the_page_before_the_row_limit() {
let db = fresh_db().await;
seed(&db, 20, 1000).await;
let page = list_messages_page(&db, SESSION, None, 20, 3200).await.expect("page");
assert_eq!(page.rows.len(), 3, "字节预算应比条数上限先生效");
assert_eq!(texts(&page), ["msg-17", "msg-18", "msg-19"], "保留的应是最新的那几条");
assert!(page.next_cursor.is_some(), "被字节预算截断必须报告还有更早");
}
#[tokio::test]
async fn oversized_single_row_still_makes_progress() {
let db = fresh_db().await;
seed(&db, 3, 5000).await;
let page = list_messages_page(&db, SESSION, None, 10, 100).await.expect("page");
assert_eq!(page.rows.len(), 1, "预算再小也得返回一条");
assert_eq!(texts(&page), ["msg-2"]);
}
#[tokio::test]
async fn short_history_returns_everything_with_no_cursor() {
let db = fresh_db().await;
seed(&db, 4, 10).await;
let page = list_messages_page(&db, SESSION, None, 100, usize::MAX).await.expect("page");
assert_eq!(page.rows.len(), 4);
assert!(page.next_cursor.is_none(), "历史已到头时不得给游标(否则前端会无限上拉)");
}
async fn seed_row(db: &SqlitePool, id: &str, role: &str, text: &str, created_at: &str) {
sqlx::query(
"INSERT INTO chat_messages (id, session_id, role, text, created_at, blocks) \
VALUES (?, ?, ?, ?, ?, NULL)",
)
.bind(id)
.bind(SESSION)
.bind(role)
.bind(text)
.bind(created_at)
.execute(db)
.await
.expect("seed row");
}
async fn blocks_of(db: &SqlitePool, id: &str) -> Option<String> {
sqlx::query_scalar("SELECT blocks FROM chat_messages WHERE id = ?")
.bind(id)
.fetch_one(db)
.await
.expect("fetch blocks")
}
async fn row_count(db: &SqlitePool) -> i64 {
sqlx::query_scalar("SELECT COUNT(*) FROM chat_messages WHERE session_id = ?")
.bind(SESSION)
.fetch_one(db)
.await
.expect("count rows")
}
fn text_entry(text: &str, blocks: &str) -> SyncMessageInput {
SyncMessageInput {
id: None,
role: "assistant".into(),
text: text.into(),
blocks: Some(blocks.into()),
}
}
fn id_entry(id: &str, text: &str, blocks: &str) -> SyncMessageInput {
SyncMessageInput {
id: Some(id.into()),
role: "assistant".into(),
text: text.into(),
blocks: Some(blocks.into()),
}
}
#[tokio::test]
async fn same_text_rows_keep_independent_blocks() {
let db = fresh_db().await;
seed_row(&db, "row-a", "assistant", "OK", "2026-08-10T00:00:01Z").await;
seed_row(&db, "row-b", "assistant", "OK", "2026-08-10T00:00:02Z").await;
let msgs = vec![text_entry("OK", "[\"first\"]"), text_entry("OK", "[\"second\"]")];
sync_messages(&db, SESSION, &msgs).await.expect("sync");
assert_eq!(blocks_of(&db, "row-a").await.as_deref(), Some("[\"first\"]"));
assert_eq!(blocks_of(&db, "row-b").await.as_deref(), Some("[\"second\"]"));
assert_eq!(row_count(&db).await, 2, "两条 payload 对应两条既有行,不得新增");
}
#[tokio::test]
async fn surplus_same_text_entry_inserts_instead_of_overwriting() {
let db = fresh_db().await;
seed_row(&db, "row-a", "assistant", "OK", "2026-08-10T00:00:01Z").await;
let msgs = vec![text_entry("OK", "[\"first\"]"), text_entry("OK", "[\"second\"]")];
sync_messages(&db, SESSION, &msgs).await.expect("sync");
assert_eq!(blocks_of(&db, "row-a").await.as_deref(), Some("[\"first\"]"));
assert_eq!(row_count(&db).await, 2, "多出的一条新增行,不覆写已消费的行");
}
#[tokio::test]
async fn row_id_path_updates_only_that_rows_blocks() {
let db = fresh_db().await;
seed_row(&db, "row-a", "assistant", "OK", "2026-08-10T00:00:01Z").await;
seed_row(&db, "row-b", "assistant", "OK", "2026-08-10T00:00:02Z").await;
sync_messages(&db, SESSION, &[id_entry("row-b", "stale frontend text", "[\"cooked\"]")])
.await
.expect("sync");
assert_eq!(blocks_of(&db, "row-b").await.as_deref(), Some("[\"cooked\"]"));
assert_eq!(blocks_of(&db, "row-a").await, None, "同 text 的另一行不得被注入");
let text: String = sqlx::query_scalar("SELECT text FROM chat_messages WHERE id = 'row-b'")
.fetch_one(&db)
.await
.expect("fetch text");
assert_eq!(text, "OK", "text 的权威在后端,不得被前端 payload 覆盖");
assert_eq!(row_count(&db).await, 2);
}
#[tokio::test]
async fn unknown_row_id_is_skipped_not_inserted() {
let db = fresh_db().await;
seed_row(&db, "row-a", "assistant", "OK", "2026-08-10T00:00:01Z").await;
sync_messages(&db, SESSION, &[id_entry("row-missing", "OK", "[\"cooked\"]")])
.await
.expect("sync");
assert_eq!(blocks_of(&db, "row-a").await, None, "不得回退到文本匹配");
assert_eq!(row_count(&db).await, 1, "不得 INSERT 重复行");
}
#[tokio::test]
async fn row_id_path_consumes_the_row_for_later_text_matches() {
let db = fresh_db().await;
seed_row(&db, "row-a", "assistant", "OK", "2026-08-10T00:00:01Z").await;
seed_row(&db, "row-b", "assistant", "OK", "2026-08-10T00:00:02Z").await;
let msgs = vec![id_entry("row-a", "OK", "[\"cooked\"]"), text_entry("OK", "[\"replay\"]")];
sync_messages(&db, SESSION, &msgs).await.expect("sync");
assert_eq!(blocks_of(&db, "row-a").await.as_deref(), Some("[\"cooked\"]"));
assert_eq!(blocks_of(&db, "row-b").await.as_deref(), Some("[\"replay\"]"));
assert_eq!(row_count(&db).await, 2);
}
#[tokio::test]
async fn text_path_inserts_messages_absent_from_db() {
let db = fresh_db().await;
let msgs = vec![text_entry("hello", "[\"a\"]"), text_entry("world", "[\"b\"]")];
sync_messages(&db, SESSION, &msgs).await.expect("sync");
assert_eq!(row_count(&db).await, 2);
let blocks: Vec<String> =
sqlx::query_scalar("SELECT blocks FROM chat_messages ORDER BY text")
.fetch_all(&db)
.await
.expect("fetch");
assert_eq!(blocks, ["[\"a\"]", "[\"b\"]"]);
}
#[tokio::test]
async fn repeated_sync_does_not_duplicate_rows() {
let db = fresh_db().await;
let msgs = vec![text_entry("hello", "[\"a\"]"), text_entry("hello", "[\"b\"]")];
sync_messages(&db, SESSION, &msgs).await.expect("first sync");
sync_messages(&db, SESSION, &msgs).await.expect("second sync");
assert_eq!(row_count(&db).await, 2, "重复 sync 不得累积行");
let blocks: Vec<String> =
sqlx::query_scalar("SELECT blocks FROM chat_messages ORDER BY created_at, id")
.fetch_all(&db)
.await
.expect("fetch");
assert_eq!(blocks, ["[\"a\"]", "[\"b\"]"], "第二轮按同顺序落到同两行");
}
#[tokio::test]
async fn row_id_path_accepts_empty_text_rows() {
let db = fresh_db().await;
seed_row(&db, "row-a", "assistant", "", "2026-08-10T00:00:01Z").await;
sync_messages(&db, SESSION, &[id_entry("row-a", "", "[\"cooked\"]")]).await.expect("sync");
assert_eq!(blocks_of(&db, "row-a").await.as_deref(), Some("[\"cooked\"]"));
assert_eq!(row_count(&db).await, 1);
}
#[tokio::test]
async fn text_path_skips_empty_text_entries() {
let db = fresh_db().await;
sync_messages(&db, SESSION, &[text_entry("", "[\"cooked\"]")]).await.expect("sync");
assert_eq!(row_count(&db).await, 0, "空 text 无 id 无法定位,不得 INSERT 空行");
}
#[test]
fn cursor_round_trips_and_rejects_malformed() {
let c = MessageCursor { created_at: "2026-08-10T00:00:00Z".into(), id: "id-1".into() };
assert_eq!(MessageCursor::parse(&c.encode()), Some(c));
assert!(MessageCursor::parse("no-separator").is_none());
assert!(MessageCursor::parse("|id").is_none());
assert!(MessageCursor::parse("ts|").is_none());
}
}