From 73d2b07d52e62b2aa8d250ad352cb4772f2a2d99 Mon Sep 17 00:00:00 2001 From: kvtodev Date: Fri, 10 Jul 2026 05:43:06 +0800 Subject: [PATCH] refactor: simplify agent memory handling and improve message flow logic --- .../module-editor-agent/src/agent/agent.rs | 4 +- .../module-editor-agent/src/agent/hook.rs | 6 +-- .../module-editor-agent/src/agent/run.rs | 39 ++++++++++--------- 3 files changed, 25 insertions(+), 24 deletions(-) diff --git a/server-rs/crates/module-editor-agent/src/agent/agent.rs b/server-rs/crates/module-editor-agent/src/agent/agent.rs index 5720ba6d9..aa000bbe0 100644 --- a/server-rs/crates/module-editor-agent/src/agent/agent.rs +++ b/server-rs/crates/module-editor-agent/src/agent/agent.rs @@ -1,10 +1,8 @@ use crate::agent::error::PromptError; use crate::agent::hook::Hook; use crate::agent::memory::AgentMemory; -use crate::agent::run::{Flow, PromptRequest}; -use crate::agent::tool::ToolOutcome; +use crate::agent::run::PromptRequest; use crate::agent::{Tool, ToolDyn}; -use std::pin::Pin; pub struct Agent, Message> { pub model: M, diff --git a/server-rs/crates/module-editor-agent/src/agent/hook.rs b/server-rs/crates/module-editor-agent/src/agent/hook.rs index bc0bbc662..5193c3d18 100644 --- a/server-rs/crates/module-editor-agent/src/agent/hook.rs +++ b/server-rs/crates/module-editor-agent/src/agent/hook.rs @@ -2,13 +2,13 @@ use crate::agent::run::{TextFlow, ToolCallFlow}; use crate::agent::tool::ToolCall; pub trait Hook: Send + Sync { - fn on_text_reply(&self, text: &str) -> TextFlow { + fn on_text_reply(&self, _text: &str) -> TextFlow { TextFlow::Continue } - fn before_tool_call(&self, tool_call: &ToolCall) -> ToolCallFlow { + fn before_tool_call(&self, _tool_call: &ToolCall) -> ToolCallFlow { ToolCallFlow::Continue } - fn after_tool_call(&self, tool_name: &str, output: &mut serde_json::Value) -> ToolCallFlow { + fn after_tool_call(&self, _tool_name: &str, _output: &mut serde_json::Value) -> ToolCallFlow { ToolCallFlow::Continue } } diff --git a/server-rs/crates/module-editor-agent/src/agent/run.rs b/server-rs/crates/module-editor-agent/src/agent/run.rs index fb0b90b1b..dd7b90fb6 100644 --- a/server-rs/crates/module-editor-agent/src/agent/run.rs +++ b/server-rs/crates/module-editor-agent/src/agent/run.rs @@ -84,15 +84,21 @@ where let max_turns = self.max_turns; Box::pin(async move { - let mut memory = agent - .memory - .unwrap_or_else(|| VecMemory::new(Vec::new()).into()); + let mut memory: Box> = match &agent.memory { + Some(m) => { + let msgs: Vec = m.get_memory().iter().cloned().collect(); + Box::new(VecMemory::new(msgs)) + } + None => Box::new(VecMemory::new(Vec::new())), + }; let mut prompt_result: Vec = Vec::new(); - let history = match agent.system_prompt { - Some(ref sp) => [sp].iter().chain(memory.iter().chain(message.into())), - None => memory.iter().chain(message.into()), - }; + let mut history: Vec = Vec::new(); + if let Some(ref sp) = agent.system_prompt { + history.push(sp.clone()); + } + history.extend(memory.get_memory().iter().cloned()); + history.push(message); for _ in 0..max_turns { // TODO perf issue for copy cost let text = agent.model.complete(&history).await?; @@ -121,8 +127,8 @@ where TextFlow::Continue => {} } } - memory.push(agent.model.build_assistant_message(&clean_text)); - prompt_result.push(Text(clean_text)); + memory.append_message(agent.model.build_assistant_message(&clean_text)); + prompt_result.push(Text(clean_text.clone())); if reply_has_end { return Ok(prompt_result); @@ -140,10 +146,7 @@ where .collect(); if tool_calls.is_empty() { - return Ok(PromptOutput { - text: clean_text, - tool_calls: vec![], - }); + return Ok(vec![Text(clean_text)]); } for (tc_id, tc) in tool_calls.iter().enumerate() { @@ -160,7 +163,7 @@ where let msg = agent .model .tool_result_message(&tc.name, "(skipped by hook)"); - memory.push(msg); + memory.append_message(msg); should_skip = true; break; } @@ -220,8 +223,8 @@ where ); let msg = agent.model.tool_result_message(&tc.name, &output_str); - memory.push(msg); - prompt_result.push(Tool(tc)) + memory.append_message(msg); + prompt_result.push(Tool(tc.clone())) } ToolOutcome::InternalError(failure) if failure.fatal => { return Err(PromptError::ToolError(failure.message)); @@ -231,14 +234,14 @@ where &tc.name, &format!("error: {}", failure.message), ); - memory.push(msg); + memory.append_message(msg); } } } } Err(_) => { // Not valid JSON — push as assistant message and continue to next _turn - memory.push(agent.model.build_assistant_message(&text)); + memory.append_message(agent.model.build_assistant_message(&text)); continue; } }