rig_core/agent/
builder.rs

1use std::{collections::HashMap, sync::Arc};
2
3use schemars::{JsonSchema, Schema, schema_for};
4
5use crate::{
6    agent::prompt_request::hooks::PromptHook,
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;
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, P = (), ToolState = NoToolConfig>
97where
98    M: CompletionModel,
99    P: PromptHook<M>,
100{
101    /// Name of the agent used for logging and debugging
102    name: Option<String>,
103    /// Agent description. Primarily useful when using sub-agents as part of an agent workflow and converting agents to other formats.
104    description: Option<String>,
105    /// Completion model (e.g.: OpenAI's gpt-3.5-turbo-1106, Cohere's command-r)
106    model: M,
107    /// System prompt
108    preamble: Option<String>,
109    /// Context documents always available to the agent
110    static_context: Vec<Document>,
111    /// Additional parameters to be passed to the model
112    additional_params: Option<serde_json::Value>,
113    /// Maximum number of tokens for the completion
114    max_tokens: Option<u64>,
115    /// List of vector store, with the sample number
116    dynamic_context: Vec<(usize, Arc<dyn VectorStoreIndexDyn + Send + Sync>)>,
117    /// Temperature of the model
118    temperature: Option<f64>,
119    /// Whether or not the underlying LLM should be forced to use a tool before providing a response.
120    tool_choice: Option<ToolChoice>,
121    /// Default maximum depth for multi-turn agent calls
122    default_max_turns: Option<usize>,
123    /// Tool configuration state (typestate pattern)
124    tool_state: ToolState,
125    /// Prompt hook
126    hook: Option<P>,
127    /// Optional JSON Schema for structured output
128    output_schema: Option<schemars::Schema>,
129    /// Optional conversation memory backend that loads/saves history per conversation id.
130    memory: Option<Arc<dyn ConversationMemory>>,
131    /// Optional default conversation id used when none is set per-request.
132    default_conversation_id: Option<String>,
133}
134
135impl<M, P, ToolState> AgentBuilder<M, P, ToolState>
136where
137    M: CompletionModel,
138    P: PromptHook<M>,
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 maximum depth that an agent will use for multi-turn.
199    pub fn default_max_turns(mut self, default_max_turns: usize) -> Self {
200        self.default_max_turns = Some(default_max_turns);
201        self
202    }
203
204    /// Set the temperature of the model
205    pub fn temperature(mut self, temperature: f64) -> Self {
206        self.temperature = Some(temperature);
207        self
208    }
209
210    /// Set the maximum number of tokens for the completion
211    pub fn max_tokens(mut self, max_tokens: u64) -> Self {
212        self.max_tokens = Some(max_tokens);
213        self
214    }
215
216    /// Set additional parameters to be passed to the model
217    pub fn additional_params(mut self, params: serde_json::Value) -> Self {
218        self.additional_params = Some(params);
219        self
220    }
221
222    /// Set the output schema for structured output. When set, providers that support
223    /// native structured outputs will constrain the model's response to match this schema.
224    pub fn output_schema<T>(mut self) -> Self
225    where
226        T: JsonSchema,
227    {
228        self.output_schema = Some(schema_for!(T));
229        self
230    }
231
232    /// 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.
233    pub fn output_schema_raw(mut self, schema: Schema) -> Self {
234        self.output_schema = Some(schema);
235        self
236    }
237
238    /// Attach a [`ConversationMemory`] backend.
239    ///
240    /// When set, the agent will automatically load prior conversation history before
241    /// each prompt and append the new turn after a successful response. A
242    /// `conversation_id` must be supplied either via [`AgentBuilder::conversation_id`]
243    /// or per-request via [`crate::agent::prompt_request::PromptRequest::conversation`].
244    /// If neither is set, memory is silently bypassed.
245    pub fn memory<B>(mut self, memory: B) -> Self
246    where
247        B: ConversationMemory + 'static,
248    {
249        self.memory = Some(Arc::new(memory));
250        self
251    }
252
253    /// Set a default conversation id used when none is provided per-request.
254    ///
255    /// Most agents are reused across users or threads; prefer setting the id
256    /// per-request via [`crate::agent::prompt_request::PromptRequest::conversation`].
257    pub fn conversation_id(mut self, id: impl Into<String>) -> Self {
258        self.default_conversation_id = Some(id.into());
259        self
260    }
261
262    /// Set the default hook for the agent.
263    ///
264    /// This hook will be used for all prompt requests unless overridden
265    /// via `.with_hook()` on the request.
266    pub fn hook<P2>(self, hook: P2) -> AgentBuilder<M, P2, ToolState>
267    where
268        P2: PromptHook<M>,
269    {
270        AgentBuilder {
271            name: self.name,
272            description: self.description,
273            model: self.model,
274            preamble: self.preamble,
275            static_context: self.static_context,
276            additional_params: self.additional_params,
277            max_tokens: self.max_tokens,
278            dynamic_context: self.dynamic_context,
279            temperature: self.temperature,
280            tool_choice: self.tool_choice,
281            default_max_turns: self.default_max_turns,
282            tool_state: self.tool_state,
283            hook: Some(hook),
284            output_schema: self.output_schema,
285            memory: self.memory,
286            default_conversation_id: self.default_conversation_id,
287        }
288    }
289}
290
291impl<M> AgentBuilder<M, (), NoToolConfig>
292where
293    M: CompletionModel,
294{
295    /// Create a new agent builder with the given model
296    pub fn new(model: M) -> Self {
297        Self {
298            name: None,
299            description: None,
300            model,
301            preamble: None,
302            static_context: vec![],
303            temperature: None,
304            max_tokens: None,
305            additional_params: None,
306            dynamic_context: vec![],
307            tool_choice: None,
308            default_max_turns: None,
309            tool_state: NoToolConfig,
310            hook: None,
311            output_schema: None,
312            memory: None,
313            default_conversation_id: None,
314        }
315    }
316}
317
318impl<M, P> AgentBuilder<M, P, NoToolConfig>
319where
320    M: CompletionModel,
321    P: PromptHook<M>,
322{
323    /// Set a pre-existing ToolServerHandle for the agent.
324    ///
325    /// After calling this method, tool-adding methods (`.tool()`, `.tools()`, etc.)
326    /// will not be available. Use this when you want to share a `ToolServer`
327    /// between multiple agents or have pre-configured tools.
328    pub fn tool_server_handle(
329        self,
330        handle: ToolServerHandle,
331    ) -> AgentBuilder<M, P, WithToolServerHandle> {
332        AgentBuilder {
333            name: self.name,
334            description: self.description,
335            model: self.model,
336            preamble: self.preamble,
337            static_context: self.static_context,
338            additional_params: self.additional_params,
339            max_tokens: self.max_tokens,
340            dynamic_context: self.dynamic_context,
341            temperature: self.temperature,
342            tool_choice: self.tool_choice,
343            default_max_turns: self.default_max_turns,
344            tool_state: WithToolServerHandle { handle },
345            hook: self.hook,
346            output_schema: self.output_schema,
347            memory: self.memory,
348            default_conversation_id: self.default_conversation_id,
349        }
350    }
351
352    /// Add a static tool to the agent.
353    ///
354    /// This transitions the builder to the `WithBuilderTools` state, where
355    /// additional tools can be added but `tool_server_handle()` is no longer available.
356    pub fn tool(self, tool: impl Tool + 'static) -> AgentBuilder<M, P, WithBuilderTools> {
357        let toolname = tool.name();
358        AgentBuilder {
359            name: self.name,
360            description: self.description,
361            model: self.model,
362            preamble: self.preamble,
363            static_context: self.static_context,
364            additional_params: self.additional_params,
365            max_tokens: self.max_tokens,
366            dynamic_context: self.dynamic_context,
367            temperature: self.temperature,
368            tool_choice: self.tool_choice,
369            default_max_turns: self.default_max_turns,
370            tool_state: WithBuilderTools {
371                static_tools: vec![toolname],
372                tools: ToolSet::from_tools(vec![tool]),
373                dynamic_tools: vec![],
374            },
375            hook: self.hook,
376            output_schema: self.output_schema,
377            memory: self.memory,
378            default_conversation_id: self.default_conversation_id,
379        }
380    }
381
382    /// Add a vector of boxed static tools to the agent.
383    ///
384    /// This is useful when you need to dynamically add static tools to the agent.
385    /// Transitions the builder to the `WithBuilderTools` state.
386    pub fn tools(self, tools: Vec<Box<dyn ToolDyn>>) -> AgentBuilder<M, P, WithBuilderTools> {
387        let static_tools = tools.iter().map(|tool| tool.name()).collect();
388        let tools = ToolSet::from_tools_boxed(tools);
389
390        AgentBuilder {
391            name: self.name,
392            description: self.description,
393            model: self.model,
394            preamble: self.preamble,
395            static_context: self.static_context,
396            additional_params: self.additional_params,
397            max_tokens: self.max_tokens,
398            dynamic_context: self.dynamic_context,
399            temperature: self.temperature,
400            tool_choice: self.tool_choice,
401            default_max_turns: self.default_max_turns,
402            hook: self.hook,
403            output_schema: self.output_schema,
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, P, 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, P, 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, P, 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, P, 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, P, 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            hook: self.hook,
503            output_schema: self.output_schema,
504            memory: self.memory,
505            default_conversation_id: self.default_conversation_id,
506            tool_state: WithBuilderTools {
507                static_tools,
508                tools: ToolSet::from_tools(toolset),
509                dynamic_tools: vec![],
510            },
511        }
512    }
513
514    /// Add some dynamic tools to the agent. On each prompt, `sample` tools from the
515    /// dynamic toolset will be inserted in the request.
516    ///
517    /// Transitions the builder to the `WithBuilderTools` state.
518    pub fn dynamic_tools(
519        self,
520        sample: usize,
521        dynamic_tools: impl VectorStoreIndexDyn + Send + Sync + 'static,
522        toolset: ToolSet,
523    ) -> AgentBuilder<M, P, WithBuilderTools> {
524        AgentBuilder {
525            name: self.name,
526            description: self.description,
527            model: self.model,
528            preamble: self.preamble,
529            static_context: self.static_context,
530            additional_params: self.additional_params,
531            max_tokens: self.max_tokens,
532            dynamic_context: self.dynamic_context,
533            temperature: self.temperature,
534            tool_choice: self.tool_choice,
535            default_max_turns: self.default_max_turns,
536            hook: self.hook,
537            output_schema: self.output_schema,
538            memory: self.memory,
539            default_conversation_id: self.default_conversation_id,
540            tool_state: WithBuilderTools {
541                static_tools: vec![],
542                tools: toolset,
543                dynamic_tools: vec![(sample, Arc::new(dynamic_tools))],
544            },
545        }
546    }
547
548    /// Build the agent with no tools configured.
549    ///
550    /// An empty `ToolServer` will be created for the agent.
551    pub fn build(self) -> Agent<M, P> {
552        let tool_server_handle = ToolServer::new().run();
553
554        Agent {
555            name: self.name,
556            description: self.description,
557            model: Arc::new(self.model),
558            preamble: self.preamble,
559            static_context: self.static_context,
560            temperature: self.temperature,
561            max_tokens: self.max_tokens,
562            additional_params: self.additional_params,
563            tool_choice: self.tool_choice,
564            dynamic_context: Arc::new(self.dynamic_context),
565            tool_server_handle,
566            default_max_turns: self.default_max_turns,
567            hook: self.hook,
568            output_schema: self.output_schema,
569            memory: self.memory,
570            default_conversation_id: self.default_conversation_id,
571        }
572    }
573}
574
575impl<M, P> AgentBuilder<M, P, WithToolServerHandle>
576where
577    M: CompletionModel,
578    P: PromptHook<M>,
579{
580    /// Build the agent using the pre-configured ToolServerHandle.
581    pub fn build(self) -> Agent<M, P> {
582        Agent {
583            name: self.name,
584            description: self.description,
585            model: Arc::new(self.model),
586            preamble: self.preamble,
587            static_context: self.static_context,
588            temperature: self.temperature,
589            max_tokens: self.max_tokens,
590            additional_params: self.additional_params,
591            tool_choice: self.tool_choice,
592            dynamic_context: Arc::new(self.dynamic_context),
593            tool_server_handle: self.tool_state.handle,
594            default_max_turns: self.default_max_turns,
595            hook: self.hook,
596            output_schema: self.output_schema,
597            memory: self.memory,
598            default_conversation_id: self.default_conversation_id,
599        }
600    }
601}
602
603impl<M, P> AgentBuilder<M, P, WithBuilderTools>
604where
605    M: CompletionModel,
606    P: PromptHook<M>,
607{
608    /// Add another static tool to the agent.
609    pub fn tool(mut self, tool: impl Tool + 'static) -> Self {
610        let toolname = tool.name();
611        self.tool_state.tools.add_tool(tool);
612        self.tool_state.static_tools.push(toolname);
613        self
614    }
615
616    /// Add a vector of boxed static tools to the agent.
617    pub fn tools(mut self, tools: Vec<Box<dyn ToolDyn>>) -> Self {
618        let toolnames: Vec<String> = tools.iter().map(|tool| tool.name()).collect();
619        let tools = ToolSet::from_tools_boxed(tools);
620        self.tool_state.tools.add_tools(tools);
621        self.tool_state.static_tools.extend(toolnames);
622        self
623    }
624
625    /// Add an array of MCP tools (from `rmcp`) to the agent, each bounded by
626    /// [`DEFAULT_MCP_TOOL_TIMEOUT`](crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT)
627    /// (see issue #1914). Use [`rmcp_tools_with_timeout`](Self::rmcp_tools_with_timeout)
628    /// to change or disable it.
629    #[cfg(feature = "rmcp")]
630    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
631    pub fn rmcp_tools(
632        self,
633        tools: Vec<rmcp::model::Tool>,
634        client: rmcp::service::ServerSink,
635    ) -> Self {
636        self.rmcp_tools_with_timeout(tools, client, crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT)
637    }
638
639    /// Add an array of MCP tools (from `rmcp`) with a per-call timeout (see
640    /// issue #1914).
641    ///
642    /// Pass a [`Duration`](std::time::Duration) to bound calls, or `None` to
643    /// disable the timeout (unbounded). On timeout a call resolves to a tool
644    /// error the agent can recover from instead of blocking forever.
645    #[cfg(feature = "rmcp")]
646    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
647    pub fn rmcp_tools_with_timeout(
648        self,
649        tools: Vec<rmcp::model::Tool>,
650        client: rmcp::service::ServerSink,
651        timeout: impl Into<Option<std::time::Duration>>,
652    ) -> Self {
653        self.add_rmcp_tools(build_rmcp_tools(tools, client, timeout.into()))
654    }
655
656    #[cfg(feature = "rmcp")]
657    fn add_rmcp_tools(mut self, built: Vec<(String, RmcpTool)>) -> Self {
658        for (name, tool) in built {
659            self.tool_state.static_tools.push(name);
660            self.tool_state.tools.add_tool(tool);
661        }
662
663        self
664    }
665
666    /// Add some dynamic tools to the agent. On each prompt, `sample` tools from the
667    /// dynamic toolset will be inserted in the request.
668    pub fn dynamic_tools(
669        mut self,
670        sample: usize,
671        dynamic_tools: impl VectorStoreIndexDyn + Send + Sync + 'static,
672        toolset: ToolSet,
673    ) -> Self {
674        self.tool_state
675            .dynamic_tools
676            .push((sample, Arc::new(dynamic_tools)));
677        self.tool_state.tools.add_tools(toolset);
678        self
679    }
680
681    /// Build the agent with the configured tools.
682    ///
683    /// A new `ToolServer` will be created containing all tools added via
684    /// `.tool()`, `.tools()`, `.dynamic_tools()`, etc.
685    pub fn build(self) -> Agent<M, P> {
686        let tool_server_handle = ToolServer::new()
687            .static_tool_names(self.tool_state.static_tools)
688            .add_tools(self.tool_state.tools)
689            .add_dynamic_tools(self.tool_state.dynamic_tools)
690            .run();
691
692        Agent {
693            name: self.name,
694            description: self.description,
695            model: Arc::new(self.model),
696            preamble: self.preamble,
697            static_context: self.static_context,
698            temperature: self.temperature,
699            max_tokens: self.max_tokens,
700            additional_params: self.additional_params,
701            tool_choice: self.tool_choice,
702            dynamic_context: Arc::new(self.dynamic_context),
703            tool_server_handle,
704            default_max_turns: self.default_max_turns,
705            hook: self.hook,
706            output_schema: self.output_schema,
707            memory: self.memory,
708            default_conversation_id: self.default_conversation_id,
709        }
710    }
711}
712
713#[cfg(test)]
714mod tests {
715    use super::*;
716    use crate::test_utils::{MockAddTool, MockCompletionModel};
717
718    #[derive(Clone)]
719    struct BuilderHook;
720
721    impl PromptHook<MockCompletionModel> for BuilderHook {}
722
723    #[test]
724    fn hook_can_be_set_after_tool_configuration() {
725        let _agent = AgentBuilder::new(MockCompletionModel::text("ok"))
726            .tool(MockAddTool)
727            .hook(BuilderHook)
728            .build();
729    }
730
731    /// The builder's shared MCP helper threads the configured timeout (default,
732    /// explicit, or `None`/disabled) onto every built tool, and the threaded
733    /// timeout actually bounds a hanging call. This covers the plumbing behind
734    /// `rmcp_tool[s]` / `rmcp_tool[s]_with_timeout` (see issue #1914).
735    #[cfg(feature = "rmcp")]
736    #[tokio::test]
737    async fn build_rmcp_tools_threads_timeout_into_built_tools() {
738        use crate::tool::ToolDyn;
739        use crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT;
740        use rmcp::model::{
741            CallToolRequestParams, CallToolResult, ClientInfo, ErrorData, Implementation,
742            ProtocolVersion, ServerCapabilities, ServerInfo, Tool,
743        };
744        use rmcp::service::RequestContext;
745        use rmcp::{RoleServer, ServerHandler, ServiceExt};
746        use std::sync::Arc;
747        use std::time::Duration;
748
749        #[derive(Clone)]
750        struct HangingServer;
751        impl ServerHandler for HangingServer {
752            fn get_info(&self) -> ServerInfo {
753                ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
754                    .with_protocol_version(ProtocolVersion::LATEST)
755                    .with_server_info(Implementation::new("builder-timeout-test", "0.1.0"))
756            }
757            async fn call_tool(
758                &self,
759                _request: CallToolRequestParams,
760                _context: RequestContext<RoleServer>,
761            ) -> Result<CallToolResult, ErrorData> {
762                std::future::pending::<Result<CallToolResult, ErrorData>>().await
763            }
764        }
765
766        fn tool(name: &str) -> Tool {
767            Tool::new(
768                name.to_string(),
769                String::new(),
770                Arc::new(serde_json::Map::new()),
771            )
772        }
773
774        let (c2s, sfc) = tokio::io::duplex(8192);
775        let (s2c, cfs) = tokio::io::duplex(8192);
776        let server_task = tokio::spawn(async move {
777            let running = HangingServer.serve((sfc, s2c)).await.expect("server start");
778            running.waiting().await.expect("server error");
779        });
780        let client = ClientInfo::default()
781            .serve((cfs, c2s))
782            .await
783            .expect("client connect");
784        let peer = client.peer().clone();
785
786        // The configured timeout (default, explicit, or disabled) is threaded
787        // onto each built tool.
788        let built_default = build_rmcp_tools(
789            vec![tool("a")],
790            peer.clone(),
791            Some(DEFAULT_MCP_TOOL_TIMEOUT),
792        );
793        assert_eq!(built_default[0].1.timeout(), Some(DEFAULT_MCP_TOOL_TIMEOUT));
794        let built_none = build_rmcp_tools(vec![tool("b")], peer.clone(), None);
795        assert_eq!(built_none[0].1.timeout(), None);
796
797        // ...and the threaded timeout actually bounds a hanging call.
798        let built = build_rmcp_tools(
799            vec![tool("hang_forever")],
800            peer,
801            Some(Duration::from_millis(200)),
802        );
803        assert_eq!(built.len(), 1);
804        assert_eq!(built[0].0, "hang_forever");
805        let timed =
806            tokio::time::timeout(Duration::from_secs(5), built[0].1.call("{}".to_string())).await;
807        let err = timed
808            .expect("built tool hung past the safety timeout")
809            .expect_err("call should time out");
810        assert!(err.to_string().contains("timed out"), "got: {err}");
811
812        drop(client);
813        server_task.abort();
814    }
815}
rig_core/agent/builder.rs

rig_core/agent/
builder.rs
