rig_core/agent/
builder.rs

1use std::{collections::HashMap, sync::Arc};
2
3use schemars::{JsonSchema, Schema, schema_for};
4
5use crate::{
6    agent::hook::{AgentHook, HookStack},
7    completion::{CompletionModel, Document},
8    memory::ConversationMemory,
9    message::ToolChoice,
10    tool::{
11        Tool, ToolDyn, ToolSet,
12        server::{ToolServer, ToolServerHandle},
13    },
14    vector_store::VectorStoreIndexDyn,
15};
16
17#[cfg(feature = "rmcp")]
18#[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
19use crate::tool::rmcp::McpTool as RmcpTool;
20
21use super::{Agent, OutputMode};
22
23/// Build [`RmcpTool`]s from MCP tool definitions, applying the given per-call
24/// timeout to each (`None` disables it; see issue #1914). Returns
25/// `(tool_name, tool)` pairs.
26#[cfg(feature = "rmcp")]
27fn build_rmcp_tools(
28    tools: Vec<rmcp::model::Tool>,
29    client: rmcp::service::ServerSink,
30    timeout: Option<std::time::Duration>,
31) -> Vec<(String, RmcpTool)> {
32    tools
33        .into_iter()
34        .map(|tool| {
35            let name = tool.name.to_string();
36            let rmcp_tool = RmcpTool::from_mcp_server(tool, client.clone()).with_timeout(timeout);
37            (name, rmcp_tool)
38        })
39        .collect()
40}
41
42/// Marker type indicating no tool configuration has been set yet.
43///
44/// This is the default state for a new `AgentBuilder`. From this state,
45/// you can either:
46/// - Add tools via `.tool()`, `.tools()`, `.dynamic_tools()`, etc. (transitions to `WithBuilderTools`)
47/// - Set a pre-existing `ToolServerHandle` via `.tool_server_handle()` (transitions to `WithToolServerHandle`)
48/// - Call `.build()` to create an agent with no tools
49#[derive(Default)]
50pub struct NoToolConfig;
51
52/// Typestate indicating a pre-existing `ToolServerHandle` has been provided.
53///
54/// In this state, tool-adding methods (`.tool()`, `.tools()`, etc.) are not available.
55/// The provided handle will be used directly when building the agent.
56pub struct WithToolServerHandle {
57    handle: ToolServerHandle,
58}
59
60/// Typestate indicating tools are being configured via the builder API.
61///
62/// In this state, you can continue adding tools via `.tool()`, `.tools()`,
63/// `.dynamic_tools()`, etc. When `.build()` is called, a new `ToolServer`
64/// will be created with all the configured tools.
65pub struct WithBuilderTools {
66    static_tools: Vec<String>,
67    tools: ToolSet,
68    dynamic_tools: Vec<(usize, Arc<dyn VectorStoreIndexDyn + Send + Sync>)>,
69}
70
71/// A builder for creating an agent
72///
73/// The builder uses a typestate pattern to enforce that tool configuration
74/// is done in a mutually exclusive way: either provide a pre-existing
75/// `ToolServerHandle`, or add tools via the builder API, but not both.
76///
77/// # Example
78/// ```no_run
79/// use rig_core::{agent::AgentBuilder, client::{CompletionClient, ProviderClient}, providers::openai};
80///
81/// # fn run() -> Result<(), Box<dyn std::error::Error>> {
82/// let openai = openai::Client::from_env()?;
83///
84/// let model = openai.completion_model(openai::GPT_5_2);
85///
86/// // Configure the agent
87/// let agent = AgentBuilder::new(model)
88///     .preamble("System prompt")
89///     .context("Context document 1")
90///     .context("Context document 2")
91///     .temperature(0.8)
92///     .build();
93/// # Ok(())
94/// # }
95/// ```
96pub struct AgentBuilder<M, ToolState = NoToolConfig>
97where
98    M: CompletionModel,
99{
100    /// Name of the agent used for logging and debugging
101    name: Option<String>,
102    /// Agent description. Primarily useful when using sub-agents as part of an agent workflow and converting agents to other formats.
103    description: Option<String>,
104    /// Completion model (e.g.: OpenAI's gpt-3.5-turbo-1106, Cohere's command-r)
105    model: M,
106    /// System prompt
107    preamble: Option<String>,
108    /// Context documents always available to the agent
109    static_context: Vec<Document>,
110    /// Additional parameters to be passed to the model
111    additional_params: Option<serde_json::Value>,
112    /// Maximum number of tokens for the completion
113    max_tokens: Option<u64>,
114    /// List of vector store, with the sample number
115    dynamic_context: Vec<(usize, Arc<dyn VectorStoreIndexDyn + Send + Sync>)>,
116    /// Temperature of the model
117    temperature: Option<f64>,
118    /// Whether or not the underlying LLM should be forced to use a tool before providing a response.
119    tool_choice: Option<ToolChoice>,
120    /// Default total model-call budget, including the initial call and retries.
121    default_max_turns: Option<usize>,
122    /// Tool configuration state (typestate pattern)
123    tool_state: ToolState,
124    /// Default hook stack applied to every prompt request from the built agent.
125    hooks: HookStack<M>,
126    /// Optional JSON Schema for structured output
127    output_schema: Option<schemars::Schema>,
128    /// How `output_schema` is enforced (tool vs native vs prompted; see #1928)
129    output_mode: OutputMode,
130    /// Optional conversation memory backend that loads/saves history per conversation id.
131    memory: Option<Arc<dyn ConversationMemory>>,
132    /// Optional default conversation id used when none is set per-request.
133    default_conversation_id: Option<String>,
134}
135
136impl<M, ToolState> AgentBuilder<M, ToolState>
137where
138    M: CompletionModel,
139{
140    /// Set the name of the agent
141    pub fn name(mut self, name: &str) -> Self {
142        self.name = Some(name.into());
143        self
144    }
145
146    /// Set the description of the agent
147    pub fn description(mut self, description: &str) -> Self {
148        self.description = Some(description.into());
149        self
150    }
151
152    /// Set the system prompt
153    pub fn preamble(mut self, preamble: &str) -> Self {
154        self.preamble = Some(preamble.into());
155        self
156    }
157
158    /// Remove the system prompt
159    pub fn without_preamble(mut self) -> Self {
160        self.preamble = None;
161        self
162    }
163
164    /// Append to the preamble of the agent
165    pub fn append_preamble(mut self, doc: &str) -> Self {
166        self.preamble = Some(format!("{}\n{}", self.preamble.unwrap_or_default(), doc));
167        self
168    }
169
170    /// Add a static context document to the agent
171    pub fn context(mut self, doc: &str) -> Self {
172        self.static_context.push(Document {
173            id: format!("static_doc_{}", self.static_context.len()),
174            text: doc.into(),
175            additional_props: HashMap::new(),
176        });
177        self
178    }
179
180    /// Add some dynamic context to the agent. On each prompt, `sample` documents from the
181    /// dynamic context will be inserted in the request.
182    pub fn dynamic_context(
183        mut self,
184        sample: usize,
185        dynamic_context: impl VectorStoreIndexDyn + Send + Sync + 'static,
186    ) -> Self {
187        self.dynamic_context
188            .push((sample, Arc::new(dynamic_context)));
189        self
190    }
191
192    /// Set the tool choice for the agent
193    pub fn tool_choice(mut self, tool_choice: ToolChoice) -> Self {
194        self.tool_choice = Some(tool_choice);
195        self
196    }
197
198    /// Set the default total model-call budget, including the initial call and
199    /// every retry or continuation. Zero permits no model calls.
200    pub fn default_max_turns(mut self, default_max_turns: usize) -> Self {
201        self.default_max_turns = Some(default_max_turns);
202        self
203    }
204
205    /// Set the temperature of the model
206    pub fn temperature(mut self, temperature: f64) -> Self {
207        self.temperature = Some(temperature);
208        self
209    }
210
211    /// Set the maximum number of tokens for the completion
212    pub fn max_tokens(mut self, max_tokens: u64) -> Self {
213        self.max_tokens = Some(max_tokens);
214        self
215    }
216
217    /// Set additional parameters to be passed to the model
218    pub fn additional_params(mut self, params: serde_json::Value) -> Self {
219        self.additional_params = Some(params);
220        self
221    }
222
223    /// Set the output schema for structured output. When set, providers that support
224    /// native structured outputs will constrain the model's response to match this schema.
225    pub fn output_schema<T>(mut self) -> Self
226    where
227        T: JsonSchema,
228    {
229        self.output_schema = Some(schema_for!(T));
230        self
231    }
232
233    /// Set the output schema for structured output. In comparison to `AgentBuilder::schema()` which requires type annotation, you can put in any schema you'd like here.
234    pub fn output_schema_raw(mut self, schema: Schema) -> Self {
235        self.output_schema = Some(schema);
236        self
237    }
238
239    /// Set how `output_schema` is enforced — [`OutputMode::Tool`] (output as a
240    /// tool call, the default when the agent has tools), [`OutputMode::Native`]
241    /// (provider structured output), or [`OutputMode::Prompted`] (see #1928).
242    /// Has no effect unless `output_schema`/`output_schema_raw` is also set.
243    pub fn output_mode(mut self, mode: OutputMode) -> Self {
244        self.output_mode = mode;
245        self
246    }
247
248    /// Attach a [`ConversationMemory`] backend.
249    ///
250    /// When set, the agent will automatically load prior conversation history before
251    /// each prompt and append the new turn after a successful response. A
252    /// `conversation_id` must be supplied either via [`AgentBuilder::conversation`]
253    /// or per-request via [`crate::agent::prompt_request::PromptRequest::conversation`].
254    /// If neither is set, memory is silently bypassed.
255    pub fn memory<B>(mut self, memory: B) -> Self
256    where
257        B: ConversationMemory + 'static,
258    {
259        self.memory = Some(Arc::new(memory));
260        self
261    }
262
263    /// Set a default conversation id used when none is provided per-request.
264    ///
265    /// Most agents are reused across users or threads; prefer setting the id
266    /// per-request via [`crate::agent::prompt_request::PromptRequest::conversation`].
267    pub fn conversation(mut self, id: impl Into<String>) -> Self {
268        self.default_conversation_id = Some(id.into());
269        self
270    }
271
272    /// Attach a default hook to the agent. Each call appends to the agent's hook
273    /// stack; hooks run for every prompt request (unless more are added per
274    /// request) in registration order. How their results compose is
275    /// event-dependent: `CompletionCall` request patches accumulate and merge,
276    /// `ToolCall`/`ToolResult` rewrites chain, and only observe-only/recovery
277    /// events use first-non-`Continue`-wins. See the
278    /// [`hook`](crate::agent::hook) module docs.
279    pub fn add_hook<H>(mut self, hook: H) -> Self
280    where
281        H: AgentHook<M> + 'static,
282    {
283        self.hooks.push(hook);
284        self
285    }
286}
287
288impl<M> AgentBuilder<M, NoToolConfig>
289where
290    M: CompletionModel,
291{
292    /// Create a new agent builder with the given model
293    pub fn new(model: M) -> Self {
294        Self {
295            name: None,
296            description: None,
297            model,
298            preamble: None,
299            static_context: vec![],
300            temperature: None,
301            max_tokens: None,
302            additional_params: None,
303            dynamic_context: vec![],
304            tool_choice: None,
305            default_max_turns: None,
306            tool_state: NoToolConfig,
307            hooks: HookStack::new(),
308            output_schema: None,
309            output_mode: OutputMode::default(),
310            memory: None,
311            default_conversation_id: None,
312        }
313    }
314}
315
316impl<M> AgentBuilder<M, NoToolConfig>
317where
318    M: CompletionModel,
319{
320    /// Set a pre-existing ToolServerHandle for the agent.
321    ///
322    /// After calling this method, tool-adding methods (`.tool()`, `.tools()`, etc.)
323    /// will not be available. Use this when you want to share a `ToolServer`
324    /// between multiple agents or have pre-configured tools.
325    pub fn tool_server_handle(
326        self,
327        handle: ToolServerHandle,
328    ) -> AgentBuilder<M, WithToolServerHandle> {
329        AgentBuilder {
330            name: self.name,
331            description: self.description,
332            model: self.model,
333            preamble: self.preamble,
334            static_context: self.static_context,
335            additional_params: self.additional_params,
336            max_tokens: self.max_tokens,
337            dynamic_context: self.dynamic_context,
338            temperature: self.temperature,
339            tool_choice: self.tool_choice,
340            default_max_turns: self.default_max_turns,
341            tool_state: WithToolServerHandle { handle },
342            hooks: self.hooks,
343            output_schema: self.output_schema,
344            output_mode: self.output_mode,
345            memory: self.memory,
346            default_conversation_id: self.default_conversation_id,
347        }
348    }
349
350    /// Add a static tool to the agent.
351    ///
352    /// This transitions the builder to the `WithBuilderTools` state, where
353    /// additional tools can be added but `tool_server_handle()` is no longer available.
354    pub fn tool(self, tool: impl Tool + 'static) -> AgentBuilder<M, WithBuilderTools> {
355        let toolname = tool.name();
356        AgentBuilder {
357            name: self.name,
358            description: self.description,
359            model: self.model,
360            preamble: self.preamble,
361            static_context: self.static_context,
362            additional_params: self.additional_params,
363            max_tokens: self.max_tokens,
364            dynamic_context: self.dynamic_context,
365            temperature: self.temperature,
366            tool_choice: self.tool_choice,
367            default_max_turns: self.default_max_turns,
368            tool_state: WithBuilderTools {
369                static_tools: vec![toolname],
370                tools: ToolSet::from_tools(vec![tool]),
371                dynamic_tools: vec![],
372            },
373            hooks: self.hooks,
374            output_schema: self.output_schema,
375            output_mode: self.output_mode,
376            memory: self.memory,
377            default_conversation_id: self.default_conversation_id,
378        }
379    }
380
381    /// Add a vector of boxed static tools to the agent.
382    ///
383    /// This is useful when you need to dynamically add static tools to the agent.
384    /// Transitions the builder to the `WithBuilderTools` state.
385    pub fn tools(self, tools: Vec<Box<dyn ToolDyn>>) -> AgentBuilder<M, WithBuilderTools> {
386        let static_tools = tools.iter().map(|tool| tool.name()).collect();
387        let tools = ToolSet::from_tools_boxed(tools);
388
389        AgentBuilder {
390            name: self.name,
391            description: self.description,
392            model: self.model,
393            preamble: self.preamble,
394            static_context: self.static_context,
395            additional_params: self.additional_params,
396            max_tokens: self.max_tokens,
397            dynamic_context: self.dynamic_context,
398            temperature: self.temperature,
399            tool_choice: self.tool_choice,
400            default_max_turns: self.default_max_turns,
401            hooks: self.hooks,
402            output_schema: self.output_schema,
403            output_mode: self.output_mode,
404            memory: self.memory,
405            default_conversation_id: self.default_conversation_id,
406            tool_state: WithBuilderTools {
407                static_tools,
408                tools,
409                dynamic_tools: vec![],
410            },
411        }
412    }
413
414    /// Add an MCP tool (from `rmcp`) to the agent, bounded by
415    /// [`DEFAULT_MCP_TOOL_TIMEOUT`](crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT)
416    /// (see issue #1914). Use [`rmcp_tool_with_timeout`](Self::rmcp_tool_with_timeout)
417    /// to change or disable it.
418    ///
419    /// Transitions the builder to the `WithBuilderTools` state.
420    #[cfg(feature = "rmcp")]
421    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
422    pub fn rmcp_tool(
423        self,
424        tool: rmcp::model::Tool,
425        client: rmcp::service::ServerSink,
426    ) -> AgentBuilder<M, WithBuilderTools> {
427        self.rmcp_tool_with_timeout(tool, client, crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT)
428    }
429
430    /// Add an MCP tool (from `rmcp`) with a per-call timeout (see issue #1914).
431    ///
432    /// Pass a [`Duration`](std::time::Duration) to bound the call, or `None` to
433    /// disable the timeout (unbounded). On timeout the call resolves to a tool
434    /// error the agent can recover from instead of blocking forever.
435    /// Transitions the builder to the `WithBuilderTools` state.
436    #[cfg(feature = "rmcp")]
437    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
438    pub fn rmcp_tool_with_timeout(
439        self,
440        tool: rmcp::model::Tool,
441        client: rmcp::service::ServerSink,
442        timeout: impl Into<Option<std::time::Duration>>,
443    ) -> AgentBuilder<M, WithBuilderTools> {
444        self.with_rmcp_toolset(build_rmcp_tools(vec![tool], client, timeout.into()))
445    }
446
447    /// Add an array of MCP tools (from `rmcp`) to the agent, each bounded by
448    /// [`DEFAULT_MCP_TOOL_TIMEOUT`](crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT)
449    /// (see issue #1914). Use [`rmcp_tools_with_timeout`](Self::rmcp_tools_with_timeout)
450    /// to change or disable it.
451    ///
452    /// Transitions the builder to the `WithBuilderTools` state.
453    #[cfg(feature = "rmcp")]
454    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
455    pub fn rmcp_tools(
456        self,
457        tools: Vec<rmcp::model::Tool>,
458        client: rmcp::service::ServerSink,
459    ) -> AgentBuilder<M, WithBuilderTools> {
460        self.rmcp_tools_with_timeout(tools, client, crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT)
461    }
462
463    /// Add an array of MCP tools (from `rmcp`) with a per-call timeout (see
464    /// issue #1914).
465    ///
466    /// Pass a [`Duration`](std::time::Duration) to bound calls, or `None` to
467    /// disable the timeout (unbounded). On timeout a call resolves to a tool
468    /// error the agent can recover from instead of blocking forever.
469    /// Transitions the builder to the `WithBuilderTools` state.
470    #[cfg(feature = "rmcp")]
471    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
472    pub fn rmcp_tools_with_timeout(
473        self,
474        tools: Vec<rmcp::model::Tool>,
475        client: rmcp::service::ServerSink,
476        timeout: impl Into<Option<std::time::Duration>>,
477    ) -> AgentBuilder<M, WithBuilderTools> {
478        self.with_rmcp_toolset(build_rmcp_tools(tools, client, timeout.into()))
479    }
480
481    /// Transition into the `WithBuilderTools` state carrying the given built
482    /// MCP tools.
483    #[cfg(feature = "rmcp")]
484    fn with_rmcp_toolset(
485        self,
486        built: Vec<(String, RmcpTool)>,
487    ) -> AgentBuilder<M, WithBuilderTools> {
488        let (static_tools, toolset): (Vec<String>, Vec<RmcpTool>) = built.into_iter().unzip();
489
490        AgentBuilder {
491            name: self.name,
492            description: self.description,
493            model: self.model,
494            preamble: self.preamble,
495            static_context: self.static_context,
496            additional_params: self.additional_params,
497            max_tokens: self.max_tokens,
498            dynamic_context: self.dynamic_context,
499            temperature: self.temperature,
500            tool_choice: self.tool_choice,
501            default_max_turns: self.default_max_turns,
502            hooks: self.hooks,
503            output_schema: self.output_schema,
504            output_mode: self.output_mode,
505            memory: self.memory,
506            default_conversation_id: self.default_conversation_id,
507            tool_state: WithBuilderTools {
508                static_tools,
509                tools: ToolSet::from_tools(toolset),
510                dynamic_tools: vec![],
511            },
512        }
513    }
514
515    /// Add some dynamic tools to the agent. On each prompt, `sample` tools from the
516    /// dynamic toolset will be inserted in the request.
517    ///
518    /// Transitions the builder to the `WithBuilderTools` state.
519    pub fn dynamic_tools(
520        self,
521        sample: usize,
522        dynamic_tools: impl VectorStoreIndexDyn + Send + Sync + 'static,
523        toolset: ToolSet,
524    ) -> AgentBuilder<M, WithBuilderTools> {
525        AgentBuilder {
526            name: self.name,
527            description: self.description,
528            model: self.model,
529            preamble: self.preamble,
530            static_context: self.static_context,
531            additional_params: self.additional_params,
532            max_tokens: self.max_tokens,
533            dynamic_context: self.dynamic_context,
534            temperature: self.temperature,
535            tool_choice: self.tool_choice,
536            default_max_turns: self.default_max_turns,
537            hooks: self.hooks,
538            output_schema: self.output_schema,
539            output_mode: self.output_mode,
540            memory: self.memory,
541            default_conversation_id: self.default_conversation_id,
542            tool_state: WithBuilderTools {
543                static_tools: vec![],
544                tools: toolset,
545                dynamic_tools: vec![(sample, Arc::new(dynamic_tools))],
546            },
547        }
548    }
549
550    /// Build the agent with no tools configured.
551    ///
552    /// An empty `ToolServer` will be created for the agent.
553    pub fn build(self) -> Agent<M> {
554        let tool_server_handle = ToolServer::new().run();
555
556        Agent {
557            name: self.name,
558            description: self.description,
559            model: Arc::new(self.model),
560            preamble: self.preamble,
561            static_context: self.static_context,
562            temperature: self.temperature,
563            max_tokens: self.max_tokens,
564            additional_params: self.additional_params,
565            tool_choice: self.tool_choice,
566            dynamic_context: Arc::new(self.dynamic_context),
567            tool_server_handle,
568            default_max_turns: self.default_max_turns,
569            hooks: self.hooks,
570            output_schema: self.output_schema,
571            output_mode: self.output_mode,
572            memory: self.memory,
573            default_conversation_id: self.default_conversation_id,
574        }
575    }
576}
577
578impl<M> AgentBuilder<M, WithToolServerHandle>
579where
580    M: CompletionModel,
581{
582    /// Build the agent using the pre-configured ToolServerHandle.
583    pub fn build(self) -> Agent<M> {
584        Agent {
585            name: self.name,
586            description: self.description,
587            model: Arc::new(self.model),
588            preamble: self.preamble,
589            static_context: self.static_context,
590            temperature: self.temperature,
591            max_tokens: self.max_tokens,
592            additional_params: self.additional_params,
593            tool_choice: self.tool_choice,
594            dynamic_context: Arc::new(self.dynamic_context),
595            tool_server_handle: self.tool_state.handle,
596            default_max_turns: self.default_max_turns,
597            hooks: self.hooks,
598            output_schema: self.output_schema,
599            output_mode: self.output_mode,
600            memory: self.memory,
601            default_conversation_id: self.default_conversation_id,
602        }
603    }
604}
605
606impl<M> AgentBuilder<M, WithBuilderTools>
607where
608    M: CompletionModel,
609{
610    /// Add another static tool to the agent.
611    pub fn tool(mut self, tool: impl Tool + 'static) -> Self {
612        let toolname = tool.name();
613        self.tool_state.tools.add_tool(tool);
614        self.tool_state.static_tools.push(toolname);
615        self
616    }
617
618    /// Add a vector of boxed static tools to the agent.
619    pub fn tools(mut self, tools: Vec<Box<dyn ToolDyn>>) -> Self {
620        let toolnames: Vec<String> = tools.iter().map(|tool| tool.name()).collect();
621        let tools = ToolSet::from_tools_boxed(tools);
622        self.tool_state.tools.add_tools(tools);
623        self.tool_state.static_tools.extend(toolnames);
624        self
625    }
626
627    /// Add an array of MCP tools (from `rmcp`) to the agent, each bounded by
628    /// [`DEFAULT_MCP_TOOL_TIMEOUT`](crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT)
629    /// (see issue #1914). Use [`rmcp_tools_with_timeout`](Self::rmcp_tools_with_timeout)
630    /// to change or disable it.
631    #[cfg(feature = "rmcp")]
632    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
633    pub fn rmcp_tools(
634        self,
635        tools: Vec<rmcp::model::Tool>,
636        client: rmcp::service::ServerSink,
637    ) -> Self {
638        self.rmcp_tools_with_timeout(tools, client, crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT)
639    }
640
641    /// Add an array of MCP tools (from `rmcp`) with a per-call timeout (see
642    /// issue #1914).
643    ///
644    /// Pass a [`Duration`](std::time::Duration) to bound calls, or `None` to
645    /// disable the timeout (unbounded). On timeout a call resolves to a tool
646    /// error the agent can recover from instead of blocking forever.
647    #[cfg(feature = "rmcp")]
648    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
649    pub fn rmcp_tools_with_timeout(
650        self,
651        tools: Vec<rmcp::model::Tool>,
652        client: rmcp::service::ServerSink,
653        timeout: impl Into<Option<std::time::Duration>>,
654    ) -> Self {
655        self.add_rmcp_tools(build_rmcp_tools(tools, client, timeout.into()))
656    }
657
658    #[cfg(feature = "rmcp")]
659    fn add_rmcp_tools(mut self, built: Vec<(String, RmcpTool)>) -> Self {
660        for (name, tool) in built {
661            self.tool_state.static_tools.push(name);
662            self.tool_state.tools.add_tool(tool);
663        }
664
665        self
666    }
667
668    /// Add some dynamic tools to the agent. On each prompt, `sample` tools from the
669    /// dynamic toolset will be inserted in the request.
670    pub fn dynamic_tools(
671        mut self,
672        sample: usize,
673        dynamic_tools: impl VectorStoreIndexDyn + Send + Sync + 'static,
674        toolset: ToolSet,
675    ) -> Self {
676        self.tool_state
677            .dynamic_tools
678            .push((sample, Arc::new(dynamic_tools)));
679        self.tool_state.tools.add_tools(toolset);
680        self
681    }
682
683    /// Build the agent with the configured tools.
684    ///
685    /// A new `ToolServer` will be created containing all tools added via
686    /// `.tool()`, `.tools()`, `.dynamic_tools()`, etc.
687    pub fn build(self) -> Agent<M> {
688        let tool_server_handle = ToolServer::new()
689            .static_tool_names(self.tool_state.static_tools)
690            .add_tools(self.tool_state.tools)
691            .add_dynamic_tools(self.tool_state.dynamic_tools)
692            .run();
693
694        Agent {
695            name: self.name,
696            description: self.description,
697            model: Arc::new(self.model),
698            preamble: self.preamble,
699            static_context: self.static_context,
700            temperature: self.temperature,
701            max_tokens: self.max_tokens,
702            additional_params: self.additional_params,
703            tool_choice: self.tool_choice,
704            dynamic_context: Arc::new(self.dynamic_context),
705            tool_server_handle,
706            default_max_turns: self.default_max_turns,
707            hooks: self.hooks,
708            output_schema: self.output_schema,
709            output_mode: self.output_mode,
710            memory: self.memory,
711            default_conversation_id: self.default_conversation_id,
712        }
713    }
714}
715
716#[cfg(test)]
717mod tests {
718    use super::*;
719    use crate::test_utils::{MockAddTool, MockCompletionModel};
720
721    #[derive(Clone)]
722    struct BuilderHook;
723
724    impl AgentHook<MockCompletionModel> for BuilderHook {}
725
726    #[test]
727    fn hook_can_be_set_after_tool_configuration() {
728        let _agent = AgentBuilder::new(MockCompletionModel::text("ok"))
729            .tool(MockAddTool)
730            .add_hook(BuilderHook)
731            .build();
732    }
733
734    /// The builder's shared MCP helper threads the configured timeout (default,
735    /// explicit, or `None`/disabled) onto every built tool, and the threaded
736    /// timeout actually bounds a hanging call. This covers the plumbing behind
737    /// `rmcp_tool[s]` / `rmcp_tool[s]_with_timeout` (see issue #1914).
738    #[cfg(feature = "rmcp")]
739    #[tokio::test]
740    async fn build_rmcp_tools_threads_timeout_into_built_tools() {
741        use crate::tool::ToolDyn;
742        use crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT;
743        use rmcp::model::{
744            CallToolRequestParams, CallToolResult, ClientInfo, ErrorData, Implementation,
745            ProtocolVersion, ServerCapabilities, ServerInfo, Tool,
746        };
747        use rmcp::service::RequestContext;
748        use rmcp::{RoleServer, ServerHandler, ServiceExt};
749        use std::sync::Arc;
750        use std::time::Duration;
751
752        #[derive(Clone)]
753        struct HangingServer;
754        impl ServerHandler for HangingServer {
755            fn get_info(&self) -> ServerInfo {
756                ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
757                    .with_protocol_version(ProtocolVersion::LATEST)
758                    .with_server_info(Implementation::new("builder-timeout-test", "0.1.0"))
759            }
760            async fn call_tool(
761                &self,
762                _request: CallToolRequestParams,
763                _context: RequestContext<RoleServer>,
764            ) -> Result<CallToolResult, ErrorData> {
765                std::future::pending::<Result<CallToolResult, ErrorData>>().await
766            }
767        }
768
769        fn tool(name: &str) -> Tool {
770            Tool::new(
771                name.to_string(),
772                String::new(),
773                Arc::new(serde_json::Map::new()),
774            )
775        }
776
777        let (c2s, sfc) = tokio::io::duplex(8192);
778        let (s2c, cfs) = tokio::io::duplex(8192);
779        let server_task = tokio::spawn(async move {
780            let running = HangingServer.serve((sfc, s2c)).await.expect("server start");
781            running.waiting().await.expect("server error");
782        });
783        let client = ClientInfo::default()
784            .serve((cfs, c2s))
785            .await
786            .expect("client connect");
787        let peer = client.peer().clone();
788
789        // The configured timeout (default, explicit, or disabled) is threaded
790        // onto each built tool.
791        let built_default = build_rmcp_tools(
792            vec![tool("a")],
793            peer.clone(),
794            Some(DEFAULT_MCP_TOOL_TIMEOUT),
795        );
796        assert_eq!(built_default[0].1.timeout(), Some(DEFAULT_MCP_TOOL_TIMEOUT));
797        let built_none = build_rmcp_tools(vec![tool("b")], peer.clone(), None);
798        assert_eq!(built_none[0].1.timeout(), None);
799
800        // ...and the threaded timeout actually bounds a hanging call.
801        let built = build_rmcp_tools(
802            vec![tool("hang_forever")],
803            peer,
804            Some(Duration::from_millis(200)),
805        );
806        assert_eq!(built.len(), 1);
807        assert_eq!(built[0].0, "hang_forever");
808        let timed =
809            tokio::time::timeout(Duration::from_secs(5), built[0].1.call("{}".to_string())).await;
810        let err = timed
811            .expect("built tool hung past the safety timeout")
812            .expect_err("call should time out");
813        assert!(err.to_string().contains("timed out"), "got: {err}");
814
815        drop(client);
816        server_task.abort();
817    }
818}
rig_core/agent/builder.rs

rig_core/agent/
builder.rs
