Skip to main content

bamboo_tools/tools/
kill_shell.rs

1use async_trait::async_trait;
2use bamboo_agent_core::{Tool, ToolCtx, ToolError, ToolOutcome, ToolResult};
3use serde::Deserialize;
4use serde_json::json;
5
6use super::bash_runtime;
7
8#[derive(Debug, Deserialize)]
9struct KillShellArgs {
10    #[serde(default)]
11    shell_id: Option<String>,
12    #[serde(default)]
13    bash_id: Option<String>,
14}
15
16impl KillShellArgs {
17    fn resolved_shell_id(&self) -> Option<&str> {
18        self.shell_id
19            .as_deref()
20            .map(str::trim)
21            .filter(|value| !value.is_empty())
22            .or_else(|| {
23                self.bash_id
24                    .as_deref()
25                    .map(str::trim)
26                    .filter(|value| !value.is_empty())
27            })
28    }
29}
30
31pub struct KillShellTool;
32
33impl KillShellTool {
34    pub fn new() -> Self {
35        Self
36    }
37}
38
39impl Default for KillShellTool {
40    fn default() -> Self {
41        Self::new()
42    }
43}
44
45#[async_trait]
46impl Tool for KillShellTool {
47    fn name(&self) -> &str {
48        "KillShell"
49    }
50
51    fn description(&self) -> &str {
52        "Kill a running background Bash shell by ID (use the bash_id returned by Bash run_in_background)"
53    }
54
55    fn parameters_schema(&self) -> serde_json::Value {
56        json!({
57            "type": "object",
58            "properties": {
59                "shell_id": {
60                    "type": "string",
61                    "description": "The ID of the background shell to kill (recommended: pass Bash's bash_id here)"
62                },
63                "bash_id": {
64                    "type": "string",
65                    "description": "Legacy alias for shell_id; use the id returned by Bash"
66                }
67            },
68            "required": ["shell_id"],
69            "additionalProperties": false
70        })
71    }
72
73    async fn invoke(
74        &self,
75        args: serde_json::Value,
76        _ctx: ToolCtx,
77    ) -> Result<ToolOutcome, ToolError> {
78        let parsed: KillShellArgs = serde_json::from_value(args)
79            .map_err(|e| ToolError::InvalidArguments(format!("Invalid KillShell args: {}", e)))?;
80
81        let shell_id = parsed.resolved_shell_id().ok_or_else(|| {
82            ToolError::InvalidArguments(
83                "KillShell requires 'shell_id' (or legacy alias 'bash_id') from Bash run_in_background".to_string(),
84            )
85        })?;
86        let shell = bash_runtime::get_shell(shell_id).ok_or_else(|| {
87            ToolError::Execution(format!(
88                "Background shell '{}' not found. Use the bash_id returned by Bash(run_in_background=true), not chat session_id.",
89                shell_id
90            ))
91        })?;
92
93        if shell.status() == "running" {
94            shell.kill().await.map_err(ToolError::Execution)?;
95        }
96        let _ = bash_runtime::remove_shell(shell_id);
97
98        Ok(ToolOutcome::Completed(ToolResult {
99            success: true,
100            result: json!({
101                "shell_id": shell_id,
102                "bash_id": shell_id,
103                "status": "killed"
104            })
105            .to_string(),
106            display_preference: Some("Collapsible".to_string()),
107            images: Vec::new(),
108        }))
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115    use crate::tools::bash::BashTool;
116    use serde_json::Value;
117
118    #[cfg(target_os = "windows")]
119    fn long_running_command() -> &'static str {
120        "powershell -NoProfile -Command \"Start-Sleep -Seconds 2\""
121    }
122
123    #[cfg(not(target_os = "windows"))]
124    fn long_running_command() -> &'static str {
125        "sleep 2"
126    }
127
128    #[tokio::test]
129    async fn kill_shell_terminates_and_removes_session() {
130        let bash = BashTool::new();
131        let spawned_out = bash
132            .invoke(
133                json!({
134                    "command": long_running_command(),
135                    "run_in_background": true
136                }),
137                ToolCtx::none("t"),
138            )
139            .await
140            .unwrap();
141        let ToolOutcome::Completed(spawned) = spawned_out else {
142            panic!("expected Completed")
143        };
144        let spawned_payload: Value = serde_json::from_str(&spawned.result).unwrap();
145        let shell_id = spawned_payload["bash_id"].as_str().unwrap().to_string();
146        assert!(super::bash_runtime::get_shell(&shell_id).is_some());
147
148        let kill = KillShellTool::new();
149        let out = kill
150            .invoke(
151                json!({
152                    "shell_id": shell_id
153                }),
154                ToolCtx::none("t"),
155            )
156            .await
157            .unwrap();
158        let ToolOutcome::Completed(result) = out else {
159            panic!("expected Completed")
160        };
161        assert!(result.success);
162
163        let payload: Value = serde_json::from_str(&result.result).unwrap();
164        let killed_id = payload["shell_id"].as_str().unwrap();
165        assert!(super::bash_runtime::get_shell(killed_id).is_none());
166    }
167
168    #[tokio::test]
169    async fn kill_shell_accepts_bash_id_alias() {
170        let bash = BashTool::new();
171        let spawned_out = bash
172            .invoke(
173                json!({
174                    "command": long_running_command(),
175                    "run_in_background": true
176                }),
177                ToolCtx::none("t"),
178            )
179            .await
180            .unwrap();
181        let ToolOutcome::Completed(spawned) = spawned_out else {
182            panic!("expected Completed")
183        };
184        let spawned_payload: Value = serde_json::from_str(&spawned.result).unwrap();
185        let shell_id = spawned_payload["bash_id"].as_str().unwrap().to_string();
186
187        let kill = KillShellTool::new();
188        let out = kill
189            .invoke(
190                json!({
191                    "bash_id": shell_id
192                }),
193                ToolCtx::none("t"),
194            )
195            .await
196            .unwrap();
197        let ToolOutcome::Completed(result) = out else {
198            panic!("expected Completed")
199        };
200        assert!(result.success);
201    }
202}