bashkit_scripted_tool/scripted_tool/
toolset.rs

1// ScriptingToolSet: higher-level wrapper around ScriptedTool that controls
2// how tools are exposed based on DiscoveryMode.
3//
4// - Exclusive (default): returns one ScriptedTool with full schemas in prompt.
5// - WithDiscovery: returns two tools — ScriptedTool (compact prompt) +
6//   DiscoverTool (discover/help only).
7
8use super::{
9    CallbackKind, RegisteredTool, ScriptedExecutionTrace, ScriptedTool, ToolArgs, ToolDef, ToolImpl,
10};
11use async_trait::async_trait;
12use bashkit::tool::{
13    Tool, ToolError, ToolRequest, ToolResponse, ToolStatus, VERSION, tool_response_schema,
14};
15use bashkit::{ExecutionLimits, ExecutionProfile};
16use std::sync::Arc;
17
18// ============================================================================
19// DiscoveryMode
20// ============================================================================
21
22/// Controls how [`ScriptingToolSet::tools`] exposes tool information to the LLM.
23#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
24pub enum DiscoveryMode {
25    /// Returns one tool with full schemas in system_prompt.
26    #[default]
27    Exclusive,
28    /// Returns two tools: a compact ScriptedTool + a DiscoverTool for
29    /// runtime schema discovery via `discover` and `help` builtins.
30    WithDiscovery,
31}
32
33// ============================================================================
34// DiscoverTool — discovery-only tool for WithDiscovery mode
35// ============================================================================
36
37/// A [`Tool`] that exposes only `discover` and `help` builtins for runtime
38/// schema discovery. Returned by [`ScriptingToolSet::tools`] in
39/// [`DiscoveryMode::WithDiscovery`] mode alongside the main script tool.
40///
41/// The LLM uses this tool to explore available commands before writing
42/// scripts for the main tool.
43///
44/// ```rust
45/// use bashkit::Tool;
46/// use bashkit_scripted_tool::{ScriptingToolSet, ToolArgs, ToolDef};
47///
48/// # tokio_test::block_on(async {
49/// let toolset = ScriptingToolSet::builder("api")
50///     .tool_fn(
51///         ToolDef::new("greet", "Greet someone").with_category("social"),
52///         |_args: &ToolArgs| Ok("hello\n".to_string()),
53///     )
54///     .with_discovery()
55///     .build();
56///
57/// let tools = toolset.tools();
58/// assert_eq!(tools.len(), 2);
59///
60/// // Second tool is the DiscoverTool
61/// let discover = &tools[1];
62/// assert!(discover.name().ends_with("_discover"));
63/// # });
64/// ```
65pub struct DiscoverTool {
66    name: String,
67    locale: String,
68    display_name: String,
69    short_desc: String,
70    inner: ScriptedTool,
71}
72
73const DISCOVER_ALLOWED_COMMANDS: &[&str] = &["discover", "help"];
74
75impl DiscoverTool {
76    fn has_forbidden_shell_syntax(input: &str) -> bool {
77        input.contains('\n')
78            || input.contains('\r')
79            || input.contains(';')
80            || input.contains('&')
81            || input.contains('|')
82            || input.contains('`')
83            || input.contains("$(")
84    }
85
86    /// Reject commands that aren't `discover` or `help`.
87    fn validate_commands(commands: &str) -> Result<(), String> {
88        if Self::has_forbidden_shell_syntax(commands) {
89            return Err(
90                "discover tool commands cannot contain shell control characters".to_string(),
91            );
92        }
93        let first_word = commands.split_whitespace().next().unwrap_or("");
94        if DISCOVER_ALLOWED_COMMANDS.contains(&first_word) {
95            Ok(())
96        } else {
97            Err("discover tool only supports: discover, help".to_string())
98        }
99    }
100
101    fn reject_response(msg: &str) -> ToolResponse {
102        ToolResponse {
103            stdout: String::new(),
104            stderr: msg.to_string(),
105            exit_code: 1,
106            error: Some(msg.to_string()),
107            ..Default::default()
108        }
109    }
110
111    /// Translate structured tool input into a [`ToolRequest`].
112    ///
113    /// Priority: `commands` (backward compat) > `all` > `query`.
114    fn resolve_request(args: &serde_json::Value) -> Result<ToolRequest, String> {
115        let obj = args.as_object().ok_or("arguments must be a JSON object")?;
116
117        let timeout_ms = obj.get("timeout_ms").and_then(|v| v.as_u64());
118
119        // Raw commands pass-through (backward compatible)
120        if let Some(commands) = obj.get("commands").and_then(|v| v.as_str()) {
121            return Ok(ToolRequest {
122                commands: commands.to_string(),
123                timeout_ms,
124            });
125        }
126
127        // Structured: all
128        if obj.get("all").and_then(|v| v.as_bool()).unwrap_or(false) {
129            return Ok(ToolRequest {
130                commands: "discover --categories".to_string(),
131                timeout_ms,
132            });
133        }
134
135        // Structured: query
136        if let Some(query) = obj.get("query").and_then(|v| v.as_str()) {
137            if Self::has_forbidden_shell_syntax(query) {
138                return Err("query contains unsupported shell control characters".to_string());
139            }
140            return Ok(ToolRequest {
141                commands: format!("discover --search {query}"),
142                timeout_ms,
143            });
144        }
145
146        Err("one of `commands`, `all`, or `query` is required".to_string())
147    }
148}
149
150#[async_trait]
151impl Tool for DiscoverTool {
152    fn name(&self) -> &str {
153        &self.name
154    }
155
156    fn display_name(&self) -> &str {
157        &self.display_name
158    }
159
160    fn short_description(&self) -> &str {
161        &self.short_desc
162    }
163
164    fn description(&self) -> &str {
165        &self.short_desc
166    }
167
168    fn help(&self) -> String {
169        format!(
170            "# {}\n\nDiscover available tool commands.\n\n## Commands\n\n- `discover --categories` — list categories\n- `discover --category <name>` — list tools in category\n- `discover --tag <tag>` — filter by tag\n- `discover --search <keyword>` — search by name/description\n- `help --list` — list all tools\n- `help <tool>` — human-readable usage\n- `help <tool> --json` — machine-readable schema\n\nAll commands support `--json` for structured output.\n",
171            self.display_name
172        )
173    }
174
175    fn system_prompt(&self) -> String {
176        format!(
177            "{}: discover available tool commands. Use `discover --categories`, `discover --search <keyword>`, `help <tool>`, `help <tool> --json`.",
178            self.name
179        )
180    }
181
182    fn locale(&self) -> &str {
183        &self.locale
184    }
185
186    fn input_schema(&self) -> serde_json::Value {
187        serde_json::json!({
188            "type": "object",
189            "properties": {
190                "query": {
191                    "type": "string",
192                    "description": "Search query to find tools (matched against names, descriptions, and categories)"
193                },
194                "all": {
195                    "type": "boolean",
196                    "description": "List all available tools grouped by category. When true, query is ignored."
197                },
198                "commands": {
199                    "type": "string",
200                    "description": "Raw bash commands (backward compatible). Takes precedence over query/all."
201                },
202                "timeout_ms": {
203                    "type": "integer",
204                    "description": "Optional timeout in milliseconds"
205                }
206            }
207        })
208    }
209
210    fn output_schema(&self) -> serde_json::Value {
211        tool_response_schema()
212    }
213
214    fn version(&self) -> &str {
215        VERSION
216    }
217
218    fn execution(
219        &self,
220        args: serde_json::Value,
221    ) -> Result<bashkit::tool::ToolExecution, bashkit::tool::ToolError> {
222        let req = Self::resolve_request(&args).map_err(ToolError::UserFacing)?;
223        if let Err(msg) = Self::validate_commands(&req.commands) {
224            return Err(ToolError::UserFacing(msg));
225        }
226        self.inner
227            .execution(serde_json::to_value(req).unwrap_or_default())
228    }
229
230    async fn execute(&self, req: ToolRequest) -> ToolResponse {
231        if let Err(msg) = Self::validate_commands(&req.commands) {
232            return Self::reject_response(&msg);
233        }
234        self.inner.execute(req).await
235    }
236
237    async fn execute_with_status(
238        &self,
239        req: ToolRequest,
240        status_callback: Box<dyn FnMut(ToolStatus) + Send>,
241    ) -> ToolResponse {
242        if let Err(msg) = Self::validate_commands(&req.commands) {
243            return Self::reject_response(&msg);
244        }
245        self.inner.execute_with_status(req, status_callback).await
246    }
247}
248
249// ============================================================================
250// ScriptingToolSetBuilder
251// ============================================================================
252
253/// Builder for [`ScriptingToolSet`].
254///
255/// ```rust
256/// use bashkit_scripted_tool::{ScriptingToolSet, ToolArgs, ToolDef};
257///
258/// let toolset = ScriptingToolSet::builder("api")
259///     .short_description("Example API")
260///     .tool_fn(
261///         ToolDef::new("ping", "Ping a host"),
262///         |_args: &ToolArgs| Ok("pong\n".to_string()),
263///     )
264///     .build();
265/// ```
266pub struct ScriptingToolSetBuilder {
267    name: String,
268    locale: String,
269    short_desc: Option<String>,
270    tools: Vec<RegisteredTool>,
271    limits: Option<ExecutionLimits>,
272    profile: Option<ExecutionProfile>,
273    env_vars: Vec<(String, String)>,
274    mode: DiscoveryMode,
275}
276
277impl ScriptingToolSetBuilder {
278    fn new(name: impl Into<String>) -> Self {
279        Self {
280            name: name.into(),
281            locale: "en-US".to_string(),
282            short_desc: None,
283            tools: Vec::new(),
284            limits: None,
285            profile: None,
286            env_vars: Vec::new(),
287            mode: DiscoveryMode::Exclusive,
288        }
289    }
290
291    /// Set locale for descriptions, help, prompts, and user-facing errors.
292    pub fn locale(mut self, locale: &str) -> Self {
293        self.locale = locale.to_string();
294        self
295    }
296
297    /// One-line description for tool listings.
298    pub fn short_description(mut self, desc: impl Into<String>) -> Self {
299        self.short_desc = Some(desc.into());
300        self
301    }
302
303    /// Register a [`ToolImpl`] (definition + exec functions).
304    pub fn tool(mut self, tool: ToolImpl) -> Self {
305        self.tools.push(RegisteredTool::from_tool_impl(tool));
306        self
307    }
308
309    /// Register a tool with its definition and synchronous exec function.
310    ///
311    /// Convenience shorthand — constructs a [`ToolImpl`] internally.
312    pub fn tool_fn(
313        mut self,
314        def: ToolDef,
315        exec: impl Fn(&ToolArgs) -> Result<String, String> + Send + Sync + 'static,
316    ) -> Self {
317        self.tools.push(RegisteredTool {
318            def,
319            callback: CallbackKind::Sync(Arc::new(exec)),
320            dry_run: None,
321        });
322        self
323    }
324
325    /// Register a tool with its definition and **async** exec function.
326    ///
327    /// Convenience shorthand — constructs a [`ToolImpl`] internally.
328    pub fn async_tool_fn<F, Fut>(mut self, def: ToolDef, exec: F) -> Self
329    where
330        F: Fn(ToolArgs) -> Fut + Send + Sync + 'static,
331        Fut: std::future::Future<Output = Result<String, String>> + Send + 'static,
332    {
333        let cb: super::AsyncToolExec = Arc::new(move |args| Box::pin(exec(args)));
334        self.tools.push(RegisteredTool {
335            def,
336            callback: CallbackKind::Async(cb),
337            dry_run: None,
338        });
339        self
340    }
341
342    /// Set execution limits for the bash interpreter.
343    pub fn limits(mut self, limits: ExecutionLimits) -> Self {
344        self.limits = Some(limits);
345        self
346    }
347
348    /// Apply a coherent execution profile to the inner scripted tool.
349    pub fn profile(mut self, profile: ExecutionProfile) -> Self {
350        self.profile = Some(profile);
351        self
352    }
353
354    /// Add an environment variable visible inside scripts.
355    pub fn env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
356        self.env_vars.push((key.into(), value.into()));
357        self
358    }
359
360    /// Switch to discovery mode: returns two tools from [`ScriptingToolSet::tools`].
361    pub fn with_discovery(mut self) -> Self {
362        self.mode = DiscoveryMode::WithDiscovery;
363        self
364    }
365
366    /// Build the [`ScriptingToolSet`].
367    pub fn build(&self) -> ScriptingToolSet {
368        let short_desc = self
369            .short_desc
370            .clone()
371            .unwrap_or_else(|| format!("ScriptingToolSet: {}", self.name));
372
373        // Inner ScriptedTool uses compact_prompt in discovery mode
374        let compact = self.mode == DiscoveryMode::WithDiscovery;
375
376        let mut builder = ScriptedTool::builder(&self.name).locale(&self.locale);
377        builder = builder.short_description(&short_desc);
378        builder = builder.compact_prompt(compact);
379
380        if let Some(profile) = &self.profile {
381            builder = builder.profile(profile.clone());
382        }
383
384        if let Some(limits) = &self.limits {
385            builder = builder.limits(limits.clone());
386        }
387        for (key, value) in &self.env_vars {
388            builder = builder.env(key, value);
389        }
390
391        // Move tools into inner ScriptedTool
392        // We need to reconstruct because ScriptedToolBuilder expects closures
393        for reg in &self.tools {
394            match &reg.callback {
395                CallbackKind::Sync(cb) => {
396                    let cb = Arc::clone(cb);
397                    builder = builder.tool_fn(reg.def.clone(), move |args: &ToolArgs| (cb)(args));
398                }
399                CallbackKind::Async(cb) => {
400                    let cb = Arc::clone(cb);
401                    builder = builder.async_tool_fn(reg.def.clone(), move |args: ToolArgs| {
402                        let cb = Arc::clone(&cb);
403                        async move { (cb)(args).await }
404                    });
405                }
406            }
407        }
408
409        ScriptingToolSet {
410            name: self.name.clone(),
411            locale: self.locale.clone(),
412            inner: builder.build(),
413            mode: self.mode,
414        }
415    }
416}
417
418// ============================================================================
419// ScriptingToolSet
420// ============================================================================
421
422/// Higher-level wrapper around [`ScriptedTool`] with mode-controlled tool exposure.
423///
424/// Use [`ScriptingToolSet::tools`] to get the tools to register with your LLM:
425///
426/// - **Exclusive** (default): Returns one tool — a [`ScriptedTool`] with full
427///   schemas in its system prompt.
428///
429/// - **WithDiscovery**: Returns two tools — a [`ScriptedTool`] with compact
430///   prompt (no schemas) + a [`DiscoverTool`] for runtime schema discovery
431///   via `discover` and `help` builtins.
432///
433/// ```rust
434/// use bashkit::{Tool, ToolRequest};
435/// use bashkit_scripted_tool::{ScriptingToolSet, ToolArgs, ToolDef};
436///
437/// # tokio_test::block_on(async {
438/// // Exclusive mode (default): one tool with full schemas
439/// let toolset = ScriptingToolSet::builder("api")
440///     .tool_fn(
441///         ToolDef::new("greet", "Greet someone")
442///             .with_schema(serde_json::json!({
443///                 "type": "object",
444///                 "properties": { "name": {"type": "string"} }
445///             })),
446///         |args: &ToolArgs| Ok(format!("hello {}\n", args.param_str("name").unwrap_or("world"))),
447///     )
448///     .build();
449///
450/// let tools = toolset.tools();
451/// assert_eq!(tools.len(), 1);
452/// assert!(tools[0].system_prompt().contains("--name <string>"));
453///
454/// let resp = tools[0].execute(ToolRequest {
455///     commands: "greet --name Alice".into(),
456///     timeout_ms: None,
457/// }).await;
458/// assert_eq!(resp.stdout.trim(), "hello Alice");
459/// # });
460/// ```
461///
462/// ```rust
463/// use bashkit::Tool;
464/// use bashkit_scripted_tool::{ScriptingToolSet, ToolArgs, ToolDef};
465///
466/// // Discovery mode: two tools
467/// let toolset = ScriptingToolSet::builder("api")
468///     .tool_fn(
469///         ToolDef::new("greet", "Greet someone"),
470///         |_args: &ToolArgs| Ok("hello\n".to_string()),
471///     )
472///     .with_discovery()
473///     .build();
474///
475/// let tools = toolset.tools();
476/// assert_eq!(tools.len(), 2);
477/// assert_eq!(tools[0].name(), "api");
478/// assert_eq!(tools[1].name(), "api_discover");
479/// ```
480pub struct ScriptingToolSet {
481    name: String,
482    locale: String,
483    inner: ScriptedTool,
484    mode: DiscoveryMode,
485}
486
487impl ScriptingToolSet {
488    /// Create a builder with the given tool name.
489    pub fn builder(name: impl Into<String>) -> ScriptingToolSetBuilder {
490        ScriptingToolSetBuilder::new(name)
491    }
492
493    /// Current discovery mode.
494    pub fn discovery_mode(&self) -> DiscoveryMode {
495        self.mode
496    }
497
498    /// Return and clear the trace from the most recent execute call.
499    pub fn take_last_execution_trace(&self) -> Option<ScriptedExecutionTrace> {
500        self.inner.take_last_execution_trace()
501    }
502
503    /// Return the tools to register with the LLM.
504    ///
505    /// - **Exclusive**: `[ScriptedTool]` — one tool with full schemas in prompt.
506    /// - **WithDiscovery**: `[ScriptedTool, DiscoverTool]` — script tool with
507    ///   compact prompt + discover tool for runtime schema exploration.
508    pub fn tools(&self) -> Vec<Box<dyn Tool>> {
509        match self.mode {
510            DiscoveryMode::Exclusive => {
511                vec![Box::new(self.inner.clone())]
512            }
513            DiscoveryMode::WithDiscovery => {
514                let discover = DiscoverTool {
515                    name: format!("{}_discover", self.name),
516                    locale: self.locale.clone(),
517                    display_name: format!("{} Discover", self.name),
518                    short_desc: format!("Discover available {} commands", self.name),
519                    inner: self.inner.clone(),
520                };
521                vec![Box::new(self.inner.clone()), Box::new(discover)]
522            }
523        }
524    }
525}
526
527// ============================================================================
528// Tests
529// ============================================================================
530
531#[cfg(test)]
532mod tests {
533    use super::*;
534    use bashkit::tool::ToolRequest;
535
536    fn make_tools() -> ScriptingToolSetBuilder {
537        ScriptingToolSet::builder("test_api")
538            .short_description("Test API")
539            .tool_fn(
540                ToolDef::new("get_user", "Fetch user by ID")
541                    .with_schema(serde_json::json!({
542                        "type": "object",
543                        "properties": {
544                            "id": {"type": "integer"}
545                        },
546                        "required": ["id"]
547                    }))
548                    .with_category("users"),
549                |args: &ToolArgs| {
550                    let id = args.param_i64("id").ok_or("missing --id")?;
551                    Ok(format!("{{\"id\":{id},\"name\":\"Alice\"}}\n"))
552                },
553            )
554            .tool_fn(
555                ToolDef::new("list_orders", "List orders for a user")
556                    .with_schema(serde_json::json!({
557                        "type": "object",
558                        "properties": {
559                            "user_id": {"type": "integer"}
560                        }
561                    }))
562                    .with_category("orders"),
563                |args: &ToolArgs| {
564                    let uid = args.param_i64("user_id").ok_or("missing --user_id")?;
565                    Ok(format!("[{{\"order_id\":1,\"user_id\":{uid}}}]\n"))
566                },
567            )
568    }
569
570    // -- Mode defaults --
571
572    #[test]
573    fn test_builder_defaults_to_exclusive() {
574        let toolset = make_tools().build();
575        assert_eq!(toolset.discovery_mode(), DiscoveryMode::Exclusive);
576    }
577
578    #[test]
579    fn test_with_discovery_switches_mode() {
580        let toolset = make_tools().with_discovery().build();
581        assert_eq!(toolset.discovery_mode(), DiscoveryMode::WithDiscovery);
582    }
583
584    // -- Exclusive mode: tools() returns 1 --
585
586    #[test]
587    fn test_exclusive_returns_one_tool() {
588        let toolset = make_tools().build();
589        let tools = toolset.tools();
590        assert_eq!(tools.len(), 1);
591        assert_eq!(tools[0].name(), "test_api");
592    }
593
594    #[test]
595    fn test_exclusive_tool_has_full_schemas() {
596        let toolset = make_tools().build();
597        let tools = toolset.tools();
598        let sp = tools[0].system_prompt();
599        assert!(sp.contains("get_user [--id <integer>]"), "prompt: {sp}");
600        assert!(
601            sp.contains("list_orders [--user_id <integer>]"),
602            "prompt: {sp}"
603        );
604    }
605
606    #[test]
607    fn test_exclusive_tool_no_discover_instructions() {
608        let toolset = make_tools().build();
609        let tools = toolset.tools();
610        let sp = tools[0].system_prompt();
611        assert!(!sp.contains("discover --categories"), "prompt: {sp}");
612    }
613
614    // -- Discovery mode: tools() returns 2 --
615
616    #[test]
617    fn test_discovery_returns_two_tools() {
618        let toolset = make_tools().with_discovery().build();
619        let tools = toolset.tools();
620        assert_eq!(tools.len(), 2);
621        assert_eq!(tools[0].name(), "test_api");
622        assert_eq!(tools[1].name(), "test_api_discover");
623    }
624
625    #[test]
626    fn test_discovery_script_tool_compact_prompt() {
627        let toolset = make_tools().with_discovery().build();
628        let tools = toolset.tools();
629        let sp = tools[0].system_prompt();
630        // Compact: no schema flags
631        assert!(!sp.contains("--id <integer>"), "prompt: {sp}");
632        assert!(!sp.contains("--user_id <integer>"), "prompt: {sp}");
633    }
634
635    #[test]
636    fn test_discovery_discover_tool_prompt() {
637        let toolset = make_tools().with_discovery().build();
638        let tools = toolset.tools();
639        let sp = tools[1].system_prompt();
640        assert!(sp.contains("discover"), "prompt: {sp}");
641        assert!(sp.contains("help"), "prompt: {sp}");
642    }
643
644    // -- Name / description --
645
646    #[test]
647    fn test_name_and_short_description() {
648        let toolset = make_tools().build();
649        let tools = toolset.tools();
650        assert_eq!(tools[0].name(), "test_api");
651        assert_eq!(tools[0].short_description(), "Test API");
652    }
653
654    #[test]
655    fn test_default_short_description() {
656        let toolset = ScriptingToolSet::builder("mytools")
657            .tool_fn(ToolDef::new("noop", "No-op"), |_: &ToolArgs| {
658                Ok("ok\n".into())
659            })
660            .build();
661        let tools = toolset.tools();
662        assert_eq!(tools[0].short_description(), "ScriptingToolSet: mytools");
663    }
664
665    // -- Execution via tools() --
666
667    #[tokio::test]
668    async fn test_execute_via_exclusive_tool() {
669        let toolset = make_tools().build();
670        let tools = toolset.tools();
671        let resp = tools[0]
672            .execute(ToolRequest {
673                commands: "get_user --id 42 | jq -r '.name'".into(),
674                timeout_ms: None,
675            })
676            .await;
677        assert_eq!(resp.exit_code, 0);
678        assert_eq!(resp.stdout.trim(), "Alice");
679    }
680
681    #[tokio::test]
682    async fn test_execute_via_discovery_script_tool() {
683        let toolset = make_tools().with_discovery().build();
684        let tools = toolset.tools();
685        let resp = tools[0]
686            .execute(ToolRequest {
687                commands: "get_user --id 1".into(),
688                timeout_ms: None,
689            })
690            .await;
691        assert_eq!(resp.exit_code, 0);
692        assert!(resp.stdout.contains("Alice"));
693    }
694
695    #[tokio::test]
696    async fn test_execute_via_discover_tool() {
697        let toolset = make_tools().with_discovery().build();
698        let tools = toolset.tools();
699        let resp = tools[1]
700            .execute(ToolRequest {
701                commands: "discover --categories".into(),
702                timeout_ms: None,
703            })
704            .await;
705        assert_eq!(resp.exit_code, 0);
706        assert!(resp.stdout.contains("users"));
707        assert!(resp.stdout.contains("orders"));
708    }
709
710    #[tokio::test]
711    async fn test_discover_tool_help_builtin() {
712        let toolset = make_tools().with_discovery().build();
713        let tools = toolset.tools();
714        let resp = tools[1]
715            .execute(ToolRequest {
716                commands: "help get_user".into(),
717                timeout_ms: None,
718            })
719            .await;
720        assert_eq!(resp.exit_code, 0);
721        assert!(resp.stdout.contains("get_user"));
722        assert!(resp.stdout.contains("--id"));
723    }
724
725    #[tokio::test]
726    async fn test_execute_with_status_via_tool() {
727        use std::sync::{Arc, Mutex};
728
729        let toolset = make_tools().build();
730        let tools = toolset.tools();
731        let phases = Arc::new(Mutex::new(Vec::new()));
732        let phases_clone = phases.clone();
733
734        let resp = tools[0]
735            .execute_with_status(
736                ToolRequest {
737                    commands: "get_user --id 1".into(),
738                    timeout_ms: None,
739                },
740                Box::new(move |status| {
741                    phases_clone
742                        .lock()
743                        .expect("lock poisoned")
744                        .push(status.phase.clone());
745                }),
746            )
747            .await;
748
749        assert_eq!(resp.exit_code, 0);
750        let phases = phases.lock().expect("lock poisoned");
751        assert!(phases.contains(&"validate".to_string()));
752        assert!(phases.contains(&"complete".to_string()));
753    }
754
755    // -- help/discover builtins work in both modes --
756
757    #[tokio::test]
758    async fn test_help_builtin_works_in_exclusive() {
759        let toolset = make_tools().build();
760        let tools = toolset.tools();
761        let resp = tools[0]
762            .execute(ToolRequest {
763                commands: "help get_user".into(),
764                timeout_ms: None,
765            })
766            .await;
767        assert_eq!(resp.exit_code, 0);
768        assert!(resp.stdout.contains("get_user"));
769        assert!(resp.stdout.contains("--id"));
770    }
771
772    #[tokio::test]
773    async fn test_discover_builtin_works_in_exclusive() {
774        let toolset = make_tools().build();
775        let tools = toolset.tools();
776        let resp = tools[0]
777            .execute(ToolRequest {
778                commands: "discover --categories".into(),
779                timeout_ms: None,
780            })
781            .await;
782        assert_eq!(resp.exit_code, 0);
783        assert!(resp.stdout.contains("users"));
784        assert!(resp.stdout.contains("orders"));
785    }
786
787    // -- env vars --
788
789    #[tokio::test]
790    async fn test_env_vars_passed_through() {
791        let toolset = ScriptingToolSet::builder("env_test")
792            .env("MY_VAR", "hello")
793            .tool_fn(ToolDef::new("noop", "No-op"), |_: &ToolArgs| {
794                Ok("ok\n".into())
795            })
796            .build();
797
798        let tools = toolset.tools();
799        let resp = tools[0]
800            .execute(ToolRequest {
801                commands: "echo $MY_VAR".into(),
802                timeout_ms: None,
803            })
804            .await;
805        assert_eq!(resp.exit_code, 0);
806        assert_eq!(resp.stdout.trim(), "hello");
807    }
808
809    // -- version / schemas --
810
811    #[test]
812    fn test_version() {
813        let toolset = make_tools().build();
814        let tools = toolset.tools();
815        assert_eq!(tools[0].version(), VERSION);
816    }
817
818    #[test]
819    fn test_schemas() {
820        let toolset = make_tools().build();
821        let tools = toolset.tools();
822        let input = tools[0].input_schema();
823        assert!(input["properties"]["commands"].is_object());
824        let output = tools[0].output_schema();
825        assert!(output["properties"]["stdout"].is_object());
826    }
827
828    #[test]
829    fn test_discover_tool_schemas() {
830        let toolset = make_tools().with_discovery().build();
831        let tools = toolset.tools();
832        let input = tools[1].input_schema();
833        assert!(input["properties"]["commands"].is_object());
834    }
835
836    // -- DiscoverTool command restriction --
837
838    #[tokio::test]
839    async fn test_discover_tool_allows_discover() {
840        let toolset = make_tools().with_discovery().build();
841        let tools = toolset.tools();
842        let resp = tools[1]
843            .execute(ToolRequest {
844                commands: "discover --categories".into(),
845                timeout_ms: None,
846            })
847            .await;
848        assert_eq!(resp.exit_code, 0);
849        assert!(resp.stdout.contains("users"));
850    }
851
852    #[tokio::test]
853    async fn test_discover_tool_allows_help() {
854        let toolset = make_tools().with_discovery().build();
855        let tools = toolset.tools();
856        let resp = tools[1]
857            .execute(ToolRequest {
858                commands: "help get_user".into(),
859                timeout_ms: None,
860            })
861            .await;
862        assert_eq!(resp.exit_code, 0);
863        assert!(resp.stdout.contains("get_user"));
864    }
865
866    #[tokio::test]
867    async fn test_discover_tool_rejects_other_commands() {
868        let toolset = make_tools().with_discovery().build();
869        let tools = toolset.tools();
870        let resp = tools[1]
871            .execute(ToolRequest {
872                commands: "get_user --id 42".into(),
873                timeout_ms: None,
874            })
875            .await;
876        assert_ne!(resp.exit_code, 0);
877        assert!(
878            resp.error
879                .as_deref()
880                .unwrap_or("")
881                .contains("discover tool only supports"),
882            "error: {:?}",
883            resp.error
884        );
885    }
886
887    #[tokio::test]
888    async fn test_discover_tool_rejects_arbitrary_bash() {
889        let toolset = make_tools().with_discovery().build();
890        let tools = toolset.tools();
891        let resp = tools[1]
892            .execute(ToolRequest {
893                commands: "echo pwned".into(),
894                timeout_ms: None,
895            })
896            .await;
897        assert_ne!(resp.exit_code, 0);
898        assert!(
899            resp.error
900                .as_deref()
901                .unwrap_or("")
902                .contains("discover tool only supports")
903        );
904    }
905
906    #[test]
907    fn test_discover_tool_execution_rejects_other_commands() {
908        let toolset = make_tools().with_discovery().build();
909        let tools = toolset.tools();
910        let args = serde_json::json!({ "commands": "get_user --id 42" });
911        let result = tools[1].execution(args);
912        match result {
913            Err(e) => assert!(
914                e.to_string().contains("discover tool only supports"),
915                "unexpected error: {e}"
916            ),
917            Ok(_) => panic!("expected error for disallowed command"),
918        }
919    }
920
921    // -- ToolImpl registration --
922
923    #[tokio::test]
924    async fn test_tool_impl_registration() {
925        let get_user = ToolImpl::new(
926            ToolDef::new("get_user", "Fetch user by ID")
927                .with_schema(serde_json::json!({
928                    "type": "object",
929                    "properties": { "id": {"type": "integer"} },
930                    "required": ["id"]
931                }))
932                .with_category("users"),
933        )
934        .with_exec_sync(|args| {
935            let id = args.param_i64("id").ok_or("missing --id")?;
936            Ok(format!("{{\"id\":{id},\"name\":\"Alice\"}}\n"))
937        });
938
939        let list_orders = ToolImpl::new(
940            ToolDef::new("list_orders", "List orders")
941                .with_schema(serde_json::json!({
942                    "type": "object",
943                    "properties": { "user_id": {"type": "integer"} }
944                }))
945                .with_category("orders"),
946        )
947        .with_exec_sync(|args| {
948            let uid = args.param_i64("user_id").ok_or("missing --user_id")?;
949            Ok(format!("[{{\"order_id\":1,\"user_id\":{uid}}}]\n"))
950        });
951
952        // Exclusive mode
953        let toolset = ScriptingToolSet::builder("api")
954            .short_description("Test API")
955            .tool(get_user.clone())
956            .tool(list_orders.clone())
957            .build();
958
959        let tools = toolset.tools();
960        assert_eq!(tools.len(), 1);
961        assert!(tools[0].system_prompt().contains("get_user"));
962        assert!(tools[0].system_prompt().contains("list_orders"));
963
964        let resp = tools[0]
965            .execute(ToolRequest {
966                commands: "get_user --id 1 | jq -r '.name'".into(),
967                timeout_ms: None,
968            })
969            .await;
970        assert_eq!(resp.stdout.trim(), "Alice");
971
972        // Discovery mode
973        let toolset = ScriptingToolSet::builder("api")
974            .tool(get_user)
975            .tool(list_orders)
976            .with_discovery()
977            .build();
978
979        let tools = toolset.tools();
980        assert_eq!(tools.len(), 2);
981        assert_eq!(tools[0].name(), "api");
982        assert_eq!(tools[1].name(), "api_discover");
983    }
984
985    // -- Structured discover input --
986
987    #[tokio::test]
988    async fn test_discover_structured_query() {
989        let toolset = make_tools().with_discovery().build();
990        let tools = toolset.tools();
991        let args = serde_json::json!({ "query": "user" });
992        let exec = tools[1].execution(args).expect("valid structured args");
993        let output = exec.execute().await.expect("execution succeeds");
994        let stdout = output.result["stdout"].as_str().unwrap_or("");
995        assert!(stdout.contains("get_user"), "stdout: {stdout}");
996    }
997
998    #[tokio::test]
999    async fn test_discover_structured_all() {
1000        let toolset = make_tools().with_discovery().build();
1001        let tools = toolset.tools();
1002        let args = serde_json::json!({ "all": true });
1003        let exec = tools[1].execution(args).expect("valid structured args");
1004        let output = exec.execute().await.expect("execution succeeds");
1005        let stdout = output.result["stdout"].as_str().unwrap_or("");
1006        assert!(stdout.contains("users"), "stdout: {stdout}");
1007        assert!(stdout.contains("orders"), "stdout: {stdout}");
1008    }
1009
1010    #[tokio::test]
1011    async fn test_discover_structured_query_rejects_injection() {
1012        let toolset = make_tools().with_discovery().build();
1013        let tools = toolset.tools();
1014        let args = serde_json::json!({ "query": "user; help get_user --json" });
1015        match tools[1].execution(args) {
1016            Err(e) => assert!(
1017                e.to_string()
1018                    .contains("query contains unsupported shell control characters"),
1019                "unexpected error: {e}"
1020            ),
1021            Ok(_) => panic!("expected injection-like query to be rejected"),
1022        }
1023    }
1024
1025    #[tokio::test]
1026    async fn test_discover_backward_compat() {
1027        let toolset = make_tools().with_discovery().build();
1028        let tools = toolset.tools();
1029        let args = serde_json::json!({ "commands": "discover --search user" });
1030        let exec = tools[1].execution(args).expect("valid commands args");
1031        let output = exec.execute().await.expect("execution succeeds");
1032        let stdout = output.result["stdout"].as_str().unwrap_or("");
1033        assert!(stdout.contains("get_user"), "stdout: {stdout}");
1034    }
1035
1036    #[tokio::test]
1037    async fn test_discover_backward_compat_rejects_shell_separators() {
1038        let toolset = make_tools().with_discovery().build();
1039        let tools = toolset.tools();
1040        let args = serde_json::json!({ "commands": "discover --search user; help get_user" });
1041        match tools[1].execution(args) {
1042            Err(e) => assert!(
1043                e.to_string()
1044                    .contains("discover tool commands cannot contain shell control characters"),
1045                "unexpected error: {e}"
1046            ),
1047            Ok(_) => panic!("expected injection-like commands to be rejected"),
1048        }
1049    }
1050
1051    #[test]
1052    fn test_discover_input_schema_is_structured() {
1053        let toolset = make_tools().with_discovery().build();
1054        let tools = toolset.tools();
1055        let schema = tools[1].input_schema();
1056        let props = &schema["properties"];
1057        assert!(props["query"].is_object(), "missing query property");
1058        assert!(props["all"].is_object(), "missing all property");
1059        assert!(props["commands"].is_object(), "missing commands property");
1060        assert!(
1061            props["timeout_ms"].is_object(),
1062            "missing timeout_ms property"
1063        );
1064    }
1065}
bashkit_scripted_tool/scripted_tool/toolset.rs

bashkit_scripted_tool/scripted_tool/
toolset.rs
