diff --git a/server-rs/crates/module-editor-agent/examples/llm_chat_agent.rs b/server-rs/crates/module-editor-agent/examples/llm_chat_agent.rs index 53eca84b6..b5c85d330 100644 --- a/server-rs/crates/module-editor-agent/examples/llm_chat_agent.rs +++ b/server-rs/crates/module-editor-agent/examples/llm_chat_agent.rs @@ -40,10 +40,6 @@ impl LlmApiAdaptor for LlmCompletionModel { } } -// --------------------------------------------------------------------------- -// 2. Tool validation hook — only allows known tools -// --------------------------------------------------------------------------- - struct ToolValidationHook { valid_names: Vec, } @@ -74,7 +70,7 @@ struct LlmChatAgentBuilder { tools: Vec>, hooks: Vec>, max_turns: usize, - memory_data: Option + Sync>>, + memory_data: Option>>, context: Option, } @@ -118,7 +114,7 @@ impl AgentBuilder for LlmChatAgentBuilder { self } - fn memory(mut self, memory: impl AgentMemory + Sync + 'static) -> Self { + fn memory(mut self, memory: impl AgentMemory + 'static) -> Self { self.memory_data = Some(Box::new(memory)); self } 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 4fe689909..5720ba6d9 100644 --- a/server-rs/crates/module-editor-agent/src/agent/agent.rs +++ b/server-rs/crates/module-editor-agent/src/agent/agent.rs @@ -12,7 +12,7 @@ pub struct Agent, Message> { pub hooks: Vec>, pub default_max_turns: usize, pub system_prompt: Option, - pub memory: Option + Sync>>, + pub memory: Option>>, pub context: Option, } @@ -43,7 +43,7 @@ where self } - pub fn memory(mut self, mem: impl AgentMemory + Sync + 'static) -> Self { + pub fn memory(mut self, mem: impl AgentMemory + 'static) -> Self { self.memory = Some(Box::new(mem)); self } @@ -62,50 +62,6 @@ where &self.tools } - pub fn call_tool<'s>( - &'s self, - name: &str, - args: serde_json::Value, - ) -> Pin> + Send + 's>> { - let tools: Vec<&Box> = self.tools.iter().collect(); - let hooks: Vec<&Box> = self.hooks.iter().collect(); - let name = name.to_string(); - let context = self.context.clone().unwrap_or_default(); - Box::pin(async move { - let result = { - let mut found = None; - for tool in tools { - if tool.tool_name() == name { - found = Some(tool.call_with_context(args, context).await); - break; - } - } - found.ok_or_else(|| format!("unknown tool: {name}"))? - }; - - match result.outcome { - ToolOutcome::Success => { - let mut json_output = result.output; - // Run after_tool_call hooks - for hook in &hooks { - match hook.after_tool_call(&name, &mut json_output) { - Flow::Stop => { - return Err("tool call output rejected by hook".to_string()); - } - Flow::Skip => { - json_output = serde_json::Value::Null; - break; - } - Flow::Continue => {} - } - } - serde_json::to_string(&json_output).map_err(|e| e.to_string()) - } - ToolOutcome::Failure(failure) if failure.fatal => Err(failure.message), - ToolOutcome::Failure(failure) => Ok(format!("error: {}", failure.message)), - } - }) - } pub fn prompt(&self, message: impl Into + Send) -> PromptRequest<'_, M, Message> where @@ -208,28 +164,4 @@ mod tests { } } } - - #[tokio::test] - async fn call_tool_returns_recoverable_failure_as_model_visible_error() { - let agent = Agent::new(TestModel).tool(FailingTool); - - let output = agent - .call_tool("fail", json!({"fatal": false})) - .await - .expect("recoverable tool failure should be model visible"); - - assert_eq!(output, "error: recoverable failure"); - } - - #[tokio::test] - async fn call_tool_propagates_fatal_failure() { - let agent = Agent::new(TestModel).tool(FailingTool); - - let error = agent - .call_tool("fail", json!({"fatal": true})) - .await - .expect_err("fatal tool failure should be propagated"); - - assert_eq!(error, "fatal failure"); - } } diff --git a/server-rs/crates/module-editor-agent/src/agent/agent_builder.rs b/server-rs/crates/module-editor-agent/src/agent/agent_builder.rs index 1d714740c..0316c2048 100644 --- a/server-rs/crates/module-editor-agent/src/agent/agent_builder.rs +++ b/server-rs/crates/module-editor-agent/src/agent/agent_builder.rs @@ -12,7 +12,7 @@ pub trait AgentBuilder> { fn tool(self, tool: impl Tool + Send + Sync + 'static) -> Self; fn add_hook(self, hook: impl Hook + 'static) -> Self; fn max_turns(self, n: usize) -> Self; - fn memory(self, memory: impl AgentMemory + Sync + 'static) -> Self; + fn memory(self, memory: impl AgentMemory + 'static) -> Self; fn context(self, context: serde_json::Value) -> Self; fn build(self) -> Agent; diff --git a/server-rs/crates/module-editor-agent/src/agent/memory.rs b/server-rs/crates/module-editor-agent/src/agent/memory.rs index 6b051bc02..3157d7387 100644 --- a/server-rs/crates/module-editor-agent/src/agent/memory.rs +++ b/server-rs/crates/module-editor-agent/src/agent/memory.rs @@ -1,4 +1,31 @@ -pub trait AgentMemory { +use serde::{Deserialize, Serialize}; + +pub trait AgentMemory: Send + Sync { fn get_memory(&self) -> &[Message]; fn append_message(&mut self, message: Message); } + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct VecMemory { + messages: Vec, +} + +impl VecMemory { + pub fn new(messages: Vec) -> Self { + Self { messages } + } + + pub fn into_inner(self) -> Vec { + self.messages + } +} + +impl AgentMemory for VecMemory { + fn get_memory(&self) -> &[Message] { + &self.messages + } + + fn append_message(&mut self, message: Message) { + self.messages.push(message); + } +} 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 e9741cf37..66e86cea7 100644 --- a/server-rs/crates/module-editor-agent/src/agent/run.rs +++ b/server-rs/crates/module-editor-agent/src/agent/run.rs @@ -165,17 +165,17 @@ where let args = tc.args.clone(); let context = self.context.clone().unwrap_or_default(); let fut = async move { - for tool in tools { - if tool.tool_name() == name { - return tool - .call_with_context(args, context.clone()) - .await; + match tools.iter().find(|tool| tool.tool_name() == name) { + Some(tool) => { + tool.call_with_context(args, context.clone()).await } + None => ToolExecutionResult::failed( + Value::Null, + ToolFailure::invalid_args(format!( + "unknown tool: {name}" + )), + ), } - ToolExecutionResult::failed( - Value::Null, - ToolFailure::invalid_args(format!("unknown tool: {name}")), - ) }; fut.await };