agent_base/engine/
builder.rs

1use std::sync::Arc;
2
3use crate::llm::{ReasoningConfig, StreamClient};
4use crate::tool::{Tool, ToolPolicy, ToolRegistry};
5use crate::types::{
6    AgentConfig, AtomicU64SessionIdGenerator, ConvertToLlmFn, ResponseFormat, RetryConfig,
7    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 StreamClient>,
19    config: AgentConfig,
20    tools: ToolRegistry,
21    approval_handler: Option<Arc<dyn ApprovalHandler>>,
22    tool_policy: Option<Arc<dyn ToolPolicy>>,
23    middlewares: Vec<MiddlewareRef>,
24    context_manager: Option<ContextWindowManager>,
25    session_store: Option<Arc<dyn SessionStore>>,
26    error_recovery: Option<Arc<dyn ToolErrorRecovery>>,
27    event_bus_capacity: usize,
28    session_id_generator: Option<Arc<dyn SessionIdGenerator>>,
29    convert_to_llm: Option<ConvertToLlmFn>,
30}
31
32impl AgentBuilder {
33    pub fn new(client: Arc<dyn StreamClient>) -> Self {
34        Self {
35            client,
36            config: AgentConfig::default(),
37            tools: ToolRegistry::default(),
38            approval_handler: None,
39            tool_policy: None,
40            middlewares: Vec::new(),
41            context_manager: None,
42            session_store: None,
43            error_recovery: None,
44            event_bus_capacity: 2048,
45            session_id_generator: None,
46            convert_to_llm: 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    /// Set a callback to transform messages before they are sent to the LLM.
61    ///
62    /// The default behavior (when `None`) is to filter out
63    /// [`ChatMessage::Custom`] variants, which most providers don't understand.
64    /// Override this to inject custom serialization logic for application-specific
65    /// message types.
66    pub fn convert_to_llm(mut self, cb: ConvertToLlmFn) -> Self {
67        self.convert_to_llm = Some(cb);
68        self
69    }
70
71    pub fn system_prompt(mut self, system_prompt: impl Into<String>) -> Self {
72        self.config.system_prompt = Some(system_prompt.into());
73        self
74    }
75
76    /// Set whether to include the reasoning content in LLM responses.
77    ///
78    /// Controls whether the `reasoning_content` field from the LLM is forwarded
79    /// to consumers (i.e., "show the thinking process").
80    /// See [`AgentConfig::enable_thought`] for the distinction from `enable_thinking()`.
81    pub fn enable_thought(mut self, enable: bool) -> Self {
82        self.config.enable_thought = enable;
83        self
84    }
85
86    pub fn reasoning(mut self, config: ReasoningConfig) -> Self {
87        self.config.reasoning = Some(config);
88        self
89    }
90
91    /// Set whether to enable the model's extended thinking / reasoning mode.
92    ///
93    /// Controls whether the model performs deep reasoning (i.e., "enable thinking mode").
94    /// See [`AgentConfig::enable_thought`] for the distinction from `enable_thought()`.
95    pub fn enable_thinking(mut self, enable: bool) -> Self {
96        let mut config = self.config.reasoning.take().unwrap_or_default();
97        config.enabled = Some(enable);
98        self.config.reasoning = Some(config);
99        self
100    }
101
102    pub fn thinking_budget(mut self, budget: u64) -> Self {
103        let mut config = self.config.reasoning.take().unwrap_or_default();
104        config.budget_tokens = Some(budget);
105        self.config.reasoning = Some(config);
106        self
107    }
108
109    pub fn tool_timeout(mut self, timeout_ms: u64) -> Self {
110        self.config.tool.tool_timeout_ms = Some(timeout_ms);
111        self
112    }
113
114    pub fn max_tool_output_chars(mut self, max_chars: usize) -> Self {
115        self.config.tool.max_tool_output_chars = Some(max_chars);
116        self
117    }
118
119    pub fn register_tool(mut self, tool: impl Tool + 'static) -> Self {
120        self.tools.register(tool);
121        self
122    }
123
124    pub fn register_tool_arc(mut self, tool: Arc<dyn Tool>) -> Self {
125        self.tools.register_arc(tool);
126        self
127    }
128
129    pub fn approval_handler(mut self, handler: Arc<dyn ApprovalHandler>) -> Self {
130        self.approval_handler = Some(handler);
131        self
132    }
133
134    pub fn tool_policy(mut self, policy: Arc<dyn ToolPolicy>) -> Self {
135        self.tool_policy = Some(policy);
136        self
137    }
138
139    pub fn middleware(mut self, mw: impl Middleware + 'static) -> Self {
140        self.middlewares.push(Arc::new(mw));
141        self
142    }
143
144    pub fn context_window(mut self, max_tokens: usize) -> Self {
145        self.context_manager = Some(ContextWindowManager::new(max_tokens));
146        self
147    }
148
149    pub fn context_window_manager(mut self, manager: ContextWindowManager) -> Self {
150        self.context_manager = Some(manager);
151        self
152    }
153
154    pub fn response_format(mut self, format: ResponseFormat) -> Self {
155        self.config.llm.response_format = Some(format);
156        self
157    }
158
159    pub fn llm_retry(mut self, retry: RetryConfig) -> Self {
160        self.config.llm.llm_retry = Some(retry);
161        self
162    }
163
164    pub fn session_store(mut self, store: Arc<dyn SessionStore>) -> Self {
165        self.session_store = Some(store);
166        self
167    }
168
169    pub fn error_recovery(mut self, recovery: Arc<dyn ToolErrorRecovery>) -> Self {
170        self.error_recovery = Some(recovery);
171        self
172    }
173
174    pub fn max_sessions(mut self, max: usize) -> Self {
175        self.config.session.max_sessions = Some(max);
176        self
177    }
178
179    pub fn max_turns_per_session(mut self, max: usize) -> Self {
180        self.config.session.max_turns_per_session = Some(max);
181        self
182    }
183
184    /// Cap the number of react-loop iterations allowed for a *single* run (one user
185    /// input). Distinct from [`Self::max_turns_per_session`], which caps turns across
186    /// the whole session. When unset, falls back to `DEFAULT_MAX_TURNS` (50).
187    pub fn execution_max_turns(mut self, max: u32) -> Self {
188        self.config.execution.max_turns = Some(max);
189        self
190    }
191
192    pub fn max_message_tokens(mut self, max: usize) -> Self {
193        self.config.session.max_message_tokens = Some(max);
194        self
195    }
196
197    pub fn tool_error_retry_prompt(mut self, prompt: impl Into<String>) -> Self {
198        self.config.tool.tool_error_retry_prompt = Some(prompt.into());
199        self
200    }
201
202    pub fn language(mut self, language: crate::types::Language) -> Self {
203        self.config.language = language;
204        self
205    }
206
207    /// Conditionally chain a builder call: apply `f` only when `value` is `Some`.
208    ///
209    /// # Example
210    /// ```ignore
211    /// let builder = AgentBuilder::new(client)
212    ///     .apply_if(config.timeout, |b, t| b.tool_timeout(t))
213    ///     .apply_if(config.max_chars, |b, c| b.max_tool_output_chars(c));
214    /// ```
215    pub fn apply_if<T>(self, value: Option<T>, f: impl FnOnce(Self, T) -> Self) -> Self {
216        match value {
217            Some(v) => f(self, v),
218            None => self,
219        }
220    }
221
222    pub fn build(self) -> crate::types::AgentResult<AgentRuntime> {
223        self.config.validate()?;
224
225        tracing::info!(
226            tool_count = self.tools.len(),
227            middleware_count = self.middlewares.len(),
228            has_approval = self.approval_handler.is_some(),
229            has_context_window = self.context_manager.is_some(),
230            "building agent runtime"
231        );
232
233        let event_bus = super::runtime::EventBus::new(self.event_bus_capacity);
234
235        let session_store = self
236            .session_store
237            .unwrap_or_else(|| Arc::new(InMemorySessionStore::new()));
238        let error_recovery = self.error_recovery.unwrap_or_else(|| Arc::new(StopOnError));
239        let session_id_generator = self
240            .session_id_generator
241            .unwrap_or_else(|| Arc::new(AtomicU64SessionIdGenerator::default()));
242
243        let session_manager = super::runtime::SessionManager::new(
244            session_id_generator,
245            session_store,
246            self.config.session.clone(),
247        );
248
249        let llm_engine = super::runtime::LlmEngine::new(self.client.clone(), event_bus.clone());
250
251        let tool_engine = super::runtime::ToolEngine::new(
252            self.tools,
253            self.approval_handler,
254            self.tool_policy,
255            error_recovery,
256            event_bus.clone(),
257        );
258
259        let runner = Arc::new(super::runtime::RuntimeCore::new(
260            self.config,
261            llm_engine,
262            tool_engine,
263            session_manager,
264            event_bus,
265            self.context_manager,
266            self.middlewares,
267            self.convert_to_llm,
268        ));
269
270        Ok(AgentRuntime { runner })
271    }
272}
273
274#[cfg(test)]
275mod tests {
276    use super::*;
277    use crate::engine::DenyAllApprovalHandler;
278    use crate::llm::{LlmClient, ReasoningEffort, StreamChunk};
279    use crate::tool::{Content, ToolContext};
280    use crate::types::{
281        AgentError, AgentResult, ApprovalRequest, ChatMessage, Language, ResponseFormat,
282    };
283    use async_trait::async_trait;
284    use futures_core::Stream;
285    use serde_json::Value;
286    use std::pin::Pin;
287
288    struct DummyClient;
289
290    #[async_trait]
291    impl LlmClient for DummyClient {
292        async fn chat(
293            &self,
294            _messages: &[ChatMessage],
295            _tools: &[Value],
296            _reasoning: Option<&ReasoningConfig>,
297            _response_format: Option<&ResponseFormat>,
298        ) -> AgentResult<Value> {
299            Ok(Value::Null)
300        }
301
302        async fn chat_stream(
303            &self,
304            _messages: &[ChatMessage],
305            _tools: &[Value],
306            _reasoning: Option<&ReasoningConfig>,
307            _response_format: Option<&ResponseFormat>,
308        ) -> AgentResult<Pin<Box<dyn Stream<Item = AgentResult<StreamChunk>> + Send>>> {
309            unimplemented!("not used in builder tests")
310        }
311
312        fn capabilities(&self) -> crate::llm::LlmCapabilities {
313            crate::llm::LlmCapabilities::default()
314        }
315    }
316
317    #[test]
318    fn execution_max_turns_writes_per_run_config() {
319        let client = crate::llm::adapt(Arc::new(DummyClient));
320        let builder = AgentBuilder::new(client).execution_max_turns(200);
321        // The `config` field is private to this module, so the test can assert directly.
322        assert_eq!(builder.config.execution.max_turns, Some(200));
323    }
324
325    #[test]
326    fn execution_max_turns_defaults_to_none() {
327        let client = crate::llm::adapt(Arc::new(DummyClient));
328        let builder = AgentBuilder::new(client);
329        assert_eq!(builder.config.execution.max_turns, None);
330    }
331
332    fn b() -> AgentBuilder {
333        AgentBuilder::new(crate::llm::adapt(Arc::new(DummyClient)))
334    }
335
336    struct NoopTool;
337
338    #[async_trait]
339    impl Tool for NoopTool {
340        fn name(&self) -> &'static str {
341            "noop"
342        }
343        fn description(&self) -> &'static str {
344            "noop tool"
345        }
346        fn schema(&self) -> Value {
347            serde_json::json!({"type": "object"})
348        }
349        async fn call(&self, _args: &Value, _ctx: &ToolContext) -> AgentResult<Vec<Content>> {
350            Ok(vec![Content::text("ok")])
351        }
352    }
353
354    struct AutoApprovePolicy;
355
356    #[async_trait]
357    impl ToolPolicy for AutoApprovePolicy {
358        async fn evaluate_approval(
359            &self,
360            _tool_name: &str,
361            _args: &Value,
362        ) -> Option<ApprovalRequest> {
363            None
364        }
365    }
366
367    struct NoopMiddleware;
368
369    impl Middleware for NoopMiddleware {}
370
371    #[test]
372    fn system_prompt_sets_config() {
373        assert_eq!(
374            b().system_prompt("be helpful")
375                .config
376                .system_prompt
377                .as_deref(),
378            Some("be helpful")
379        );
380    }
381
382    #[test]
383    fn enable_thought_sets_config() {
384        assert!(b().enable_thought(true).config.enable_thought);
385    }
386
387    #[test]
388    fn reasoning_sets_config() {
389        let rc = ReasoningConfig {
390            enabled: Some(true),
391            budget_tokens: Some(64),
392            effort: Some(ReasoningEffort::Medium),
393        };
394        let builder = b().reasoning(rc);
395        let got = builder.config.reasoning.as_ref().unwrap();
396        assert_eq!(got.enabled, Some(true));
397        assert_eq!(got.budget_tokens, Some(64));
398        assert!(matches!(got.effort.as_ref(), Some(ReasoningEffort::Medium)));
399    }
400
401    #[test]
402    fn enable_thinking_and_budget_set_reasoning() {
403        let builder = b().enable_thinking(true).thinking_budget(128);
404        let got = builder.config.reasoning.as_ref().unwrap();
405        assert_eq!(got.enabled, Some(true));
406        assert_eq!(got.budget_tokens, Some(128));
407    }
408
409    #[test]
410    fn tool_limits_set_config() {
411        let builder = b().tool_timeout(5_000).max_tool_output_chars(1_024);
412        assert_eq!(builder.config.tool.tool_timeout_ms, Some(5_000));
413        assert_eq!(builder.config.tool.max_tool_output_chars, Some(1_024));
414    }
415
416    #[test]
417    fn register_tool_adds_to_registry() {
418        assert_eq!(b().register_tool(NoopTool).tools.len(), 1);
419    }
420
421    #[test]
422    fn approval_handler_and_tool_policy_are_set() {
423        let builder = b()
424            .approval_handler(Arc::new(DenyAllApprovalHandler))
425            .tool_policy(Arc::new(AutoApprovePolicy));
426        assert!(builder.approval_handler.is_some());
427        assert!(builder.tool_policy.is_some());
428    }
429
430    #[test]
431    fn middleware_and_context_window_are_set() {
432        let builder = b().middleware(NoopMiddleware).context_window(8_000);
433        assert_eq!(builder.middlewares.len(), 1);
434        assert!(builder.context_manager.is_some());
435    }
436
437    #[test]
438    fn response_format_and_retry_set_config() {
439        let builder = b()
440            .response_format(ResponseFormat::JsonObject)
441            .llm_retry(RetryConfig::default().max_retries(5));
442        assert!(builder.config.llm.response_format.is_some());
443        assert_eq!(
444            builder.config.llm.llm_retry.as_ref().unwrap().max_retries,
445            5
446        );
447    }
448
449    #[test]
450    fn session_store_and_error_recovery_are_set() {
451        let builder = b()
452            .session_store(Arc::new(InMemorySessionStore::new()))
453            .error_recovery(Arc::new(StopOnError));
454        assert!(builder.session_store.is_some());
455        assert!(builder.error_recovery.is_some());
456    }
457
458    #[test]
459    fn session_limits_set_config() {
460        let builder = b()
461            .max_sessions(10)
462            .max_turns_per_session(20)
463            .max_message_tokens(30);
464        assert_eq!(builder.config.session.max_sessions, Some(10));
465        assert_eq!(builder.config.session.max_turns_per_session, Some(20));
466        assert_eq!(builder.config.session.max_message_tokens, Some(30));
467    }
468
469    #[test]
470    fn tool_error_retry_prompt_and_language_set_config() {
471        let builder = b()
472            .tool_error_retry_prompt("try again")
473            .language(Language::Zh);
474        assert_eq!(
475            builder.config.tool.tool_error_retry_prompt.as_deref(),
476            Some("try again")
477        );
478        assert_eq!(builder.config.language, Language::Zh);
479    }
480
481    #[test]
482    fn apply_if_applies_when_some_and_skips_when_none() {
483        let applied = b().apply_if(Some(3_000_u64), |b, t| b.tool_timeout(t));
484        assert_eq!(applied.config.tool.tool_timeout_ms, Some(3_000));
485
486        let skipped = b().apply_if(None, |b, t| b.tool_timeout(t));
487        assert_eq!(skipped.config.tool.tool_timeout_ms, None);
488    }
489
490    #[test]
491    fn build_ok_with_defaults() {
492        assert!(b().build().is_ok());
493    }
494
495    #[test]
496    fn build_err_on_invalid_config() {
497        assert!(matches!(
498            b().execution_max_turns(0).build(),
499            Err(AgentError::ConfigError(_))
500        ));
501    }
502}
agent_base/engine/builder.rs

agent_base/engine/
builder.rs
