agent_base/engine/
builder.rs

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

agent_base/engine/
builder.rs
