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
66#[cfg(test)]
67mod tests {
68    use super::*;
69    use serde_json::json;
70
71    /// A minimal mock tool for testing delegation.
72    struct MockInnerTool;
73
74    #[async_trait]
75    impl ToolFunction for MockInnerTool {
76        fn name(&self) -> &str {
77            "slow_operation"
78        }
79        fn description(&self) -> &str {
80            "Performs a slow operation"
81        }
82        fn parameters(&self) -> Option<serde_json::Value> {
83            Some(json!({
84                "type": "object",
85                "properties": {
86                    "task_id": { "type": "string" }
87                }
88            }))
89        }
90        async fn call(&self, args: serde_json::Value) -> Result<serde_json::Value, ToolError> {
91            let task_id = args
92                .get("task_id")
93                .and_then(|v| v.as_str())
94                .unwrap_or("unknown");
95            Ok(json!({ "status": "completed", "task_id": task_id }))
96        }
97    }
98
99    #[test]
100    fn description_is_augmented() {
101        let inner = Arc::new(MockInnerTool);
102        let tool = LongRunningFunctionTool::new(inner);
103
104        let desc = tool.description();
105        assert!(
106            desc.starts_with("Performs a slow operation"),
107            "should start with the inner description, got: {desc}"
108        );
109        assert!(
110            desc.contains(LONG_RUNNING_INSTRUCTION),
111            "should contain the long-running instruction, got: {desc}"
112        );
113        assert!(
114            desc.contains('\n'),
115            "inner description and instruction should be separated by a newline"
116        );
117    }
118
119    #[test]
120    fn name_delegates_to_inner() {
121        let inner = Arc::new(MockInnerTool);
122        let tool = LongRunningFunctionTool::new(inner);
123        assert_eq!(tool.name(), "slow_operation");
124    }
125
126    #[test]
127    fn parameters_delegates_to_inner() {
128        let inner = Arc::new(MockInnerTool);
129        let tool = LongRunningFunctionTool::new(inner);
130
131        let params = tool.parameters().expect("should have parameters");
132        assert!(params["properties"]["task_id"].is_object());
133    }
134
135    #[tokio::test]
136    async fn call_delegates_to_inner() {
137        let inner = Arc::new(MockInnerTool);
138        let tool = LongRunningFunctionTool::new(inner);
139
140        let result = tool
141            .call(json!({ "task_id": "abc-123" }))
142            .await
143            .expect("call should succeed");
144        assert_eq!(result["status"], "completed");
145        assert_eq!(result["task_id"], "abc-123");
146    }
147
148    #[test]
149    fn is_long_running_returns_true() {
150        let inner = Arc::new(MockInnerTool);
151        let tool = LongRunningFunctionTool::new(inner);
152        assert!(tool.is_long_running());
153    }
154
155    #[test]
156    fn description_format_is_correct() {
157        let inner = Arc::new(MockInnerTool);
158        let tool = LongRunningFunctionTool::new(inner);
159
160        let expected = format!("Performs a slow operation\n{}", LONG_RUNNING_INSTRUCTION);
161        assert_eq!(tool.description(), expected);
162    }
163}
gemini_adk_rs/tools/long_running.rs

gemini_adk_rs/tools/
long_running.rs
