gemini_adk_rs/tools/
long_running.rs

1//! Long-running function tool wrapper.
2//!
3//! Wraps any [`ToolFunction`] and marks it as long-running by appending an
4//! instruction to the tool description that tells the LLM not to re-invoke
5//! the tool while it is still pending.
6
7use std::sync::Arc;
8
9use async_trait::async_trait;
10
11use crate::error::ToolError;
12use crate::tool::ToolFunction;
13
14/// Instruction appended to the tool description for long-running tools.
15const LONG_RUNNING_INSTRUCTION: &str = "NOTE: This is a long-running operation. \
16    Do not call this tool again if it has already returned some intermediate or pending status.";
17
18/// Wraps a [`ToolFunction`] and marks it as long-running.
19///
20/// The wrapper appends the long-running instruction to the inner tool's
21/// description so the LLM knows not to re-invoke it while a previous call
22/// is still in progress. All other trait methods delegate directly to the
23/// inner tool.
24pub struct LongRunningFunctionTool {
25    inner: Arc<dyn ToolFunction>,
26    /// Cached description: inner description + "\n" + LONG_RUNNING_INSTRUCTION.
27    augmented_description: String,
28}
29
30impl LongRunningFunctionTool {
31    /// Create a new `LongRunningFunctionTool` wrapping the given inner tool.
32    pub fn new(inner: Arc<dyn ToolFunction>) -> Self {
33        let augmented_description =
34            format!("{}\n{}", inner.description(), LONG_RUNNING_INSTRUCTION);
35        Self {
36            inner,
37            augmented_description,
38        }
39    }
40
41    /// Returns `true` — this tool is always considered long-running.
42    pub fn is_long_running(&self) -> bool {
43        true
44    }
45}
46
47#[async_trait]
48impl ToolFunction for LongRunningFunctionTool {
49    fn name(&self) -> &str {
50        self.inner.name()
51    }
52
53    fn description(&self) -> &str {
54        &self.augmented_description
55    }
56
57    fn parameters(&self) -> Option<serde_json::Value> {
58        self.inner.parameters()
59    }
60
61    async fn call(&self, args: serde_json::Value) -> Result<serde_json::Value, ToolError> {
62        self.inner.call(args).await
63    }
64
65    async fn call_with_context(
66        &self,
67        args: serde_json::Value,
68        ctx: crate::tool::ToolContext,
69    ) -> Result<serde_json::Value, ToolError> {
70        self.inner.call_with_context(args, ctx).await
71    }
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77    use serde_json::json;
78
79    /// A minimal mock tool for testing delegation.
80    struct MockInnerTool;
81
82    #[async_trait]
83    impl ToolFunction for MockInnerTool {
84        fn name(&self) -> &str {
85            "slow_operation"
86        }
87        fn description(&self) -> &str {
88            "Performs a slow operation"
89        }
90        fn parameters(&self) -> Option<serde_json::Value> {
91            Some(json!({
92                "type": "object",
93                "properties": {
94                    "task_id": { "type": "string" }
95                }
96            }))
97        }
98        async fn call(&self, args: serde_json::Value) -> Result<serde_json::Value, ToolError> {
99            let task_id = args
100                .get("task_id")
101                .and_then(|v| v.as_str())
102                .unwrap_or("unknown");
103            Ok(json!({ "status": "completed", "task_id": task_id }))
104        }
105    }
106
107    #[test]
108    fn description_is_augmented() {
109        let inner = Arc::new(MockInnerTool);
110        let tool = LongRunningFunctionTool::new(inner);
111
112        let desc = tool.description();
113        assert!(
114            desc.starts_with("Performs a slow operation"),
115            "should start with the inner description, got: {desc}"
116        );
117        assert!(
118            desc.contains(LONG_RUNNING_INSTRUCTION),
119            "should contain the long-running instruction, got: {desc}"
120        );
121        assert!(
122            desc.contains('\n'),
123            "inner description and instruction should be separated by a newline"
124        );
125    }
126
127    #[test]
128    fn name_delegates_to_inner() {
129        let inner = Arc::new(MockInnerTool);
130        let tool = LongRunningFunctionTool::new(inner);
131        assert_eq!(tool.name(), "slow_operation");
132    }
133
134    #[test]
135    fn parameters_delegates_to_inner() {
136        let inner = Arc::new(MockInnerTool);
137        let tool = LongRunningFunctionTool::new(inner);
138
139        let params = tool.parameters().expect("should have parameters");
140        assert!(params["properties"]["task_id"].is_object());
141    }
142
143    #[tokio::test]
144    async fn call_delegates_to_inner() {
145        let inner = Arc::new(MockInnerTool);
146        let tool = LongRunningFunctionTool::new(inner);
147
148        let result = tool
149            .call(json!({ "task_id": "abc-123" }))
150            .await
151            .expect("call should succeed");
152        assert_eq!(result["status"], "completed");
153        assert_eq!(result["task_id"], "abc-123");
154    }
155
156    #[test]
157    fn is_long_running_returns_true() {
158        let inner = Arc::new(MockInnerTool);
159        let tool = LongRunningFunctionTool::new(inner);
160        assert!(tool.is_long_running());
161    }
162
163    #[test]
164    fn description_format_is_correct() {
165        let inner = Arc::new(MockInnerTool);
166        let tool = LongRunningFunctionTool::new(inner);
167
168        let expected = format!("Performs a slow operation\n{LONG_RUNNING_INSTRUCTION}");
169        assert_eq!(tool.description(), expected);
170    }
171}
gemini_adk_rs/tools/long_running.rs

gemini_adk_rs/tools/
long_running.rs
