fix: 用户隔离完整性 + 聊天去重
1. 启动时不加载历史消息(等第一条消息确定身份后按需恢复) 2. MemoryStore 改为 HashMap<user_id, Vec<mem>>,按用户隔离缓存 3. session.current_user_id 与外部 mutex 同步 4. chat_records.message_id 加 UNIQUE 索引 + ON CONFLICT DO NOTHING
This commit is contained in:
@@ -27,3 +27,6 @@ CREATE INDEX IF NOT EXISTS idx_chat_records_user_created
|
|||||||
|
|
||||||
CREATE INDEX IF NOT EXISTS idx_chat_records_account_created
|
CREATE INDEX IF NOT EXISTS idx_chat_records_account_created
|
||||||
ON chat_records (account_id, created_at DESC);
|
ON chat_records (account_id, created_at DESC);
|
||||||
|
|
||||||
|
CREATE UNIQUE INDEX IF NOT EXISTS idx_chat_records_message_id
|
||||||
|
ON chat_records (message_id) WHERE message_id != '';
|
||||||
|
|||||||
@@ -0,0 +1,3 @@
|
|||||||
|
-- 聊天记录去重索引
|
||||||
|
CREATE UNIQUE INDEX IF NOT EXISTS idx_chat_records_message_id
|
||||||
|
ON chat_records (message_id) WHERE message_id != '';
|
||||||
+22
-33
@@ -1,22 +1,22 @@
|
|||||||
use sqlx::PgPool;
|
use sqlx::PgPool;
|
||||||
|
use std::collections::HashMap;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use tokio::sync::Mutex;
|
use tokio::sync::Mutex;
|
||||||
|
|
||||||
/// 长期记忆管理器(内存 + PostgreSQL 双写)
|
/// 长期记忆管理器(内存 + PostgreSQL 双写,按用户隔离)
|
||||||
pub struct MemoryStore {
|
pub struct MemoryStore {
|
||||||
pool: Option<Arc<PgPool>>,
|
pool: Option<Arc<PgPool>>,
|
||||||
cache: Arc<Mutex<Vec<String>>>,
|
cache: Arc<Mutex<HashMap<String, Vec<String>>>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl MemoryStore {
|
impl MemoryStore {
|
||||||
pub fn new(pool: Option<Arc<PgPool>>) -> Self {
|
pub fn new(pool: Option<Arc<PgPool>>) -> Self {
|
||||||
Self {
|
Self {
|
||||||
pool,
|
pool,
|
||||||
cache: Arc::new(Mutex::new(Vec::new())),
|
cache: Arc::new(Mutex::new(HashMap::new())),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// 从数据库加载记忆到缓存
|
|
||||||
pub async fn load(&self, user_id: &str) {
|
pub async fn load(&self, user_id: &str) {
|
||||||
let pool = match self.pool.as_ref() {
|
let pool = match self.pool.as_ref() {
|
||||||
Some(p) => p.clone(),
|
Some(p) => p.clone(),
|
||||||
@@ -30,33 +30,28 @@ impl MemoryStore {
|
|||||||
.fetch_all(pool.as_ref())
|
.fetch_all(pool.as_ref())
|
||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
let mut cache = self.cache.lock().await;
|
let mems: Vec<String> = rows.into_iter().map(|(c,)| c).collect();
|
||||||
*cache = rows.into_iter().map(|(c,)| c).collect();
|
self.cache.lock().await.insert(uid.to_string(), mems);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// 读取记忆
|
pub async fn read_for(&self, user_id: &str) -> String {
|
||||||
pub async fn read(&self) -> String {
|
|
||||||
let mems = self.cache.lock().await;
|
|
||||||
if mems.is_empty() { "暂无长期记忆".to_string() } else {
|
|
||||||
mems.iter().enumerate().map(|(i, m)| format!("{}. {}", i+1, m)).collect::<Vec<_>>().join("\n")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// 按用户读取记忆(缓存按 user_id 过滤由调用方保证,这里简化为全量读取)
|
|
||||||
pub async fn read_for(&self, _user_id: &str) -> String { self.read().await }
|
|
||||||
|
|
||||||
/// 写入记忆
|
|
||||||
pub async fn write(&self, content: &str) -> String { self.write_for("default", content).await }
|
|
||||||
|
|
||||||
/// 按用户写入
|
|
||||||
pub async fn write_for(&self, user_id: &str, content: &str) -> String {
|
|
||||||
self.cache.lock().await.push(content.to_string());
|
|
||||||
let uid = if user_id.is_empty() { "default" } else { user_id };
|
let uid = if user_id.is_empty() { "default" } else { user_id };
|
||||||
|
let cache = self.cache.lock().await;
|
||||||
|
match cache.get(uid) {
|
||||||
|
Some(m) if !m.is_empty() => m.iter().enumerate()
|
||||||
|
.map(|(i, v)| format!("{}. {}", i + 1, v))
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
.join("\n"),
|
||||||
|
_ => "暂无长期记忆".to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn write_for(&self, user_id: &str, content: &str) -> String {
|
||||||
|
let uid = if user_id.is_empty() { "default" } else { user_id };
|
||||||
|
self.cache.lock().await.entry(uid.to_string()).or_default().push(content.to_string());
|
||||||
if let Some(ref pool) = self.pool {
|
if let Some(ref pool) = self.pool {
|
||||||
let _ = sqlx::query(
|
let _ = sqlx::query("INSERT INTO user_memories (user_id, content) VALUES ($1, $2)")
|
||||||
"INSERT INTO user_memories (user_id, content) VALUES ($1, $2)",
|
|
||||||
)
|
|
||||||
.bind(uid)
|
.bind(uid)
|
||||||
.bind(content)
|
.bind(content)
|
||||||
.execute(pool.as_ref())
|
.execute(pool.as_ref())
|
||||||
@@ -73,7 +68,6 @@ pub async fn read_summaries(
|
|||||||
let s = session.lock().await;
|
let s = session.lock().await;
|
||||||
let mut lines: Vec<String> = Vec::new();
|
let mut lines: Vec<String> = Vec::new();
|
||||||
|
|
||||||
// 内存中的摘要
|
|
||||||
for sum in &s.summaries {
|
for sum in &s.summaries {
|
||||||
let reason = match sum.reason {
|
let reason = match sum.reason {
|
||||||
super::types::SummaryReason::Overflow => "上下文压缩",
|
super::types::SummaryReason::Overflow => "上下文压缩",
|
||||||
@@ -87,7 +81,6 @@ pub async fn read_summaries(
|
|||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
// 数据库中的历史摘要
|
|
||||||
if let Some(pool) = &s.db_pool {
|
if let Some(pool) = &s.db_pool {
|
||||||
if let Ok(db_summaries) = crate::db::models::load_summaries(pool, 5).await {
|
if let Ok(db_summaries) = crate::db::models::load_summaries(pool, 5).await {
|
||||||
for text in db_summaries {
|
for text in db_summaries {
|
||||||
@@ -98,9 +91,5 @@ pub async fn read_summaries(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if lines.is_empty() {
|
if lines.is_empty() { "暂无历史摘要".to_string() } else { lines.join("\n---\n") }
|
||||||
"暂无历史摘要".to_string()
|
|
||||||
} else {
|
|
||||||
lines.join("\n---\n")
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -70,6 +70,7 @@ pub async fn insert_chat_record(
|
|||||||
r#"
|
r#"
|
||||||
INSERT INTO chat_records (direction, user_id, account_id, text, source, context_token, message_id)
|
INSERT INTO chat_records (direction, user_id, account_id, text, source, context_token, message_id)
|
||||||
VALUES ($1, $2, $3, $4, $5, $6, $7)
|
VALUES ($1, $2, $3, $4, $5, $6, $7)
|
||||||
|
ON CONFLICT (message_id) DO NOTHING
|
||||||
RETURNING id
|
RETURNING id
|
||||||
"#,
|
"#,
|
||||||
)
|
)
|
||||||
|
|||||||
+5
-4
@@ -306,14 +306,11 @@ async fn cmd_listen(
|
|||||||
// 创建 Conversation
|
// 创建 Conversation
|
||||||
let mut conv = match Conversation::new(cfg) {
|
let mut conv = match Conversation::new(cfg) {
|
||||||
Ok(c) => {
|
Ok(c) => {
|
||||||
// 设置 db_pool 并加载历史消息
|
// 设置 db_pool(启动时不加载消息,等第一条消息到达后再按需恢复)
|
||||||
if let Some(db) = database {
|
if let Some(db) = database {
|
||||||
let session = c.session();
|
let session = c.session();
|
||||||
{
|
|
||||||
let mut s = session.lock().await;
|
let mut s = session.lock().await;
|
||||||
s.db_pool = Some(Arc::new(db.pool().clone()));
|
s.db_pool = Some(Arc::new(db.pool().clone()));
|
||||||
s.load_recent_messages(20).await;
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
info!("LLM 已初始化: {} / {}", c.provider_name(), c.model());
|
info!("LLM 已初始化: {} / {}", c.provider_name(), c.model());
|
||||||
c
|
c
|
||||||
@@ -541,6 +538,10 @@ async fn listen_loop(
|
|||||||
|
|
||||||
// 更新当前用户(审批流程需要)
|
// 更新当前用户(审批流程需要)
|
||||||
*current_user_id.lock().await = from.to_string();
|
*current_user_id.lock().await = from.to_string();
|
||||||
|
// 同步到 session(executor 从 session 读取 user_id)
|
||||||
|
if let Some(conv) = &conversation {
|
||||||
|
conv.session().lock().await.current_user_id = from.to_string();
|
||||||
|
}
|
||||||
|
|
||||||
// LLM 回复(异步执行,避免审批阻塞主循环)
|
// LLM 回复(异步执行,避免审批阻塞主循环)
|
||||||
if enable_llm {
|
if enable_llm {
|
||||||
|
|||||||
Reference in New Issue
Block a user