From 476a555f225f34572b1484bee370b2e57219c1c9 Mon Sep 17 00:00:00 2001 From: kvtodev Date: Thu, 9 Jul 2026 21:55:23 +0800 Subject: [PATCH] refactor: rename flow enums and improve tool interaction logic --- .../module-editor-agent/src/agent/hook.rs | 24 ++- .../module-editor-agent/src/agent/run.rs | 141 +++++++++--------- .../module-editor-agent/src/agent/tool.rs | 12 +- 3 files changed, 86 insertions(+), 91 deletions(-) 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 205c63e77..bc0bbc662 100644 --- a/server-rs/crates/module-editor-agent/src/agent/hook.rs +++ b/server-rs/crates/module-editor-agent/src/agent/hook.rs @@ -1,20 +1,14 @@ -use crate::agent::run::Flow; +use crate::agent::run::{TextFlow, ToolCallFlow}; use crate::agent::tool::ToolCall; pub trait Hook: Send + Sync { - fn before_tool_call(&self, tool_call: &ToolCall) -> Flow; - - /// Called after a tool call completes, before the JSON output is serialized to a string. - /// The `output` value can be modified in place. - /// Return `Flow::Stop` to abort the agent loop, `Flow::Skip` to discard this result, - /// or `Flow::Continue` to proceed normally. - fn after_tool_call(&self, _tool_name: &str, _output: &mut serde_json::Value) -> Flow { - Flow::Continue - } -} - -impl Hook for () { - fn before_tool_call(&self, _tool_call: &ToolCall) -> Flow { - Flow::Continue + fn on_text_reply(&self, text: &str) -> TextFlow { + TextFlow::Continue + } + 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 { + 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 5b254867c..15beacdf4 100644 --- a/server-rs/crates/module-editor-agent/src/agent/run.rs +++ b/server-rs/crates/module-editor-agent/src/agent/run.rs @@ -2,24 +2,29 @@ use crate::agent::agent::Agent; use crate::agent::agent::LlmApiAdaptor; use crate::agent::error::PromptError; use crate::agent::hook::Hook; +use crate::agent::memory::{AgentMemory, VecMemory}; use crate::agent::tool::{ToolCall, ToolDyn, ToolExecutionResult, ToolFailure, ToolOutcome}; use serde::Deserialize; use serde_json::Value; use std::pin::Pin; +use crate::agent::run::PromptOutput::{Text, Tool}; #[derive(Debug, Clone)] -pub struct PromptOutput { - pub text: String, - pub tool_calls: Vec, +pub enum PromptOutput { + Text(String), + Tool(ToolCall), } -/// Flow control for tool execution within the agent loop. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum Flow { +pub enum TextFlow { + Continue, + Stop, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ToolCallFlow { Continue, - /// Skip this tool call but continue processing other responses. Skip, - /// Stop the agent execution entirely. Stop, } @@ -69,45 +74,35 @@ where M: LlmApiAdaptor + Send + Sync + 'a, Message: Send + Sync + Clone + 'a, { - type Output = Result; + type Output = Result, PromptError>; type IntoFuture = Pin + Send + 'a>>; fn into_future(self) -> Self::IntoFuture { let agent = self.agent; let message = self.message; - let per_prompt_system_prompt = self.system_prompt; let extra_hooks = self.hooks; let max_turns = self.max_turns; Box::pin(async move { - let mut memory: Vec = Vec::new(); + let mut memory = agent + .memory + .unwrap_or_else(|| VecMemory::new(Vec::new()).into()); + let mut prompt_result: Vec = Vec::new(); - // Pre-populate from agent's memory backend if available - if let Some(ref mem_backend) = agent.memory { - memory.extend(mem_backend.get_memory().iter().cloned()); - } + let history = match agent.system_prompt { + Some(ref sp) => [sp].iter().chain(memory.iter().chain(message.into())), + None => memory.iter().chain(message.into()), + }; + for _ in 0..max_turns { + // TODO perf issue for copy cost + let text = agent.model.complete(&history).await?; - // Effective system prompt: per-prompt override > agent default - match per_prompt_system_prompt { - Some(sp) => memory.push(sp), - None => { - if let Some(ref sp) = agent.system_prompt { - memory.push(sp.clone()); - } - } - } - memory.push(message); - - for turn in 0..max_turns { - // Delegate completion to the model (rig's CompletionModel pattern) - let text = agent.model.complete(&memory).await?; - - // Try to parse the LLM reply as JSON (handle markdown fences) + // Try to parse the LLM reply as JSON (handle Markdown fences) let cleaned = clean_json_response(&text); match serde_json::from_str::(&cleaned) { Ok(json_resp) => { - // Check for marker — stop the turn immediately + // Check for marker — stop the _turn immediately let reply_has_end = json_resp.reply_text.contains(""); let clean_text = json_resp .reply_text @@ -115,7 +110,19 @@ where .trim() .to_string(); + // Run on_text_reply hooks + for hook in agent.hooks.iter().chain(extra_hooks.iter()) { + match hook.on_text_reply(&clean_text) { + TextFlow::Stop => { + return Err(PromptError::ToolError( + "text reply rejected by hook".to_string(), + )); + } + TextFlow::Continue => {} + } + } memory.push(agent.model.build_assistant_message(&clean_text)); + prompt_result.push(Text(clean_text)); if reply_has_end { return Ok(PromptOutput { @@ -129,7 +136,7 @@ where .into_iter() .enumerate() .map(|(idx, tc)| ToolCall { - id: format!("call_{turn}_{idx}"), + id: format!("{idx}"), name: tc.tool_name, args: tc.args, }) @@ -143,20 +150,28 @@ where } for (tc_id, tc) in tool_calls.iter().enumerate() { - match run_hooks(&agent.hooks, &extra_hooks, tc) { - Flow::Stop => { - return Err(PromptError::ToolError( - "tool call rejected by hook".to_string(), - )); + // inline run_hooks: before_tool_call hook + let mut should_skip = false; + for hook in agent.hooks.iter().chain(extra_hooks.iter()) { + match hook.before_tool_call(tc) { + ToolCallFlow::Stop => { + return Err(PromptError::ToolError( + "tool call rejected by hook".to_string(), + )); + } + ToolCallFlow::Skip => { + let msg = agent + .model + .tool_result_message(&tc.name, "(skipped by hook)"); + memory.push(msg); + should_skip = true; + break; + } + ToolCallFlow::Continue => {} } - Flow::Skip => { - let msg = agent - .model - .tool_result_message(&tc.name, "(skipped by hook)"); - memory.push(msg); - continue; - } - Flow::Continue => {} + } + if should_skip { + continue; } let result = { @@ -181,21 +196,21 @@ where }; match result.outcome { - ToolOutcome::Success => { + ToolOutcome::InternalOk => { let mut json_output = result.output; // Run after_tool_call hooks to allow output modification for hook in agent.hooks.iter().chain(extra_hooks.iter()) { match hook.after_tool_call(&tc.name, &mut json_output) { - Flow::Stop => { + ToolCallFlow::Stop => { return Err(PromptError::ToolError( - "tool call output caused this turn stopped by hook".to_string(), + "tool call output caused this _turn stopped by hook".to_string(), )); } - Flow::Skip => { + ToolCallFlow::Skip => { json_output = serde_json::json!({"message":"tool call is ignored by hook"}); break; } - Flow::Continue => {} + ToolCallFlow::Continue => {} } } let arg_json = serde_json::to_string(&tc.args) @@ -209,11 +224,12 @@ where let msg = agent.model.tool_result_message(&tc.name, &output_str); memory.push(msg); + prompt_result.push(Tool(tc)) } - ToolOutcome::Failure(failure) if failure.fatal => { - return Err(PromptError::InternalError(failure.message)); + ToolOutcome::InternalError(failure) if failure.fatal => { + return Err(PromptError::ToolError(failure.message)); } - ToolOutcome::Failure(failure) => { + ToolOutcome::InternalError(failure) => { let msg = agent.model.tool_result_message( &tc.name, &format!("error: {}", failure.message), @@ -224,7 +240,7 @@ where } } Err(_) => { - // Not valid JSON — push as assistant message and continue to next turn + // Not valid JSON — push as assistant message and continue to next _turn memory.push(agent.model.build_assistant_message(&text)); continue; } @@ -236,21 +252,6 @@ where } } -fn run_hooks( - agent_hooks: &[Box], - extra_hooks: &[Box], - tool_call: &ToolCall, -) -> Flow { - for hook in agent_hooks.iter().chain(extra_hooks.iter()) { - match hook.before_tool_call(tool_call) { - Flow::Stop => return Flow::Stop, - Flow::Skip => return Flow::Skip, - Flow::Continue => {} - } - } - Flow::Continue -} - #[derive(Deserialize)] struct LlmJsonResponse { reply_text: String, @@ -262,7 +263,7 @@ struct LlmJsonResponse { struct LlmToolCallRequest { tool_name: String, #[serde(default)] - args: serde_json::Value, + args: Value, } pub(crate) fn clean_json_response(text: &str) -> String { diff --git a/server-rs/crates/module-editor-agent/src/agent/tool.rs b/server-rs/crates/module-editor-agent/src/agent/tool.rs index fd893fcf8..f671475b6 100644 --- a/server-rs/crates/module-editor-agent/src/agent/tool.rs +++ b/server-rs/crates/module-editor-agent/src/agent/tool.rs @@ -84,8 +84,8 @@ impl ToolFailure { #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub enum ToolOutcome { - Success, - Failure(ToolFailure), + InternalOk, + InternalError(ToolFailure), } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -98,21 +98,21 @@ impl ToolExecutionResult { pub fn success(output: serde_json::Value) -> Self { Self { output, - outcome: ToolOutcome::Success, + outcome: ToolOutcome::InternalOk, } } pub fn failed(output: serde_json::Value, failure: ToolFailure) -> Self { Self { output, - outcome: ToolOutcome::Failure(failure), + outcome: ToolOutcome::InternalError(failure), } } pub fn failure(&self) -> Option<&ToolFailure> { match &self.outcome { - ToolOutcome::Success => None, - ToolOutcome::Failure(failure) => Some(failure), + ToolOutcome::InternalOk => None, + ToolOutcome::InternalError(failure) => Some(failure), } }