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

agent_base/engine/
builder.rs
