agent_base/engine/
builder.rs

1use std::collections::HashSet;
2use std::sync::Arc;
3
4use crate::llm::{LlmClient, ReasoningConfig};
5use crate::tool::{Tool, ToolPolicy, ToolRegistry};
6use crate::types::{
7    AgentConfig, AtomicU64SessionIdGenerator, ResponseFormat, RetryConfig, SessionIdGenerator,
8};
9
10use super::AgentRuntime;
11use super::approval::ApprovalHandler;
12use super::context::ContextWindowManager;
13use super::middleware::{Middleware, MiddlewareRef};
14use super::recovery::{StopOnError, ToolErrorRecovery};
15use super::session_store::{InMemorySessionStore, SessionStore};
16
17pub struct AgentBuilder {
18    client: Arc<dyn LlmClient>,
19    config: AgentConfig,
20    tools: ToolRegistry,
21    tool_names: HashSet<String>,
22    approval_handler: Option<Arc<dyn ApprovalHandler>>,
23    tool_policy: Option<Arc<dyn ToolPolicy>>,
24    middlewares: Vec<MiddlewareRef>,
25    context_manager: Option<ContextWindowManager>,
26    session_store: Option<Arc<dyn SessionStore>>,
27    error_recovery: Option<Arc<dyn ToolErrorRecovery>>,
28    event_bus_capacity: usize,
29    session_id_generator: Option<Arc<dyn SessionIdGenerator>>,
30}
31
32impl AgentBuilder {
33    pub fn new(client: Arc<dyn LlmClient>) -> Self {
34        Self {
35            client,
36            config: AgentConfig::default(),
37            tools: ToolRegistry::default(),
38            tool_names: HashSet::new(),
39            approval_handler: None,
40            tool_policy: None,
41            middlewares: Vec::new(),
42            context_manager: None,
43            session_store: None,
44            error_recovery: None,
45            event_bus_capacity: 2048,
46            session_id_generator: None,
47        }
48    }
49
50    pub fn event_bus_capacity(mut self, capacity: usize) -> Self {
51        self.event_bus_capacity = capacity;
52        self
53    }
54
55    pub fn session_id_generator(mut self, generator: Arc<dyn SessionIdGenerator>) -> Self {
56        self.session_id_generator = Some(generator);
57        self
58    }
59
60    pub fn system_prompt(mut self, system_prompt: impl Into<String>) -> Self {
61        self.config.system_prompt = Some(system_prompt.into());
62        self
63    }
64
65    /// Set whether to include the reasoning content in LLM responses.
66    ///
67    /// Controls whether the `reasoning_content` field from the LLM is forwarded
68    /// to consumers (i.e., "show the thinking process").
69    /// See [`AgentConfig::enable_thought`] for the distinction from `enable_thinking()`.
70    pub fn enable_thought(mut self, enable: bool) -> Self {
71        self.config.enable_thought = enable;
72        self
73    }
74
75    pub fn reasoning(mut self, config: ReasoningConfig) -> Self {
76        self.config.reasoning = Some(config);
77        self
78    }
79
80    /// Set whether to enable the model's extended thinking / reasoning mode.
81    ///
82    /// Controls whether the model performs deep reasoning (i.e., "enable thinking mode").
83    /// See [`AgentConfig::enable_thought`] for the distinction from `enable_thought()`.
84    pub fn enable_thinking(mut self, enable: bool) -> Self {
85        let mut config = self.config.reasoning.take().unwrap_or_default();
86        config.enabled = Some(enable);
87        self.config.reasoning = Some(config);
88        self
89    }
90
91    pub fn thinking_budget(mut self, budget: u64) -> Self {
92        let mut config = self.config.reasoning.take().unwrap_or_default();
93        config.budget_tokens = Some(budget);
94        self.config.reasoning = Some(config);
95        self
96    }
97
98    pub fn tool_timeout(mut self, timeout_ms: u64) -> Self {
99        self.config.tool.tool_timeout_ms = Some(timeout_ms);
100        self
101    }
102
103    pub fn max_tool_output_chars(mut self, max_chars: usize) -> Self {
104        self.config.tool.max_tool_output_chars = Some(max_chars);
105        self
106    }
107
108    pub fn register_tool(mut self, tool: impl Tool + 'static) -> Self {
109        self.tools.register(tool);
110        self
111    }
112
113    pub fn register_tool_arc(mut self, tool: Arc<dyn Tool>) -> Self {
114        self.tools.register_arc(tool);
115        self
116    }
117
118    pub fn approval_handler(mut self, handler: Arc<dyn ApprovalHandler>) -> Self {
119        self.approval_handler = Some(handler);
120        self
121    }
122
123    pub fn tool_policy(mut self, policy: Arc<dyn ToolPolicy>) -> Self {
124        self.tool_policy = Some(policy);
125        self
126    }
127
128    pub fn middleware(mut self, mw: impl Middleware + 'static) -> Self {
129        self.middlewares.push(Arc::new(mw));
130        self
131    }
132
133    pub fn context_window(mut self, max_tokens: usize) -> Self {
134        self.context_manager = Some(ContextWindowManager::new(max_tokens));
135        self
136    }
137
138    pub fn context_window_manager(mut self, manager: ContextWindowManager) -> Self {
139        self.context_manager = Some(manager);
140        self
141    }
142
143    pub fn response_format(mut self, format: ResponseFormat) -> Self {
144        self.config.llm.response_format = Some(format);
145        self
146    }
147
148    pub fn llm_retry(mut self, retry: RetryConfig) -> Self {
149        self.config.llm.llm_retry = Some(retry);
150        self
151    }
152
153    pub fn session_store(mut self, store: Arc<dyn SessionStore>) -> Self {
154        self.session_store = Some(store);
155        self
156    }
157
158    pub fn error_recovery(mut self, recovery: Arc<dyn ToolErrorRecovery>) -> Self {
159        self.error_recovery = Some(recovery);
160        self
161    }
162
163    pub fn max_sessions(mut self, max: usize) -> Self {
164        self.config.session.max_sessions = Some(max);
165        self
166    }
167
168    pub fn max_turns_per_session(mut self, max: usize) -> Self {
169        self.config.session.max_turns_per_session = Some(max);
170        self
171    }
172
173    /// Cap the number of react-loop iterations allowed for a *single* run (one user
174    /// input). Distinct from [`Self::max_turns_per_session`], which caps turns across
175    /// the whole session. When unset, falls back to `DEFAULT_MAX_TURNS` (50).
176    pub fn execution_max_turns(mut self, max: u32) -> Self {
177        self.config.execution.max_turns = Some(max);
178        self
179    }
180
181    pub fn max_message_tokens(mut self, max: usize) -> Self {
182        self.config.session.max_message_tokens = Some(max);
183        self
184    }
185
186    pub fn tool_error_retry_prompt(mut self, prompt: impl Into<String>) -> Self {
187        self.config.tool.tool_error_retry_prompt = Some(prompt.into());
188        self
189    }
190
191    pub fn language(mut self, language: crate::types::Language) -> Self {
192        self.config.language = language;
193        self
194    }
195
196    /// Conditionally chain a builder call: apply `f` only when `value` is `Some`.
197    ///
198    /// # Example
199    /// ```ignore
200    /// let builder = AgentBuilder::new(client)
201    ///     .apply_if(config.timeout, |b, t| b.tool_timeout(t))
202    ///     .apply_if(config.max_chars, |b, c| b.max_tool_output_chars(c));
203    /// ```
204    pub fn apply_if<T>(self, value: Option<T>, f: impl FnOnce(Self, T) -> Self) -> Self {
205        match value {
206            Some(v) => f(self, v),
207            None => self,
208        }
209    }
210
211    pub fn build(self) -> crate::types::AgentResult<AgentRuntime> {
212        self.config.validate()?;
213
214        tracing::info!(
215            tool_count = self.tools.len(),
216            middleware_count = self.middlewares.len(),
217            has_approval = self.approval_handler.is_some(),
218            has_context_window = self.context_manager.is_some(),
219            "building agent runtime"
220        );
221
222        let event_bus = super::runtime::EventBus::new(self.event_bus_capacity);
223
224        // Inject EventBus into framework-provided tools that need it
225        // (e.g., UpdatePlanTool via FrameworkTool::set_event_bus)
226        self.tools.inject_event_bus(&event_bus);
227
228        let session_store = self
229            .session_store
230            .unwrap_or_else(|| Arc::new(InMemorySessionStore::new()));
231        let error_recovery = self.error_recovery.unwrap_or_else(|| Arc::new(StopOnError));
232        let session_id_generator = self
233            .session_id_generator
234            .unwrap_or_else(|| Arc::new(AtomicU64SessionIdGenerator::default()));
235
236        let session_manager = super::runtime::SessionManager::new(
237            session_id_generator,
238            session_store,
239            self.config.session.clone(),
240        );
241
242        let llm_engine = super::runtime::LlmEngine::new(self.client.clone(), event_bus.clone());
243
244        let tool_engine = super::runtime::ToolEngine::new(
245            self.tools,
246            self.approval_handler,
247            self.tool_policy,
248            error_recovery,
249            event_bus.clone(),
250        );
251
252        let runner = Arc::new(super::runtime::RuntimeCore::new(
253            self.config,
254            llm_engine,
255            tool_engine,
256            session_manager,
257            event_bus,
258            self.context_manager,
259            self.middlewares,
260        ));
261
262        Ok(AgentRuntime { runner })
263    }
264}
265
266#[cfg(test)]
267mod tests {
268    use super::*;
269    use crate::llm::StreamChunk;
270    use crate::types::{AgentResult, ChatMessage, ResponseFormat};
271    use async_trait::async_trait;
272    use futures_core::Stream;
273    use serde_json::Value;
274    use std::pin::Pin;
275
276    struct DummyClient;
277
278    #[async_trait]
279    impl LlmClient for DummyClient {
280        async fn chat(
281            &self,
282            _messages: &[ChatMessage],
283            _tools: &[Value],
284            _reasoning: Option<&ReasoningConfig>,
285            _response_format: Option<&ResponseFormat>,
286        ) -> AgentResult<Value> {
287            Ok(Value::Null)
288        }
289
290        async fn chat_stream(
291            &self,
292            _messages: &[ChatMessage],
293            _tools: &[Value],
294            _reasoning: Option<&ReasoningConfig>,
295            _response_format: Option<&ResponseFormat>,
296        ) -> AgentResult<Pin<Box<dyn Stream<Item = AgentResult<StreamChunk>> + Send>>> {
297            unimplemented!("not used in builder tests")
298        }
299
300        fn capabilities(&self) -> crate::llm::LlmCapabilities {
301            crate::llm::LlmCapabilities::default()
302        }
303    }
304
305    #[test]
306    fn execution_max_turns_writes_per_run_config() {
307        let client: Arc<dyn LlmClient> = Arc::new(DummyClient);
308        let builder = AgentBuilder::new(client).execution_max_turns(200);
309        // The `config` field is private to this module, so the test can assert directly.
310        assert_eq!(builder.config.execution.max_turns, Some(200));
311    }
312
313    #[test]
314    fn execution_max_turns_defaults_to_none() {
315        let client: Arc<dyn LlmClient> = Arc::new(DummyClient);
316        let builder = AgentBuilder::new(client);
317        assert_eq!(builder.config.execution.max_turns, None);
318    }
319}
agent_base/engine/builder.rs

agent_base/engine/
builder.rs
