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::{AgentHook, CompletionCall, CompletionCallAction, HookContext, RequestPatch},
13    completion::{CompletionModel, Document},
14    tool::{
15        DynamicTool, PortableDynamicTool, Tool, ToolSet,
16        server::{ToolServer, ToolServerHandle},
17    },
18};
19
20use super::{Agent, ModelHandle, OutputMode, completion::AgentConfig};
21
22struct DynamicContext<I> {
23    samples: usize,
24    index: I,
25}
26
27impl<I> AgentHook for DynamicContext<I>
28where
29    I: VectorStoreIndexDyn,
30{
31    async fn on_completion_call(
32        &self,
33        _ctx: &HookContext,
34        event: CompletionCall<'_>,
35    ) -> CompletionCallAction {
36        let query = event.prompt.rag_text().or_else(|| {
37            event
38                .history
39                .iter()
40                .rev()
41                .find_map(|message| message.rag_text())
42        });
43        let Some(query) = query else {
44            return CompletionCallAction::continue_run();
45        };
46
47        let request = VectorSearchRequest::builder()
48            .query(query)
49            .samples(self.samples as u64)
50            .build();
51        match self.index.top_n(request).await {
52            Ok(results) => CompletionCallAction::patch(RequestPatch::new().extra_context(
53                results.into_iter().map(|(_, id, value)| Document {
54                    id,
55                    text:
56                        serde_json::to_string_pretty(&value).unwrap_or_else(|_| value.to_string()),
57                    additional_props: Default::default(),
58                }),
59            )),
60            Err(error) => {
61                CompletionCallAction::stop(format!("failed to retrieve dynamic context: {error}"))
62            }
63        }
64    }
65}
66
67/// Marker type indicating no tool configuration has been set yet.
68///
69/// This is the default state for a new `AgentBuilder`. From this state,
70/// you can either:
71/// - Add tools via `.tool()`, `.dynamic_tool()`, `.dynamic_tools()`, or
72///   `.retrieved_tools()` (transitions to `WithBuilderTools`)
73/// - Set a pre-existing `ToolServerHandle` via `.tool_server_handle()` (transitions to `WithToolServerHandle`)
74/// - Call `.build()` to create an agent with no tools
75#[derive(Default)]
76pub struct NoToolConfig;
77
78/// Typestate indicating a pre-existing `ToolServerHandle` has been provided.
79///
80/// In this state, tool-adding methods (`.tool()`, `.dynamic_tool()`, etc.) are not available.
81/// The provided handle will be used directly when building the agent.
82pub struct WithToolServerHandle {
83    handle: ToolServerHandle,
84}
85
86/// Typestate indicating tools are being configured via the builder API.
87///
88/// In this state, you can continue adding tools via `.tool()`,
89/// `.dynamic_tool()`, `.dynamic_tools()`, and `.retrieved_tools()`. When
90/// `.build()` is called, a new `ToolServer`
91/// will be created with all the configured tools.
92pub struct WithBuilderTools(ToolServer);
93
94/// A builder for creating an agent
95///
96/// The builder uses a typestate pattern to enforce that tool configuration
97/// is done in a mutually exclusive way: either provide a pre-existing
98/// `ToolServerHandle`, or add tools via the builder API, but not both.
99///
100/// # Example
101/// ```no_run
102/// use rig_agent::AgentBuilder;
103/// use rig_core::{client::{CompletionClient, ProviderClient}, providers::openai};
104///
105/// # fn run() -> Result<(), Box<dyn std::error::Error>> {
106/// let openai = openai::Client::from_env()?;
107///
108/// let model = openai.completion_model(openai::GPT_5_2);
109///
110/// // Configure the agent
111/// let agent = AgentBuilder::new(model)
112///     .preamble("System prompt")
113///     .context("Context document 1")
114///     .context("Context document 2")
115///     .temperature(0.8)
116///     .build();
117/// # Ok(())
118/// # }
119/// ```
120pub struct AgentBuilder<ToolState = NoToolConfig> {
121    /// Everything the built [`Agent`] carries unchanged.
122    config: AgentConfig,
123    /// Tool configuration state (typestate pattern)
124    tool_state: ToolState,
125}
126
127impl<ToolState> AgentBuilder<ToolState> {
128    /// Set the name of the agent
129    pub fn name(mut self, name: &str) -> Self {
130        self.config.name = Some(name.into());
131        self
132    }
133
134    /// Set the description of the agent
135    pub fn description(mut self, description: &str) -> Self {
136        self.config.description = Some(description.into());
137        self
138    }
139
140    /// Set the system prompt
141    pub fn preamble(mut self, preamble: &str) -> Self {
142        self.config.preamble = Some(preamble.into());
143        self
144    }
145
146    /// Remove the system prompt
147    pub fn without_preamble(mut self) -> Self {
148        self.config.preamble = None;
149        self
150    }
151
152    /// Append to the preamble of the agent
153    pub fn append_preamble(mut self, doc: &str) -> Self {
154        self.config.preamble = Some(format!(
155            "{}\n{}",
156            self.config.preamble.unwrap_or_default(),
157            doc
158        ));
159        self
160    }
161
162    /// Add a static context document to the agent
163    pub fn context(mut self, doc: &str) -> Self {
164        self.config.static_context.push(Document {
165            id: format!("static_doc_{}", self.config.static_context.len()),
166            text: doc.into(),
167            additional_props: HashMap::new(),
168        });
169        self
170    }
171
172    /// Add dynamic context retrieved from a vector store on every model call.
173    ///
174    /// This is a convenience wrapper around an internal completion-call hook.
175    /// The hook searches with the current prompt's first text part, falling back
176    /// to the latest textual history message, and appends the retrieved documents
177    /// to the request after static context. Retrieval and injected documents
178    /// follow registration order relative to application hooks, so register a
179    /// stop policy before this helper when it should prevent retrieval. A
180    /// retrieval failure stops the run before provider I/O.
181    pub fn dynamic_context<I>(self, samples: usize, index: I) -> Self
182    where
183        I: VectorStoreIndexDyn + 'static,
184    {
185        self.add_hook(DynamicContext { samples, index })
186    }
187
188    /// Set the tool choice for the agent
189    pub fn tool_choice(mut self, tool_choice: ToolChoice) -> Self {
190        self.config.tool_choice = Some(tool_choice);
191        self
192    }
193
194    /// Set the default total model-call budget, including the initial call and
195    /// every retry or continuation. Zero permits no model calls.
196    pub fn default_max_turns(mut self, default_max_turns: usize) -> Self {
197        self.config.max_turns = default_max_turns;
198        self
199    }
200
201    /// Set the temperature of the model
202    pub fn temperature(mut self, temperature: f64) -> Self {
203        self.config.temperature = Some(temperature);
204        self
205    }
206
207    /// Set the maximum number of tokens for the completion
208    pub fn max_tokens(mut self, max_tokens: u64) -> Self {
209        self.config.max_tokens = Some(max_tokens);
210        self
211    }
212
213    /// Set additional parameters to be passed to the model
214    pub fn additional_params(mut self, params: serde_json::Value) -> Self {
215        self.config.additional_params = Some(params);
216        self
217    }
218
219    /// Opt in or out of recording sensitive request, response, and tool content
220    /// on GenAI telemetry spans for requests made by this agent.
221    ///
222    /// Defaults to `false`. Enabling this can expose prompts, retrieved context,
223    /// tool results, model responses, and other sensitive or high-cardinality data
224    /// through OpenTelemetry span attributes, which can increase observability
225    /// backend storage and query costs. Only enable it when content telemetry is
226    /// acceptable for this agent. Structural metadata and token usage remain
227    /// available when this is disabled.
228    pub fn record_content_telemetry(mut self, enabled: bool) -> Self {
229        self.config.record_telemetry_content = enabled;
230        self
231    }
232
233    /// Set the output schema for structured output. When set, providers that support
234    /// native structured outputs will constrain the model's response to match this schema.
235    pub fn output_schema<T>(mut self) -> Self
236    where
237        T: JsonSchema,
238    {
239        self.config.output_schema = Some(schema_for!(T));
240        self
241    }
242
243    /// 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.
244    pub fn output_schema_raw(mut self, schema: Schema) -> Self {
245        self.config.output_schema = Some(schema);
246        self
247    }
248
249    /// Set how `output_schema` is enforced — [`OutputMode::Tool`] (output as a
250    /// tool call, the default when the agent has tools), [`OutputMode::Native`]
251    /// (provider structured output), or [`OutputMode::Prompted`] (see #1928).
252    /// Has no effect unless `output_schema`/`output_schema_raw` is also set.
253    pub fn output_mode(mut self, mode: OutputMode) -> Self {
254        self.config.output_mode = mode;
255        self
256    }
257
258    /// Attach a [`ConversationMemory`] backend.
259    ///
260    /// When set, the agent will automatically load prior conversation history before
261    /// each prompt and append the new turn after a successful response. A
262    /// `conversation_id` must be supplied either via [`AgentBuilder::conversation`]
263    /// or per-request via [`crate::agent::prompt_request::PromptRequest::conversation`].
264    /// If neither is set, memory is silently bypassed.
265    pub fn memory<B>(mut self, memory: B) -> Self
266    where
267        B: ConversationMemory + 'static,
268    {
269        self.config.memory = Some(Arc::new(memory));
270        self
271    }
272
273    /// Set a default conversation id used when none is provided per-request.
274    ///
275    /// Most agents are reused across users or threads; prefer setting the id
276    /// per-request via [`crate::agent::prompt_request::PromptRequest::conversation`].
277    pub fn conversation(mut self, id: impl Into<String>) -> Self {
278        self.config.conversation_id = Some(id.into());
279        self
280    }
281
282    /// Attach a default hook to the agent. Each call appends to the agent's hook
283    /// stack; hooks run for every prompt request (unless more are added per
284    /// request) in registration order. How their results compose is
285    /// event-dependent: model selections and `ToolCall`/`ToolResult` rewrites
286    /// chain, `CompletionCall` request patches accumulate and merge, while
287    /// model-turn steering and observe-only/recovery events use
288    /// first-non-`Continue`-wins. See the [`hook`](crate::agent::hook) module
289    /// docs.
290    pub fn add_hook<H>(mut self, hook: H) -> Self
291    where
292        H: AgentHook + 'static,
293    {
294        self.config.hooks.push(hook);
295        self
296    }
297
298    /// Carry the configuration into a builder with a new tool state.
299    fn with_tool_state<S>(self, tool_state: S) -> AgentBuilder<S> {
300        AgentBuilder {
301            config: self.config,
302            tool_state,
303        }
304    }
305
306    /// Assemble the [`Agent`], resolving the tool server handle from the final
307    /// tool state.
308    fn build_agent(self, handle: impl FnOnce(ToolState) -> ToolServerHandle) -> Agent {
309        Agent {
310            tool_server_handle: handle(self.tool_state),
311            config: self.config,
312        }
313    }
314}
315
316impl AgentBuilder<NoToolConfig> {
317    /// Create a new agent builder with the given model.
318    ///
319    /// The typed model is erased once, here, into a [`ModelHandle`]; the built
320    /// [`Agent`] carries no model type parameter.
321    pub fn new<M>(model: M) -> Self
322    where
323        M: CompletionModel + 'static,
324    {
325        Self::from_model_handle(ModelHandle::new(model))
326    }
327
328    /// Create an agent builder from an already-erased runtime model handle.
329    pub fn from_model_handle(model: ModelHandle) -> Self {
330        Self {
331            config: AgentConfig::new(model),
332            tool_state: NoToolConfig,
333        }
334    }
335}
336
337impl AgentBuilder<NoToolConfig> {
338    /// Set a pre-existing ToolServerHandle for the agent.
339    ///
340    /// After calling this method, tool-adding methods (`.tool()`, `.dynamic_tool()`, etc.)
341    /// will not be available. Use this when you want to share a `ToolServer`
342    /// between multiple agents or have pre-configured tools.
343    pub fn tool_server_handle(
344        self,
345        handle: ToolServerHandle,
346    ) -> AgentBuilder<WithToolServerHandle> {
347        self.with_tool_state(WithToolServerHandle { handle })
348    }
349
350    /// Transition into the `WithBuilderTools` state with no tools yet; every
351    /// tool-adding method below is the `WithBuilderTools` method after this
352    /// one-way step.
353    fn into_tool_builder(self) -> AgentBuilder<WithBuilderTools> {
354        self.with_tool_state(WithBuilderTools(ToolServer::new()))
355    }
356
357    /// Add a static tool to the agent.
358    ///
359    /// This transitions the builder to the `WithBuilderTools` state, where
360    /// additional tools can be added but `tool_server_handle()` is no longer available.
361    pub fn tool<T>(self, tool: T) -> AgentBuilder<WithBuilderTools>
362    where
363        T: Tool + 'static,
364    {
365        self.into_tool_builder().tool(tool)
366    }
367
368    /// Add an MCP tool (from `rmcp`) to the agent, bounded by
369    /// [`DEFAULT_MCP_TOOL_TIMEOUT`](crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT)
370    /// (see issue #1914). Use [`rmcp_tool_with_timeout`](Self::rmcp_tool_with_timeout)
371    /// to change or disable it.
372    ///
373    /// Transitions the builder to the `WithBuilderTools` state.
374    #[cfg(all(feature = "rmcp", not(target_family = "wasm")))]
375    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
376    pub fn rmcp_tool(
377        self,
378        tool: rmcp::model::Tool,
379        client: rmcp::service::ServerSink,
380    ) -> AgentBuilder<WithBuilderTools> {
381        self.rmcp_tool_with_timeout(tool, client, crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT)
382    }
383
384    /// Add an MCP tool (from `rmcp`) with a per-call timeout (see issue #1914).
385    ///
386    /// Pass a [`Duration`](std::time::Duration) to bound the call, or `None` to
387    /// disable the timeout (unbounded). On timeout the call resolves to a tool
388    /// error the agent can recover from instead of blocking forever.
389    /// Transitions the builder to the `WithBuilderTools` state.
390    #[cfg(all(feature = "rmcp", not(target_family = "wasm")))]
391    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
392    pub fn rmcp_tool_with_timeout(
393        self,
394        tool: rmcp::model::Tool,
395        client: rmcp::service::ServerSink,
396        timeout: impl Into<Option<std::time::Duration>>,
397    ) -> AgentBuilder<WithBuilderTools> {
398        self.rmcp_tools_with_timeout(vec![tool], client, timeout)
399    }
400
401    /// Build the agent with no tools configured.
402    ///
403    /// An empty `ToolServer` will be created for the agent.
404    pub fn build(self) -> Agent {
405        self.build_agent(|_| ToolServer::new().run())
406    }
407}
408
409/// Generate the `NoToolConfig` tool methods that transition into the
410/// `WithBuilderTools` state by forwarding verbatim through
411/// [`AgentBuilder::into_tool_builder`] to the `WithBuilderTools` method of the
412/// same name. Doc comments live at each invocation; `tool` (generic over the
413/// tool type) and the single-tool rmcp helpers stay hand-written above.
414macro_rules! forward_into_tool_builder {
415    ($( $(#[$attr:meta])* $name:ident ( $($arg:ident : $ty:ty),* $(,)? ) );* $(;)?) => {
416        impl AgentBuilder<NoToolConfig> {
417            $(
418                $(#[$attr])*
419                pub fn $name(self, $($arg: $ty),*) -> AgentBuilder<WithBuilderTools> {
420                    self.into_tool_builder().$name($($arg),*)
421                }
422            )*
423        }
424    };
425}
426
427forward_into_tool_builder! {
428    /// Add one runtime-defined tool to the agent.
429    dynamic_tool(tool: DynamicTool);
430
431    /// Add one context-free dynamic tool through the classic registry adapter.
432    portable_dynamic_tool(tool: PortableDynamicTool);
433
434    /// Add runtime-defined tools to the agent.
435    ///
436    /// This is useful when tool definitions and callbacks are constructed at runtime.
437    /// Transitions the builder to the `WithBuilderTools` state.
438    dynamic_tools(tools: Vec<DynamicTool>);
439
440    /// Add an array of MCP tools (from `rmcp`) to the agent, each bounded by
441    /// [`DEFAULT_MCP_TOOL_TIMEOUT`](crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT)
442    /// (see issue #1914). Use [`rmcp_tools_with_timeout`](Self::rmcp_tools_with_timeout)
443    /// to change or disable it.
444    ///
445    /// Transitions the builder to the `WithBuilderTools` state.
446    #[cfg(all(feature = "rmcp", not(target_family = "wasm")))]
447    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
448    rmcp_tools(tools: Vec<rmcp::model::Tool>, client: rmcp::service::ServerSink);
449
450    /// Add an array of MCP tools (from `rmcp`) with a per-call timeout (see
451    /// issue #1914).
452    ///
453    /// Pass a [`Duration`](std::time::Duration) to bound calls, or `None` to
454    /// disable the timeout (unbounded). On timeout a call resolves to a tool
455    /// error the agent can recover from instead of blocking forever.
456    /// Transitions the builder to the `WithBuilderTools` state.
457    #[cfg(all(feature = "rmcp", not(target_family = "wasm")))]
458    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
459    rmcp_tools_with_timeout(
460        tools: Vec<rmcp::model::Tool>,
461        client: rmcp::service::ServerSink,
462        timeout: impl Into<Option<std::time::Duration>>
463    );
464
465    /// Configure tools retrieved from a vector index for each prompt.
466    ///
467    /// Transitions the builder to the `WithBuilderTools` state.
468    retrieved_tools(
469        sample: usize,
470        index: impl VectorStoreIndexDyn + Send + Sync + 'static,
471        toolset: ToolSet
472    );
473}
474
475impl AgentBuilder<WithToolServerHandle> {
476    /// Build the agent using the pre-configured ToolServerHandle.
477    pub fn build(self) -> Agent {
478        self.build_agent(|state| state.handle)
479    }
480}
481
482impl AgentBuilder<WithBuilderTools> {
483    /// Configure the [`ToolServer`] the builder is accumulating tools into. Every
484    /// tool-adding method here is one of its registrations, so registration
485    /// semantics live in exactly one place.
486    fn map_server(self, register: impl FnOnce(ToolServer) -> ToolServer) -> Self {
487        let Self { config, tool_state } = self;
488        Self {
489            config,
490            tool_state: WithBuilderTools(register(tool_state.0)),
491        }
492    }
493
494    /// Add another static tool to the agent.
495    pub fn tool<T>(self, tool: T) -> Self
496    where
497        T: Tool + 'static,
498    {
499        self.map_server(|server| server.tool(tool))
500    }
501
502    /// Add one runtime-defined tool to the agent.
503    pub fn dynamic_tool(self, tool: DynamicTool) -> Self {
504        self.map_server(|server| server.dynamic_tool(tool))
505    }
506
507    /// Add one context-free dynamic tool through the classic registry adapter.
508    pub fn portable_dynamic_tool(self, tool: PortableDynamicTool) -> Self {
509        self.map_server(|server| server.portable_dynamic_tool(tool))
510    }
511
512    /// Add runtime-defined tools to the agent.
513    pub fn dynamic_tools(self, tools: Vec<DynamicTool>) -> Self {
514        self.map_server(|server| server.dynamic_tools(tools))
515    }
516
517    /// Add an array of MCP tools (from `rmcp`) to the agent, each bounded by
518    /// [`DEFAULT_MCP_TOOL_TIMEOUT`](crate::tool::rmcp::DEFAULT_MCP_TOOL_TIMEOUT)
519    /// (see issue #1914). Use [`rmcp_tools_with_timeout`](Self::rmcp_tools_with_timeout)
520    /// to change or disable it.
521    #[cfg(all(feature = "rmcp", not(target_family = "wasm")))]
522    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
523    pub fn rmcp_tools(
524        self,
525        tools: Vec<rmcp::model::Tool>,
526        client: rmcp::service::ServerSink,
527    ) -> Self {
528        self.map_server(|server| server.rmcp_tools(tools, client))
529    }
530
531    /// Add an array of MCP tools (from `rmcp`) with a per-call timeout (see
532    /// issue #1914).
533    ///
534    /// Pass a [`Duration`](std::time::Duration) to bound calls, or `None` to
535    /// disable the timeout (unbounded). On timeout a call resolves to a tool
536    /// error the agent can recover from instead of blocking forever.
537    #[cfg(all(feature = "rmcp", not(target_family = "wasm")))]
538    #[cfg_attr(docsrs, doc(cfg(feature = "rmcp")))]
539    pub fn rmcp_tools_with_timeout(
540        self,
541        tools: Vec<rmcp::model::Tool>,
542        client: rmcp::service::ServerSink,
543        timeout: impl Into<Option<std::time::Duration>>,
544    ) -> Self {
545        self.map_server(|server| server.rmcp_tools_with_timeout(tools, client, timeout))
546    }
547
548    /// Configure tools retrieved from a vector index for each prompt.
549    pub fn retrieved_tools(
550        self,
551        sample: usize,
552        index: impl VectorStoreIndexDyn + Send + Sync + 'static,
553        toolset: ToolSet,
554    ) -> Self {
555        self.map_server(|server| server.retrieved_tools(sample, index, toolset))
556    }
557
558    /// Build the agent with the configured tools.
559    ///
560    /// A new `ToolServer` will be created containing all tools added via
561    /// `.tool()`, `.dynamic_tool()`, `.dynamic_tools()`, and
562    /// `.retrieved_tools()`.
563    pub fn build(self) -> Agent {
564        self.build_agent(|state| state.0.run())
565    }
566}
567#[cfg(test)]
568mod tests {
569    use super::*;
570    use crate::test_utils::{MockAddTool, MockCompletionModel, MockSubtractTool, MockToolIndex};
571    use crate::tool::{ToolContext, ToolExecutionError};
572
573    #[derive(Clone)]
574    struct BuilderHook;
575
576    impl AgentHook for BuilderHook {}
577
578    #[test]
579    fn hook_can_be_set_after_tool_configuration() {
580        let _agent = AgentBuilder::new(MockCompletionModel::text("ok"))
581            .tool(MockAddTool)
582            .add_hook(BuilderHook)
583            .build();
584    }
585
586    struct NamedTool;
587
588    impl NamedTool {
589        fn new() -> Self {
590            Self
591        }
592    }
593
594    impl Tool for NamedTool {
595        const NAME: &'static str = "registered_named";
596        type Error = rig::tool::ToolExecutionError;
597        type Args = serde_json::Value;
598        type Output = String;
599
600        fn description(&self) -> String {
601            "uses its canonical name".to_string()
602        }
603
604        fn parameters(&self) -> serde_json::Value {
605            serde_json::json!({"type": "object", "properties": {}})
606        }
607
608        async fn call(
609            &self,
610            _context: &mut ToolContext,
611            _args: Self::Args,
612        ) -> Result<Self::Output, ToolExecutionError> {
613            Ok("ok".to_string())
614        }
615    }
616
617    #[tokio::test]
618    async fn typed_tool_builder_paths_advertise_canonical_name() {
619        for agent in [
620            AgentBuilder::new(MockCompletionModel::text("ok"))
621                .tool(NamedTool::new())
622                .build(),
623            AgentBuilder::new(MockCompletionModel::text("ok"))
624                .tool(MockAddTool)
625                .tool(NamedTool::new())
626                .build(),
627        ] {
628            let definitions = agent.tool_server_handle.get_tool_defs(None).await.unwrap();
629            assert!(
630                definitions
631                    .iter()
632                    .any(|definition| definition.name == NamedTool::NAME),
633                "the provider definitions dropped the canonical tool name"
634            );
635
636            let mut context = ToolContext::new();
637            let result = agent
638                .tool_server_handle
639                .execute(NamedTool::NAME, "{}", &mut context)
640                .await;
641            assert!(result.is_success());
642            assert_eq!(result.output().as_text(), Some("ok"));
643        }
644    }
645
646    #[tokio::test]
647    async fn retrieved_tools_are_exposed_only_for_prompted_retrieval() {
648        let retrieval_only = AgentBuilder::new(MockCompletionModel::text("ok"))
649            .retrieved_tools(
650                1,
651                MockToolIndex::new(["add"]),
652                ToolSet::from_tools(vec![MockAddTool]),
653            )
654            .build();
655        assert!(
656            retrieval_only
657                .tool_server_handle
658                .get_tool_defs(None)
659                .await
660                .unwrap()
661                .is_empty()
662        );
663
664        let agent = AgentBuilder::new(MockCompletionModel::text("ok"))
665            .tool(MockSubtractTool)
666            .retrieved_tools(
667                1,
668                MockToolIndex::new(["add"]),
669                ToolSet::from_tools(vec![MockAddTool]),
670            )
671            .build();
672
673        let always = agent.tool_server_handle.get_tool_defs(None).await.unwrap();
674        assert_eq!(
675            always
676                .iter()
677                .map(|definition| definition.name.as_str())
678                .collect::<Vec<_>>(),
679            vec!["subtract"]
680        );
681
682        let with_retrieval = agent
683            .tool_server_handle
684            .get_tool_defs(Some("add two numbers".to_string()))
685            .await
686            .unwrap();
687        assert_eq!(
688            with_retrieval
689                .iter()
690                .map(|definition| definition.name.as_str())
691                .collect::<Vec<_>>(),
692            vec!["add", "subtract"]
693        );
694    }
695
696    /// The builder's MCP path registers every requested tool against the shared
697    /// client and threads the configured timeout onto each of them, so a hanging
698    /// call is bounded instead of blocking forever. This covers the plumbing
699    /// behind `rmcp_tool[s]` / `rmcp_tool[s]_with_timeout` (see issue #1914).
700    #[cfg(all(feature = "rmcp", not(target_family = "wasm")))]
701    #[tokio::test]
702    async fn builder_rmcp_tools_thread_timeout_into_registered_tools() {
703        use crate::tool::rmcp::{DEFAULT_MCP_TOOL_TIMEOUT, McpTool as RmcpTool};
704        use crate::tool::{ToolContext, ToolErrorKind};
705        use rmcp::model::{
706            CallToolRequestParams, CallToolResult, ClientInfo, ErrorData, Implementation,
707            ProtocolVersion, ServerCapabilities, ServerInfo, Tool,
708        };
709        use rmcp::service::RequestContext;
710        use rmcp::{RoleServer, ServerHandler, ServiceExt};
711        use std::sync::Arc;
712        use std::time::Duration;
713
714        #[derive(Clone)]
715        struct HangingServer;
716        impl ServerHandler for HangingServer {
717            fn get_info(&self) -> ServerInfo {
718                ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
719                    .with_protocol_version(ProtocolVersion::LATEST)
720                    .with_server_info(Implementation::new("builder-timeout-test", "0.1.0"))
721            }
722            async fn call_tool(
723                &self,
724                _request: CallToolRequestParams,
725                _context: RequestContext<RoleServer>,
726            ) -> Result<CallToolResult, ErrorData> {
727                std::future::pending::<Result<CallToolResult, ErrorData>>().await
728            }
729        }
730
731        fn tool(name: &str) -> Tool {
732            Tool::new(
733                name.to_string(),
734                String::new(),
735                Arc::new(serde_json::Map::new()),
736            )
737        }
738
739        let (c2s, sfc) = tokio::io::duplex(8192);
740        let (s2c, cfs) = tokio::io::duplex(8192);
741        let server_task = tokio::spawn(async move {
742            let running = HangingServer.serve((sfc, s2c)).await.expect("server start");
743            running.waiting().await.expect("server error");
744        });
745        let client = ClientInfo::default()
746            .serve((cfs, c2s))
747            .await
748            .expect("client connect");
749        let peer = client.peer().clone();
750
751        // The default the plural builders pass, and a disabled timeout, both
752        // reach the built tool verbatim.
753        let built = RmcpTool::from_mcp_server(tool("a"), peer.clone());
754        assert_eq!(built.timeout(), Some(DEFAULT_MCP_TOOL_TIMEOUT));
755        assert_eq!(built.with_timeout(None).timeout(), None);
756
757        // Every requested tool is registered against the shared client...
758        let agent = AgentBuilder::new(MockCompletionModel::text("ok"))
759            .rmcp_tools(vec![tool("a"), tool("b")], peer.clone())
760            .build();
761        let definitions = agent.tool_server_handle.get_tool_defs(None).await.unwrap();
762        assert_eq!(
763            definitions
764                .iter()
765                .map(|definition| definition.name.as_str())
766                .collect::<Vec<_>>(),
767            vec!["a", "b"]
768        );
769
770        // ...and the configured timeout actually bounds a hanging call.
771        let agent = AgentBuilder::new(MockCompletionModel::text("ok"))
772            .rmcp_tools_with_timeout(vec![tool("hang_forever")], peer, Duration::from_millis(200))
773            .build();
774        let timed = tokio::time::timeout(Duration::from_secs(5), async {
775            let mut context = ToolContext::new();
776            agent
777                .tool_server_handle
778                .execute("hang_forever", "{}", &mut context)
779                .await
780        })
781        .await;
782        let result = timed.expect("registered tool hung past the safety timeout");
783        assert!(result.is_error_kind(ToolErrorKind::Timeout));
784        assert!(result.output().render().contains("timed out"));
785
786        drop(client);
787        server_task.abort();
788    }
789}
rig_agent/agent/builder.rs

rig_agent/agent/
builder.rs
