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:
+26
-37
@@ -1,22 +1,22 @@
|
||||
use sqlx::PgPool;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
/// 长期记忆管理器(内存 + PostgreSQL 双写)
|
||||
/// 长期记忆管理器(内存 + PostgreSQL 双写,按用户隔离)
|
||||
pub struct MemoryStore {
|
||||
pool: Option<Arc<PgPool>>,
|
||||
cache: Arc<Mutex<Vec<String>>>,
|
||||
cache: Arc<Mutex<HashMap<String, Vec<String>>>>,
|
||||
}
|
||||
|
||||
impl MemoryStore {
|
||||
pub fn new(pool: Option<Arc<PgPool>>) -> Self {
|
||||
Self {
|
||||
pool,
|
||||
cache: Arc::new(Mutex::new(Vec::new())),
|
||||
cache: Arc::new(Mutex::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// 从数据库加载记忆到缓存
|
||||
pub async fn load(&self, user_id: &str) {
|
||||
let pool = match self.pool.as_ref() {
|
||||
Some(p) => p.clone(),
|
||||
@@ -30,37 +30,32 @@ impl MemoryStore {
|
||||
.fetch_all(pool.as_ref())
|
||||
.await
|
||||
{
|
||||
let mut cache = self.cache.lock().await;
|
||||
*cache = rows.into_iter().map(|(c,)| c).collect();
|
||||
let mems: Vec<String> = rows.into_iter().map(|(c,)| c).collect();
|
||||
self.cache.lock().await.insert(uid.to_string(), mems);
|
||||
}
|
||||
}
|
||||
|
||||
/// 读取记忆
|
||||
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());
|
||||
pub async fn read_for(&self, user_id: &str) -> String {
|
||||
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 {
|
||||
let _ = sqlx::query(
|
||||
"INSERT INTO user_memories (user_id, content) VALUES ($1, $2)",
|
||||
)
|
||||
.bind(uid)
|
||||
.bind(content)
|
||||
.execute(pool.as_ref())
|
||||
.await;
|
||||
let _ = sqlx::query("INSERT INTO user_memories (user_id, content) VALUES ($1, $2)")
|
||||
.bind(uid)
|
||||
.bind(content)
|
||||
.execute(pool.as_ref())
|
||||
.await;
|
||||
}
|
||||
"已添加长期记忆".to_string()
|
||||
}
|
||||
@@ -73,7 +68,6 @@ pub async fn read_summaries(
|
||||
let s = session.lock().await;
|
||||
let mut lines: Vec<String> = Vec::new();
|
||||
|
||||
// 内存中的摘要
|
||||
for sum in &s.summaries {
|
||||
let reason = match sum.reason {
|
||||
super::types::SummaryReason::Overflow => "上下文压缩",
|
||||
@@ -87,7 +81,6 @@ pub async fn read_summaries(
|
||||
));
|
||||
}
|
||||
|
||||
// 数据库中的历史摘要
|
||||
if let Some(pool) = &s.db_pool {
|
||||
if let Ok(db_summaries) = crate::db::models::load_summaries(pool, 5).await {
|
||||
for text in db_summaries {
|
||||
@@ -98,9 +91,5 @@ pub async fn read_summaries(
|
||||
}
|
||||
}
|
||||
|
||||
if lines.is_empty() {
|
||||
"暂无历史摘要".to_string()
|
||||
} else {
|
||||
lines.join("\n---\n")
|
||||
}
|
||||
if lines.is_empty() { "暂无历史摘要".to_string() } else { lines.join("\n---\n") }
|
||||
}
|
||||
|
||||
@@ -70,6 +70,7 @@ pub async fn insert_chat_record(
|
||||
r#"
|
||||
INSERT INTO chat_records (direction, user_id, account_id, text, source, context_token, message_id)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7)
|
||||
ON CONFLICT (message_id) DO NOTHING
|
||||
RETURNING id
|
||||
"#,
|
||||
)
|
||||
|
||||
+7
-6
@@ -306,14 +306,11 @@ async fn cmd_listen(
|
||||
// 创建 Conversation
|
||||
let mut conv = match Conversation::new(cfg) {
|
||||
Ok(c) => {
|
||||
// 设置 db_pool 并加载历史消息
|
||||
// 设置 db_pool(启动时不加载消息,等第一条消息到达后再按需恢复)
|
||||
if let Some(db) = database {
|
||||
let session = c.session();
|
||||
{
|
||||
let mut s = session.lock().await;
|
||||
s.db_pool = Some(Arc::new(db.pool().clone()));
|
||||
s.load_recent_messages(20).await;
|
||||
}
|
||||
let mut s = session.lock().await;
|
||||
s.db_pool = Some(Arc::new(db.pool().clone()));
|
||||
}
|
||||
info!("LLM 已初始化: {} / {}", c.provider_name(), c.model());
|
||||
c
|
||||
@@ -541,6 +538,10 @@ async fn listen_loop(
|
||||
|
||||
// 更新当前用户(审批流程需要)
|
||||
*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 回复(异步执行,避免审批阻塞主循环)
|
||||
if enable_llm {
|
||||
|
||||
Reference in New Issue
Block a user