From 1daae3ec4551c18345fdd3a915cf31841f9d4602 Mon Sep 17 00:00:00 2001 From: YayoiNanoka <1159066485@qq.com> Date: Wed, 1 Jul 2026 20:20:56 +0800 Subject: [PATCH 1/7] feat: add agent loader abstraction --- src/agent/agent.rs | 56 +++++++++++++++++++- src/agent/custom_agent.rs | 14 ++--- src/agent/loader.rs | 88 ++++++++++++++++++++++++++++++++ src/agent/mod.rs | 4 +- src/llm_agent/agent.rs | 5 +- src/llm_agent/config.rs | 4 +- src/runner/runner.rs | 25 ++++----- src/workflow/loop_agent.rs | 22 +++++--- src/workflow/sequential_agent.rs | 24 ++++++--- tests/agent_loader.rs | 54 ++++++++++++++++++++ tests/multi_agent.rs | 4 +- 11 files changed, 260 insertions(+), 40 deletions(-) create mode 100644 src/agent/loader.rs create mode 100644 tests/agent_loader.rs diff --git a/src/agent/agent.rs b/src/agent/agent.rs index 7297969..c0b15c9 100644 --- a/src/agent/agent.rs +++ b/src/agent/agent.rs @@ -1,3 +1,5 @@ +use std::{collections::HashSet, sync::Arc}; + use crate::{Result, event::Event}; use super::{AsyncEventStream, EventStream, InvocationContext, event_stream_from_result}; @@ -7,7 +9,7 @@ pub trait Agent: Send + Sync { fn description(&self) -> &str; - fn sub_agents(&self) -> &[Box] { + fn sub_agents(&self) -> &[Arc] { &[] } @@ -36,3 +38,55 @@ pub fn find_agent<'a>(agent: &'a dyn Agent, name: &str) -> Option<&'a dyn Agent> None } + +pub trait IntoAgentArc { + fn into_agent_arc(self) -> Arc; +} + +impl IntoAgentArc for Arc { + fn into_agent_arc(self) -> Arc { + self + } +} + +impl IntoAgentArc for Arc +where + T: Agent + 'static, +{ + fn into_agent_arc(self) -> Arc { + self + } +} + +impl IntoAgentArc for Box { + fn into_agent_arc(self) -> Arc { + Arc::from(self) + } +} + +impl IntoAgentArc for Box +where + T: Agent + 'static, +{ + fn into_agent_arc(self) -> Arc { + Arc::new(*self) + } +} + +pub fn validate_agent_names(agent: &dyn Agent) -> Result<()> { + let mut seen = HashSet::new(); + validate_agent_names_inner(agent, &mut seen) +} + +fn validate_agent_names_inner(agent: &dyn Agent, seen: &mut HashSet) -> Result<()> { + let name = agent.name().to_string(); + if !seen.insert(name.clone()) { + return Err(crate::Error::msg(format!("duplicate agent name: {name}"))); + } + + for sub_agent in agent.sub_agents() { + validate_agent_names_inner(sub_agent.as_ref(), seen)?; + } + + Ok(()) +} diff --git a/src/agent/custom_agent.rs b/src/agent/custom_agent.rs index 277867c..f54f60e 100644 --- a/src/agent/custom_agent.rs +++ b/src/agent/custom_agent.rs @@ -1,12 +1,14 @@ +use std::sync::Arc; + use crate::{Result, event::Event}; -use super::{Agent, InvocationContext}; +use super::{Agent, IntoAgentArc, InvocationContext}; pub struct CustomAgent { name: String, description: String, handler: Box Result> + Send + Sync>, - sub_agents: Vec>, + sub_agents: Vec>, } impl CustomAgent { @@ -23,12 +25,12 @@ impl CustomAgent { } } - pub fn with_sub_agent(mut self, sub_agent: Box) -> Self { - self.sub_agents.push(sub_agent); + pub fn with_sub_agent(mut self, sub_agent: impl IntoAgentArc) -> Self { + self.sub_agents.push(sub_agent.into_agent_arc()); self } - pub fn with_sub_agents(mut self, sub_agents: Vec>) -> Self { + pub fn with_sub_agents(mut self, sub_agents: Vec>) -> Self { self.sub_agents = sub_agents; self } @@ -43,7 +45,7 @@ impl Agent for CustomAgent { &self.description } - fn sub_agents(&self) -> &[Box] { + fn sub_agents(&self) -> &[Arc] { &self.sub_agents } diff --git a/src/agent/loader.rs b/src/agent/loader.rs new file mode 100644 index 0000000..15c42c0 --- /dev/null +++ b/src/agent/loader.rs @@ -0,0 +1,88 @@ +use std::{collections::HashMap, sync::Arc}; + +use crate::{Error, Result}; + +use super::{Agent, validate_agent_names}; + +pub trait AgentLoader: Send + Sync { + fn list_agents(&self) -> Result>; + + fn load_agent(&self, name: &str) -> Result>; + + fn root_agent(&self) -> Result>; +} + +pub struct SingleAgentLoader { + root_agent: Arc, +} + +impl SingleAgentLoader { + pub fn new(root_agent: Arc) -> Self { + Self { root_agent } + } +} + +impl AgentLoader for SingleAgentLoader { + fn list_agents(&self) -> Result> { + Ok(vec![self.root_agent.name().to_string()]) + } + + fn load_agent(&self, name: &str) -> Result> { + if self.root_agent.name() == name { + Ok(self.root_agent.clone()) + } else { + Err(Error::msg(format!("agent not found: {name}"))) + } + } + + fn root_agent(&self) -> Result> { + Ok(self.root_agent.clone()) + } +} + +pub struct MultiAgentLoader { + root_agent: Arc, + agents: HashMap>, +} + +impl MultiAgentLoader { + pub fn new(root_agent: Arc, agents: Vec>) -> Result { + validate_agent_names(root_agent.as_ref())?; + + let mut by_name = HashMap::new(); + for agent in agents { + let name = agent.name().to_string(); + if by_name.insert(name.clone(), agent).is_some() { + return Err(Error::msg(format!("duplicate agent name: {name}"))); + } + } + + if !by_name.contains_key(root_agent.name()) { + by_name.insert(root_agent.name().to_string(), root_agent.clone()); + } + + Ok(Self { + root_agent, + agents: by_name, + }) + } +} + +impl AgentLoader for MultiAgentLoader { + fn list_agents(&self) -> Result> { + let mut names = self.agents.keys().cloned().collect::>(); + names.sort(); + Ok(names) + } + + fn load_agent(&self, name: &str) -> Result> { + self.agents + .get(name) + .cloned() + .ok_or_else(|| Error::msg(format!("agent not found: {name}"))) + } + + fn root_agent(&self) -> Result> { + Ok(self.root_agent.clone()) + } +} diff --git a/src/agent/mod.rs b/src/agent/mod.rs index 0b97c13..b921214 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -1,11 +1,13 @@ mod agent; mod context; mod custom_agent; +mod loader; mod module; mod stream; -pub use agent::{Agent, find_agent}; +pub use agent::{Agent, IntoAgentArc, find_agent, validate_agent_names}; pub use context::InvocationContext; pub use custom_agent::CustomAgent; +pub use loader::{AgentLoader, MultiAgentLoader, SingleAgentLoader}; pub use module::MODULE; pub use stream::{AsyncEventStream, EventStream, event_stream_from_result}; diff --git a/src/llm_agent/agent.rs b/src/llm_agent/agent.rs index 0eb3637..a1b18cf 100644 --- a/src/llm_agent/agent.rs +++ b/src/llm_agent/agent.rs @@ -1,5 +1,6 @@ use std::{ collections::HashMap, + sync::Arc, sync::atomic::{AtomicU64, Ordering}, time::SystemTime, }; @@ -31,7 +32,7 @@ pub struct LlmAgent { system_prompt: String, model: Box, tools: Vec>, - sub_agents: Vec>, + sub_agents: Vec>, max_steps: usize, include_contents: IncludeContents, callbacks: LlmAgentCallbacks, @@ -433,7 +434,7 @@ impl Agent for LlmAgent { &self.description } - fn sub_agents(&self) -> &[Box] { + fn sub_agents(&self) -> &[Arc] { &self.sub_agents } diff --git a/src/llm_agent/config.rs b/src/llm_agent/config.rs index 6a033fa..66c2ed9 100644 --- a/src/llm_agent/config.rs +++ b/src/llm_agent/config.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use crate::{agent::Agent, model::Model, tool::Tool}; use super::{IncludeContents, LlmAgentCallbacks}; @@ -8,7 +10,7 @@ pub struct LlmAgentConfig { pub system_prompt: String, pub model: Box, pub tools: Vec>, - pub sub_agents: Vec>, + pub sub_agents: Vec>, pub max_steps: usize, pub include_contents: IncludeContents, pub callbacks: LlmAgentCallbacks, diff --git a/src/runner/runner.rs b/src/runner/runner.rs index efe77d2..baec0b2 100644 --- a/src/runner/runner.rs +++ b/src/runner/runner.rs @@ -9,7 +9,8 @@ use futures_util::StreamExt; use crate::{ Error, Result, agent::{ - Agent, AsyncEventStream, EventStream, InvocationContext, event_stream_from_result, + Agent, AsyncEventStream, EventStream, IntoAgentArc, InvocationContext, + event_stream_from_result, find_agent, }, artifact::{ArtifactFacade, ArtifactService, InMemoryArtifactService}, @@ -27,7 +28,7 @@ static NEXT_ID: AtomicU64 = AtomicU64::new(1); pub struct Runner { pub app_name: String, - pub root_agent: Box, + pub root_agent: Arc, pub session_service: S, pub auto_create_session: bool, pub artifact_service: Arc, @@ -39,13 +40,13 @@ pub struct Runner { impl Runner { pub fn new( app_name: impl Into, - root_agent: Box, + root_agent: impl IntoAgentArc, session_service: S, auto_create_session: bool, ) -> Self { Self { app_name: app_name.into(), - root_agent, + root_agent: root_agent.into_agent_arc(), session_service, auto_create_session, artifact_service: Arc::new(InMemoryArtifactService::new()), @@ -57,14 +58,14 @@ impl Runner { pub fn with_artifact_service( app_name: impl Into, - root_agent: Box, + root_agent: impl IntoAgentArc, session_service: S, auto_create_session: bool, artifact_service: Arc, ) -> Self { Self { app_name: app_name.into(), - root_agent, + root_agent: root_agent.into_agent_arc(), session_service, auto_create_session, artifact_service, @@ -76,14 +77,14 @@ impl Runner { pub fn with_memory_service( app_name: impl Into, - root_agent: Box, + root_agent: impl IntoAgentArc, session_service: S, auto_create_session: bool, memory_service: Arc, ) -> Self { Self { app_name: app_name.into(), - root_agent, + root_agent: root_agent.into_agent_arc(), session_service, auto_create_session, artifact_service: Arc::new(InMemoryArtifactService::new()), @@ -95,14 +96,14 @@ impl Runner { pub fn with_plugin_manager( app_name: impl Into, - root_agent: Box, + root_agent: impl IntoAgentArc, session_service: S, auto_create_session: bool, plugin_manager: PluginManager, ) -> Self { Self { app_name: app_name.into(), - root_agent, + root_agent: root_agent.into_agent_arc(), session_service, auto_create_session, artifact_service: Arc::new(InMemoryArtifactService::new()), @@ -114,7 +115,7 @@ impl Runner { pub fn with_services( app_name: impl Into, - root_agent: Box, + root_agent: impl IntoAgentArc, session_service: S, auto_create_session: bool, artifact_service: Arc, @@ -122,7 +123,7 @@ impl Runner { ) -> Self { Self { app_name: app_name.into(), - root_agent, + root_agent: root_agent.into_agent_arc(), session_service, auto_create_session, artifact_service, diff --git a/src/workflow/loop_agent.rs b/src/workflow/loop_agent.rs index ee8e37e..28dcebf 100644 --- a/src/workflow/loop_agent.rs +++ b/src/workflow/loop_agent.rs @@ -1,6 +1,8 @@ +use std::sync::Arc; + use crate::{ Result, - agent::{Agent, InvocationContext}, + agent::{Agent, IntoAgentArc, InvocationContext}, event::Event, }; @@ -10,21 +12,27 @@ pub struct LoopAgent { name: String, description: String, max_iterations: usize, - sub_agents: Vec>, + sub_agents: Vec>, } impl LoopAgent { - pub fn new( + pub fn new( name: impl Into, description: impl Into, max_iterations: usize, - sub_agents: Vec>, - ) -> Self { + sub_agents: Vec, + ) -> Self + where + T: IntoAgentArc, + { Self { name: name.into(), description: description.into(), max_iterations, - sub_agents, + sub_agents: sub_agents + .into_iter() + .map(IntoAgentArc::into_agent_arc) + .collect(), } } } @@ -38,7 +46,7 @@ impl Agent for LoopAgent { &self.description } - fn sub_agents(&self) -> &[Box] { + fn sub_agents(&self) -> &[Arc] { &self.sub_agents } diff --git a/src/workflow/sequential_agent.rs b/src/workflow/sequential_agent.rs index 5f22f5a..14e5ccf 100644 --- a/src/workflow/sequential_agent.rs +++ b/src/workflow/sequential_agent.rs @@ -1,6 +1,8 @@ +use std::sync::Arc; + use crate::{ Result, - agent::{Agent, InvocationContext}, + agent::{Agent, IntoAgentArc, InvocationContext}, event::Event, session::Session, }; @@ -8,19 +10,25 @@ use crate::{ pub struct SequentialAgent { name: String, description: String, - sub_agents: Vec>, + sub_agents: Vec>, } impl SequentialAgent { - pub fn new( + pub fn new( name: impl Into, description: impl Into, - sub_agents: Vec>, - ) -> Self { + sub_agents: Vec, + ) -> Self + where + T: IntoAgentArc, + { Self { name: name.into(), description: description.into(), - sub_agents, + sub_agents: sub_agents + .into_iter() + .map(IntoAgentArc::into_agent_arc) + .collect(), } } } @@ -34,7 +42,7 @@ impl Agent for SequentialAgent { &self.description } - fn sub_agents(&self) -> &[Box] { + fn sub_agents(&self) -> &[Arc] { &self.sub_agents } @@ -44,7 +52,7 @@ impl Agent for SequentialAgent { } pub(crate) fn run_sub_agents_once( - sub_agents: &[Box], + sub_agents: &[Arc], ctx: InvocationContext, ) -> Result> { let mut session_snapshot = ctx.session_snapshot.clone(); diff --git a/tests/agent_loader.rs b/tests/agent_loader.rs new file mode 100644 index 0000000..97f2d2f --- /dev/null +++ b/tests/agent_loader.rs @@ -0,0 +1,54 @@ +use std::sync::Arc; + +use r_agent::{ + agent::{AgentLoader, CustomAgent, MultiAgentLoader, SingleAgentLoader, validate_agent_names}, + runner::Runner, + session::InMemorySessionService, +}; + +#[test] +fn single_agent_loader_returns_root_agent() { + let root = Arc::new(CustomAgent::new("root", "Root agent.", |_| Ok(Vec::new()))); + let loader = SingleAgentLoader::new(root.clone()); + + assert_eq!(loader.list_agents().unwrap(), vec!["root".to_string()]); + assert_eq!(loader.root_agent().unwrap().name(), "root"); + assert_eq!(loader.load_agent("root").unwrap().name(), "root"); +} + +#[test] +fn multi_agent_loader_loads_agents_by_name() { + let root = Arc::new(CustomAgent::new("root", "Root agent.", |_| Ok(Vec::new()))); + let helper = Arc::new(CustomAgent::new("helper", "Helper agent.", |_| { + Ok(Vec::new()) + })); + let loader = MultiAgentLoader::new(root.clone(), vec![root, helper]).unwrap(); + + assert_eq!( + loader.list_agents().unwrap(), + vec!["helper".to_string(), "root".to_string()] + ); + assert_eq!(loader.root_agent().unwrap().name(), "root"); + assert_eq!(loader.load_agent("helper").unwrap().name(), "helper"); +} + +#[test] +fn duplicate_agent_names_are_rejected() { + let child_a = CustomAgent::new("child", "Child A.", |_| Ok(Vec::new())); + let child_b = CustomAgent::new("child", "Child B.", |_| Ok(Vec::new())); + let root = CustomAgent::new("root", "Root agent.", |_| Ok(Vec::new())) + .with_sub_agent(Box::new(child_a)) + .with_sub_agent(Box::new(child_b)); + + let err = validate_agent_names(&root).unwrap_err(); + + assert_eq!(err.to_string(), "duplicate agent name: child"); +} + +#[test] +fn runner_accepts_arc_agent() { + let root = Arc::new(CustomAgent::new("root", "Root agent.", |_| Ok(Vec::new()))); + let runner = Runner::new("app", root, InMemorySessionService::new(), true); + + assert_eq!(runner.root_agent.name(), "root"); +} diff --git a/tests/multi_agent.rs b/tests/multi_agent.rs index 98cfae6..36c9406 100644 --- a/tests/multi_agent.rs +++ b/tests/multi_agent.rs @@ -1,4 +1,4 @@ -use std::time::SystemTime; +use std::{sync::Arc, time::SystemTime}; use serde_json::json; @@ -153,7 +153,7 @@ fn llm_agent_can_transfer_to_sub_agent_through_transfer_tool() { system_prompt: "Route to a sub agent when useful.".to_string(), model: Box::new(model), tools: Vec::new(), - sub_agents: vec![Box::new(child)], + sub_agents: vec![Arc::new(child)], max_steps: 4, include_contents: IncludeContents::Default, callbacks: LlmAgentCallbacks::default(), From 37f295926a0b2324c32e38dfe77b96aad01271d3 Mon Sep 17 00:00:00 2001 From: YayoiNanoka <1159066485@qq.com> Date: Wed, 1 Jul 2026 20:25:53 +0800 Subject: [PATCH 2/7] feat: add model factory and generation config --- examples/openai_compatible_demo.rs | 1 + src/llm_agent/agent.rs | 3 +++ src/llm_agent/config.rs | 7 +++++- src/main.rs | 1 + src/model/factory.rs | 29 ++++++++++++++++++++++ src/model/generate_config.rs | 18 ++++++++++++++ src/model/llm.rs | 5 +++- src/model/mod.rs | 4 ++++ src/model/openai_compatible/convert.rs | 15 ++++++++++++ tests/artifact.rs | 1 + tests/async_streaming.rs | 2 ++ tests/callbacks.rs | 8 +++++++ tests/content_event.rs | 1 + tests/instruction_history.rs | 1 + tests/llm_agent.rs | 3 +++ tests/long_running_tool.rs | 1 + tests/memory.rs | 1 + tests/model_factory.rs | 33 ++++++++++++++++++++++++++ tests/multi_agent.rs | 1 + tests/openai_compatible_model.rs | 28 +++++++++++++++++++++- tests/openai_compatible_real.rs | 1 + tests/streaming.rs | 1 + tests/telemetry.rs | 1 + tests/tool_confirmation.rs | 1 + tests/tool_loop.rs | 2 ++ 25 files changed, 166 insertions(+), 3 deletions(-) create mode 100644 src/model/factory.rs create mode 100644 src/model/generate_config.rs create mode 100644 tests/model_factory.rs diff --git a/examples/openai_compatible_demo.rs b/examples/openai_compatible_demo.rs index 4f3f8a9..e9d3cd3 100644 --- a/examples/openai_compatible_demo.rs +++ b/examples/openai_compatible_demo.rs @@ -60,6 +60,7 @@ fn main() -> Result<()> { description: "A tool-using assistant.".to_string(), system_prompt: "Use tools when useful, then answer the user.".to_string(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(GetTimeTool)], sub_agents: Vec::new(), max_steps: 4, diff --git a/src/llm_agent/agent.rs b/src/llm_agent/agent.rs index a1b18cf..5cc003d 100644 --- a/src/llm_agent/agent.rs +++ b/src/llm_agent/agent.rs @@ -31,6 +31,7 @@ pub struct LlmAgent { description: String, system_prompt: String, model: Box, + generate_content_config: Option, tools: Vec>, sub_agents: Vec>, max_steps: usize, @@ -65,6 +66,7 @@ impl LlmAgent { description: config.description, system_prompt: config.system_prompt, model: config.model, + generate_content_config: config.generate_content_config, tools, sub_agents: config.sub_agents, max_steps: config.max_steps, @@ -83,6 +85,7 @@ impl LlmAgent { system_prompt: self.render_system_prompt(session_snapshot)?, contents: self.build_contents(ctx, session_snapshot), tools: self.tool_declarations(), + generate_content_config: self.generate_content_config.clone(), }) } diff --git a/src/llm_agent/config.rs b/src/llm_agent/config.rs index 66c2ed9..d7ccb2d 100644 --- a/src/llm_agent/config.rs +++ b/src/llm_agent/config.rs @@ -1,6 +1,10 @@ use std::sync::Arc; -use crate::{agent::Agent, model::Model, tool::Tool}; +use crate::{ + agent::Agent, + model::{GenerateContentConfig, Model}, + tool::Tool, +}; use super::{IncludeContents, LlmAgentCallbacks}; @@ -9,6 +13,7 @@ pub struct LlmAgentConfig { pub description: String, pub system_prompt: String, pub model: Box, + pub generate_content_config: Option, pub tools: Vec>, pub sub_agents: Vec>, pub max_steps: usize, diff --git a/src/main.rs b/src/main.rs index 2ba4cec..d4eeb31 100644 --- a/src/main.rs +++ b/src/main.rs @@ -60,6 +60,7 @@ fn main() -> Result<()> { description: "A DeepSeek-backed tool-using assistant.".to_string(), system_prompt: "Use tools when useful, then answer the user.".to_string(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(GetTimeTool)], sub_agents: Vec::new(), max_steps: 4, diff --git a/src/model/factory.rs b/src/model/factory.rs new file mode 100644 index 0000000..5a4f8a5 --- /dev/null +++ b/src/model/factory.rs @@ -0,0 +1,29 @@ +use crate::{Error, Result}; + +use super::{Model, ModelConfig, OpenAiCompatibleModel}; + +pub trait ModelFactory: Send + Sync { + fn create(&self, config: ModelConfig) -> Result>; +} + +#[derive(Debug, Default, Clone)] +pub struct DefaultModelFactory; + +impl DefaultModelFactory { + pub fn new() -> Self { + Self + } +} + +impl ModelFactory for DefaultModelFactory { + fn create(&self, config: ModelConfig) -> Result> { + match config.protocol.as_str() { + "openai-compatible" | "openai-chat-completions" => { + Ok(Box::new(OpenAiCompatibleModel::new(config))) + } + protocol => Err(Error::msg(format!( + "unsupported model protocol: {protocol}" + ))), + } + } +} diff --git a/src/model/generate_config.rs b/src/model/generate_config.rs new file mode 100644 index 0000000..7afef2a --- /dev/null +++ b/src/model/generate_config.rs @@ -0,0 +1,18 @@ +#[derive(Debug, Clone, PartialEq)] +pub struct GenerateContentConfig { + pub temperature: Option, + pub top_p: Option, + pub max_output_tokens: Option, + pub stop_sequences: Vec, +} + +impl Default for GenerateContentConfig { + fn default() -> Self { + Self { + temperature: None, + top_p: None, + max_output_tokens: None, + stop_sequences: Vec::new(), + } + } +} diff --git a/src/model/llm.rs b/src/model/llm.rs index e330bd6..055b8fe 100644 --- a/src/model/llm.rs +++ b/src/model/llm.rs @@ -4,16 +4,19 @@ use futures_util::Stream; use crate::{content::Content, tool::ToolDeclaration}; +use super::GenerateContentConfig; + pub type LlmResponseStream<'a> = Box> + 'a>; pub type AsyncLlmResponseStream<'a> = Pin> + Send + 'a>>; -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq)] pub struct LlmRequest { pub model: String, pub system_prompt: Option, pub contents: Vec, pub tools: Vec, + pub generate_content_config: Option, } #[derive(Debug, Clone, PartialEq, Eq)] diff --git a/src/model/mod.rs b/src/model/mod.rs index edc7af0..0041865 100644 --- a/src/model/mod.rs +++ b/src/model/mod.rs @@ -1,12 +1,16 @@ mod config; +mod factory; mod fake_model; +mod generate_config; mod llm; mod model; mod module; pub mod openai_compatible; pub use config::ModelConfig; +pub use factory::{DefaultModelFactory, ModelFactory}; pub use fake_model::FakeModel; +pub use generate_config::GenerateContentConfig; pub use llm::{AsyncLlmResponseStream, LlmRequest, LlmResponse, ModelUsage}; pub use llm::{LlmResponseStream, llm_response_stream_from_result}; pub use model::Model; diff --git a/src/model/openai_compatible/convert.rs b/src/model/openai_compatible/convert.rs index c3512fb..781714e 100644 --- a/src/model/openai_compatible/convert.rs +++ b/src/model/openai_compatible/convert.rs @@ -44,6 +44,21 @@ pub fn request_json(request: &LlmRequest) -> Value { ); } + if let Some(config) = &request.generate_content_config { + if let Some(temperature) = config.temperature { + body["temperature"] = json!(temperature); + } + if let Some(top_p) = config.top_p { + body["top_p"] = json!(top_p); + } + if let Some(max_output_tokens) = config.max_output_tokens { + body["max_tokens"] = json!(max_output_tokens); + } + if !config.stop_sequences.is_empty() { + body["stop"] = json!(config.stop_sequences); + } + } + body } diff --git a/tests/artifact.rs b/tests/artifact.rs index 74ef367..d606645 100644 --- a/tests/artifact.rs +++ b/tests/artifact.rs @@ -132,6 +132,7 @@ fn tool_context_save_artifact_records_event_delta() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(tool)], sub_agents: Vec::new(), max_steps: 4, diff --git a/tests/async_streaming.rs b/tests/async_streaming.rs index 2756be3..05dfb74 100644 --- a/tests/async_streaming.rs +++ b/tests/async_streaming.rs @@ -98,6 +98,7 @@ async fn llm_agent_async_stream_yields_model_partials_before_final() { description: "llm".to_string(), system_prompt: String::new(), model: Box::new(model), + generate_content_config: None, tools: vec![], sub_agents: vec![], max_steps: 4, @@ -181,6 +182,7 @@ async fn async_tool_loop_waits_for_final_tool_call() { description: "llm".to_string(), system_prompt: String::new(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(tool)], sub_agents: vec![], max_steps: 4, diff --git a/tests/callbacks.rs b/tests/callbacks.rs index 01a68b5..5a4614a 100644 --- a/tests/callbacks.rs +++ b/tests/callbacks.rs @@ -128,6 +128,7 @@ fn before_agent_callback_can_return_content_and_skip_agent_flow() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: Vec::new(), sub_agents: Vec::new(), max_steps: 1, @@ -164,6 +165,7 @@ fn after_agent_callback_can_append_agent_event() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: Vec::new(), sub_agents: Vec::new(), max_steps: 1, @@ -208,6 +210,7 @@ fn before_model_callback_can_return_cached_response_and_skip_model() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: Vec::new(), sub_agents: Vec::new(), max_steps: 1, @@ -248,6 +251,7 @@ fn after_model_callback_can_replace_model_response() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: Vec::new(), sub_agents: Vec::new(), max_steps: 1, @@ -279,6 +283,7 @@ fn on_model_error_callback_can_recover_with_response() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(FailingModel), + generate_content_config: None, tools: Vec::new(), sub_agents: Vec::new(), max_steps: 1, @@ -323,6 +328,7 @@ fn before_tool_callback_can_return_result_and_skip_tool() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(CountingWeatherTool { calls: tool_calls.clone(), })], @@ -380,6 +386,7 @@ fn after_tool_callback_can_replace_tool_result() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(CountingWeatherTool { calls: tool_calls.clone(), })], @@ -435,6 +442,7 @@ fn on_tool_error_callback_can_recover_with_result() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(FailingTool)], sub_agents: Vec::new(), max_steps: 4, diff --git a/tests/content_event.rs b/tests/content_event.rs index fde1dd7..1380dc2 100644 --- a/tests/content_event.rs +++ b/tests/content_event.rs @@ -71,6 +71,7 @@ fn llm_request_and_response_use_content_as_the_model_boundary() { system_prompt: Some("You are helpful.".to_string()), contents: vec![user_content.clone()], tools: Vec::new(), + generate_content_config: None, }; let response = LlmResponse { content: Some(model_content.clone()), diff --git a/tests/instruction_history.rs b/tests/instruction_history.rs index 17e4bf1..c4c98a1 100644 --- a/tests/instruction_history.rs +++ b/tests/instruction_history.rs @@ -83,6 +83,7 @@ fn agent_with_prompt( description: "A helpful assistant.".to_string(), system_prompt: system_prompt.to_string(), model: Box::new(model), + generate_content_config: None, tools: Vec::new(), sub_agents: Vec::new(), max_steps: 1, diff --git a/tests/llm_agent.rs b/tests/llm_agent.rs index 1e606fa..88c7bcf 100644 --- a/tests/llm_agent.rs +++ b/tests/llm_agent.rs @@ -63,6 +63,7 @@ fn fake_model_returns_configured_response_and_records_request() { system_prompt: Some("You are helpful.".to_string()), contents: vec![text_content(Role::User, "hello")], tools: Vec::new(), + generate_content_config: None, }; let actual = model.generate(request.clone()).unwrap(); @@ -85,6 +86,7 @@ fn llm_agent_builds_request_from_session_snapshot_and_returns_final_event() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: Vec::new(), sub_agents: Vec::new(), max_steps: 1, @@ -139,6 +141,7 @@ fn runner_persists_llm_agent_response() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(FakeModel::new("fake-model", response.clone())), + generate_content_config: None, tools: Vec::new(), sub_agents: Vec::new(), max_steps: 1, diff --git a/tests/long_running_tool.rs b/tests/long_running_tool.rs index 51adfd3..0d0f317 100644 --- a/tests/long_running_tool.rs +++ b/tests/long_running_tool.rs @@ -79,6 +79,7 @@ fn long_running_agent(model: FakeModel) -> LlmAgent { description: "A helpful assistant.".to_string(), system_prompt: String::new(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(SandboxJobTool)], sub_agents: Vec::new(), max_steps: 4, diff --git a/tests/memory.rs b/tests/memory.rs index f7998c7..2e873ed 100644 --- a/tests/memory.rs +++ b/tests/memory.rs @@ -121,6 +121,7 @@ fn load_memory_tool_reads_memory_through_runner_context() -> Result<()> { description: "A helpful assistant.".to_string(), system_prompt: "Use memory when needed.".to_string(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(LoadMemoryTool::new())], sub_agents: Vec::new(), max_steps: 4, diff --git a/tests/model_factory.rs b/tests/model_factory.rs new file mode 100644 index 0000000..dbb9d06 --- /dev/null +++ b/tests/model_factory.rs @@ -0,0 +1,33 @@ +use r_agent::model::{DefaultModelFactory, ModelConfig, ModelFactory}; + +#[test] +fn default_model_factory_creates_openai_compatible_model() { + let model = DefaultModelFactory::new() + .create(ModelConfig { + provider: "deepseek".to_string(), + protocol: "openai-compatible".to_string(), + model: "deepseek-chat".to_string(), + base_url: Some("https://api.deepseek.com/v1".to_string()), + api_key_env: Some("DEEPSEEK_API_KEY".to_string()), + }) + .unwrap(); + + assert_eq!(model.name(), "deepseek-chat"); +} + +#[test] +fn default_model_factory_rejects_unknown_protocol() { + let result = DefaultModelFactory::new().create(ModelConfig { + provider: "example".to_string(), + protocol: "unknown-protocol".to_string(), + model: "example-model".to_string(), + base_url: None, + api_key_env: None, + }); + let err = match result { + Ok(_) => panic!("expected unsupported protocol error"), + Err(err) => err, + }; + + assert_eq!(err.to_string(), "unsupported model protocol: unknown-protocol"); +} diff --git a/tests/multi_agent.rs b/tests/multi_agent.rs index 36c9406..dd1ec57 100644 --- a/tests/multi_agent.rs +++ b/tests/multi_agent.rs @@ -152,6 +152,7 @@ fn llm_agent_can_transfer_to_sub_agent_through_transfer_tool() { description: "Root agent.".to_string(), system_prompt: "Route to a sub agent when useful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: Vec::new(), sub_agents: vec![Arc::new(child)], max_steps: 4, diff --git a/tests/openai_compatible_model.rs b/tests/openai_compatible_model.rs index ac30eba..a810cf9 100644 --- a/tests/openai_compatible_model.rs +++ b/tests/openai_compatible_model.rs @@ -9,7 +9,7 @@ use serde_json::json; use r_agent::{ content::{Content, Part, Role, ToolCall, ToolResponse}, model::{ - LlmRequest, Model, ModelConfig, ModelUsage, OpenAiCompatibleModel, + GenerateContentConfig, LlmRequest, Model, ModelConfig, ModelUsage, OpenAiCompatibleModel, openai_compatible::{request_json, response_from_json}, }, tool::ToolDeclaration, @@ -58,6 +58,7 @@ fn converts_internal_request_to_openai_compatible_json() { })), output_schema: None, }], + generate_content_config: None, }; let actual = request_json(&request); @@ -110,6 +111,29 @@ fn converts_internal_request_to_openai_compatible_json() { ); } +#[test] +fn request_json_includes_generate_content_config() { + let request = LlmRequest { + model: "test-model".to_string(), + system_prompt: None, + contents: vec![text_content(Role::User, "hello")], + tools: Vec::new(), + generate_content_config: Some(GenerateContentConfig { + temperature: Some(0.2), + top_p: Some(0.9), + max_output_tokens: Some(1024), + stop_sequences: vec!["END".to_string()], + }), + }; + + let actual = request_json(&request); + + assert_eq!(actual["temperature"], json!(0.2)); + assert_eq!(actual["top_p"], json!(0.9)); + assert_eq!(actual["max_tokens"], json!(1024)); + assert_eq!(actual["stop"], json!(["END"])); +} + #[test] fn converts_openai_compatible_text_response_to_llm_response() { let response = response_from_json(&json!({ @@ -263,6 +287,7 @@ fn generate_posts_chat_completion_request_and_parses_response() { system_prompt: None, contents: vec![text_content(Role::User, "hello")], tools: Vec::new(), + generate_content_config: None, }) .unwrap(); @@ -289,6 +314,7 @@ fn generate_errors_when_api_key_env_is_missing() { system_prompt: None, contents: vec![text_content(Role::User, "hello")], tools: Vec::new(), + generate_content_config: None, }) .unwrap_err(); diff --git a/tests/openai_compatible_real.rs b/tests/openai_compatible_real.rs index 6de7a82..5e2b2da 100644 --- a/tests/openai_compatible_real.rs +++ b/tests/openai_compatible_real.rs @@ -39,6 +39,7 @@ fn real_openai_compatible_text_generation_when_enabled() { system_prompt: Some("Reply with exactly: pong".to_string()), contents: vec![text_content(Role::User, "ping")], tools: Vec::new(), + generate_content_config: None, }) .unwrap(); diff --git a/tests/streaming.rs b/tests/streaming.rs index ead7ad5..9df7206 100644 --- a/tests/streaming.rs +++ b/tests/streaming.rs @@ -111,6 +111,7 @@ fn llm_agent_turns_streaming_model_responses_into_partial_and_final_events() -> description: "A streaming assistant.".to_string(), system_prompt: "You stream.".to_string(), model: Box::new(model), + generate_content_config: None, tools: Vec::new(), sub_agents: Vec::new(), max_steps: 1, diff --git a/tests/telemetry.rs b/tests/telemetry.rs index 81e232f..42e1e49 100644 --- a/tests/telemetry.rs +++ b/tests/telemetry.rs @@ -40,6 +40,7 @@ fn agent(model: FakeModel) -> LlmAgent { description: "A helpful assistant.".to_string(), system_prompt: String::new(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(tool)], sub_agents: Vec::new(), max_steps: 4, diff --git a/tests/tool_confirmation.rs b/tests/tool_confirmation.rs index ae0c66d..17247d9 100644 --- a/tests/tool_confirmation.rs +++ b/tests/tool_confirmation.rs @@ -79,6 +79,7 @@ fn agent_with_confirmed_delete_tool(calls: Arc, model: FakeModel) - description: "A helpful assistant.".to_string(), system_prompt: String::new(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(delete_tool)], sub_agents: Vec::new(), max_steps: 4, diff --git a/tests/tool_loop.rs b/tests/tool_loop.rs index ab7bbc8..6185c16 100644 --- a/tests/tool_loop.rs +++ b/tests/tool_loop.rs @@ -75,6 +75,7 @@ fn llm_agent_runs_tool_loop_until_final_response() { description: "A helpful assistant.".to_string(), system_prompt: "You are helpful.".to_string(), model: Box::new(model), + generate_content_config: None, tools: vec![Box::new(WeatherTool)], sub_agents: Vec::new(), max_steps: 4, @@ -155,6 +156,7 @@ fn llm_agent_errors_when_tool_loop_exceeds_max_steps() { usage: None, }, )), + generate_content_config: None, tools: vec![Box::new(WeatherTool)], sub_agents: Vec::new(), max_steps: 1, From 0addf8ccace197bda3621614142f6448480cd08d Mon Sep 17 00:00:00 2001 From: YayoiNanoka <1159066485@qq.com> Date: Wed, 1 Jul 2026 20:31:48 +0800 Subject: [PATCH 3/7] feat: add dynamic toolsets --- examples/openai_compatible_demo.rs | 1 + src/llm_agent/agent.rs | 96 ++++++++++++---- src/llm_agent/config.rs | 3 +- src/main.rs | 1 + src/tool/mod.rs | 6 +- src/tool/tool.rs | 36 ++++++ src/tool/toolset.rs | 111 ++++++++++++++++++ tests/artifact.rs | 1 + tests/async_streaming.rs | 2 + tests/callbacks.rs | 8 ++ tests/instruction_history.rs | 1 + tests/llm_agent.rs | 2 + tests/long_running_tool.rs | 1 + tests/memory.rs | 1 + tests/multi_agent.rs | 1 + tests/streaming.rs | 1 + tests/telemetry.rs | 1 + tests/tool_confirmation.rs | 1 + tests/tool_loop.rs | 2 + tests/toolset.rs | 179 +++++++++++++++++++++++++++++ 20 files changed, 432 insertions(+), 23 deletions(-) create mode 100644 src/tool/toolset.rs create mode 100644 tests/toolset.rs diff --git a/examples/openai_compatible_demo.rs b/examples/openai_compatible_demo.rs index e9d3cd3..35e6240 100644 --- a/examples/openai_compatible_demo.rs +++ b/examples/openai_compatible_demo.rs @@ -62,6 +62,7 @@ fn main() -> Result<()> { model: Box::new(model), generate_content_config: None, tools: vec![Box::new(GetTimeTool)], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, diff --git a/src/llm_agent/agent.rs b/src/llm_agent/agent.rs index 5cc003d..832c1be 100644 --- a/src/llm_agent/agent.rs +++ b/src/llm_agent/agent.rs @@ -16,9 +16,10 @@ use crate::{ session::Session, telemetry::span_attr, tool::{ - REQUEST_CONFIRMATION_TOOL_NAME, TRANSFER_TO_AGENT_TOOL_NAME, Tool, ToolConfirmation, - ToolContext, ToolDeclaration, TransferToAgentTool, confirmation_request_args, - original_tool_call_from_confirmation_args, tool_confirmation_from_request_args, + IntoToolArc, REQUEST_CONFIRMATION_TOOL_NAME, TRANSFER_TO_AGENT_TOOL_NAME, Tool, + ToolConfirmation, ToolContext, ToolDeclaration, Toolset, ToolsetContext, + TransferToAgentTool, confirmation_request_args, original_tool_call_from_confirmation_args, + tool_confirmation_from_request_args, }, }; @@ -32,7 +33,8 @@ pub struct LlmAgent { system_prompt: String, model: Box, generate_content_config: Option, - tools: Vec>, + tools: Vec>, + toolsets: Vec>, sub_agents: Vec>, max_steps: usize, include_contents: IncludeContents, @@ -46,7 +48,11 @@ struct ConfirmedToolCall { impl LlmAgent { pub fn new(config: LlmAgentConfig) -> Self { - let mut tools = config.tools; + let mut tools = config + .tools + .into_iter() + .map(IntoToolArc::into_tool_arc) + .collect::>(); let transfer_targets = config .sub_agents .iter() @@ -58,7 +64,7 @@ impl LlmAgent { .iter() .any(|tool| tool.name() == TRANSFER_TO_AGENT_TOOL_NAME) { - tools.push(Box::new(TransferToAgentTool::new(transfer_targets))); + tools.push(Arc::new(TransferToAgentTool::new(transfer_targets))); } Self { @@ -68,6 +74,7 @@ impl LlmAgent { model: config.model, generate_content_config: config.generate_content_config, tools, + toolsets: config.toolsets, sub_agents: config.sub_agents, max_steps: config.max_steps, include_contents: config.include_contents, @@ -79,12 +86,13 @@ impl LlmAgent { &self, ctx: &InvocationContext, session_snapshot: &Session, + visible_tools: &[Arc], ) -> Result { Ok(LlmRequest { model: self.model.name().to_string(), system_prompt: self.render_system_prompt(session_snapshot)?, contents: self.build_contents(ctx, session_snapshot), - tools: self.tool_declarations(), + tools: self.tool_declarations(visible_tools), generate_content_config: self.generate_content_config.clone(), }) } @@ -113,8 +121,30 @@ impl LlmAgent { )?)) } - fn tool_declarations(&self) -> Vec { - self.tools + fn resolve_tools( + &self, + ctx: &InvocationContext, + session_snapshot: &Session, + ) -> Result>> { + let toolset_ctx = ToolsetContext::new(ctx, session_snapshot); + let mut tools = self.tools.clone(); + for toolset in &self.toolsets { + tools.extend(toolset.tools(&toolset_ctx)?); + } + + let mut names = std::collections::HashSet::new(); + for tool in &tools { + let name = tool.name().to_string(); + if !names.insert(name.clone()) { + return Err(Error::msg(format!("duplicate tool name: {name}"))); + } + } + + Ok(tools) + } + + fn tool_declarations(&self, tools: &[Arc]) -> Vec { + tools .iter() .map(|tool| ToolDeclaration { name: tool.name().to_string(), @@ -125,8 +155,8 @@ impl LlmAgent { .collect() } - fn tool_for_name(&self, name: &str) -> Option<&dyn Tool> { - self.tools + fn tool_for_name<'a>(&self, tools: &'a [Arc], name: &str) -> Option<&'a dyn Tool> { + tools .iter() .find(|tool| tool.name() == name) .map(|tool| tool.as_ref()) @@ -463,9 +493,15 @@ impl Agent for LlmAgent { let tool_event = if confirmed_tool_call.confirmation.confirmed == Some(false) { self.rejected_tool_response_event(&ctx, &tool_call) } else { + let visible_tools = self.resolve_tools(&ctx, &session_snapshot)?; let tool = self - .tool_for_name(&tool_call.name) - .ok_or_else(|| Error::msg(format!("tool not found: {}", tool_call.name)))?; + .tool_for_name(&visible_tools, &tool_call.name) + .ok_or_else(|| { + Error::msg(format!( + "tool not visible in current invocation: {}", + tool_call.name + )) + })?; let (result, actions) = self.run_tool_with_callbacks( &ctx, &session_snapshot, @@ -496,7 +532,8 @@ impl Agent for LlmAgent { continue; } - let request = self.build_request(&ctx, &session_snapshot)?; + let visible_tools = self.resolve_tools(&ctx, &session_snapshot)?; + let request = self.build_request(&ctx, &session_snapshot, &visible_tools)?; let responses = self.generate_stream_with_callbacks(&ctx, request)?; for response in responses { @@ -526,8 +563,13 @@ impl Agent for LlmAgent { for tool_call in tool_calls { let tool = self - .tool_for_name(&tool_call.name) - .ok_or_else(|| Error::msg(format!("tool not found: {}", tool_call.name)))?; + .tool_for_name(&visible_tools, &tool_call.name) + .ok_or_else(|| { + Error::msg(format!( + "tool not visible in current invocation: {}", + tool_call.name + )) + })?; if let Some(confirmation) = tool.confirmation(&tool_call.args) { let confirmation_event = self.confirmation_request_event(&ctx, &tool_call, confirmation); @@ -593,9 +635,15 @@ impl Agent for LlmAgent { let tool_event = if confirmed_tool_call.confirmation.confirmed == Some(false) { self.rejected_tool_response_event(&ctx, &tool_call) } else { + let visible_tools = self.resolve_tools(&ctx, &session_snapshot)?; let tool = self - .tool_for_name(&tool_call.name) - .ok_or_else(|| Error::msg(format!("tool not found: {}", tool_call.name)))?; + .tool_for_name(&visible_tools, &tool_call.name) + .ok_or_else(|| { + Error::msg(format!( + "tool not visible in current invocation: {}", + tool_call.name + )) + })?; let (result, actions) = self.run_tool_with_callbacks( &ctx, &session_snapshot, @@ -627,7 +675,8 @@ impl Agent for LlmAgent { continue; } - let mut request = self.build_request(&ctx, &session_snapshot)?; + let visible_tools = self.resolve_tools(&ctx, &session_snapshot)?; + let mut request = self.build_request(&ctx, &session_snapshot, &visible_tools)?; let short_circuit_response = if let Some(response) = ctx.plugin_manager.run_before_model(&ctx, &mut request)? { Some(response) @@ -715,8 +764,13 @@ impl Agent for LlmAgent { for tool_call in tool_calls { let tool = self - .tool_for_name(&tool_call.name) - .ok_or_else(|| Error::msg(format!("tool not found: {}", tool_call.name)))?; + .tool_for_name(&visible_tools, &tool_call.name) + .ok_or_else(|| { + Error::msg(format!( + "tool not visible in current invocation: {}", + tool_call.name + )) + })?; if let Some(confirmation) = tool.confirmation(&tool_call.args) { let confirmation_event = self.confirmation_request_event(&ctx, &tool_call, confirmation); diff --git a/src/llm_agent/config.rs b/src/llm_agent/config.rs index d7ccb2d..92d462f 100644 --- a/src/llm_agent/config.rs +++ b/src/llm_agent/config.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use crate::{ agent::Agent, model::{GenerateContentConfig, Model}, - tool::Tool, + tool::{Tool, Toolset}, }; use super::{IncludeContents, LlmAgentCallbacks}; @@ -15,6 +15,7 @@ pub struct LlmAgentConfig { pub model: Box, pub generate_content_config: Option, pub tools: Vec>, + pub toolsets: Vec>, pub sub_agents: Vec>, pub max_steps: usize, pub include_contents: IncludeContents, diff --git a/src/main.rs b/src/main.rs index d4eeb31..dd1c3a6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -62,6 +62,7 @@ fn main() -> Result<()> { model: Box::new(model), generate_content_config: None, tools: vec![Box::new(GetTimeTool)], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, diff --git a/src/tool/mod.rs b/src/tool/mod.rs index e9f6b20..3ffe97a 100644 --- a/src/tool/mod.rs +++ b/src/tool/mod.rs @@ -5,6 +5,7 @@ mod function_tool; mod module; mod schema; mod tool; +mod toolset; mod transfer_tool; pub use confirmation::{ @@ -18,5 +19,8 @@ pub use context::ToolContext; pub use declaration::ToolDeclaration; pub use function_tool::FunctionTool; pub use module::MODULE; -pub use tool::Tool; +pub use tool::{IntoToolArc, Tool}; +pub use toolset::{ + StaticToolset, ToolPredicate, Toolset, ToolsetContext, allowed_tools_predicate, filter_toolset, +}; pub use transfer_tool::{TRANSFER_TO_AGENT_TOOL_NAME, TransferToAgentTool}; diff --git a/src/tool/tool.rs b/src/tool/tool.rs index 7f81e56..122d8a4 100644 --- a/src/tool/tool.rs +++ b/src/tool/tool.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use serde_json::Value; use crate::Result; @@ -25,3 +27,37 @@ pub trait Tool: Send + Sync { fn run(&self, ctx: &mut ToolContext, args: Value) -> Result; } + +pub trait IntoToolArc { + fn into_tool_arc(self) -> Arc; +} + +impl IntoToolArc for Arc { + fn into_tool_arc(self) -> Arc { + self + } +} + +impl IntoToolArc for Arc +where + T: Tool + 'static, +{ + fn into_tool_arc(self) -> Arc { + self + } +} + +impl IntoToolArc for Box { + fn into_tool_arc(self) -> Arc { + Arc::from(self) + } +} + +impl IntoToolArc for Box +where + T: Tool + 'static, +{ + fn into_tool_arc(self) -> Arc { + Arc::new(*self) + } +} diff --git a/src/tool/toolset.rs b/src/tool/toolset.rs new file mode 100644 index 0000000..2d1af13 --- /dev/null +++ b/src/tool/toolset.rs @@ -0,0 +1,111 @@ +use std::{collections::HashSet, sync::Arc}; + +use crate::{Result, agent::InvocationContext, session::Session}; + +use super::Tool; + +#[derive(Debug, Clone)] +pub struct ToolsetContext { + pub app_name: String, + pub user_id: String, + pub session_id: String, + pub invocation_id: String, + pub agent_name: String, + pub session_snapshot: Session, +} + +impl ToolsetContext { + pub fn new(ctx: &InvocationContext, session_snapshot: &Session) -> Self { + Self { + app_name: ctx.app_name.clone(), + user_id: ctx.user_id.clone(), + session_id: ctx.session_id.clone(), + invocation_id: ctx.invocation_id.clone(), + agent_name: ctx.agent_name.clone(), + session_snapshot: session_snapshot.clone(), + } + } + + pub fn for_test() -> Self { + Self { + app_name: "test_app".to_string(), + user_id: "test_user".to_string(), + session_id: "test_session".to_string(), + invocation_id: "test_invocation".to_string(), + agent_name: "test_agent".to_string(), + session_snapshot: Session::new("test_app", "test_user", "test_session"), + } + } +} + +pub trait Toolset: Send + Sync { + fn name(&self) -> &str; + + fn tools(&self, ctx: &ToolsetContext) -> Result>>; +} + +pub type ToolPredicate = Arc bool + Send + Sync>; + +pub fn allowed_tools_predicate(allowed_tools: I) -> ToolPredicate +where + I: IntoIterator, + S: Into, +{ + let allowed = allowed_tools + .into_iter() + .map(Into::into) + .collect::>(); + + Arc::new(move |_ctx, tool| allowed.contains(tool.name())) +} + +pub fn filter_toolset( + toolset: Arc, + predicate: ToolPredicate, +) -> Arc { + Arc::new(FilterToolset { toolset, predicate }) +} + +pub struct StaticToolset { + name: String, + tools: Vec>, +} + +impl StaticToolset { + pub fn new(name: impl Into, tools: Vec>) -> Self { + Self { + name: name.into(), + tools, + } + } +} + +impl Toolset for StaticToolset { + fn name(&self) -> &str { + &self.name + } + + fn tools(&self, _ctx: &ToolsetContext) -> Result>> { + Ok(self.tools.clone()) + } +} + +struct FilterToolset { + toolset: Arc, + predicate: ToolPredicate, +} + +impl Toolset for FilterToolset { + fn name(&self) -> &str { + self.toolset.name() + } + + fn tools(&self, ctx: &ToolsetContext) -> Result>> { + Ok(self + .toolset + .tools(ctx)? + .into_iter() + .filter(|tool| (self.predicate)(ctx, tool.as_ref())) + .collect()) + } +} diff --git a/tests/artifact.rs b/tests/artifact.rs index d606645..2f74afe 100644 --- a/tests/artifact.rs +++ b/tests/artifact.rs @@ -134,6 +134,7 @@ fn tool_context_save_artifact_records_event_delta() { model: Box::new(model), generate_content_config: None, tools: vec![Box::new(tool)], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, diff --git a/tests/async_streaming.rs b/tests/async_streaming.rs index 05dfb74..5605810 100644 --- a/tests/async_streaming.rs +++ b/tests/async_streaming.rs @@ -100,6 +100,7 @@ async fn llm_agent_async_stream_yields_model_partials_before_final() { model: Box::new(model), generate_content_config: None, tools: vec![], + toolsets: Vec::new(), sub_agents: vec![], max_steps: 4, include_contents: IncludeContents::Default, @@ -184,6 +185,7 @@ async fn async_tool_loop_waits_for_final_tool_call() { model: Box::new(model), generate_content_config: None, tools: vec![Box::new(tool)], + toolsets: Vec::new(), sub_agents: vec![], max_steps: 4, include_contents: IncludeContents::Default, diff --git a/tests/callbacks.rs b/tests/callbacks.rs index 5a4614a..7d3181c 100644 --- a/tests/callbacks.rs +++ b/tests/callbacks.rs @@ -130,6 +130,7 @@ fn before_agent_callback_can_return_content_and_skip_agent_flow() { model: Box::new(model), generate_content_config: None, tools: Vec::new(), + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents: IncludeContents::Default, @@ -167,6 +168,7 @@ fn after_agent_callback_can_append_agent_event() { model: Box::new(model), generate_content_config: None, tools: Vec::new(), + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents: IncludeContents::Default, @@ -212,6 +214,7 @@ fn before_model_callback_can_return_cached_response_and_skip_model() { model: Box::new(model), generate_content_config: None, tools: Vec::new(), + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents: IncludeContents::Default, @@ -253,6 +256,7 @@ fn after_model_callback_can_replace_model_response() { model: Box::new(model), generate_content_config: None, tools: Vec::new(), + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents: IncludeContents::Default, @@ -285,6 +289,7 @@ fn on_model_error_callback_can_recover_with_response() { model: Box::new(FailingModel), generate_content_config: None, tools: Vec::new(), + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents: IncludeContents::Default, @@ -332,6 +337,7 @@ fn before_tool_callback_can_return_result_and_skip_tool() { tools: vec![Box::new(CountingWeatherTool { calls: tool_calls.clone(), })], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, @@ -390,6 +396,7 @@ fn after_tool_callback_can_replace_tool_result() { tools: vec![Box::new(CountingWeatherTool { calls: tool_calls.clone(), })], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, @@ -444,6 +451,7 @@ fn on_tool_error_callback_can_recover_with_result() { model: Box::new(model), generate_content_config: None, tools: vec![Box::new(FailingTool)], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, diff --git a/tests/instruction_history.rs b/tests/instruction_history.rs index c4c98a1..56e0b9c 100644 --- a/tests/instruction_history.rs +++ b/tests/instruction_history.rs @@ -85,6 +85,7 @@ fn agent_with_prompt( model: Box::new(model), generate_content_config: None, tools: Vec::new(), + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents, diff --git a/tests/llm_agent.rs b/tests/llm_agent.rs index 88c7bcf..41344fa 100644 --- a/tests/llm_agent.rs +++ b/tests/llm_agent.rs @@ -88,6 +88,7 @@ fn llm_agent_builds_request_from_session_snapshot_and_returns_final_event() { model: Box::new(model), generate_content_config: None, tools: Vec::new(), + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents: IncludeContents::Default, @@ -143,6 +144,7 @@ fn runner_persists_llm_agent_response() { model: Box::new(FakeModel::new("fake-model", response.clone())), generate_content_config: None, tools: Vec::new(), + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents: IncludeContents::Default, diff --git a/tests/long_running_tool.rs b/tests/long_running_tool.rs index 0d0f317..ae8ea40 100644 --- a/tests/long_running_tool.rs +++ b/tests/long_running_tool.rs @@ -81,6 +81,7 @@ fn long_running_agent(model: FakeModel) -> LlmAgent { model: Box::new(model), generate_content_config: None, tools: vec![Box::new(SandboxJobTool)], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, diff --git a/tests/memory.rs b/tests/memory.rs index 2e873ed..fa9a138 100644 --- a/tests/memory.rs +++ b/tests/memory.rs @@ -123,6 +123,7 @@ fn load_memory_tool_reads_memory_through_runner_context() -> Result<()> { model: Box::new(model), generate_content_config: None, tools: vec![Box::new(LoadMemoryTool::new())], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, diff --git a/tests/multi_agent.rs b/tests/multi_agent.rs index dd1ec57..229a61c 100644 --- a/tests/multi_agent.rs +++ b/tests/multi_agent.rs @@ -154,6 +154,7 @@ fn llm_agent_can_transfer_to_sub_agent_through_transfer_tool() { model: Box::new(model), generate_content_config: None, tools: Vec::new(), + toolsets: Vec::new(), sub_agents: vec![Arc::new(child)], max_steps: 4, include_contents: IncludeContents::Default, diff --git a/tests/streaming.rs b/tests/streaming.rs index 9df7206..dd1e868 100644 --- a/tests/streaming.rs +++ b/tests/streaming.rs @@ -113,6 +113,7 @@ fn llm_agent_turns_streaming_model_responses_into_partial_and_final_events() -> model: Box::new(model), generate_content_config: None, tools: Vec::new(), + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents: IncludeContents::Default, diff --git a/tests/telemetry.rs b/tests/telemetry.rs index 42e1e49..1179b67 100644 --- a/tests/telemetry.rs +++ b/tests/telemetry.rs @@ -42,6 +42,7 @@ fn agent(model: FakeModel) -> LlmAgent { model: Box::new(model), generate_content_config: None, tools: vec![Box::new(tool)], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, diff --git a/tests/tool_confirmation.rs b/tests/tool_confirmation.rs index 17247d9..6a1ce3e 100644 --- a/tests/tool_confirmation.rs +++ b/tests/tool_confirmation.rs @@ -81,6 +81,7 @@ fn agent_with_confirmed_delete_tool(calls: Arc, model: FakeModel) - model: Box::new(model), generate_content_config: None, tools: vec![Box::new(delete_tool)], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, diff --git a/tests/tool_loop.rs b/tests/tool_loop.rs index 6185c16..4ac7f34 100644 --- a/tests/tool_loop.rs +++ b/tests/tool_loop.rs @@ -77,6 +77,7 @@ fn llm_agent_runs_tool_loop_until_final_response() { model: Box::new(model), generate_content_config: None, tools: vec![Box::new(WeatherTool)], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 4, include_contents: IncludeContents::Default, @@ -158,6 +159,7 @@ fn llm_agent_errors_when_tool_loop_exceeds_max_steps() { )), generate_content_config: None, tools: vec![Box::new(WeatherTool)], + toolsets: Vec::new(), sub_agents: Vec::new(), max_steps: 1, include_contents: IncludeContents::Default, diff --git a/tests/toolset.rs b/tests/toolset.rs new file mode 100644 index 0000000..793977e --- /dev/null +++ b/tests/toolset.rs @@ -0,0 +1,179 @@ +use std::time::SystemTime; +use std::sync::Arc; + +use serde_json::{Value, json}; + +use r_agent::{ + Result, + content::{Content, Part, Role, ToolCall}, + event::{Event, EventActions}, + llm_agent::{IncludeContents, LlmAgent, LlmAgentCallbacks, LlmAgentConfig}, + model::{FakeModel, LlmResponse}, + runner::Runner, + session::{InMemorySessionService, SessionService}, + tool::{ + FunctionTool, StaticToolset, Tool, ToolContext, Toolset, ToolsetContext, + allowed_tools_predicate, filter_toolset, + }, +}; + +fn text_content(role: Role, text: &str) -> Content { + Content { + role, + parts: vec![Part::Text(text.to_string())], + } +} + +fn tool_call_content(name: &str) -> Content { + Content { + role: Role::Model, + parts: vec![Part::ToolCall(ToolCall { + id: "call-1".to_string(), + name: name.to_string(), + args: json!({}), + })], + } +} + +fn named_tool(name: &str) -> Arc { + Arc::new(FunctionTool::new( + name, + format!("Runs {name}."), + |_ctx: &mut ToolContext, _args: Value| Ok(json!({ "ok": true })), + )) +} + +struct StateToolset { + required_key: String, + tool: Arc, +} + +impl Toolset for StateToolset { + fn name(&self) -> &str { + "state_toolset" + } + + fn tools(&self, ctx: &ToolsetContext) -> Result>> { + if ctx.session_snapshot.state.contains_key(&self.required_key) { + Ok(vec![self.tool.clone()]) + } else { + Ok(Vec::new()) + } + } +} + +#[test] +fn dynamic_toolset_exposes_tools_from_current_session_state() { + let model = FakeModel::new( + "fake-model", + LlmResponse { + content: Some(text_content(Role::Model, "done")), + partial: false, + usage: None, + }, + ); + let model_handle = model.clone(); + let agent = LlmAgent::new(LlmAgentConfig { + name: "assistant".to_string(), + description: "Assistant.".to_string(), + system_prompt: String::new(), + model: Box::new(model), + generate_content_config: None, + tools: Vec::new(), + toolsets: vec![Arc::new(StateToolset { + required_key: "project_id".to_string(), + tool: named_tool("project_search"), + })], + sub_agents: Vec::new(), + max_steps: 2, + include_contents: IncludeContents::Default, + callbacks: LlmAgentCallbacks::default(), + }); + let session_service = InMemorySessionService::new(); + session_service.create("app", "user-1", "session-1").unwrap(); + let mut actions = EventActions::default(); + actions + .state_delta + .insert("project_id".to_string(), json!("p1")); + session_service + .append_event( + "app", + "user-1", + "session-1", + Event { + id: "state-event".to_string(), + invocation_id: "setup".to_string(), + branch: None, + author: "setup".to_string(), + timestamp: SystemTime::now(), + content: None, + actions, + partial: false, + }, + ) + .unwrap(); + let runner = Runner::new("app", Box::new(agent), session_service, true); + + runner + .run( + "user-1", + "session-1", + text_content(Role::User, "search project"), + ) + .unwrap(); + + let requests = model_handle.requests(); + assert_eq!(requests[0].tools.len(), 1); + assert_eq!(requests[0].tools[0].name, "project_search"); +} + +#[test] +fn filter_toolset_uses_predicate_to_limit_visible_tools() { + let source = Arc::new(StaticToolset::new( + "all_tools", + vec![named_tool("read_file"), named_tool("write_file")], + )); + let filtered = filter_toolset(source, allowed_tools_predicate(["read_file"])); + let ctx = ToolsetContext::for_test(); + + let tools = filtered.tools(&ctx).unwrap(); + + assert_eq!(tools.len(), 1); + assert_eq!(tools[0].name(), "read_file"); +} + +#[test] +fn llm_agent_rejects_tool_call_that_was_not_visible_to_model() { + let model = FakeModel::new( + "fake-model", + LlmResponse { + content: Some(tool_call_content("hidden_tool")), + partial: false, + usage: None, + }, + ); + let agent = LlmAgent::new(LlmAgentConfig { + name: "assistant".to_string(), + description: "Assistant.".to_string(), + system_prompt: String::new(), + model: Box::new(model), + generate_content_config: None, + tools: Vec::new(), + toolsets: Vec::new(), + sub_agents: Vec::new(), + max_steps: 2, + include_contents: IncludeContents::Default, + callbacks: LlmAgentCallbacks::default(), + }); + let runner = Runner::new("app", Box::new(agent), InMemorySessionService::new(), true); + + let err = runner + .run( + "user-1", + "session-1", + text_content(Role::User, "call hidden tool"), + ) + .unwrap_err(); + + assert_eq!(err.to_string(), "tool not visible in current invocation: hidden_tool"); +} From e202aac5fca088cf595e4c1b8381bebdc091c4bb Mon Sep 17 00:00:00 2001 From: YayoiNanoka <1159066485@qq.com> Date: Wed, 1 Jul 2026 20:35:07 +0800 Subject: [PATCH 4/7] feat: add TOML agent loader --- Cargo.lock | 82 +++++++++++++ Cargo.toml | 1 + src/agent/mod.rs | 2 + src/agent/toml_loader.rs | 228 +++++++++++++++++++++++++++++++++++++ src/tool/mod.rs | 2 + src/tool/registry.rs | 57 ++++++++++ tests/toml_agent_loader.rs | 127 +++++++++++++++++++++ 7 files changed, 499 insertions(+) create mode 100644 src/agent/toml_loader.rs create mode 100644 src/tool/registry.rs create mode 100644 tests/toml_agent_loader.rs diff --git a/Cargo.lock b/Cargo.lock index d4a5cfa..bd897b1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -186,6 +186,12 @@ dependencies = [ "syn", ] +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + [[package]] name = "find-msvc-tools" version = "0.1.9" @@ -288,6 +294,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + [[package]] name = "http" version = "1.4.2" @@ -520,6 +532,16 @@ dependencies = [ "icu_properties", ] +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown", +] + [[package]] name = "ipnet" version = "2.12.0" @@ -733,6 +755,7 @@ dependencies = [ "serde_json", "tokio", "tokio-stream", + "toml", "tower", "tower-http", "ureq", @@ -930,6 +953,15 @@ dependencies = [ "serde_core", ] +[[package]] +name = "serde_spanned" +version = "0.6.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf41e0cfaf7226dca15e8197172c295a782857fcb97fad1808a166870dee75a3" +dependencies = [ + "serde", +] + [[package]] name = "serde_urlencoded" version = "0.7.1" @@ -1118,6 +1150,47 @@ dependencies = [ "tokio", ] +[[package]] +name = "toml" +version = "0.8.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1beb996b9d83529a9e75c17a1686767d148d70663143c7854d8b4a09ced362" +dependencies = [ + "serde", + "serde_spanned", + "toml_datetime", + "toml_edit", +] + +[[package]] +name = "toml_datetime" +version = "0.6.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22cddaf88f4fbc13c51aebbf5f8eceb5c7c5a9da2ac40a13519eb5b0a0e8f11c" +dependencies = [ + "serde", +] + +[[package]] +name = "toml_edit" +version = "0.22.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41fe8c660ae4257887cf66394862d21dbca4a6ddd26f04a3560410406a2f819a" +dependencies = [ + "indexmap", + "serde", + "serde_spanned", + "toml_datetime", + "toml_write", + "winnow", +] + +[[package]] +name = "toml_write" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d99f8c9a7727884afe522e9bd5edbfc91a3312b36a77b5fb8926e4c31a41801" + [[package]] name = "tower" version = "0.5.3" @@ -1506,6 +1579,15 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "winnow" +version = "0.7.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945" +dependencies = [ + "memchr", +] + [[package]] name = "wit-bindgen" version = "0.57.1" diff --git a/Cargo.toml b/Cargo.toml index 52563c4..0062f59 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,6 +13,7 @@ serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" tokio = { version = "1", features = ["macros", "rt-multi-thread", "sync"] } tokio-stream = "0.1" +toml = "0.8" tower = { version = "0.5", features = ["util"] } tower-http = { version = "0.6", features = ["cors"] } ureq = { version = "2", default-features = false, features = ["tls"] } diff --git a/src/agent/mod.rs b/src/agent/mod.rs index b921214..377c6da 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -4,6 +4,7 @@ mod custom_agent; mod loader; mod module; mod stream; +mod toml_loader; pub use agent::{Agent, IntoAgentArc, find_agent, validate_agent_names}; pub use context::InvocationContext; @@ -11,3 +12,4 @@ pub use custom_agent::CustomAgent; pub use loader::{AgentLoader, MultiAgentLoader, SingleAgentLoader}; pub use module::MODULE; pub use stream::{AsyncEventStream, EventStream, event_stream_from_result}; +pub use toml_loader::TomlAgentLoader; diff --git a/src/agent/toml_loader.rs b/src/agent/toml_loader.rs new file mode 100644 index 0000000..f694214 --- /dev/null +++ b/src/agent/toml_loader.rs @@ -0,0 +1,228 @@ +use std::{collections::HashMap, fs, path::Path, sync::Arc}; + +use serde::Deserialize; + +use crate::{ + Error, Result, + llm_agent::{IncludeContents, LlmAgent, LlmAgentCallbacks, LlmAgentConfig}, + model::{DefaultModelFactory, GenerateContentConfig, ModelConfig, ModelFactory}, + tool::{StaticToolset, ToolRegistry, ToolsetRegistry}, +}; + +use super::{Agent, AgentLoader, MultiAgentLoader}; + +pub struct TomlAgentLoader { + inner: MultiAgentLoader, +} + +impl TomlAgentLoader { + pub fn from_path( + path: impl AsRef, + model_factory: Arc, + tool_registry: ToolRegistry, + ) -> Result { + let source = fs::read_to_string(path) + .map_err(|err| Error::msg(format!("failed to read agent config: {err}")))?; + Self::from_str( + &source, + model_factory, + tool_registry, + ToolsetRegistry::new(), + ) + } + + pub fn from_str( + source: &str, + model_factory: Arc, + tool_registry: ToolRegistry, + toolset_registry: ToolsetRegistry, + ) -> Result { + let config = toml::from_str::(source) + .map_err(|err| Error::msg(format!("invalid agent config TOML: {err}")))?; + let mut builder = TomlAgentBuilder { + config, + model_factory, + tool_registry, + toolset_registry, + built_agents: HashMap::new(), + }; + + let root = builder.build_agent(&builder.config.root_agent.clone())?; + let agents = builder.built_agents.values().cloned().collect::>(); + Ok(Self { + inner: MultiAgentLoader::new(root, agents)?, + }) + } + + pub fn from_path_with_defaults(path: impl AsRef) -> Result { + Self::from_path( + path, + Arc::new(DefaultModelFactory::new()), + ToolRegistry::new(), + ) + } +} + +impl AgentLoader for TomlAgentLoader { + fn list_agents(&self) -> Result> { + self.inner.list_agents() + } + + fn load_agent(&self, name: &str) -> Result> { + self.inner.load_agent(name) + } + + fn root_agent(&self) -> Result> { + self.inner.root_agent() + } +} + +struct TomlAgentBuilder { + config: AdkTomlConfig, + model_factory: Arc, + tool_registry: ToolRegistry, + toolset_registry: ToolsetRegistry, + built_agents: HashMap>, +} + +impl TomlAgentBuilder { + fn build_agent(&mut self, id: &str) -> Result> { + if let Some(agent) = self.built_agents.get(id) { + return Ok(agent.clone()); + } + + let config = self + .config + .agents + .get(id) + .cloned() + .ok_or_else(|| Error::msg(format!("agent config not found: {id}")))?; + + if config.agent_type != "llm" { + return Err(Error::msg(format!( + "unsupported agent type for {id}: {}", + config.agent_type + ))); + } + + let model_config = self + .config + .models + .get(&config.model) + .cloned() + .ok_or_else(|| Error::msg(format!("model config not found: {}", config.model)))?; + let model = self.model_factory.create(model_config.into())?; + let sub_agents = config + .sub_agents + .iter() + .map(|name| self.build_agent(name)) + .collect::>>()?; + let static_tools = config + .tools + .iter() + .map(|name| self.tool_registry.get(name)) + .collect::>>()?; + let mut toolsets = config + .toolsets + .iter() + .map(|name| self.toolset_registry.get(name)) + .collect::>>()?; + if !static_tools.is_empty() { + toolsets.push(Arc::new(StaticToolset::new( + format!("{id}_static_tools"), + static_tools, + ))); + } + + let agent = Arc::new(LlmAgent::new(LlmAgentConfig { + name: config.name.unwrap_or_else(|| id.to_string()), + description: config.description, + system_prompt: config.system_prompt.unwrap_or_default(), + model, + generate_content_config: config.generate_content_config.map(Into::into), + tools: Vec::new(), + toolsets, + sub_agents, + max_steps: config.max_steps.unwrap_or(8), + include_contents: parse_include_contents(config.include_contents.as_deref())?, + callbacks: LlmAgentCallbacks::default(), + })); + self.built_agents.insert(id.to_string(), agent.clone()); + Ok(agent) + } +} + +fn parse_include_contents(value: Option<&str>) -> Result { + match value.unwrap_or("default") { + "default" => Ok(IncludeContents::Default), + "none" => Ok(IncludeContents::None), + other => Err(Error::msg(format!("unsupported include_contents: {other}"))), + } +} + +#[derive(Debug, Clone, Deserialize)] +struct AdkTomlConfig { + root_agent: String, + models: HashMap, + agents: HashMap, +} + +#[derive(Debug, Clone, Deserialize)] +struct TomlModelConfig { + provider: String, + protocol: String, + model: String, + base_url: Option, + api_key_env: Option, +} + +impl From for ModelConfig { + fn from(config: TomlModelConfig) -> Self { + Self { + provider: config.provider, + protocol: config.protocol, + model: config.model, + base_url: config.base_url, + api_key_env: config.api_key_env, + } + } +} + +#[derive(Debug, Clone, Deserialize)] +struct TomlAgentConfig { + #[serde(rename = "type")] + agent_type: String, + name: Option, + description: String, + model: String, + system_prompt: Option, + include_contents: Option, + max_steps: Option, + #[serde(default)] + tools: Vec, + #[serde(default)] + toolsets: Vec, + #[serde(default)] + sub_agents: Vec, + generate_content_config: Option, +} + +#[derive(Debug, Clone, Deserialize)] +struct TomlGenerateContentConfig { + temperature: Option, + top_p: Option, + max_output_tokens: Option, + #[serde(default)] + stop_sequences: Vec, +} + +impl From for GenerateContentConfig { + fn from(config: TomlGenerateContentConfig) -> Self { + Self { + temperature: config.temperature, + top_p: config.top_p, + max_output_tokens: config.max_output_tokens, + stop_sequences: config.stop_sequences, + } + } +} diff --git a/src/tool/mod.rs b/src/tool/mod.rs index 3ffe97a..1077f32 100644 --- a/src/tool/mod.rs +++ b/src/tool/mod.rs @@ -3,6 +3,7 @@ mod context; mod declaration; mod function_tool; mod module; +mod registry; mod schema; mod tool; mod toolset; @@ -19,6 +20,7 @@ pub use context::ToolContext; pub use declaration::ToolDeclaration; pub use function_tool::FunctionTool; pub use module::MODULE; +pub use registry::{ToolRegistry, ToolsetRegistry}; pub use tool::{IntoToolArc, Tool}; pub use toolset::{ StaticToolset, ToolPredicate, Toolset, ToolsetContext, allowed_tools_predicate, filter_toolset, diff --git a/src/tool/registry.rs b/src/tool/registry.rs new file mode 100644 index 0000000..92a31df --- /dev/null +++ b/src/tool/registry.rs @@ -0,0 +1,57 @@ +use std::{collections::HashMap, sync::Arc}; + +use crate::{Error, Result}; + +use super::{Tool, Toolset}; + +#[derive(Default, Clone)] +pub struct ToolRegistry { + tools: HashMap>, +} + +impl ToolRegistry { + pub fn new() -> Self { + Self::default() + } + + pub fn register(&mut self, tool: Arc) -> Result<()> { + let name = tool.name().to_string(); + if self.tools.insert(name.clone(), tool).is_some() { + return Err(Error::msg(format!("duplicate tool name: {name}"))); + } + Ok(()) + } + + pub fn get(&self, name: &str) -> Result> { + self.tools + .get(name) + .cloned() + .ok_or_else(|| Error::msg(format!("unknown tool in config: {name}"))) + } +} + +#[derive(Default, Clone)] +pub struct ToolsetRegistry { + toolsets: HashMap>, +} + +impl ToolsetRegistry { + pub fn new() -> Self { + Self::default() + } + + pub fn register(&mut self, toolset: Arc) -> Result<()> { + let name = toolset.name().to_string(); + if self.toolsets.insert(name.clone(), toolset).is_some() { + return Err(Error::msg(format!("duplicate toolset name: {name}"))); + } + Ok(()) + } + + pub fn get(&self, name: &str) -> Result> { + self.toolsets + .get(name) + .cloned() + .ok_or_else(|| Error::msg(format!("unknown toolset in config: {name}"))) + } +} diff --git a/tests/toml_agent_loader.rs b/tests/toml_agent_loader.rs new file mode 100644 index 0000000..ab9bd1e --- /dev/null +++ b/tests/toml_agent_loader.rs @@ -0,0 +1,127 @@ +use std::{ + fs, + sync::Arc, + time::{SystemTime, UNIX_EPOCH}, +}; + +use serde_json::{Value, json}; + +use r_agent::{ + agent::{AgentLoader, TomlAgentLoader}, + model::DefaultModelFactory, + tool::{FunctionTool, ToolContext, ToolRegistry}, +}; + +fn write_config(contents: &str) -> std::path::PathBuf { + let path = std::env::temp_dir().join(format!( + "r_agent_config_{}.toml", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + fs::write(&path, contents).unwrap(); + path +} + +fn registry() -> ToolRegistry { + let mut registry = ToolRegistry::new(); + registry + .register(Arc::new(FunctionTool::new( + "get_weather", + "Gets weather.", + |_ctx: &mut ToolContext, _args: Value| Ok(json!({ "forecast": "sunny" })), + ))) + .unwrap(); + registry +} + +#[test] +fn toml_agent_loader_loads_llm_agent_tree() { + let path = write_config( + r#" +root_agent = "assistant" + +[models.fake_openai] +provider = "test" +protocol = "openai-compatible" +model = "test-model" +base_url = "http://localhost:11434/v1" +api_key_env = "TEST_API_KEY" + +[agents.assistant] +type = "llm" +description = "Main assistant." +model = "fake_openai" +system_prompt = "You are helpful." +include_contents = "default" +max_steps = 4 +tools = ["get_weather"] +sub_agents = ["researcher"] + +[agents.assistant.generate_content_config] +temperature = 0.2 +top_p = 0.9 +max_output_tokens = 128 +stop_sequences = ["END"] + +[agents.researcher] +type = "llm" +description = "Research assistant." +model = "fake_openai" +system_prompt = "Research carefully." +include_contents = "none" +max_steps = 2 +tools = [] +sub_agents = [] +"#, + ); + + let loader = + TomlAgentLoader::from_path(&path, Arc::new(DefaultModelFactory::new()), registry()) + .unwrap(); + + assert_eq!( + loader.list_agents().unwrap(), + vec!["assistant".to_string(), "researcher".to_string()] + ); + let root = loader.root_agent().unwrap(); + assert_eq!(root.name(), "assistant"); + assert_eq!(root.sub_agents()[0].name(), "researcher"); +} + +#[test] +fn toml_agent_loader_rejects_unknown_tool() { + let path = write_config( + r#" +root_agent = "assistant" + +[models.fake_openai] +provider = "test" +protocol = "openai-compatible" +model = "test-model" +base_url = "http://localhost:11434/v1" + +[agents.assistant] +type = "llm" +description = "Main assistant." +model = "fake_openai" +system_prompt = "You are helpful." +include_contents = "default" +max_steps = 4 +tools = ["missing_tool"] +sub_agents = [] +"#, + ); + + let err = match TomlAgentLoader::from_path( + &path, + Arc::new(DefaultModelFactory::new()), + ToolRegistry::new(), + ) { + Ok(_) => panic!("expected unknown tool error"), + Err(err) => err, + }; + + assert_eq!(err.to_string(), "unknown tool in config: missing_tool"); +} From fb2a957319a5561c58f7a8e0765ef985c142e512 Mon Sep 17 00:00:00 2001 From: YayoiNanoka <1159066485@qq.com> Date: Wed, 1 Jul 2026 20:37:36 +0800 Subject: [PATCH 5/7] feat: add scoped session state --- src/session/in_memory.rs | 89 +++++++++++++++++++++++++++++-- src/session/mod.rs | 5 +- src/session/state.rs | 25 +++++++++ tests/scoped_state.rs | 110 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 224 insertions(+), 5 deletions(-) create mode 100644 tests/scoped_state.rs diff --git a/src/session/in_memory.rs b/src/session/in_memory.rs index 7474f49..15aed81 100644 --- a/src/session/in_memory.rs +++ b/src/session/in_memory.rs @@ -3,13 +3,17 @@ use std::{ sync::{Arc, RwLock}, }; +use serde_json::Map; + use crate::{Error, Result, event::Event}; -use super::{Session, SessionService}; +use super::{Session, SessionService, State, StateScope, state_scope}; #[derive(Debug, Clone, Default)] pub struct InMemorySessionService { sessions: Arc>>, + app_states: Arc>>, + user_states: Arc>>, } impl InMemorySessionService { @@ -33,12 +37,32 @@ impl SessionService for InMemorySessionService { fn get(&self, app_name: &str, user_id: &str, session_id: &str) -> Result { let key = SessionKey::new(app_name, user_id, session_id); - self.sessions + let mut session = self + .sessions .read() .expect("session lock poisoned") .get(&key) .cloned() - .ok_or_else(|| Error::msg("session not found")) + .ok_or_else(|| Error::msg("session not found"))?; + + if let Some(app_state) = self + .app_states + .read() + .expect("app state lock poisoned") + .get(app_name) + { + session.state.extend(app_state.clone()); + } + if let Some(user_state) = self + .user_states + .read() + .expect("user state lock poisoned") + .get(&UserKey::new(app_name, user_id)) + { + session.state.extend(user_state.clone()); + } + + Ok(session) } fn list(&self, app_name: &str, user_id: &str) -> Result> { @@ -72,14 +96,56 @@ impl SessionService for InMemorySessionService { event: Event, ) -> Result<()> { let key = SessionKey::new(app_name, user_id, session_id); + let mut event = event; let mut sessions = self.sessions.write().expect("session lock poisoned"); let session = sessions .get_mut(&key) .ok_or_else(|| Error::msg("session not found"))?; let state_delta = event.actions.state_delta.clone(); + let mut session_delta = State::new(); + let mut app_delta = State::new(); + let mut user_delta = State::new(); + let mut persisted_state_delta = Map::new(); + + for (key, value) in state_delta { + match state_scope(&key) { + StateScope::App => { + app_delta.insert(key.clone(), value.clone()); + persisted_state_delta.insert(key, value); + } + StateScope::User => { + user_delta.insert(key.clone(), value.clone()); + persisted_state_delta.insert(key, value); + } + StateScope::Session => { + session_delta.insert(key.clone(), value.clone()); + persisted_state_delta.insert(key, value); + } + StateScope::Temp => {} + } + } + + if !app_delta.is_empty() { + self.app_states + .write() + .expect("app state lock poisoned") + .entry(app_name.to_string()) + .or_default() + .extend(app_delta); + } + if !user_delta.is_empty() { + self.user_states + .write() + .expect("user state lock poisoned") + .entry(UserKey::new(app_name, user_id)) + .or_default() + .extend(user_delta); + } + + event.actions.state_delta = persisted_state_delta; session.events.push(event); - session.state.extend(state_delta); + session.state.extend(session_delta); Ok(()) } @@ -92,6 +158,21 @@ struct SessionKey { session_id: String, } +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct UserKey { + app_name: String, + user_id: String, +} + +impl UserKey { + fn new(app_name: &str, user_id: &str) -> Self { + Self { + app_name: app_name.to_string(), + user_id: user_id.to_string(), + } + } +} + impl SessionKey { fn new(app_name: &str, user_id: &str, session_id: &str) -> Self { Self { diff --git a/src/session/mod.rs b/src/session/mod.rs index 22d1c70..51607da 100644 --- a/src/session/mod.rs +++ b/src/session/mod.rs @@ -8,4 +8,7 @@ pub use in_memory::InMemorySessionService; pub use module::MODULE; pub use service::SessionService; pub use session::Session; -pub use state::State; +pub use state::{ + APP_STATE_PREFIX, SESSION_STATE_PREFIX, State, StateScope, TEMP_STATE_PREFIX, + USER_STATE_PREFIX, state_scope, +}; diff --git a/src/session/state.rs b/src/session/state.rs index 9d5c9e3..3944b7c 100644 --- a/src/session/state.rs +++ b/src/session/state.rs @@ -3,3 +3,28 @@ use std::collections::HashMap; use serde_json::Value; pub type State = HashMap; + +pub const APP_STATE_PREFIX: &str = "app:"; +pub const USER_STATE_PREFIX: &str = "user:"; +pub const TEMP_STATE_PREFIX: &str = "temp:"; +pub const SESSION_STATE_PREFIX: &str = ""; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StateScope { + App, + User, + Temp, + Session, +} + +pub fn state_scope(key: &str) -> StateScope { + if key.starts_with(APP_STATE_PREFIX) { + StateScope::App + } else if key.starts_with(USER_STATE_PREFIX) { + StateScope::User + } else if key.starts_with(TEMP_STATE_PREFIX) { + StateScope::Temp + } else { + StateScope::Session + } +} diff --git a/tests/scoped_state.rs b/tests/scoped_state.rs new file mode 100644 index 0000000..f2f1181 --- /dev/null +++ b/tests/scoped_state.rs @@ -0,0 +1,110 @@ +use std::time::SystemTime; + +use serde_json::json; + +use r_agent::{ + event::{Event, EventActions}, + session::{ + APP_STATE_PREFIX, InMemorySessionService, SESSION_STATE_PREFIX, SessionService, + TEMP_STATE_PREFIX, USER_STATE_PREFIX, + }, +}; + +fn event_with_delta(id: &str, state_delta: Vec<(&str, serde_json::Value)>) -> Event { + let mut actions = EventActions::default(); + for (key, value) in state_delta { + actions.state_delta.insert(key.to_string(), value); + } + + Event { + id: id.to_string(), + invocation_id: "invocation-1".to_string(), + branch: None, + author: "agent".to_string(), + timestamp: SystemTime::now(), + content: None, + actions, + partial: false, + } +} + +#[test] +fn app_state_is_shared_across_users_and_sessions() { + let service = InMemorySessionService::new(); + service.create("app", "user-1", "s1").unwrap(); + service.create("app", "user-2", "s2").unwrap(); + + service + .append_event( + "app", + "user-1", + "s1", + event_with_delta("event-1", vec![("app:timezone", json!("UTC"))]), + ) + .unwrap(); + + let session = service.get("app", "user-2", "s2").unwrap(); + + assert_eq!(session.state.get("app:timezone"), Some(&json!("UTC"))); +} + +#[test] +fn user_state_is_shared_only_for_same_user() { + let service = InMemorySessionService::new(); + service.create("app", "user-1", "s1").unwrap(); + service.create("app", "user-1", "s2").unwrap(); + service.create("app", "user-2", "s3").unwrap(); + + service + .append_event( + "app", + "user-1", + "s1", + event_with_delta("event-1", vec![("user:name", json!("Yayoi"))]), + ) + .unwrap(); + + let same_user = service.get("app", "user-1", "s2").unwrap(); + let other_user = service.get("app", "user-2", "s3").unwrap(); + + assert_eq!(same_user.state.get("user:name"), Some(&json!("Yayoi"))); + assert_eq!(other_user.state.get("user:name"), None); +} + +#[test] +fn temp_state_is_not_persisted_to_state_or_event_delta() { + let service = InMemorySessionService::new(); + service.create("app", "user-1", "s1").unwrap(); + + service + .append_event( + "app", + "user-1", + "s1", + event_with_delta( + "event-1", + vec![ + ("temp:scratch", json!("discard")), + ("session_value", json!("keep")), + ], + ), + ) + .unwrap(); + + let session = service.get("app", "user-1", "s1").unwrap(); + + assert_eq!(session.state.get("temp:scratch"), None); + assert_eq!(session.state.get("session_value"), Some(&json!("keep"))); + assert!(!session.events[0] + .actions + .state_delta + .contains_key("temp:scratch")); +} + +#[test] +fn state_prefix_constants_match_google_adk_style() { + assert_eq!(APP_STATE_PREFIX, "app:"); + assert_eq!(USER_STATE_PREFIX, "user:"); + assert_eq!(TEMP_STATE_PREFIX, "temp:"); + assert_eq!(SESSION_STATE_PREFIX, ""); +} From 471194a0db684767061fc4145960061f91c88237 Mon Sep 17 00:00:00 2001 From: YayoiNanoka <1159066485@qq.com> Date: Wed, 1 Jul 2026 20:42:24 +0800 Subject: [PATCH 6/7] feat: add database session service --- Cargo.lock | 88 +++++++- Cargo.toml | 1 + src/content/content.rs | 4 +- src/content/part.rs | 4 +- src/content/role.rs | 4 +- src/content/tool_call.rs | 3 +- src/content/tool_response.rs | 3 +- src/event/actions.rs | 3 +- src/event/event.rs | 4 +- src/session/database.rs | 59 ++++++ src/session/mod.rs | 4 + src/session/sqlite.rs | 397 +++++++++++++++++++++++++++++++++++ tests/database_session.rs | 105 +++++++++ 13 files changed, 671 insertions(+), 8 deletions(-) create mode 100644 src/session/database.rs create mode 100644 src/session/sqlite.rs create mode 100644 tests/database_session.rs diff --git a/Cargo.lock b/Cargo.lock index bd897b1..1be1df2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,18 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "once_cell", + "version_check", + "zerocopy", +] + [[package]] name = "android_system_properties" version = "0.1.5" @@ -192,6 +204,18 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "fallible-iterator" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649" + +[[package]] +name = "fallible-streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" + [[package]] name = "find-msvc-tools" version = "0.1.9" @@ -294,12 +318,30 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" +dependencies = [ + "ahash", +] + [[package]] name = "hashbrown" version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +[[package]] +name = "hashlink" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ba4ff7128dee98c7dc9794b6a411377e1404dba1c97deb8d1a55297bd25d8af" +dependencies = [ + "hashbrown 0.14.5", +] + [[package]] name = "http" version = "1.4.2" @@ -539,7 +581,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" dependencies = [ "equivalent", - "hashbrown", + "hashbrown 0.17.1", ] [[package]] @@ -571,6 +613,17 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "libsqlite3-sys" +version = "0.30.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e99fb7a497b1e3339bc746195567ed8d3e24945ecd636e3619d20b9de9e9149" +dependencies = [ + "cc", + "pkg-config", + "vcpkg", +] + [[package]] name = "litemap" version = "0.8.2" @@ -645,6 +698,12 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + [[package]] name = "potential_utf" version = "0.1.5" @@ -751,6 +810,7 @@ dependencies = [ "chrono", "futures-util", "reqwest", + "rusqlite", "serde", "serde_json", "tokio", @@ -845,6 +905,20 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rusqlite" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7753b721174eb8ff87a9a0e799e2d7bc3749323e773db92e0984debb00019d6e" +dependencies = [ + "bitflags", + "fallible-iterator", + "fallible-streaming-iterator", + "hashlink", + "libsqlite3-sys", + "smallvec", +] + [[package]] name = "rustc-hash" version = "2.1.2" @@ -1308,6 +1382,18 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "vcpkg" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + [[package]] name = "want" version = "0.3.1" diff --git a/Cargo.toml b/Cargo.toml index 0062f59..f396fd0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,7 @@ axum = "0.7" chrono = { version = "0.4", default-features = false, features = ["std", "clock", "serde"] } futures-util = "0.3" reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "rustls-tls"] } +rusqlite = { version = "0.32", features = ["bundled"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" tokio = { version = "1", features = ["macros", "rt-multi-thread", "sync"] } diff --git a/src/content/content.rs b/src/content/content.rs index 640ea64..18c6826 100644 --- a/src/content/content.rs +++ b/src/content/content.rs @@ -1,6 +1,8 @@ +use serde::{Deserialize, Serialize}; + use super::{Part, Role}; -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct Content { pub role: Role, pub parts: Vec, diff --git a/src/content/part.rs b/src/content/part.rs index f26db20..0860c4c 100644 --- a/src/content/part.rs +++ b/src/content/part.rs @@ -1,6 +1,8 @@ +use serde::{Deserialize, Serialize}; + use super::{ToolCall, ToolResponse}; -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub enum Part { Text(String), ToolCall(ToolCall), diff --git a/src/content/role.rs b/src/content/role.rs index 9b2557f..5345a6c 100644 --- a/src/content/role.rs +++ b/src/content/role.rs @@ -1,4 +1,6 @@ -#[derive(Debug, Clone, PartialEq, Eq)] +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub enum Role { User, Model, diff --git a/src/content/tool_call.rs b/src/content/tool_call.rs index 602e780..f8bba4e 100644 --- a/src/content/tool_call.rs +++ b/src/content/tool_call.rs @@ -1,6 +1,7 @@ +use serde::{Deserialize, Serialize}; use serde_json::Value; -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct ToolCall { pub id: String, pub name: String, diff --git a/src/content/tool_response.rs b/src/content/tool_response.rs index 4be2082..e2e193a 100644 --- a/src/content/tool_response.rs +++ b/src/content/tool_response.rs @@ -1,6 +1,7 @@ +use serde::{Deserialize, Serialize}; use serde_json::Value; -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct ToolResponse { pub call_id: String, pub name: String, diff --git a/src/event/actions.rs b/src/event/actions.rs index aef14c9..25ce4c8 100644 --- a/src/event/actions.rs +++ b/src/event/actions.rs @@ -1,10 +1,11 @@ use std::collections::HashMap; +use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use crate::tool::ToolConfirmation; -#[derive(Debug, Clone, PartialEq, Eq, Default)] +#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)] pub struct EventActions { pub state_delta: Map, pub artifact_delta: HashMap, diff --git a/src/event/event.rs b/src/event/event.rs index 954c735..d22b9c5 100644 --- a/src/event/event.rs +++ b/src/event/event.rs @@ -1,10 +1,12 @@ use std::time::SystemTime; +use serde::{Deserialize, Serialize}; + use crate::content::{Content, Part, Role}; use super::{EventActions, EventKind}; -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct Event { pub id: String, pub invocation_id: String, diff --git a/src/session/database.rs b/src/session/database.rs new file mode 100644 index 0000000..e47c708 --- /dev/null +++ b/src/session/database.rs @@ -0,0 +1,59 @@ +use crate::{Result, event::Event}; + +use super::{Session, SessionService}; + +pub trait SessionStore: Send + Sync { + fn create(&self, app_name: &str, user_id: &str, session_id: &str) -> Result; + + fn get(&self, app_name: &str, user_id: &str, session_id: &str) -> Result; + + fn list(&self, app_name: &str, user_id: &str) -> Result>; + + fn delete(&self, app_name: &str, user_id: &str, session_id: &str) -> Result<()>; + + fn append_event( + &self, + app_name: &str, + user_id: &str, + session_id: &str, + event: Event, + ) -> Result<()>; +} + +pub struct DatabaseSessionService { + store: S, +} + +impl DatabaseSessionService { + pub fn new(store: S) -> Result { + Ok(Self { store }) + } +} + +impl SessionService for DatabaseSessionService { + fn create(&self, app_name: &str, user_id: &str, session_id: &str) -> Result { + self.store.create(app_name, user_id, session_id) + } + + fn get(&self, app_name: &str, user_id: &str, session_id: &str) -> Result { + self.store.get(app_name, user_id, session_id) + } + + fn list(&self, app_name: &str, user_id: &str) -> Result> { + self.store.list(app_name, user_id) + } + + fn delete(&self, app_name: &str, user_id: &str, session_id: &str) -> Result<()> { + self.store.delete(app_name, user_id, session_id) + } + + fn append_event( + &self, + app_name: &str, + user_id: &str, + session_id: &str, + event: Event, + ) -> Result<()> { + self.store.append_event(app_name, user_id, session_id, event) + } +} diff --git a/src/session/mod.rs b/src/session/mod.rs index 51607da..ed8c9d6 100644 --- a/src/session/mod.rs +++ b/src/session/mod.rs @@ -1,13 +1,17 @@ +mod database; mod in_memory; mod module; mod service; mod session; +mod sqlite; mod state; +pub use database::{DatabaseSessionService, SessionStore}; pub use in_memory::InMemorySessionService; pub use module::MODULE; pub use service::SessionService; pub use session::Session; +pub use sqlite::SqliteSessionStore; pub use state::{ APP_STATE_PREFIX, SESSION_STATE_PREFIX, State, StateScope, TEMP_STATE_PREFIX, USER_STATE_PREFIX, state_scope, diff --git a/src/session/sqlite.rs b/src/session/sqlite.rs new file mode 100644 index 0000000..25d68fe --- /dev/null +++ b/src/session/sqlite.rs @@ -0,0 +1,397 @@ +use std::{ + path::Path, + sync::{Arc, Mutex}, + time::{SystemTime, UNIX_EPOCH}, +}; + +use rusqlite::{Connection, OptionalExtension, params}; +use serde_json::Map; + +use crate::{Error, Result, event::Event}; + +use super::{Session, SessionStore, State, StateScope, state_scope}; + +#[derive(Clone)] +pub struct SqliteSessionStore { + connection: Arc>, +} + +impl SqliteSessionStore { + pub fn open(path: impl AsRef) -> Result { + let connection = Connection::open(path) + .map_err(|err| Error::msg(format!("failed to open sqlite session store: {err}")))?; + let store = Self { + connection: Arc::new(Mutex::new(connection)), + }; + store.migrate()?; + Ok(store) + } + + fn migrate(&self) -> Result<()> { + let connection = self + .connection + .lock() + .map_err(|_| Error::msg("sqlite session store lock poisoned"))?; + connection + .execute_batch( + r#" + PRAGMA foreign_keys = ON; + + CREATE TABLE IF NOT EXISTS sessions ( + app_name TEXT NOT NULL, + user_id TEXT NOT NULL, + session_id TEXT NOT NULL, + state_json TEXT NOT NULL, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + PRIMARY KEY (app_name, user_id, session_id) + ); + + CREATE TABLE IF NOT EXISTS events ( + app_name TEXT NOT NULL, + user_id TEXT NOT NULL, + session_id TEXT NOT NULL, + event_index INTEGER NOT NULL, + event_id TEXT NOT NULL, + event_json TEXT NOT NULL, + PRIMARY KEY (app_name, user_id, session_id, event_index), + FOREIGN KEY (app_name, user_id, session_id) + REFERENCES sessions(app_name, user_id, session_id) + ON DELETE CASCADE + ); + + CREATE TABLE IF NOT EXISTS app_states ( + app_name TEXT PRIMARY KEY, + state_json TEXT NOT NULL + ); + + CREATE TABLE IF NOT EXISTS user_states ( + app_name TEXT NOT NULL, + user_id TEXT NOT NULL, + state_json TEXT NOT NULL, + PRIMARY KEY (app_name, user_id) + ); + "#, + ) + .map_err(|err| Error::msg(format!("failed to migrate sqlite session store: {err}")))?; + Ok(()) + } +} + +impl SessionStore for SqliteSessionStore { + fn create(&self, app_name: &str, user_id: &str, session_id: &str) -> Result { + let now = unix_millis(SystemTime::now())?; + let state = State::new(); + let state_json = state_to_json(&state)?; + let connection = self + .connection + .lock() + .map_err(|_| Error::msg("sqlite session store lock poisoned"))?; + let inserted = connection + .execute( + r#" + INSERT OR IGNORE INTO sessions + (app_name, user_id, session_id, state_json, created_at, updated_at) + VALUES (?1, ?2, ?3, ?4, ?5, ?6) + "#, + params![app_name, user_id, session_id, state_json, now, now], + ) + .map_err(|err| Error::msg(format!("failed to create sqlite session: {err}")))?; + if inserted == 0 { + return Err(Error::msg("session already exists")); + } + + Ok(Session::new(app_name, user_id, session_id)) + } + + fn get(&self, app_name: &str, user_id: &str, session_id: &str) -> Result { + let connection = self + .connection + .lock() + .map_err(|_| Error::msg("sqlite session store lock poisoned"))?; + load_session(&connection, app_name, user_id, session_id) + } + + fn list(&self, app_name: &str, user_id: &str) -> Result> { + let connection = self + .connection + .lock() + .map_err(|_| Error::msg("sqlite session store lock poisoned"))?; + let mut statement = connection + .prepare( + r#" + SELECT session_id + FROM sessions + WHERE app_name = ?1 AND user_id = ?2 + ORDER BY session_id + "#, + ) + .map_err(|err| Error::msg(format!("failed to list sqlite sessions: {err}")))?; + let ids = statement + .query_map(params![app_name, user_id], |row| row.get::<_, String>(0)) + .map_err(|err| Error::msg(format!("failed to list sqlite sessions: {err}")))? + .collect::, _>>() + .map_err(|err| Error::msg(format!("failed to read sqlite sessions: {err}")))?; + + ids.into_iter() + .map(|session_id| load_session(&connection, app_name, user_id, &session_id)) + .collect() + } + + fn delete(&self, app_name: &str, user_id: &str, session_id: &str) -> Result<()> { + let connection = self + .connection + .lock() + .map_err(|_| Error::msg("sqlite session store lock poisoned"))?; + let deleted = connection + .execute( + "DELETE FROM sessions WHERE app_name = ?1 AND user_id = ?2 AND session_id = ?3", + params![app_name, user_id, session_id], + ) + .map_err(|err| Error::msg(format!("failed to delete sqlite session: {err}")))?; + if deleted == 0 { + return Err(Error::msg("session not found")); + } + Ok(()) + } + + fn append_event( + &self, + app_name: &str, + user_id: &str, + session_id: &str, + event: Event, + ) -> Result<()> { + let now = unix_millis(event.timestamp)?; + let mut event = event; + let mut connection = self + .connection + .lock() + .map_err(|_| Error::msg("sqlite session store lock poisoned"))?; + let tx = connection + .transaction() + .map_err(|err| Error::msg(format!("failed to start sqlite transaction: {err}")))?; + + let session_state_json = tx + .query_row( + "SELECT state_json FROM sessions WHERE app_name = ?1 AND user_id = ?2 AND session_id = ?3", + params![app_name, user_id, session_id], + |row| row.get::<_, String>(0), + ) + .optional() + .map_err(|err| Error::msg(format!("failed to load sqlite session: {err}")))? + .ok_or_else(|| Error::msg("session not found"))?; + let mut session_state = state_from_json(&session_state_json)?; + let mut app_state = load_scoped_state(&tx, "app_states", app_name, None)?; + let mut user_state = load_scoped_state(&tx, "user_states", app_name, Some(user_id))?; + let mut persisted_state_delta = Map::new(); + + for (key, value) in event.actions.state_delta.clone() { + match state_scope(&key) { + StateScope::App => { + app_state.insert(key.clone(), value.clone()); + persisted_state_delta.insert(key, value); + } + StateScope::User => { + user_state.insert(key.clone(), value.clone()); + persisted_state_delta.insert(key, value); + } + StateScope::Session => { + session_state.insert(key.clone(), value.clone()); + persisted_state_delta.insert(key, value); + } + StateScope::Temp => {} + } + } + event.actions.state_delta = persisted_state_delta; + + save_app_state(&tx, app_name, &app_state)?; + save_user_state(&tx, app_name, user_id, &user_state)?; + + let event_index = tx + .query_row( + r#" + SELECT COALESCE(MAX(event_index), -1) + 1 + FROM events + WHERE app_name = ?1 AND user_id = ?2 AND session_id = ?3 + "#, + params![app_name, user_id, session_id], + |row| row.get::<_, i64>(0), + ) + .map_err(|err| Error::msg(format!("failed to allocate event index: {err}")))?; + tx.execute( + r#" + INSERT INTO events + (app_name, user_id, session_id, event_index, event_id, event_json) + VALUES (?1, ?2, ?3, ?4, ?5, ?6) + "#, + params![ + app_name, + user_id, + session_id, + event_index, + event.id, + event_to_json(&event)? + ], + ) + .map_err(|err| Error::msg(format!("failed to insert sqlite event: {err}")))?; + tx.execute( + r#" + UPDATE sessions + SET state_json = ?4, updated_at = ?5 + WHERE app_name = ?1 AND user_id = ?2 AND session_id = ?3 + "#, + params![ + app_name, + user_id, + session_id, + state_to_json(&session_state)?, + now + ], + ) + .map_err(|err| Error::msg(format!("failed to update sqlite session state: {err}")))?; + tx.commit() + .map_err(|err| Error::msg(format!("failed to commit sqlite transaction: {err}")))?; + Ok(()) + } +} + +fn load_session( + connection: &Connection, + app_name: &str, + user_id: &str, + session_id: &str, +) -> Result { + let state_json = connection + .query_row( + "SELECT state_json FROM sessions WHERE app_name = ?1 AND user_id = ?2 AND session_id = ?3", + params![app_name, user_id, session_id], + |row| row.get::<_, String>(0), + ) + .optional() + .map_err(|err| Error::msg(format!("failed to load sqlite session: {err}")))? + .ok_or_else(|| Error::msg("session not found"))?; + let mut session = Session::new(app_name, user_id, session_id); + session.state = state_from_json(&state_json)?; + session + .state + .extend(load_scoped_state(connection, "app_states", app_name, None)?); + session + .state + .extend(load_scoped_state(connection, "user_states", app_name, Some(user_id))?); + session.events = load_events(connection, app_name, user_id, session_id)?; + Ok(session) +} + +fn load_events( + connection: &Connection, + app_name: &str, + user_id: &str, + session_id: &str, +) -> Result> { + let mut statement = connection + .prepare( + r#" + SELECT event_json + FROM events + WHERE app_name = ?1 AND user_id = ?2 AND session_id = ?3 + ORDER BY event_index + "#, + ) + .map_err(|err| Error::msg(format!("failed to load sqlite events: {err}")))?; + statement + .query_map(params![app_name, user_id, session_id], |row| { + row.get::<_, String>(0) + }) + .map_err(|err| Error::msg(format!("failed to load sqlite events: {err}")))? + .map(|row| { + row.map_err(|err| Error::msg(format!("failed to read sqlite event: {err}"))) + .and_then(|json| event_from_json(&json)) + }) + .collect() +} + +fn load_scoped_state( + connection: &Connection, + table: &str, + app_name: &str, + user_id: Option<&str>, +) -> Result { + let state_json = match user_id { + Some(user_id) => connection + .query_row( + &format!("SELECT state_json FROM {table} WHERE app_name = ?1 AND user_id = ?2"), + params![app_name, user_id], + |row| row.get::<_, String>(0), + ) + .optional(), + None => connection + .query_row( + &format!("SELECT state_json FROM {table} WHERE app_name = ?1"), + params![app_name], + |row| row.get::<_, String>(0), + ) + .optional(), + } + .map_err(|err| Error::msg(format!("failed to load sqlite scoped state: {err}")))?; + + state_json + .map(|json| state_from_json(&json)) + .unwrap_or_else(|| Ok(State::new())) +} + +fn save_app_state(connection: &Connection, app_name: &str, state: &State) -> Result<()> { + connection + .execute( + r#" + INSERT INTO app_states (app_name, state_json) + VALUES (?1, ?2) + ON CONFLICT(app_name) DO UPDATE SET state_json = excluded.state_json + "#, + params![app_name, state_to_json(state)?], + ) + .map_err(|err| Error::msg(format!("failed to save app state: {err}")))?; + Ok(()) +} + +fn save_user_state( + connection: &Connection, + app_name: &str, + user_id: &str, + state: &State, +) -> Result<()> { + connection + .execute( + r#" + INSERT INTO user_states (app_name, user_id, state_json) + VALUES (?1, ?2, ?3) + ON CONFLICT(app_name, user_id) DO UPDATE SET state_json = excluded.state_json + "#, + params![app_name, user_id, state_to_json(state)?], + ) + .map_err(|err| Error::msg(format!("failed to save user state: {err}")))?; + Ok(()) +} + +fn state_to_json(state: &State) -> Result { + serde_json::to_string(state).map_err(|err| Error::msg(format!("failed to encode state: {err}"))) +} + +fn state_from_json(json: &str) -> Result { + serde_json::from_str(json).map_err(|err| Error::msg(format!("failed to decode state: {err}"))) +} + +fn event_to_json(event: &Event) -> Result { + serde_json::to_string(event).map_err(|err| Error::msg(format!("failed to encode event: {err}"))) +} + +fn event_from_json(json: &str) -> Result { + serde_json::from_str(json).map_err(|err| Error::msg(format!("failed to decode event: {err}"))) +} + +fn unix_millis(time: SystemTime) -> Result { + let duration = time + .duration_since(UNIX_EPOCH) + .map_err(|err| Error::msg(format!("timestamp before unix epoch: {err}")))?; + Ok(duration.as_millis() as i64) +} diff --git a/tests/database_session.rs b/tests/database_session.rs new file mode 100644 index 0000000..1bc999e --- /dev/null +++ b/tests/database_session.rs @@ -0,0 +1,105 @@ +use std::{ + fs, + time::{SystemTime, UNIX_EPOCH}, +}; + +use serde_json::json; + +use r_agent::{ + event::{Event, EventActions}, + session::{DatabaseSessionService, SessionService, SqliteSessionStore}, +}; + +fn sqlite_path() -> std::path::PathBuf { + std::env::temp_dir().join(format!( + "r_agent_session_{}.sqlite", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() + )) +} + +fn event_with_delta(id: &str, state_delta: Vec<(&str, serde_json::Value)>) -> Event { + let mut actions = EventActions::default(); + for (key, value) in state_delta { + actions.state_delta.insert(key.to_string(), value); + } + + Event { + id: id.to_string(), + invocation_id: "invocation-1".to_string(), + branch: None, + author: "agent".to_string(), + timestamp: SystemTime::UNIX_EPOCH, + content: None, + actions, + partial: false, + } +} + +#[test] +fn database_session_service_persists_events_and_state() { + let path = sqlite_path(); + let service = DatabaseSessionService::new(SqliteSessionStore::open(&path).unwrap()) + .unwrap(); + service.create("app", "user-1", "session-1").unwrap(); + service + .append_event( + "app", + "user-1", + "session-1", + event_with_delta("event-1", vec![("answer", json!(42))]), + ) + .unwrap(); + drop(service); + + let reopened = DatabaseSessionService::new(SqliteSessionStore::open(&path).unwrap()) + .unwrap(); + let session = reopened.get("app", "user-1", "session-1").unwrap(); + + assert_eq!(session.events.len(), 1); + assert_eq!(session.events[0].id, "event-1"); + assert_eq!(session.state.get("answer"), Some(&json!(42))); + + fs::remove_file(path).unwrap(); +} + +#[test] +fn database_session_service_applies_scoped_state_and_discards_temp_state() { + let path = sqlite_path(); + let service = DatabaseSessionService::new(SqliteSessionStore::open(&path).unwrap()) + .unwrap(); + service.create("app", "user-1", "s1").unwrap(); + service.create("app", "user-2", "s2").unwrap(); + service + .append_event( + "app", + "user-1", + "s1", + event_with_delta( + "event-1", + vec![ + ("app:timezone", json!("UTC")), + ("user:name", json!("Yayoi")), + ("temp:scratch", json!("discard")), + ], + ), + ) + .unwrap(); + + let same_user = service.get("app", "user-1", "s1").unwrap(); + let other_user = service.get("app", "user-2", "s2").unwrap(); + + assert_eq!(same_user.state.get("app:timezone"), Some(&json!("UTC"))); + assert_eq!(same_user.state.get("user:name"), Some(&json!("Yayoi"))); + assert_eq!(same_user.state.get("temp:scratch"), None); + assert_eq!(other_user.state.get("app:timezone"), Some(&json!("UTC"))); + assert_eq!(other_user.state.get("user:name"), None); + assert!(!same_user.events[0] + .actions + .state_delta + .contains_key("temp:scratch")); + + fs::remove_file(path).unwrap(); +} From f0807a222fe906f8cd660d59769cf6b3dec93336 Mon Sep 17 00:00:00 2001 From: YayoiNanoka <1159066485@qq.com> Date: Wed, 1 Jul 2026 20:43:00 +0800 Subject: [PATCH 7/7] style: format phase 1 changes --- src/runner/runner.rs | 3 +-- src/session/database.rs | 3 ++- src/session/sqlite.rs | 9 ++++++--- src/tool/toolset.rs | 5 +---- tests/database_session.rs | 19 +++++++++---------- tests/model_factory.rs | 17 ++++++++++------- tests/scoped_state.rs | 10 ++++++---- tests/toml_agent_loader.rs | 8 ++++---- tests/toolset.rs | 11 ++++++++--- 9 files changed, 47 insertions(+), 38 deletions(-) diff --git a/src/runner/runner.rs b/src/runner/runner.rs index baec0b2..d86a6ac 100644 --- a/src/runner/runner.rs +++ b/src/runner/runner.rs @@ -10,8 +10,7 @@ use crate::{ Error, Result, agent::{ Agent, AsyncEventStream, EventStream, IntoAgentArc, InvocationContext, - event_stream_from_result, - find_agent, + event_stream_from_result, find_agent, }, artifact::{ArtifactFacade, ArtifactService, InMemoryArtifactService}, content::Content, diff --git a/src/session/database.rs b/src/session/database.rs index e47c708..5835a5a 100644 --- a/src/session/database.rs +++ b/src/session/database.rs @@ -54,6 +54,7 @@ impl SessionService for DatabaseSessionService { session_id: &str, event: Event, ) -> Result<()> { - self.store.append_event(app_name, user_id, session_id, event) + self.store + .append_event(app_name, user_id, session_id, event) } } diff --git a/src/session/sqlite.rs b/src/session/sqlite.rs index 25d68fe..9efdd2a 100644 --- a/src/session/sqlite.rs +++ b/src/session/sqlite.rs @@ -276,9 +276,12 @@ fn load_session( session .state .extend(load_scoped_state(connection, "app_states", app_name, None)?); - session - .state - .extend(load_scoped_state(connection, "user_states", app_name, Some(user_id))?); + session.state.extend(load_scoped_state( + connection, + "user_states", + app_name, + Some(user_id), + )?); session.events = load_events(connection, app_name, user_id, session_id)?; Ok(session) } diff --git a/src/tool/toolset.rs b/src/tool/toolset.rs index 2d1af13..258128d 100644 --- a/src/tool/toolset.rs +++ b/src/tool/toolset.rs @@ -59,10 +59,7 @@ where Arc::new(move |_ctx, tool| allowed.contains(tool.name())) } -pub fn filter_toolset( - toolset: Arc, - predicate: ToolPredicate, -) -> Arc { +pub fn filter_toolset(toolset: Arc, predicate: ToolPredicate) -> Arc { Arc::new(FilterToolset { toolset, predicate }) } diff --git a/tests/database_session.rs b/tests/database_session.rs index 1bc999e..8510012 100644 --- a/tests/database_session.rs +++ b/tests/database_session.rs @@ -41,8 +41,7 @@ fn event_with_delta(id: &str, state_delta: Vec<(&str, serde_json::Value)>) -> Ev #[test] fn database_session_service_persists_events_and_state() { let path = sqlite_path(); - let service = DatabaseSessionService::new(SqliteSessionStore::open(&path).unwrap()) - .unwrap(); + let service = DatabaseSessionService::new(SqliteSessionStore::open(&path).unwrap()).unwrap(); service.create("app", "user-1", "session-1").unwrap(); service .append_event( @@ -54,8 +53,7 @@ fn database_session_service_persists_events_and_state() { .unwrap(); drop(service); - let reopened = DatabaseSessionService::new(SqliteSessionStore::open(&path).unwrap()) - .unwrap(); + let reopened = DatabaseSessionService::new(SqliteSessionStore::open(&path).unwrap()).unwrap(); let session = reopened.get("app", "user-1", "session-1").unwrap(); assert_eq!(session.events.len(), 1); @@ -68,8 +66,7 @@ fn database_session_service_persists_events_and_state() { #[test] fn database_session_service_applies_scoped_state_and_discards_temp_state() { let path = sqlite_path(); - let service = DatabaseSessionService::new(SqliteSessionStore::open(&path).unwrap()) - .unwrap(); + let service = DatabaseSessionService::new(SqliteSessionStore::open(&path).unwrap()).unwrap(); service.create("app", "user-1", "s1").unwrap(); service.create("app", "user-2", "s2").unwrap(); service @@ -96,10 +93,12 @@ fn database_session_service_applies_scoped_state_and_discards_temp_state() { assert_eq!(same_user.state.get("temp:scratch"), None); assert_eq!(other_user.state.get("app:timezone"), Some(&json!("UTC"))); assert_eq!(other_user.state.get("user:name"), None); - assert!(!same_user.events[0] - .actions - .state_delta - .contains_key("temp:scratch")); + assert!( + !same_user.events[0] + .actions + .state_delta + .contains_key("temp:scratch") + ); fs::remove_file(path).unwrap(); } diff --git a/tests/model_factory.rs b/tests/model_factory.rs index dbb9d06..3bf8593 100644 --- a/tests/model_factory.rs +++ b/tests/model_factory.rs @@ -18,16 +18,19 @@ fn default_model_factory_creates_openai_compatible_model() { #[test] fn default_model_factory_rejects_unknown_protocol() { let result = DefaultModelFactory::new().create(ModelConfig { - provider: "example".to_string(), - protocol: "unknown-protocol".to_string(), - model: "example-model".to_string(), - base_url: None, - api_key_env: None, - }); + provider: "example".to_string(), + protocol: "unknown-protocol".to_string(), + model: "example-model".to_string(), + base_url: None, + api_key_env: None, + }); let err = match result { Ok(_) => panic!("expected unsupported protocol error"), Err(err) => err, }; - assert_eq!(err.to_string(), "unsupported model protocol: unknown-protocol"); + assert_eq!( + err.to_string(), + "unsupported model protocol: unknown-protocol" + ); } diff --git a/tests/scoped_state.rs b/tests/scoped_state.rs index f2f1181..66b96cb 100644 --- a/tests/scoped_state.rs +++ b/tests/scoped_state.rs @@ -95,10 +95,12 @@ fn temp_state_is_not_persisted_to_state_or_event_delta() { assert_eq!(session.state.get("temp:scratch"), None); assert_eq!(session.state.get("session_value"), Some(&json!("keep"))); - assert!(!session.events[0] - .actions - .state_delta - .contains_key("temp:scratch")); + assert!( + !session.events[0] + .actions + .state_delta + .contains_key("temp:scratch") + ); } #[test] diff --git a/tests/toml_agent_loader.rs b/tests/toml_agent_loader.rs index ab9bd1e..b9b114a 100644 --- a/tests/toml_agent_loader.rs +++ b/tests/toml_agent_loader.rs @@ -28,10 +28,10 @@ fn registry() -> ToolRegistry { let mut registry = ToolRegistry::new(); registry .register(Arc::new(FunctionTool::new( - "get_weather", - "Gets weather.", - |_ctx: &mut ToolContext, _args: Value| Ok(json!({ "forecast": "sunny" })), - ))) + "get_weather", + "Gets weather.", + |_ctx: &mut ToolContext, _args: Value| Ok(json!({ "forecast": "sunny" })), + ))) .unwrap(); registry } diff --git a/tests/toolset.rs b/tests/toolset.rs index 793977e..4895dd5 100644 --- a/tests/toolset.rs +++ b/tests/toolset.rs @@ -1,5 +1,5 @@ -use std::time::SystemTime; use std::sync::Arc; +use std::time::SystemTime; use serde_json::{Value, json}; @@ -90,7 +90,9 @@ fn dynamic_toolset_exposes_tools_from_current_session_state() { callbacks: LlmAgentCallbacks::default(), }); let session_service = InMemorySessionService::new(); - session_service.create("app", "user-1", "session-1").unwrap(); + session_service + .create("app", "user-1", "session-1") + .unwrap(); let mut actions = EventActions::default(); actions .state_delta @@ -175,5 +177,8 @@ fn llm_agent_rejects_tool_call_that_was_not_visible_to_model() { ) .unwrap_err(); - assert_eq!(err.to_string(), "tool not visible in current invocation: hidden_tool"); + assert_eq!( + err.to_string(), + "tool not visible in current invocation: hidden_tool" + ); }