diff --git a/src/storage/mod.rs b/src/storage/mod.rs index 493675c..ad9ae7b 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -1581,10 +1581,19 @@ impl SessionStore { } } - /// 获取指定话题的消息数量(动态计算,确保准确) + /// 获取指定话题的消息数量。 + /// + /// 使用 `SELECT COUNT(*)` 在数据库侧计数,避免将所有消息 + /// (含 content、tool_calls_json 等大字段反序列化)加载到内存。 + /// 查询命中 `idx_messages_topic_seq(topic_id, seq)` 索引。 pub fn get_topic_message_count(&self, topic_id: &str) -> Result { - self.load_messages_for_topic(topic_id, None) - .map(|msgs| msgs.len()) + 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, StorageError> { diff --git a/src/storage/tests.rs b/src/storage/tests.rs index 1cefb3f..ef4553e 100644 --- a/src/storage/tests.rs +++ b/src/storage/tests.rs @@ -687,3 +687,41 @@ fn test_scheduler_job_roundtrip_and_runtime_update() { assert_eq!(fetched.run_count, 1); assert_eq!(fetched.completed_at, Some(1_700_000_000_100)); } + +#[test] +fn test_get_topic_message_count_uses_count_query() { + let store = SessionStore::in_memory().unwrap(); + let session = store.create_cli_session(Some("topic-count")).unwrap(); + let topic = store + .create_topic(&session.id, "topic-1", None) + .unwrap(); + + // 初始计数为 0 + assert_eq!(store.get_topic_message_count(&topic.id).unwrap(), 0); + + // 追加 3 条带 topic_id 的消息 + for content in ["m1", "m2", "m3"] { + store + .append_message_with_topic( + &session.id, + Some(&topic.id), + &ChatMessage::user(content), + ) + .unwrap(); + } + + // 计数应为 3,且不需要加载消息内容 + assert_eq!(store.get_topic_message_count(&topic.id).unwrap(), 3); + + // 另一个 topic 的计数应为 0(隔离验证) + let other_topic = store + .create_topic(&session.id, "topic-2", None) + .unwrap(); + assert_eq!(store.get_topic_message_count(&other_topic.id).unwrap(), 0); + + // 不存在的 topic_id 返回 0 + assert_eq!( + store.get_topic_message_count("topic:nonexistent").unwrap(), + 0 + ); +}