Compare commits

..

No commits in common. "7c48a0f7f932762967eb711d859b2aa7900a2906" and "495c8cdc7e8065e61940acc40fae33374671ace6" have entirely different histories.

6 changed files with 145 additions and 201 deletions

View File

@ -61,8 +61,13 @@ impl SchedulerMaintenanceService {
self.session_manager.cleanup_expired_sessions().await self.session_manager.cleanup_expired_sessions().await
} }
async fn run_memory_maintenance(&self) -> Result<Vec<MemoryMaintenanceScopeResult>, AgentError> { async fn run_memory_maintenance(
self.session_manager.run_memory_maintenance_for_all_scopes().await &self,
updated_since: Option<i64>,
) -> Result<Vec<MemoryMaintenanceScopeResult>, AgentError> {
self.session_manager
.run_memory_maintenance_for_all_scopes(updated_since)
.await
} }
} }
@ -72,8 +77,11 @@ impl MaintenanceExecutor for SchedulerMaintenanceService {
self.cleanup_sessions().await self.cleanup_sessions().await
} }
async fn run_memory_maintenance_for_all_scopes(&self) -> anyhow::Result<Vec<MaintenanceRunSummary>> { async fn run_memory_maintenance_for_all_scopes(
self.run_memory_maintenance() &self,
updated_since: Option<i64>,
) -> anyhow::Result<Vec<MaintenanceRunSummary>> {
self.run_memory_maintenance(updated_since)
.await .await
.map(|results| { .map(|results| {
results results

View File

@ -220,10 +220,22 @@ impl MemoryMaintenanceService {
pub(crate) async fn run_for_all_scopes( pub(crate) async fn run_for_all_scopes(
&self, &self,
updated_since: Option<i64>,
) -> Result<Vec<MemoryMaintenanceScopeResult>, AgentError> { ) -> Result<Vec<MemoryMaintenanceScopeResult>, AgentError> {
let scope_keys = self.store.list_memory_scope_keys("user").map_err(|err| { let scope_keys = if let Some(cutoff) = updated_since {
self.store
.list_memory_scope_keys_updated_since("user", cutoff)
.map_err(|err| {
AgentError::Other(format!(
"list memory scope keys updated since error: {}",
err
))
})?
} else {
self.store.list_memory_scope_keys("user").map_err(|err| {
AgentError::Other(format!("list memory scope keys error: {}", err)) AgentError::Other(format!("list memory scope keys error: {}", err))
})?; })?
};
let mut results = Vec::new(); let mut results = Vec::new();
for scope_key in scope_keys { for scope_key in scope_keys {

View File

@ -32,8 +32,9 @@ impl MemoryMaintenanceCoordinator {
pub(crate) async fn run_for_all_scopes( pub(crate) async fn run_for_all_scopes(
&self, &self,
updated_since: Option<i64>,
) -> Result<Vec<MemoryMaintenanceScopeResult>, AgentError> { ) -> Result<Vec<MemoryMaintenanceScopeResult>, AgentError> {
self.service()?.run_for_all_scopes().await self.service()?.run_for_all_scopes(updated_since).await
} }
fn service(&self) -> Result<MemoryMaintenanceService, AgentError> { fn service(&self) -> Result<MemoryMaintenanceService, AgentError> {

View File

@ -4,63 +4,23 @@ use std::path::{Path, PathBuf};
use crate::agent::AgentError; use crate::agent::AgentError;
pub(crate) const DEFAULT_AGENT_PROMPT: &str = include_str!("default_agent_prompt.md"); pub(crate) const DEFAULT_AGENT_PROMPT: &str = include_str!("default_agent_prompt.md");
pub(crate) const MANAGED_AGENT_MEMORY_BLOCK_START: &str = "<!-- PICOBOT_MANAGED_MEMORY:START -->";
#[derive(Clone)] pub(crate) const MANAGED_AGENT_MEMORY_BLOCK_END: &str = "<!-- PICOBOT_MANAGED_MEMORY:END -->";
struct PromptSource { pub(crate) const MANAGED_AGENT_MEMORY_TITLE: &str = "## 用户记忆摘要";
path: PathBuf,
default_content: Option<&'static str>,
}
pub(crate) fn load_agent_prompt() -> Result<Option<String>, AgentError> { pub(crate) fn load_agent_prompt() -> Result<Option<String>, AgentError> {
load_prompt_from_sources(&prompt_sources()?) let path = agent_prompt_path()?;
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)
.map_err(|err| AgentError::Other(format!("create agent prompt dir error: {}", err)))?;
} }
pub(crate) fn upsert_managed_agent_memory_summary(markdown_body: &str) -> Result<(), AgentError> { if !path.exists() {
persist_memory_summary(&memory_summary_path()?, markdown_body) write_agent_prompt(&path, DEFAULT_AGENT_PROMPT)?;
} }
fn prompt_sources() -> Result<Vec<PromptSource>, AgentError> { let content = fs::read_to_string(&path)
Ok(vec![ .map_err(|err| AgentError::Other(format!("read agent prompt file error: {}", err)))?;
PromptSource {
path: agent_prompt_path()?,
default_content: Some(DEFAULT_AGENT_PROMPT),
},
PromptSource {
path: memory_summary_path()?,
default_content: None,
},
])
}
fn load_prompt_from_sources(sources: &[PromptSource]) -> Result<Option<String>, AgentError> {
let mut fragments = Vec::with_capacity(sources.len());
for source in sources {
if let Some(fragment) = read_prompt_fragment(source)? {
fragments.push(fragment);
}
}
if fragments.is_empty() {
return Ok(None);
}
Ok(Some(fragments.join("\n\n")))
}
fn read_prompt_fragment(source: &PromptSource) -> Result<Option<String>, AgentError> {
ensure_parent_dir(&source.path)?;
if !source.path.exists() {
if let Some(default_content) = source.default_content {
write_prompt_file(&source.path, default_content)?;
} else {
return Ok(None);
}
}
let content = fs::read_to_string(&source.path)
.map_err(|err| AgentError::Other(format!("read prompt file error: {}", err)))?;
let trimmed = content.trim(); let trimmed = content.trim();
if trimmed.is_empty() { if trimmed.is_empty() {
return Ok(None); return Ok(None);
@ -69,151 +29,121 @@ fn read_prompt_fragment(source: &PromptSource) -> Result<Option<String>, AgentEr
Ok(Some(trimmed.to_string())) Ok(Some(trimmed.to_string()))
} }
fn persist_memory_summary(path: &Path, markdown_body: &str) -> Result<(), AgentError> { pub(crate) fn upsert_managed_agent_memory_summary(markdown_body: &str) -> Result<(), AgentError> {
let trimmed = markdown_body.trim(); let path = agent_prompt_path()?;
if trimmed.is_empty() { let existing = if path.exists() {
return Ok(()); fs::read_to_string(&path)
.map_err(|err| AgentError::Other(format!("read agent prompt file error: {}", err)))?
} else {
DEFAULT_AGENT_PROMPT.to_string()
};
let updated = upsert_managed_agent_memory_block(&existing, markdown_body);
write_agent_prompt(&path, &updated)
} }
write_prompt_file(path, trimmed) pub(crate) fn upsert_managed_agent_memory_block(existing: &str, markdown_body: &str) -> String {
let managed_block = render_managed_agent_memory_block(markdown_body);
if let (Some(start), Some(end)) = (
existing.find(MANAGED_AGENT_MEMORY_BLOCK_START),
existing.find(MANAGED_AGENT_MEMORY_BLOCK_END),
) {
let end = end + MANAGED_AGENT_MEMORY_BLOCK_END.len();
let mut updated = String::new();
updated.push_str(existing[..start].trim_end());
updated.push_str("\n\n");
updated.push_str(&managed_block);
updated.push_str("\n\n");
updated.push_str(existing[end..].trim_start());
return updated.trim().to_string() + "\n";
} }
fn write_prompt_file(path: &Path, content: &str) -> Result<(), AgentError> { if let Some(reply_rules_index) = existing.find("## 回复规则") {
ensure_parent_dir(path)?; let mut updated = String::new();
updated.push_str(existing[..reply_rules_index].trim_end());
let normalized = content.trim_end().to_string() + "\n"; updated.push_str("\n\n");
let temp_path = path.with_extension("md.tmp"); updated.push_str(&managed_block);
fs::write(&temp_path, normalized) updated.push_str("\n\n");
.map_err(|err| AgentError::Other(format!("write prompt temp file error: {}", err)))?; updated.push_str(existing[reply_rules_index..].trim_start());
fs::rename(&temp_path, path) return updated.trim().to_string() + "\n";
.map_err(|err| AgentError::Other(format!("replace prompt file error: {}", err)))?;
Ok(())
} }
fn ensure_parent_dir(path: &Path) -> Result<(), AgentError> { let mut updated = existing.trim_end().to_string();
if !updated.is_empty() {
updated.push_str("\n\n");
}
updated.push_str(&managed_block);
updated.push('\n');
updated
}
fn render_managed_agent_memory_block(markdown_body: &str) -> String {
format!(
"{MANAGED_AGENT_MEMORY_BLOCK_START}\n{MANAGED_AGENT_MEMORY_TITLE}\n\n{}\n{MANAGED_AGENT_MEMORY_BLOCK_END}",
markdown_body.trim()
)
}
fn write_agent_prompt(path: &Path, content: &str) -> Result<(), AgentError> {
if let Some(parent) = path.parent() { if let Some(parent) = path.parent() {
fs::create_dir_all(parent) fs::create_dir_all(parent)
.map_err(|err| AgentError::Other(format!("create agent prompt dir error: {}", err)))?; .map_err(|err| AgentError::Other(format!("create agent prompt dir error: {}", err)))?;
} }
let temp_path = path.with_extension("md.tmp");
fs::write(&temp_path, content)
.map_err(|err| AgentError::Other(format!("write agent prompt temp file error: {}", err)))?;
fs::rename(&temp_path, path)
.map_err(|err| AgentError::Other(format!("replace agent prompt file error: {}", err)))?;
Ok(()) Ok(())
} }
fn agent_prompt_path() -> Result<PathBuf, AgentError> { fn agent_prompt_path() -> Result<PathBuf, AgentError> {
Ok(agent_dir_path()?.join("AGENT.md"))
}
fn memory_summary_path() -> Result<PathBuf, AgentError> {
Ok(agent_dir_path()?.join("MEMORY_SUMMARY.md"))
}
fn agent_dir_path() -> Result<PathBuf, AgentError> {
let home = dirs::home_dir() let home = dirs::home_dir()
.ok_or_else(|| AgentError::Other("home directory not found".to_string()))?; .ok_or_else(|| AgentError::Other("home directory not found".to_string()))?;
Ok(home.join(".picobot").join("agent")) Ok(home.join(".picobot").join("agent").join("AGENT.md"))
} }
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use tempfile::tempdir;
#[test] #[test]
fn test_load_prompt_from_sources_aggregates_multiple_fragments_in_order() { fn test_upsert_managed_agent_memory_block_inserts_before_reply_rules() {
let temp = tempdir().unwrap(); let original =
let agent_path = temp.path().join("AGENT.md"); "# PicoBot 代理配置\n\n## 身份\n- 你是 PicoBot。\n\n## 回复规则\n- 使用中文回复。\n";
let memory_path = temp.path().join("MEMORY_SUMMARY.md"); let updated = upsert_managed_agent_memory_block(
original,
"### 用户事实\n- 用户在做AI产品\n\n### 用户偏好\n- 偏好简洁表达",
);
write_prompt_file(&agent_path, "# Agent\n静态规则").unwrap(); let managed_pos = updated.find(MANAGED_AGENT_MEMORY_BLOCK_START).unwrap();
write_prompt_file(&memory_path, "## 用户记忆摘要\n- 偏好简洁").unwrap(); let reply_rules_pos = updated.find("## 回复规则").unwrap();
assert!(managed_pos < reply_rules_pos);
let prompt = load_prompt_from_sources(&[ assert!(updated.contains(MANAGED_AGENT_MEMORY_TITLE));
PromptSource { assert!(updated.contains("用户在做AI产品"));
path: agent_path.clone(), assert!(updated.contains("偏好简洁表达"));
default_content: Some(DEFAULT_AGENT_PROMPT),
},
PromptSource {
path: memory_path.clone(),
default_content: None,
},
])
.unwrap()
.unwrap();
assert_eq!(prompt, "# Agent\n静态规则\n\n## 用户记忆摘要\n- 偏好简洁");
} }
#[test] #[test]
fn test_load_prompt_from_sources_ignores_missing_optional_source() { fn test_upsert_managed_agent_memory_block_replaces_existing_block() {
let temp = tempdir().unwrap(); let original = format!(
let agent_path = temp.path().join("AGENT.md"); "# PicoBot\n\n{MANAGED_AGENT_MEMORY_BLOCK_START}\n{MANAGED_AGENT_MEMORY_TITLE}\n\nold\n{MANAGED_AGENT_MEMORY_BLOCK_END}\n\n## 回复规则\n- 简洁。\n"
let memory_path = temp.path().join("MEMORY_SUMMARY.md"); );
write_prompt_file(&agent_path, "# Agent\n静态规则").unwrap(); let updated = upsert_managed_agent_memory_block(&original, "new");
let prompt = load_prompt_from_sources(&[ assert!(updated.contains("new"));
PromptSource { assert!(!updated.contains("old"));
path: agent_path.clone(), assert_eq!(updated.matches(MANAGED_AGENT_MEMORY_BLOCK_START).count(), 1);
default_content: Some(DEFAULT_AGENT_PROMPT), assert_eq!(updated.matches(MANAGED_AGENT_MEMORY_BLOCK_END).count(), 1);
},
PromptSource {
path: memory_path.clone(),
default_content: None,
},
])
.unwrap()
.unwrap();
assert_eq!(prompt, "# Agent\n静态规则");
assert!(!memory_path.exists());
} }
#[test] #[test]
fn test_load_prompt_from_sources_creates_default_agent_prompt() { fn test_upsert_managed_agent_memory_block_trims_summary_body() {
let temp = tempdir().unwrap(); let updated = upsert_managed_agent_memory_block("# PicoBot\n", "\n\nsummary\n\n");
let agent_path = temp.path().join("AGENT.md");
let prompt = load_prompt_from_sources(&[PromptSource { assert!(updated.contains("\n\nsummary\n"));
path: agent_path.clone(), assert!(!updated.contains("\n\nsummary\n\n\n"));
default_content: Some("# Default Agent\n规则"),
}])
.unwrap()
.unwrap();
assert_eq!(prompt, "# Default Agent\n规则");
assert_eq!(fs::read_to_string(&agent_path).unwrap(), "# Default Agent\n规则\n");
}
#[test]
fn test_load_prompt_from_sources_returns_none_when_all_sources_empty() {
let temp = tempdir().unwrap();
let agent_path = temp.path().join("AGENT.md");
let memory_path = temp.path().join("MEMORY_SUMMARY.md");
write_prompt_file(&agent_path, " ").unwrap();
write_prompt_file(&memory_path, "\n\n").unwrap();
let prompt = load_prompt_from_sources(&[
PromptSource {
path: agent_path.clone(),
default_content: None,
},
PromptSource {
path: memory_path.clone(),
default_content: None,
},
])
.unwrap();
assert!(prompt.is_none());
}
#[test]
fn test_upsert_managed_agent_memory_summary_writes_trimmed_markdown() {
let temp = tempdir().unwrap();
let memory_path = temp.path().join("MEMORY_SUMMARY.md");
persist_memory_summary(&memory_path, "\n## 用户记忆摘要\n- 偏好简洁\n\n").unwrap();
assert_eq!(fs::read_to_string(&memory_path).unwrap(), "## 用户记忆摘要\n- 偏好简洁\n");
} }
} }

View File

@ -420,8 +420,11 @@ impl SessionManager {
pub(crate) async fn run_memory_maintenance_for_all_scopes( pub(crate) async fn run_memory_maintenance_for_all_scopes(
&self, &self,
updated_since: Option<i64>,
) -> Result<Vec<MemoryMaintenanceScopeResult>, AgentError> { ) -> Result<Vec<MemoryMaintenanceScopeResult>, AgentError> {
self.memory_maintenance.run_for_all_scopes().await self.memory_maintenance
.run_for_all_scopes(updated_since)
.await
} }
/// 确保 session 存在且未超时,超时则重建 /// 确保 session 存在且未超时,超时则重建
@ -1236,34 +1239,18 @@ mod tests {
} }
#[tokio::test] #[tokio::test]
async fn test_run_memory_maintenance_for_all_scopes_scans_all_scopes_even_without_recent_updates() { async fn test_run_memory_maintenance_for_all_scopes_returns_empty_when_no_recent_updates() {
let mock_response_content = serde_json::to_string(&json!({
"user_facts": ["用户在做AI产品"],
"preferences": [],
"behavior_patterns": [],
"merges": [],
"conflicts": [],
"low_value_ids": [],
"managed_markdown": "### 用户事实\n- 用户在做AI产品"
}))
.unwrap();
let base_url =
start_mock_openai_server_with_content(Some(mock_response_content.clone())).await;
let provider_config = LLMProviderConfig { let provider_config = LLMProviderConfig {
provider_type: "openai".to_string(), provider_type: "openai".to_string(),
name: "maintenance-provider".to_string(), name: "maintenance-provider".to_string(),
base_url, base_url: "http://localhost".to_string(),
api_key: "test-key".to_string(), api_key: "test-key".to_string(),
extra_headers: HashMap::new(), extra_headers: HashMap::new(),
model_id: "maintenance-model".to_string(), model_id: "maintenance-model".to_string(),
temperature: Some(0.0), temperature: Some(0.0),
max_tokens: Some(256), max_tokens: Some(256),
context_window_tokens: None, context_window_tokens: None,
model_extra: HashMap::from([( model_extra: HashMap::new(),
"mock_response_content".to_string(),
json!(mock_response_content),
)]),
max_tool_iterations: 1, max_tool_iterations: 1,
llm_timeout_secs: 30, llm_timeout_secs: 30,
tool_result_max_chars: 20_000, tool_result_max_chars: 20_000,
@ -1281,7 +1268,7 @@ mod tests {
) )
.unwrap(); .unwrap();
session_manager let memory = session_manager
.store() .store()
.put_memory(&crate::storage::MemoryUpsert { .put_memory(&crate::storage::MemoryUpsert {
scope_kind: "user".to_string(), scope_kind: "user".to_string(),
@ -1299,13 +1286,11 @@ mod tests {
.unwrap(); .unwrap();
let results = session_manager let results = session_manager
.run_memory_maintenance_for_all_scopes() .run_memory_maintenance_for_all_scopes(Some(memory.updated_at + 1))
.await .await
.unwrap(); .unwrap();
assert_eq!(results.len(), 1); assert!(results.is_empty());
assert_eq!(results[0].scope_key, "feishu:user-1");
assert!(results[0].output.user_facts.contains(&"用户在做AI产品".to_string()));
} }
#[tokio::test] #[tokio::test]
@ -1358,7 +1343,7 @@ mod tests {
} }
let results = session_manager let results = session_manager
.run_memory_maintenance_for_all_scopes() .run_memory_maintenance_for_all_scopes(None)
.await .await
.unwrap(); .unwrap();

View File

@ -52,7 +52,10 @@ pub trait AgentTaskExecutor: Send + Sync {
pub trait MaintenanceExecutor: Send + Sync { pub trait MaintenanceExecutor: Send + Sync {
async fn cleanup_expired_sessions(&self) -> usize; async fn cleanup_expired_sessions(&self) -> usize;
async fn run_memory_maintenance_for_all_scopes(&self) -> anyhow::Result<Vec<MaintenanceRunSummary>>; async fn run_memory_maintenance_for_all_scopes(
&self,
updated_since: Option<i64>,
) -> anyhow::Result<Vec<MaintenanceRunSummary>>;
} }
pub struct Scheduler { pub struct Scheduler {
@ -649,7 +652,9 @@ async fn execute_internal_event(
Ok(()) Ok(())
} }
"memory_maintenance" => { "memory_maintenance" => {
let results = maintenance_executor.run_memory_maintenance_for_all_scopes().await?; let results = maintenance_executor
.run_memory_maintenance_for_all_scopes(job.last_fired_at)
.await?;
for result in &results { for result in &results {
tracing::info!( tracing::info!(
job_id = %job.id, job_id = %job.id,
@ -1034,7 +1039,10 @@ mod tests {
0 0
} }
async fn run_memory_maintenance_for_all_scopes(&self) -> anyhow::Result<Vec<MaintenanceRunSummary>> { async fn run_memory_maintenance_for_all_scopes(
&self,
_updated_since: Option<i64>,
) -> anyhow::Result<Vec<MaintenanceRunSummary>> {
Ok(Vec::new()) Ok(Vec::new())
} }
} }