From 5a0c018ee73d0f8b992e30fe4659067d79120b0d Mon Sep 17 00:00:00 2001 From: ooodc <549496103@qq.com> Date: Tue, 12 May 2026 21:04:17 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=9B=B4=E6=96=B0=E5=86=85=E5=AD=98?= =?UTF-8?q?=E7=BB=B4=E6=8A=A4=E9=80=BB=E8=BE=91=EF=BC=8C=E8=B0=83=E6=95=B4?= =?UTF-8?q?=E8=BF=94=E5=9B=9E=E7=B1=BB=E5=9E=8B=E4=B8=BAOption=E4=BB=A5?= =?UTF-8?q?=E5=A4=84=E7=90=86=E6=97=A0=E7=BB=93=E6=9E=9C=E6=83=85=E5=86=B5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/gateway/agent_task_executor.rs | 2 +- src/gateway/memory_maintenance.rs | 143 ++++++++++-------- src/gateway/memory_maintenance_coordinator.rs | 2 +- .../memory_maintenance_step2_system_prompt.md | 1 - src/gateway/session.rs | 33 ++-- 5 files changed, 95 insertions(+), 86 deletions(-) diff --git a/src/gateway/agent_task_executor.rs b/src/gateway/agent_task_executor.rs index a562d77..3955fd5 100644 --- a/src/gateway/agent_task_executor.rs +++ b/src/gateway/agent_task_executor.rs @@ -61,7 +61,7 @@ impl SchedulerMaintenanceService { self.session_manager.cleanup_expired_sessions().await } - async fn run_memory_maintenance(&self) -> Result, AgentError> { + async fn run_memory_maintenance(&self) -> Result, AgentError> { self.session_manager.run_memory_maintenance_for_all_scopes().await } } diff --git a/src/gateway/memory_maintenance.rs b/src/gateway/memory_maintenance.rs index a36aa03..a333f90 100644 --- a/src/gateway/memory_maintenance.rs +++ b/src/gateway/memory_maintenance.rs @@ -64,7 +64,6 @@ pub(crate) struct MemoryOrganizationOutput { #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub(crate) struct MemorySummaryInput { - pub(crate) scope_key: String, pub(crate) organized_memories: Vec, } @@ -217,6 +216,7 @@ impl MemoryMaintenanceService { Ok(output) } + #[cfg_attr(not(test), allow(dead_code))] pub(crate) async fn run_for_scope( &self, scope_key: &str, @@ -264,7 +264,6 @@ impl MemoryMaintenanceService { })?; let input = MemorySummaryInput { - scope_key: scope_key.to_string(), organized_memories: remaining_memories .iter() .map(|m| OrganizedMemory { @@ -353,15 +352,20 @@ impl MemoryMaintenanceService { pub(crate) async fn run_for_all_scopes( &self, - ) -> Result, AgentError> { + ) -> Result, AgentError> { let scope_keys = self.store.list_memory_scope_keys("user").map_err(|err| { AgentError::Other(format!("list memory scope keys error: {}", err)) })?; - let mut results = Vec::new(); - for scope_key in scope_keys { - let result = match self.run_for_scope(&scope_key).await { - Ok(Some(result)) => result, + if scope_keys.is_empty() { + return Ok(None); + } + + // 步骤1：逐个 scope 整理记忆（merge/delete），但不生成摘要 + let mut all_outputs = Vec::new(); + for scope_key in &scope_keys { + match self.run_organize_for_scope(scope_key).await { + Ok(Some(output)) => all_outputs.push((scope_key.clone(), output)), Ok(None) => continue, Err(error) if is_recoverable_maintenance_scope_error(&error) => { tracing::warn!( @@ -372,23 +376,78 @@ impl MemoryMaintenanceService { continue; } Err(error) => return Err(error), - }; - - results.push(result); + } } - let combined_markdown = combine_managed_memory_markdown( - &results + if all_outputs.is_empty() { + return Ok(None); + } + + // 步骤2：收集所有 scope 整理后的剩余记忆 + let mut all_remaining_memories = Vec::new(); + for (scope_key, _) in &all_outputs { + let memories = self + .store + .list_memories_for_scope("user", scope_key) + .map_err(|err| { + AgentError::Other(format!( + "list remaining memories for scope {} error: {}", + scope_key, err + )) + })?; + all_remaining_memories.extend(memories); + } + + // 步骤3：统一生成一个摘要 + let managed_markdown = if all_remaining_memories.is_empty() { + String::new() + } else { + self.generate_summary("all", &all_remaining_memories).await? + }; + + if !managed_markdown.is_empty() { + upsert_managed_agent_memory_summary(&managed_markdown)?; + } + + // 合并所有输出用于返回 + let combined_output = MemoryOrganizationOutput { + merges: all_outputs .iter() - .map(|result| result.managed_markdown.clone()) - .collect::>(), - ); + .flat_map(|(_, o)| o.merges.clone()) + .collect(), + conflicts: all_outputs + .iter() + .flat_map(|(_, o)| o.conflicts.clone()) + .collect(), + low_value_ids: all_outputs + .iter() + .flat_map(|(_, o)| o.low_value_ids.clone()) + .collect(), + }; - if !combined_markdown.is_empty() { - upsert_managed_agent_memory_summary(&combined_markdown)?; - } + Ok(Some(MemoryMaintenanceScopeResult { + scope_key: "all".to_string(), + output: combined_output, + managed_markdown, + })) + } - Ok(results) + /// 仅执行整理步骤（organize + apply），不生成摘要 + async fn run_organize_for_scope( + &self, + scope_key: &str, + ) -> Result, AgentError> { + let Some(plan) = self.build_plan_for_scope(scope_key)? else { + return Ok(None); + }; + + // 步骤1：整理记忆 + let organize_output = self.organize_plan(scope_key, &plan).await?; + + // 应用整理结果（merge和delete） + apply_memory_maintenance_output(self.store.as_ref(), scope_key, &plan, &organize_output)?; + + Ok(Some(organize_output)) } } @@ -512,47 +571,6 @@ pub(crate) fn extract_json_object(content: &str) -> Option<&str> { None } -pub(crate) fn combine_managed_memory_markdown(chunks: &[String]) -> String { - let normalized_chunks = chunks - .iter() - .map(|chunk| chunk.trim()) - .filter(|chunk| !chunk.is_empty()) - .collect::>(); - - let mut combined = Vec::new(); - for (index, chunk) in normalized_chunks.iter().enumerate() { - let chunk_lines = chunk - .lines() - .map(str::trim) - .filter(|line| !line.is_empty()) - .collect::>(); - - let is_subset_of_other = - normalized_chunks - .iter() - .enumerate() - .any(|(other_index, other)| { - if index == other_index { - return false; - } - - let other_lines = other - .lines() - .map(str::trim) - .filter(|line| !line.is_empty()) - .collect::>(); - - chunk_lines.len() < other_lines.len() && chunk_lines.is_subset(&other_lines) - }); - - if !is_subset_of_other && !combined.iter().any(|existing: &String| existing == chunk) { - combined.push((*chunk).to_string()); - } - } - - combined.join("\n\n") -} - pub(crate) fn apply_memory_maintenance_output( store: &SessionStore, scope_key: &str, @@ -651,3 +669,8 @@ fn preview_text(content: &str, max_chars: usize) -> String { } preview.replace('\n', "\\n") } + +#[cfg(test)] +mod tests { + use super::*; +} diff --git a/src/gateway/memory_maintenance_coordinator.rs b/src/gateway/memory_maintenance_coordinator.rs index 1e5f547..1c1a09f 100644 --- a/src/gateway/memory_maintenance_coordinator.rs +++ b/src/gateway/memory_maintenance_coordinator.rs @@ -32,7 +32,7 @@ impl MemoryMaintenanceCoordinator { pub(crate) async fn run_for_all_scopes( &self, - ) -> Result, AgentError> { + ) -> Result, AgentError> { self.service()?.run_for_all_scopes().await } diff --git a/src/gateway/memory_maintenance_step2_system_prompt.md b/src/gateway/memory_maintenance_step2_system_prompt.md index 9e7e64e..23ac0c4 100644 --- a/src/gateway/memory_maintenance_step2_system_prompt.md +++ b/src/gateway/memory_maintenance_step2_system_prompt.md @@ -3,7 +3,6 @@ 你的任务是基于整理后的用户记忆生成结构化的 Markdown 摘要。 输入格式： -- scope_key: 用户的唯一标识 - organized_memories: 整理后的记忆列表，每个包含 namespace、memory_key、content 输出要求： diff --git a/src/gateway/session.rs b/src/gateway/session.rs index 702ac3f..725d688 100644 --- a/src/gateway/session.rs +++ b/src/gateway/session.rs @@ -20,7 +20,7 @@ use super::execution::should_display_message_to_user; #[cfg(test)] use super::memory_maintenance::{ MemoryMaintenanceMerge, apply_memory_maintenance_output, build_memory_maintenance_plan, - combine_managed_memory_markdown, extract_json_object, is_recoverable_maintenance_llm_error, + extract_json_object, is_recoverable_maintenance_llm_error, strip_json_code_fence, }; use super::memory_maintenance::{MemoryMaintenanceScopeResult, MemoryOrganizationOutput}; @@ -427,7 +427,7 @@ impl SessionManager { pub(crate) async fn run_memory_maintenance_for_all_scopes( &self, - ) -> Result, AgentError> { + ) -> Result, AgentError> { self.memory_maintenance.run_for_all_scopes().await } @@ -1317,15 +1317,16 @@ mod tests { }) .unwrap(); - let results = session_manager + let result = session_manager .run_memory_maintenance_for_all_scopes() .await .unwrap(); - assert_eq!(results.len(), 1); - assert_eq!(results[0].scope_key, "feishu:user-1"); + assert!(result.is_some()); + let result = result.unwrap(); + assert_eq!(result.scope_key, "all"); // 由于步骤2需要新的提示词和输入格式，这里只验证基本功能 - assert!(!results[0].managed_markdown.is_empty()); + assert!(!result.managed_markdown.is_empty()); } #[tokio::test] @@ -1378,12 +1379,13 @@ mod tests { .unwrap(); } - let results = session_manager + let result = session_manager .run_memory_maintenance_for_all_scopes() .await .unwrap(); - assert!(results.is_empty()); + // 当遇到可恢复错误时，没有整理任何记忆，返回 None + assert!(result.is_none()); } #[test] @@ -1460,21 +1462,6 @@ mod tests { assert_eq!(all_memories[0].content, "用户主要在做AI产品设计与实现"); } - #[test] - fn test_combine_managed_memory_markdown_prefers_richer_summary_over_subset() { - let combined = combine_managed_memory_markdown(&[ - "### 用户事实\n- 用户在做AI产品\n\n### 用户偏好\n- 偏好简洁表达".to_string(), - "- 用户在做AI产品".to_string(), - "### 用户事实\n- 用户名为区德成，昵称DC。".to_string(), - ]); - - assert!( - combined.contains("### 用户事实\n- 用户在做AI产品\n\n### 用户偏好\n- 偏好简洁表达") - ); - assert!(combined.contains("### 用户事实\n- 用户名为区德成，昵称DC。")); - assert_eq!(combined.matches("- 用户在做AI产品").count(), 1); - } - #[test] fn test_should_display_message_to_user_hides_completed_tool_results_by_default() { let completed = ChatMessage::tool("call-1", "calculator", "2");