Skip to main content

bamboo_tools/tools/
kill_shell.rs

1use async_trait::async_trait;
2use bamboo_agent_core::{Tool, ToolError, 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 execute(&self, args: serde_json::Value) -> Result<ToolResult, ToolError> {
74        let parsed: KillShellArgs = serde_json::from_value(args)
75            .map_err(|e| ToolError::InvalidArguments(format!("Invalid KillShell args: {}", e)))?;
76
77        let shell_id = parsed.resolved_shell_id().ok_or_else(|| {
78            ToolError::InvalidArguments(
79                "KillShell requires 'shell_id' (or legacy alias 'bash_id') from Bash run_in_background".to_string(),
80            )
81        })?;
82        let shell = bash_runtime::get_shell(shell_id).ok_or_else(|| {
83            ToolError::Execution(format!(
84                "Background shell '{}' not found. Use the bash_id returned by Bash(run_in_background=true), not chat session_id.",
85                shell_id
86            ))
87        })?;
88
89        if shell.status() == "running" {
90            shell.kill().await.map_err(ToolError::Execution)?;
91        }
92        let _ = bash_runtime::remove_shell(shell_id);
93
94        Ok(ToolResult {
95            success: true,
96            result: json!({
97                "shell_id": shell_id,
98                "bash_id": shell_id,
99                "status": "killed"
100            })
101            .to_string(),
102            display_preference: Some("Collapsible".to_string()),
103        })
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110    use crate::tools::bash::BashTool;
111    use serde_json::Value;
112
113    #[cfg(target_os = "windows")]
114    fn long_running_command() -> &'static str {
115        "powershell -NoProfile -Command \"Start-Sleep -Seconds 2\""
116    }
117
118    #[cfg(not(target_os = "windows"))]
119    fn long_running_command() -> &'static str {
120        "sleep 2"
121    }
122
123    #[tokio::test]
124    async fn kill_shell_terminates_and_removes_session() {
125        let bash = BashTool::new();
126        let spawned = bash
127            .execute(json!({
128                "command": long_running_command(),
129                "run_in_background": true
130            }))
131            .await
132            .unwrap();
133        let spawned_payload: Value = serde_json::from_str(&spawned.result).unwrap();
134        let shell_id = spawned_payload["bash_id"].as_str().unwrap().to_string();
135        assert!(super::bash_runtime::get_shell(&shell_id).is_some());
136
137        let kill = KillShellTool::new();
138        let result = kill
139            .execute(json!({
140                "shell_id": shell_id
141            }))
142            .await
143            .unwrap();
144        assert!(result.success);
145
146        let payload: Value = serde_json::from_str(&result.result).unwrap();
147        let killed_id = payload["shell_id"].as_str().unwrap();
148        assert!(super::bash_runtime::get_shell(killed_id).is_none());
149    }
150
151    #[tokio::test]
152    async fn kill_shell_accepts_bash_id_alias() {
153        let bash = BashTool::new();
154        let spawned = bash
155            .execute(json!({
156                "command": long_running_command(),
157                "run_in_background": true
158            }))
159            .await
160            .unwrap();
161        let spawned_payload: Value = serde_json::from_str(&spawned.result).unwrap();
162        let shell_id = spawned_payload["bash_id"].as_str().unwrap().to_string();
163
164        let kill = KillShellTool::new();
165        let result = kill
166            .execute(json!({
167                "bash_id": shell_id
168            }))
169            .await
170            .unwrap();
171        assert!(result.success);
172    }
173}