refactor: simplify agent memory handling and improve message flow logic

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<M: LlmApiAdaptor<Message>, Message> {
     pub model: M,

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
     }
 }

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<dyn AgentMemory<Message>> = match &agent.memory {
+                Some(m) => {
+                    let msgs: Vec<Message> = m.get_memory().iter().cloned().collect();
+                    Box::new(VecMemory::new(msgs))
+                }
+                None => Box::new(VecMemory::new(Vec::new())),
+            };
             let mut prompt_result: Vec<PromptOutput> = 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<Message> = 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;
                     }
                 }