rig_agent/agent/
builder.rs

1use std::{collections::HashMap, sync::Arc};
2
3use schemars::{JsonSchema, Schema, schema_for};
4
5use rig_core::{
6    memory::ConversationMemory,
7    message::ToolChoice,
8    vector_store::{VectorSearchRequest, VectorStoreIndexDyn},
9};
10
11use crate::{
12    agent::hook::{
13        AgentHook, CompletionCall, CompletionCallAction, HookContext, HookStack, RequestPatch,
14    },
15    completion::{CompletionModel, Document},
16    tool::{
17        DynamicTool, PortableDynamicTool, Tool, ToolSet,
18        server::{ToolServer, ToolServerHandle},
19    },
20};
21
22#[cfg(all(feature = "rmcp", not(target_family = "wasm")))]
23#[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
24use crate::tool::rmcp::McpTool as RmcpTool;
25
26use super::{Agent, OutputMode};
27
28struct DynamicContext<I> {
29    samples: usize,
30    index: I,
31}
32
33impl<I> AgentHook for DynamicContext<I>
34where
35    I: VectorStoreIndexDyn,
36{
37    async fn on_completion_call(
38        &self,
39        _ctx: &HookContext,
40        event: CompletionCall<'_>,
41    ) -> CompletionCallAction {
42        let query = event.prompt.rag_text().or_else(|| {
43            event
44                .history
45                .iter()
46                .rev()
47                .find_map(|message| message.rag_text())
48        });
49        let Some(query) = query else {
50            return CompletionCallAction::continue_run();
51        };
52
53        let request = VectorSearchRequest::builder()
54            .query(query)
55            .samples(self.samples as u64)
56            .build();
57        match self.index.top_n(request).await {
58            Ok(results) => CompletionCallAction::patch(RequestPatch::new().extra_context(
59                results.into_iter().map(|(_, id, value)| Document {
60                    id,
61                    text:
62                        serde_json::to_string_pretty(&value).unwrap_or_else(|_| value.to_string()),
63                    additional_props: Default::default(),
64                }),
65            )),
66            Err(error) => {
67                CompletionCallAction::stop(format!("failed to retrieve dynamic context: {error}"))
68            }
69        }
70    }
71}
72
73/// Build [`RmcpTool`]s from MCP tool definitions, applying the given per-call
74/// timeout to each (`None` disables it; see issue #1914). Returns
75/// `(tool_name, tool)` pairs.
76#[cfg(all(feature = "rmcp", not(target_family = "wasm")))]
77fn build_rmcp_tools(
78    tools: Vec<rmcp::model::Tool>,
79    client: rmcp::service::ServerSink,
80    timeout: Option<std::time::Duration>,
81) -> Vec<(String, RmcpTool)> {
82    tools
83        .into_iter()
84        .map(|tool| {
85            let name = tool.name.to_string();
86            let rmcp_tool = RmcpTool::from_mcp_server(tool, client.clone()).with_timeout(timeout);
87            (name, rmcp_tool)
88        })
89        .collect()
90}
91
92/// Marker type indicating no tool configuration has been set yet.
93///
94/// This is the default state for a new `AgentBuilder`. From this state,
95/// you can either:
96/// - Add tools via `.tool()`, `.dynamic_tool()`, `.dynamic_tools()`, or
97///   `.retrieved_tools()` (transitions to `WithBuilderTools`)
98/// - Set a pre-existing `ToolServerHandle` via `.tool_server_handle()` (transitions to `WithToolServerHandle`)
99/// - Call `.build()` to create an agent with no tools
100#[derive(Default)]
101pub struct NoToolConfig;
102
103/// Typestate indicating a pre-existing `ToolServerHandle` has been provided.
104///
105/// In this state, tool-adding methods (`.tool()`, `.dynamic_tool()`, etc.) are not available.
106/// The provided handle will be used directly when building the agent.
107pub struct WithToolServerHandle {
108    handle: ToolServerHandle,
109}
110
111/// Typestate indicating tools are being configured via the builder API.
112///
113/// In this state, you can continue adding tools via `.tool()`,
114/// `.dynamic_tool()`, `.dynamic_tools()`, and `.retrieved_tools()`. When
115/// `.build()` is called, a new `ToolServer`
116/// will be created with all the configured tools.
117pub struct WithBuilderTools {
118    tools: ToolSet,
119    retrieval_indexes: Vec<(usize, Arc<dyn VectorStoreIndexDyn + Send + Sync>)>,
120}
121
122/// A builder for creating an agent
123///
124/// The builder uses a typestate pattern to enforce that tool configuration
125/// is done in a mutually exclusive way: either provide a pre-existing
126/// `ToolServerHandle`, or add tools via the builder API, but not both.
127///
128/// # Example
129/// ```no_run
130/// use rig_agent::AgentBuilder;
131/// use rig_core::{client::{CompletionClient, ProviderClient}, providers::openai};
132///
133/// # fn run() -> Result<(), Box<dyn std::error::Error>> {
134/// let openai = openai::Client::from_env()?;
135///
136/// let model = openai.completion_model(openai::GPT_5_2);
137///
138/// // Configure the agent
139/// let agent = AgentBuilder::new(model)
140///     .preamble("System prompt")
141///     .context("Context document 1")
142///     .context("Context document 2")
143///     .temperature(0.8)
144///     .build();
145/// # Ok(())
146/// # }
147/// ```
148pub struct AgentBuilder<M, ToolState = NoToolConfig>
149where
150    M: CompletionModel,
151{
152    /// Name of the agent used for logging and debugging
153    name: Option<String>,
154    /// Agent description. Primarily useful when using sub-agents as part of an agent workflow and converting agents to other formats.
155    description: Option<String>,
156    /// Completion model (e.g.: OpenAI's gpt-3.5-turbo-1106, Cohere's command-r)
157    model: M,
158    /// System prompt
159    preamble: Option<String>,
160    /// Context documents always available to the agent
161    static_context: Vec<Document>,
162    /// Additional parameters to be passed to the model
163    additional_params: Option<serde_json::Value>,
164    /// Whether to record sensitive request, response, and tool content on telemetry spans.
165    record_telemetry_content: bool,
166    /// Maximum number of tokens for the completion
167    max_tokens: Option<u64>,
168    /// Temperature of the model
169    temperature: Option<f64>,
170    /// Whether or not the underlying LLM should be forced to use a tool before providing a response.
171    tool_choice: Option<ToolChoice>,
172    /// Default total model-call budget, including the initial call and retries.
173    default_max_turns: Option<usize>,
174    /// Tool configuration state (typestate pattern)
175    tool_state: ToolState,
176    /// Default hook stack applied to every prompt request from the built agent.
177    hooks: HookStack,
178    /// Optional JSON Schema for structured output
179    output_schema: Option<schemars::Schema>,
180    /// How `output_schema` is enforced (tool vs native vs prompted; see #1928)
181    output_mode: OutputMode,
182    /// Optional conversation memory backend that loads/saves history per conversation id.
183    memory: Option<Arc<dyn ConversationMemory>>,
184    /// Optional default conversation id used when none is set per-request.
185    default_conversation_id: Option<String>,
186}
187
188impl<M, ToolState> AgentBuilder<M, ToolState>
189where
190    M: CompletionModel,
191{
192    /// Set the name of the agent
193    pub fn name(mut self, name: &str) -> Self {
194        self.name = Some(name.into());
195        self
196    }
197
198    /// Set the description of the agent
199    pub fn description(mut self, description: &str) -> Self {
200        self.description = Some(description.into());
201        self
202    }
203
204    /// Set the system prompt
205    pub fn preamble(mut self, preamble: &str) -> Self {
206        self.preamble = Some(preamble.into());
207        self
208    }
209
210    /// Remove the system prompt
211    pub fn without_preamble(mut self) -> Self {
212        self.preamble = None;
213        self
214    }
215
216    /// Append to the preamble of the agent
217    pub fn append_preamble(mut self, doc: &str) -> Self {
218        self.preamble = Some(format!("{}\n{}", self.preamble.unwrap_or_default(), doc));
219        self
220    }
221
222    /// Add a static context document to the agent
223    pub fn context(mut self, doc: &str) -> Self {
224        self.static_context.push(Document {
225            id: format!("static_doc_{}", self.static_context.len()),
226            text: doc.into(),
227            additional_props: HashMap::new(),
228        });
229        self
230    }
231
232    /// Add dynamic context retrieved from a vector store on every model call.
233    ///
234    /// This is a convenience wrapper around an internal completion-call hook.
235    /// The hook searches with the current prompt's first text part, falling back
236    /// to the latest textual history message, and appends the retrieved documents
237    /// to the request after static context. Retrieval and injected documents
238    /// follow registration order relative to application hooks, so register a
239    /// stop policy before this helper when it should prevent retrieval. A
240    /// retrieval failure stops the run before provider I/O.
241    pub fn dynamic_context<I>(self, samples: usize, index: I) -> Self
242    where
243        I: VectorStoreIndexDyn + 'static,
244    {
245        self.add_hook(DynamicContext { samples, index })
246    }
247
248    /// Set the tool choice for the agent
249    pub fn tool_choice(mut self, tool_choice: ToolChoice) -> Self {
250        self.tool_choice = Some(tool_choice);
251        self
252    }
253
254    /// Set the default total model-call budget, including the initial call and
255    /// every retry or continuation. Zero permits no model calls.
256    pub fn default_max_turns(mut self, default_max_turns: usize) -> Self {
257        self.default_max_turns = Some(default_max_turns);
258        self
259    }
260
261    /// Set the temperature of the model
262    pub fn temperature(mut self, temperature: f64) -> Self {
263        self.temperature = Some(temperature);
264        self
265    }
266
267    /// Set the maximum number of tokens for the completion
268    pub fn max_tokens(mut self, max_tokens: u64) -> Self {
269        self.max_tokens = Some(max_tokens);
270        self
271    }
272
273    /// Set additional parameters to be passed to the model
274    pub fn additional_params(mut self, params: serde_json::Value) -> Self {
275        self.additional_params = Some(params);
276        self
277    }
278
279    /// Opt in or out of recording sensitive request, response, and tool content
280    /// on GenAI telemetry spans for requests made by this agent.
281    ///
282    /// Defaults to `false`. Enabling this can expose prompts, retrieved context,
283    /// tool results, model responses, and other sensitive or high-cardinality data
284    /// through OpenTelemetry span attributes, which can increase observability
285    /// backend storage and query costs. Only enable it when content telemetry is
286    /// acceptable for this agent. Structural metadata and token usage remain
287    /// available when this is disabled.
288    pub fn record_content_telemetry(mut self, enabled: bool) -> Self {
289        self.record_telemetry_content = enabled;
290        self
291    }
292
293    /// Set the output schema for structured output. When set, providers that support
294    /// native structured outputs will constrain the model's response to match this schema.
295    pub fn output_schema<T>(mut self) -> Self
296    where
297        T: JsonSchema,
298    {
299        self.output_schema = Some(schema_for!(T));
300        self
301    }
302
303    /// 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.
304    pub fn output_schema_raw(mut self, schema: Schema) -> Self {
305        self.output_schema = Some(schema);
306        self
307    }
308
309    /// Set how `output_schema` is enforced — [`OutputMode::Tool`] (output as a
310    /// tool call, the default when the agent has tools), [`OutputMode::Native`]
311    /// (provider structured output), or [`OutputMode::Prompted`] (see #1928).
312    /// Has no effect unless `output_schema`/`output_schema_raw` is also set.
313    pub fn output_mode(mut self, mode: OutputMode) -> Self {
314        self.output_mode = mode;
315        self
316    }
317
318    /// Attach a [`ConversationMemory`] backend.
319    ///
320    /// When set, the agent will automatically load prior conversation history before
321    /// each prompt and append the new turn after a successful response. A
322    /// `conversation_id` must be supplied either via [`AgentBuilder::conversation`]
323    /// or per-request via [`crate::agent::prompt_request::PromptRequest::conversation`].
324    /// If neither is set, memory is silently bypassed.
325    pub fn memory<B>(mut self, memory: B) -> Self
326    where
327        B: ConversationMemory + 'static,
328    {
329        self.memory = Some(Arc::new(memory));
330        self
331    }
332
333    /// Set a default conversation id used when none is provided per-request.
334    ///
335    /// Most agents are reused across users or threads; prefer setting the id
336    /// per-request via [`crate::agent::prompt_request::PromptRequest::conversation`].
337    pub fn conversation(mut self, id: impl Into<String>) -> Self {
338        self.default_conversation_id = Some(id.into());
339        self
340    }
341
342    /// Attach a default hook to the agent. Each call appends to the agent's hook
343    /// stack; hooks run for every prompt request (unless more are added per
344    /// request) in registration order. How their results compose is
345    /// event-dependent: `CompletionCall` request patches accumulate and merge,
346    /// `ToolCall`/`ToolResult` rewrites chain, while model-turn steering and
347    /// observe-only/recovery events use first-non-`Continue`-wins. See the
348    /// [`hook`](crate::agent::hook) module docs.
349    pub fn add_hook<H>(mut self, hook: H) -> Self
350    where
351        H: AgentHook + 'static,
352    {
353        self.hooks.push(hook);
354        self
355    }
356}
357
358impl<M> AgentBuilder<M, NoToolConfig>
359where
360    M: CompletionModel,
361{
362    /// Create a new agent builder with the given model
363    pub fn new(model: M) -> Self {
364        Self {
365            name: None,
366            description: None,
367            model,
368            preamble: None,
369            static_context: vec![],
370            temperature: None,
371            max_tokens: None,
372            additional_params: None,
373            record_telemetry_content: false,
374            tool_choice: None,
375            default_max_turns: None,
376            tool_state: NoToolConfig,
377            hooks: HookStack::new(),
378            output_schema: None,
379            output_mode: OutputMode::default(),
380            memory: None,
381            default_conversation_id: None,
382        }
383    }
384}
385
386impl<M> AgentBuilder<M, NoToolConfig>
387where
388    M: CompletionModel,
389{
390    /// Set a pre-existing ToolServerHandle for the agent.
391    ///
392    /// After calling this method, tool-adding methods (`.tool()`, `.dynamic_tool()`, etc.)
393    /// will not be available. Use this when you want to share a `ToolServer`
394    /// between multiple agents or have pre-configured tools.
395    pub fn tool_server_handle(
396        self,
397        handle: ToolServerHandle,
398    ) -> AgentBuilder<M, WithToolServerHandle> {
399        AgentBuilder {
400            name: self.name,
401            description: self.description,
402            model: self.model,
403            preamble: self.preamble,
404            static_context: self.static_context,
405            additional_params: self.additional_params,
406            record_telemetry_content: self.record_telemetry_content,
407            max_tokens: self.max_tokens,
408            temperature: self.temperature,
409            tool_choice: self.tool_choice,
410            default_max_turns: self.default_max_turns,
411            tool_state: WithToolServerHandle { handle },
412            hooks: self.hooks,
413            output_schema: self.output_schema,
414            output_mode: self.output_mode,
415            memory: self.memory,
416            default_conversation_id: self.default_conversation_id,
417        }
418    }
419
420    /// Add a static tool to the agent.
421    ///
422    /// This transitions the builder to the `WithBuilderTools` state, where
423    /// additional tools can be added but `tool_server_handle()` is no longer available.
424    pub fn tool<T>(self, tool: T) -> AgentBuilder<M, WithBuilderTools>
425    where
426        T: Tool + 'static,
427    {
428        let mut tools = ToolSet::default();
429        tools.add_tool(tool);
430        AgentBuilder {
431            name: self.name,
432            description: self.description,
433            model: self.model,
434            preamble: self.preamble,
435            static_context: self.static_context,
436            additional_params: self.additional_params,
437            record_telemetry_content: self.record_telemetry_content,
438            max_tokens: self.max_tokens,
439            temperature: self.temperature,
440            tool_choice: self.tool_choice,
441            default_max_turns: self.default_max_turns,
442            tool_state: WithBuilderTools {
443                tools,
444                retrieval_indexes: vec![],
445            },
446            hooks: self.hooks,
447            output_schema: self.output_schema,
448            output_mode: self.output_mode,
449            memory: self.memory,
450            default_conversation_id: self.default_conversation_id,
451        }
452    }
453
454    /// Add one runtime-defined tool to the agent.
455    pub fn dynamic_tool(self, tool: DynamicTool) -> AgentBuilder<M, WithBuilderTools> {
456        self.dynamic_tools(vec![tool])
457    }
458
459    /// Add one context-free dynamic tool through the classic registry adapter.
460    pub fn portable_dynamic_tool(
461        self,
462        tool: PortableDynamicTool,
463    ) -> AgentBuilder<M, WithBuilderTools> {
464        self.dynamic_tool(DynamicTool::from_portable(tool))
465    }
466
467    /// Add runtime-defined tools to the agent.
468    ///
469    /// This is useful when tool definitions and callbacks are constructed at runtime.
470    /// Transitions the builder to the `WithBuilderTools` state.
471    pub fn dynamic_tools(self, tools: Vec<DynamicTool>) -> AgentBuilder<M, WithBuilderTools> {
472        let tools = ToolSet::from_dynamic_tools(tools);
473
474        AgentBuilder {
475            name: self.name,
476            description: self.description,
477            model: self.model,
478            preamble: self.preamble,
479            static_context: self.static_context,
480            additional_params: self.additional_params,
481            record_telemetry_content: self.record_telemetry_content,
482            max_tokens: self.max_tokens,
483            temperature: self.temperature,
484            tool_choice: self.tool_choice,
485            default_max_turns: self.default_max_turns,
486            hooks: self.hooks,
487            output_schema: self.output_schema,
488            output_mode: self.output_mode,
489            memory: self.memory,
490            default_conversation_id: self.default_conversation_id,
491            tool_state: WithBuilderTools {
492                tools,
493                retrieval_indexes: vec![],
494            },
495        }
496    }
497
498    /// Add an MCP tool (from `rmcp`) to the agent, bounded by
499    /// [`DEFAULT_MCP_TOOL_TIMEOUT`](crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT)
500    /// (see issue #1914). Use [`rmcp_tool_with_timeout`](Self::rmcp_tool_with_timeout)
501    /// to change or disable it.
502    ///
503    /// Transitions the builder to the `WithBuilderTools` state.
504    #[cfg(all(feature = "rmcp", not(target_family = "wasm")))]
505    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
506    pub fn rmcp_tool(
507        self,
508        tool: rmcp::model::Tool,
509        client: rmcp::service::ServerSink,
510    ) -> AgentBuilder<M, WithBuilderTools> {
511        self.rmcp_tool_with_timeout(tool, client, crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT)
512    }
513
514    /// Add an MCP tool (from `rmcp`) with a per-call timeout (see issue #1914).
515    ///
516    /// Pass a [`Duration`](std::time::Duration) to bound the call, or `None` to
517    /// disable the timeout (unbounded). On timeout the call resolves to a tool
518    /// error the agent can recover from instead of blocking forever.
519    /// Transitions the builder to the `WithBuilderTools` state.
520    #[cfg(all(feature = "rmcp", not(target_family = "wasm")))]
521    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
522    pub fn rmcp_tool_with_timeout(
523        self,
524        tool: rmcp::model::Tool,
525        client: rmcp::service::ServerSink,
526        timeout: impl Into<Option<std::time::Duration>>,
527    ) -> AgentBuilder<M, WithBuilderTools> {
528        self.with_rmcp_toolset(build_rmcp_tools(vec![tool], client, timeout.into()))
529    }
530
531    /// Add an array of MCP tools (from `rmcp`) to the agent, each bounded by
532    /// [`DEFAULT_MCP_TOOL_TIMEOUT`](crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT)
533    /// (see issue #1914). Use [`rmcp_tools_with_timeout`](Self::rmcp_tools_with_timeout)
534    /// to change or disable it.
535    ///
536    /// Transitions the builder to the `WithBuilderTools` state.
537    #[cfg(all(feature = "rmcp", not(target_family = "wasm")))]
538    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
539    pub fn rmcp_tools(
540        self,
541        tools: Vec<rmcp::model::Tool>,
542        client: rmcp::service::ServerSink,
543    ) -> AgentBuilder<M, WithBuilderTools> {
544        self.rmcp_tools_with_timeout(tools, client, crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT)
545    }
546
547    /// Add an array of MCP tools (from `rmcp`) with a per-call timeout (see
548    /// issue #1914).
549    ///
550    /// Pass a [`Duration`](std::time::Duration) to bound calls, or `None` to
551    /// disable the timeout (unbounded). On timeout a call resolves to a tool
552    /// error the agent can recover from instead of blocking forever.
553    /// Transitions the builder to the `WithBuilderTools` state.
554    #[cfg(all(feature = "rmcp", not(target_family = "wasm")))]
555    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
556    pub fn rmcp_tools_with_timeout(
557        self,
558        tools: Vec<rmcp::model::Tool>,
559        client: rmcp::service::ServerSink,
560        timeout: impl Into<Option<std::time::Duration>>,
561    ) -> AgentBuilder<M, WithBuilderTools> {
562        self.with_rmcp_toolset(build_rmcp_tools(tools, client, timeout.into()))
563    }
564
565    /// Transition into the `WithBuilderTools` state carrying the given built
566    /// MCP tools.
567    #[cfg(all(feature = "rmcp", not(target_family = "wasm")))]
568    fn with_rmcp_toolset(
569        self,
570        built: Vec<(String, RmcpTool)>,
571    ) -> AgentBuilder<M, WithBuilderTools> {
572        AgentBuilder {
573            name: self.name,
574            description: self.description,
575            model: self.model,
576            preamble: self.preamble,
577            static_context: self.static_context,
578            additional_params: self.additional_params,
579            record_telemetry_content: self.record_telemetry_content,
580            max_tokens: self.max_tokens,
581            temperature: self.temperature,
582            tool_choice: self.tool_choice,
583            default_max_turns: self.default_max_turns,
584            hooks: self.hooks,
585            output_schema: self.output_schema,
586            output_mode: self.output_mode,
587            memory: self.memory,
588            default_conversation_id: self.default_conversation_id,
589            tool_state: WithBuilderTools {
590                tools: {
591                    let mut set = ToolSet::default();
592                    for (_, tool) in built {
593                        set.add_erased(std::sync::Arc::new(tool));
594                    }
595                    set
596                },
597                retrieval_indexes: vec![],
598            },
599        }
600    }
601
602    /// Configure tools retrieved from a vector index for each prompt.
603    ///
604    /// Transitions the builder to the `WithBuilderTools` state.
605    pub fn retrieved_tools(
606        self,
607        sample: usize,
608        index: impl VectorStoreIndexDyn + Send + Sync + 'static,
609        toolset: ToolSet,
610    ) -> AgentBuilder<M, WithBuilderTools> {
611        let mut tools = ToolSet::default();
612        tools.add_retrievable_tools(toolset);
613        AgentBuilder {
614            name: self.name,
615            description: self.description,
616            model: self.model,
617            preamble: self.preamble,
618            static_context: self.static_context,
619            additional_params: self.additional_params,
620            record_telemetry_content: self.record_telemetry_content,
621            max_tokens: self.max_tokens,
622            temperature: self.temperature,
623            tool_choice: self.tool_choice,
624            default_max_turns: self.default_max_turns,
625            hooks: self.hooks,
626            output_schema: self.output_schema,
627            output_mode: self.output_mode,
628            memory: self.memory,
629            default_conversation_id: self.default_conversation_id,
630            tool_state: WithBuilderTools {
631                tools,
632                retrieval_indexes: vec![(sample, Arc::new(index))],
633            },
634        }
635    }
636
637    /// Build the agent with no tools configured.
638    ///
639    /// An empty `ToolServer` will be created for the agent.
640    pub fn build(self) -> Agent<M> {
641        let tool_server_handle = ToolServer::new().run();
642
643        Agent {
644            name: self.name,
645            description: self.description,
646            model: Arc::new(self.model),
647            preamble: self.preamble,
648            static_context: self.static_context,
649            temperature: self.temperature,
650            max_tokens: self.max_tokens,
651            additional_params: self.additional_params,
652            record_telemetry_content: self.record_telemetry_content,
653            tool_choice: self.tool_choice,
654            tool_server_handle,
655            default_max_turns: self.default_max_turns,
656            hooks: self.hooks,
657            output_schema: self.output_schema,
658            output_mode: self.output_mode,
659            memory: self.memory,
660            default_conversation_id: self.default_conversation_id,
661        }
662    }
663}
664
665impl<M> AgentBuilder<M, WithToolServerHandle>
666where
667    M: CompletionModel,
668{
669    /// Build the agent using the pre-configured ToolServerHandle.
670    pub fn build(self) -> Agent<M> {
671        Agent {
672            name: self.name,
673            description: self.description,
674            model: Arc::new(self.model),
675            preamble: self.preamble,
676            static_context: self.static_context,
677            temperature: self.temperature,
678            max_tokens: self.max_tokens,
679            additional_params: self.additional_params,
680            record_telemetry_content: self.record_telemetry_content,
681            tool_choice: self.tool_choice,
682            tool_server_handle: self.tool_state.handle,
683            default_max_turns: self.default_max_turns,
684            hooks: self.hooks,
685            output_schema: self.output_schema,
686            output_mode: self.output_mode,
687            memory: self.memory,
688            default_conversation_id: self.default_conversation_id,
689        }
690    }
691}
692
693impl<M> AgentBuilder<M, WithBuilderTools>
694where
695    M: CompletionModel,
696{
697    /// Add another static tool to the agent.
698    pub fn tool<T>(mut self, tool: T) -> Self
699    where
700        T: Tool + 'static,
701    {
702        self.tool_state.tools.add_tool(tool);
703        self
704    }
705
706    /// Add one runtime-defined tool to the agent.
707    pub fn dynamic_tool(mut self, tool: DynamicTool) -> Self {
708        self.tool_state.tools.add_dynamic_tool(tool);
709        self
710    }
711
712    /// Add one context-free dynamic tool through the classic registry adapter.
713    pub fn portable_dynamic_tool(mut self, tool: PortableDynamicTool) -> Self {
714        self.tool_state.tools.add_portable_dynamic_tool(tool);
715        self
716    }
717
718    /// Add runtime-defined tools to the agent.
719    pub fn dynamic_tools(mut self, tools: Vec<DynamicTool>) -> Self {
720        let tools = ToolSet::from_dynamic_tools(tools);
721        self.tool_state.tools.add_tools(tools);
722        self
723    }
724
725    /// Add an array of MCP tools (from `rmcp`) to the agent, each bounded by
726    /// [`DEFAULT_MCP_TOOL_TIMEOUT`](crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT)
727    /// (see issue #1914). Use [`rmcp_tools_with_timeout`](Self::rmcp_tools_with_timeout)
728    /// to change or disable it.
729    #[cfg(all(feature = "rmcp", not(target_family = "wasm")))]
730    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
731    pub fn rmcp_tools(
732        self,
733        tools: Vec<rmcp::model::Tool>,
734        client: rmcp::service::ServerSink,
735    ) -> Self {
736        self.rmcp_tools_with_timeout(tools, client, crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT)
737    }
738
739    /// Add an array of MCP tools (from `rmcp`) with a per-call timeout (see
740    /// issue #1914).
741    ///
742    /// Pass a [`Duration`](std::time::Duration) to bound calls, or `None` to
743    /// disable the timeout (unbounded). On timeout a call resolves to a tool
744    /// error the agent can recover from instead of blocking forever.
745    #[cfg(all(feature = "rmcp", not(target_family = "wasm")))]
746    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
747    pub fn rmcp_tools_with_timeout(
748        self,
749        tools: Vec<rmcp::model::Tool>,
750        client: rmcp::service::ServerSink,
751        timeout: impl Into<Option<std::time::Duration>>,
752    ) -> Self {
753        self.add_rmcp_tools(build_rmcp_tools(tools, client, timeout.into()))
754    }
755
756    #[cfg(all(feature = "rmcp", not(target_family = "wasm")))]
757    fn add_rmcp_tools(mut self, built: Vec<(String, RmcpTool)>) -> Self {
758        for (_, tool) in built {
759            self.tool_state.tools.add_erased(std::sync::Arc::new(tool));
760        }
761
762        self
763    }
764
765    /// Configure tools retrieved from a vector index for each prompt.
766    pub fn retrieved_tools(
767        mut self,
768        sample: usize,
769        index: impl VectorStoreIndexDyn + Send + Sync + 'static,
770        toolset: ToolSet,
771    ) -> Self {
772        self.tool_state
773            .retrieval_indexes
774            .push((sample, Arc::new(index)));
775        self.tool_state.tools.add_retrievable_tools(toolset);
776        self
777    }
778
779    /// Build the agent with the configured tools.
780    ///
781    /// A new `ToolServer` will be created containing all tools added via
782    /// `.tool()`, `.dynamic_tool()`, `.dynamic_tools()`, and
783    /// `.retrieved_tools()`.
784    pub fn build(self) -> Agent<M> {
785        let tool_server_handle = ToolServer::new()
786            .add_tools(self.tool_state.tools)
787            .add_retrieval_indexes(self.tool_state.retrieval_indexes)
788            .run();
789
790        Agent {
791            name: self.name,
792            description: self.description,
793            model: Arc::new(self.model),
794            preamble: self.preamble,
795            static_context: self.static_context,
796            temperature: self.temperature,
797            max_tokens: self.max_tokens,
798            additional_params: self.additional_params,
799            record_telemetry_content: self.record_telemetry_content,
800            tool_choice: self.tool_choice,
801            tool_server_handle,
802            default_max_turns: self.default_max_turns,
803            hooks: self.hooks,
804            output_schema: self.output_schema,
805            output_mode: self.output_mode,
806            memory: self.memory,
807            default_conversation_id: self.default_conversation_id,
808        }
809    }
810}
811#[cfg(test)]
812mod tests {
813    use super::*;
814    use crate::test_utils::{MockAddTool, MockCompletionModel, MockSubtractTool, MockToolIndex};
815    use crate::tool::{ToolContext, ToolExecutionError};
816
817    #[derive(Clone)]
818    struct BuilderHook;
819
820    impl AgentHook for BuilderHook {}
821
822    #[test]
823    fn hook_can_be_set_after_tool_configuration() {
824        let _agent = AgentBuilder::new(MockCompletionModel::text("ok"))
825            .tool(MockAddTool)
826            .add_hook(BuilderHook)
827            .build();
828    }
829
830    struct NamedTool;
831
832    impl NamedTool {
833        fn new() -> Self {
834            Self
835        }
836    }
837
838    impl Tool for NamedTool {
839        const NAME: &'static str = "registered_named";
840        type Error = rig::tool::ToolExecutionError;
841        type Args = serde_json::Value;
842        type Output = String;
843
844        fn description(&self) -> String {
845            "uses its canonical name".to_string()
846        }
847
848        fn parameters(&self) -> serde_json::Value {
849            serde_json::json!({"type": "object", "properties": {}})
850        }
851
852        async fn call(
853            &self,
854            _context: &mut ToolContext,
855            _args: Self::Args,
856        ) -> Result<Self::Output, ToolExecutionError> {
857            Ok("ok".to_string())
858        }
859    }
860
861    #[tokio::test]
862    async fn typed_tool_builder_paths_advertise_canonical_name() {
863        for agent in [
864            AgentBuilder::new(MockCompletionModel::text("ok"))
865                .tool(NamedTool::new())
866                .build(),
867            AgentBuilder::new(MockCompletionModel::text("ok"))
868                .tool(MockAddTool)
869                .tool(NamedTool::new())
870                .build(),
871        ] {
872            let definitions = agent.tool_server_handle.get_tool_defs(None).await.unwrap();
873            assert!(
874                definitions
875                    .iter()
876                    .any(|definition| definition.name == NamedTool::NAME),
877                "the provider definitions dropped the canonical tool name"
878            );
879
880            let mut context = ToolContext::new();
881            let result = agent
882                .tool_server_handle
883                .execute(NamedTool::NAME, "{}", &mut context)
884                .await;
885            assert!(result.is_success());
886            assert_eq!(result.output().as_text(), Some("ok"));
887        }
888    }
889
890    #[tokio::test]
891    async fn retrieved_tools_are_exposed_only_for_prompted_retrieval() {
892        let retrieval_only = AgentBuilder::new(MockCompletionModel::text("ok"))
893            .retrieved_tools(
894                1,
895                MockToolIndex::new(["add"]),
896                ToolSet::from_tools(vec![MockAddTool]),
897            )
898            .build();
899        assert!(
900            retrieval_only
901                .tool_server_handle
902                .get_tool_defs(None)
903                .await
904                .unwrap()
905                .is_empty()
906        );
907
908        let agent = AgentBuilder::new(MockCompletionModel::text("ok"))
909            .tool(MockSubtractTool)
910            .retrieved_tools(
911                1,
912                MockToolIndex::new(["add"]),
913                ToolSet::from_tools(vec![MockAddTool]),
914            )
915            .build();
916
917        let always = agent.tool_server_handle.get_tool_defs(None).await.unwrap();
918        assert_eq!(
919            always
920                .iter()
921                .map(|definition| definition.name.as_str())
922                .collect::<Vec<_>>(),
923            vec!["subtract"]
924        );
925
926        let with_retrieval = agent
927            .tool_server_handle
928            .get_tool_defs(Some("add two numbers".to_string()))
929            .await
930            .unwrap();
931        assert_eq!(
932            with_retrieval
933                .iter()
934                .map(|definition| definition.name.as_str())
935                .collect::<Vec<_>>(),
936            vec!["add", "subtract"]
937        );
938    }
939
940    /// The builder's shared MCP helper threads the configured timeout (default,
941    /// explicit, or `None`/disabled) onto every built tool, and the threaded
942    /// timeout actually bounds a hanging call. This covers the plumbing behind
943    /// `rmcp_tool[s]` / `rmcp_tool[s]_with_timeout` (see issue #1914).
944    #[cfg(all(feature = "rmcp", not(target_family = "wasm")))]
945    #[tokio::test]
946    async fn build_rmcp_tools_threads_timeout_into_built_tools() {
947        use crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT;
948        use crate::tool::{ToolContext, ToolErrorKind, server::ToolServer};
949        use rmcp::model::{
950            CallToolRequestParams, CallToolResult, ClientInfo, ErrorData, Implementation,
951            ProtocolVersion, ServerCapabilities, ServerInfo, Tool,
952        };
953        use rmcp::service::RequestContext;
954        use rmcp::{RoleServer, ServerHandler, ServiceExt};
955        use std::sync::Arc;
956        use std::time::Duration;
957
958        #[derive(Clone)]
959        struct HangingServer;
960        impl ServerHandler for HangingServer {
961            fn get_info(&self) -> ServerInfo {
962                ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
963                    .with_protocol_version(ProtocolVersion::LATEST)
964                    .with_server_info(Implementation::new("builder-timeout-test", "0.1.0"))
965            }
966            async fn call_tool(
967                &self,
968                _request: CallToolRequestParams,
969                _context: RequestContext<RoleServer>,
970            ) -> Result<CallToolResult, ErrorData> {
971                std::future::pending::<Result<CallToolResult, ErrorData>>().await
972            }
973        }
974
975        fn tool(name: &str) -> Tool {
976            Tool::new(
977                name.to_string(),
978                String::new(),
979                Arc::new(serde_json::Map::new()),
980            )
981        }
982
983        let (c2s, sfc) = tokio::io::duplex(8192);
984        let (s2c, cfs) = tokio::io::duplex(8192);
985        let server_task = tokio::spawn(async move {
986            let running = HangingServer.serve((sfc, s2c)).await.expect("server start");
987            running.waiting().await.expect("server error");
988        });
989        let client = ClientInfo::default()
990            .serve((cfs, c2s))
991            .await
992            .expect("client connect");
993        let peer = client.peer().clone();
994
995        // The configured timeout (default, explicit, or disabled) is threaded
996        // onto each built tool.
997        let built_default = build_rmcp_tools(
998            vec![tool("a")],
999            peer.clone(),
1000            Some(DEFAULT_MCP_TOOL_TIMEOUT),
1001        );
1002        assert_eq!(built_default[0].1.timeout(), Some(DEFAULT_MCP_TOOL_TIMEOUT));
1003        let built_none = build_rmcp_tools(vec![tool("b")], peer.clone(), None);
1004        assert_eq!(built_none[0].1.timeout(), None);
1005
1006        // ...and the threaded timeout actually bounds a hanging call.
1007        let built = build_rmcp_tools(
1008            vec![tool("hang_forever")],
1009            peer,
1010            Some(Duration::from_millis(200)),
1011        );
1012        assert_eq!(built.len(), 1);
1013        assert_eq!(built[0].0, "hang_forever");
1014        let handle = ToolServer::new().run();
1015        handle
1016            .add_erased_tool(Arc::new(built.into_iter().next().unwrap().1))
1017            .await;
1018        let timed = tokio::time::timeout(Duration::from_secs(5), async {
1019            let mut context = ToolContext::new();
1020            handle.execute("hang_forever", "{}", &mut context).await
1021        })
1022        .await;
1023        let result = timed.expect("built tool hung past the safety timeout");
1024        assert!(result.is_error_kind(ToolErrorKind::Timeout));
1025        assert!(result.output().render().contains("timed out"));
1026
1027        drop(client);
1028        server_task.abort();
1029    }
1030}
rig_agent/agent/builder.rs

rig_agent/agent/
builder.rs
