rig_core/agent/
tool.rs

1use crate::{
2    agent::Agent,
3    completion::{CompletionModel, Prompt, PromptError},
4    tool::{Tool, ToolCallExtensions},
5};
6use schemars::{JsonSchema, schema_for};
7use serde::{Deserialize, Serialize};
8use serde_json::json;
9
10#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
11pub struct AgentToolArgs {
12    /// The prompt for the agent to call.
13    prompt: String,
14}
15
16impl<M: CompletionModel + 'static> Tool for Agent<M> {
17    const NAME: &'static str = "agent_tool";
18
19    type Error = PromptError;
20    type Args = AgentToolArgs;
21    type Output = String;
22
23    fn description(&self) -> String {
24        format!(
25            "
26            Prompt a sub-agent to do a task for you.
27
28            Agent name: {name}
29            Agent description: {description}
30            Agent system prompt: {sysprompt}
31            ",
32            name = self.name(),
33            description = self.description.clone().unwrap_or_default(),
34            sysprompt = self.preamble.clone().unwrap_or_default()
35        )
36    }
37
38    fn parameters(&self) -> serde_json::Value {
39        json!(schema_for!(AgentToolArgs))
40    }
41
42    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
43        self.prompt(args.prompt).await
44    }
45
46    /// Propagate the caller's [`ToolCallExtensions`] into the sub-agent run, so the
47    /// inner agent's own tools observe them too (sub-agent delegation / A2A
48    /// chains). Without this, a sub-agent invoked as a tool would start with
49    /// empty extensions.
50    async fn call_with_extensions(
51        &self,
52        args: Self::Args,
53        extensions: &ToolCallExtensions,
54    ) -> Result<Self::Output, Self::Error> {
55        self.prompt(args.prompt)
56            .tool_extensions(extensions.clone())
57            .await
58    }
59
60    fn name(&self) -> String {
61        self.name.clone().unwrap_or_else(|| Self::NAME.to_string())
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68    use crate::agent::AgentBuilder;
69    use crate::test_utils::{MockCompletionModel, MockExtensionsProbeTool, MockTurn, SessionId};
70
71    /// A `ToolCallExtensions` set on the outer run propagates into a sub-agent
72    /// invoked as a tool, so the inner agent's own tools observe it.
73    #[tokio::test]
74    async fn context_propagates_into_sub_agent() {
75        // Inner agent: calls a context-probing tool, then answers.
76        let probe = MockExtensionsProbeTool::default();
77        let inner_model = MockCompletionModel::new([
78            MockTurn::tool_call("c1", "context_probe", json!({})),
79            MockTurn::text("inner done"),
80        ]);
81        let inner = AgentBuilder::new(inner_model)
82            .name("researcher")
83            .tool(probe.clone())
84            .build();
85
86        // Outer agent: delegates to the inner agent (registered as the
87        // "researcher" tool), then answers.
88        let outer_model = MockCompletionModel::new([
89            MockTurn::tool_call("c2", "researcher", json!({"prompt": "do research"})),
90            MockTurn::text("outer done"),
91        ]);
92        let outer = AgentBuilder::new(outer_model).tool(inner).build();
93
94        let mut extensions = ToolCallExtensions::new();
95        extensions.insert(SessionId("abc-123".to_string()));
96
97        let out = outer
98            .prompt("start")
99            .tool_extensions(extensions)
100            .max_turns(5)
101            .await
102            .expect("run succeeds");
103
104        assert_eq!(out, "outer done");
105        assert_eq!(probe.observed().as_deref(), Some("session:abc-123"));
106    }
107}
rig_core/agent/tool.rs

rig_core/agent/
tool.rs
