Skip to main content

stakpak_shared/
container.rs

1use std::collections::HashMap;
2use std::net::TcpListener;
3use std::process::Command;
4
5// ── Stakpak agent container constants ──────────────────────────────────────
6
7/// The container image used for sandboxed agent sessions.
8/// Override at runtime with STAKPAK_AGENT_IMAGE env var for local testing.
9pub fn stakpak_agent_image() -> String {
10    std::env::var("STAKPAK_AGENT_IMAGE")
11        .unwrap_or_else(|_| format!("ghcr.io/stakpak/agent:v{}", env!("CARGO_PKG_VERSION")))
12}
13
14/// Default volume mounts for the stakpak agent container.
15///
16/// Single source of truth for every path the container needs.
17/// Used by `WardenConfig::readonly_profile()`, `prepare_volumes()`,
18/// and `build_dynamic_subagent_command()`.
19pub fn stakpak_agent_default_mounts() -> Vec<String> {
20    vec![
21        // Stakpak config & credentials
22        "~/.stakpak/config.toml:/home/agent/.stakpak/config.toml:ro".to_string(),
23        "~/.stakpak/auth.toml:/home/agent/.stakpak/auth.toml:ro".to_string(),
24        "~/.stakpak/data/local.db:/home/agent/.stakpak/data/local.db".to_string(),
25        "~/.agent-board/data.db:/home/agent/.agent-board/data.db".to_string(),
26        // Working directory
27        "./:/agent:ro".to_string(),
28        "./.stakpak:/agent/.stakpak".to_string(),
29        // AWS — config read-only, SSO/STS cache writable for token refresh
30        "~/.aws/config:/home/agent/.aws/config:ro".to_string(),
31        "~/.aws/credentials:/home/agent/.aws/credentials:ro".to_string(),
32        "~/.aws/sso:/home/agent/.aws/sso".to_string(),
33        "~/.aws/cli:/home/agent/.aws/cli".to_string(),
34        // GCP — credential files read-only, cache/logs/db writable for gcloud to function
35        "~/.config/gcloud/active_config:/home/agent/.config/gcloud/active_config:ro".to_string(),
36        "~/.config/gcloud/configurations:/home/agent/.config/gcloud/configurations:ro".to_string(),
37        "~/.config/gcloud/application_default_credentials.json:/home/agent/.config/gcloud/application_default_credentials.json:ro".to_string(),
38        "~/.config/gcloud/credentials.db:/home/agent/.config/gcloud/credentials.db:ro".to_string(),
39        "~/.config/gcloud/access_tokens.db:/home/agent/.config/gcloud/access_tokens.db:ro".to_string(),
40        "~/.config/gcloud/logs:/home/agent/.config/gcloud/logs".to_string(),
41        "~/.config/gcloud/cache:/home/agent/.config/gcloud/cache".to_string(),
42        // Azure — config read-only, MSAL token cache and session writable
43        "~/.azure/config:/home/agent/.azure/config:ro".to_string(),
44        "~/.azure/clouds.config:/home/agent/.azure/clouds.config:ro".to_string(),
45        "~/.azure/azureProfile.json:/home/agent/.azure/azureProfile.json:ro".to_string(),
46        "~/.azure/msal_token_cache.json:/home/agent/.azure/msal_token_cache.json".to_string(),
47        "~/.azure/msal_http_cache.bin:/home/agent/.azure/msal_http_cache.bin".to_string(),
48        "~/.azure/logs:/home/agent/.azure/logs".to_string(),
49        // DigitalOcean & Kubernetes
50        "~/.digitalocean:/home/agent/.digitalocean:ro".to_string(),
51        "~/.kube:/home/agent/.kube:ro".to_string(),
52        // SSH — config and keys read-only (useful for host aliases and remote connections)
53        "~/.ssh:/home/agent/.ssh:ro".to_string(),
54        // Aqua tool cache (named volume — persists downloaded CLIs across runs)
55        "stakpak-aqua-cache:/home/agent/.local/share/aquaproj-aqua".to_string(),
56    ]
57}
58
59/// Expand `~` to `$HOME` in a volume mount string.
60pub fn expand_volume_path(volume: &str) -> String {
61    if (volume.starts_with("~/") || volume.starts_with("~:"))
62        && let Ok(home_dir) = std::env::var("HOME")
63    {
64        return volume.replacen("~", &home_dir, 1);
65    }
66    volume.to_string()
67}
68
69/// Check whether the host-side part of a volume mount is a Docker named volume
70/// (as opposed to a bind mount path).
71///
72/// Named volumes don't start with `/`, `.`, or `~` and contain no `/`.
73pub fn is_named_volume(host_part: &str) -> bool {
74    !host_part.starts_with('/')
75        && !host_part.starts_with('.')
76        && !host_part.starts_with('~')
77        && !host_part.contains('/')
78}
79
80/// Pre-create any Docker named volumes found in [`stakpak_agent_default_mounts`].
81///
82/// Running `docker volume create` is idempotent and prevents a race condition
83/// when multiple sandbox containers first-use the same named volume in parallel.
84pub fn ensure_named_volumes_exist() {
85    for vol in stakpak_agent_default_mounts() {
86        let host_part = vol.split(':').next().unwrap_or(&vol);
87        if is_named_volume(host_part) {
88            let _ = Command::new("docker")
89                .args(["volume", "create", host_part])
90                .stdout(std::process::Stdio::null())
91                .stderr(std::process::Stdio::null())
92                .status();
93        }
94    }
95}
96
97#[derive(Debug, Clone)]
98pub struct ContainerConfig {
99    pub image: String,
100    pub env_vars: HashMap<String, String>,
101    pub ports: Vec<String>,       // Format: "host_port:container_port"
102    pub extra_hosts: Vec<String>, // Format: "host:ip"
103    pub volumes: Vec<String>,     // Format: "host_path:container_path"
104}
105
106pub fn find_available_port() -> Option<u16> {
107    match TcpListener::bind("0.0.0.0:0") {
108        Ok(listener) => listener.local_addr().ok().map(|addr| addr.port()),
109        Err(_) => None,
110    }
111}
112
113/// Checks if Docker is installed and accessible
114pub fn is_docker_available() -> bool {
115    Command::new("docker")
116        .arg("--version")
117        .output()
118        .map(|output| output.status.success())
119        .unwrap_or(false)
120}
121
122/// Checks if a Docker image exists locally
123pub fn image_exists_locally(image: &str) -> Result<bool, String> {
124    let output = Command::new("docker")
125        .args(["images", "-q", image])
126        .output()
127        .map_err(|e| format!("Failed to execute docker images command: {}", e))?;
128
129    if output.status.success() {
130        let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
131        Ok(!stdout.is_empty())
132    } else {
133        let stderr = String::from_utf8_lossy(&output.stderr).to_string();
134        Err(format!("Docker images command failed: {}", stderr))
135    }
136}
137
138/// The platform warden uses for all containers.
139/// Warden hardcodes `--platform linux/amd64` on every `docker run` / `docker create`,
140/// so all image operations must target the same platform for consistency.
141pub const WARDEN_PLATFORM: &str = "linux/amd64";
142
143/// Checks if a Docker image for the warden platform (linux/amd64) exists locally.
144///
145/// Unlike [`image_exists_locally`], this uses `docker image inspect --platform`
146/// to verify the correct architecture variant is cached. On Apple Silicon a plain
147/// `docker images -q` would match a cached arm64 image, giving a false positive.
148pub fn warden_image_exists_locally(image: &str) -> bool {
149    Command::new("docker")
150        .args(["image", "inspect", "--platform", WARDEN_PLATFORM, image])
151        .stdout(std::process::Stdio::null())
152        .stderr(std::process::Stdio::null())
153        .status()
154        .map(|s| s.success())
155        .unwrap_or(false)
156}
157
158/// Pull a Docker image for the warden platform (linux/amd64) with visible progress.
159///
160/// Inherits stdout/stderr so the caller sees Docker's native progress bars
161/// (layer downloads, extraction, etc.). Returns an error if the pull fails.
162pub fn pull_warden_image(image: &str) -> Result<(), String> {
163    let status = Command::new("docker")
164        .args(["pull", "--platform", WARDEN_PLATFORM, image])
165        .stdout(std::process::Stdio::inherit())
166        .stderr(std::process::Stdio::inherit())
167        .status()
168        .map_err(|e| format!("Failed to run docker pull: {e}"))?;
169
170    if status.success() {
171        Ok(())
172    } else {
173        Err(format!(
174            "Failed to pull image '{image}' for platform {WARDEN_PLATFORM}. \
175             Check your network connection and that the image exists."
176        ))
177    }
178}
179
180pub fn run_container_detached(config: ContainerConfig) -> Result<String, String> {
181    let mut cmd = Command::new("docker");
182
183    cmd.arg("run").arg("-d").arg("--rm");
184
185    // Add ports
186    for port_mapping in &config.ports {
187        cmd.arg("-p").arg(port_mapping);
188    }
189
190    // Add environment variables
191    for (key, value) in &config.env_vars {
192        cmd.arg("-e").arg(format!("{}={}", key, value));
193    }
194
195    // Add extra hosts
196    for host_mapping in &config.extra_hosts {
197        cmd.arg("--add-host").arg(host_mapping);
198    }
199
200    // Add volumes
201    for volume_mapping in &config.volumes {
202        cmd.arg("-v").arg(volume_mapping);
203    }
204
205    // Add image
206    cmd.arg(&config.image);
207
208    let output = cmd
209        .output()
210        .map_err(|e| format!("Failed to execute docker command: {}", e))?;
211
212    if output.status.success() {
213        let container_id = String::from_utf8_lossy(&output.stdout).trim().to_string();
214        Ok(container_id)
215    } else {
216        let stderr = String::from_utf8_lossy(&output.stderr).to_string();
217        Err(format!("Docker command failed: {}", stderr))
218    }
219}
220
221pub fn stop_container(container_id: &str) -> Result<(), String> {
222    let output = Command::new("docker")
223        .arg("stop")
224        .arg(container_id)
225        .output()
226        .map_err(|e| format!("Failed to execute docker stop: {}", e))?;
227
228    if output.status.success() {
229        Ok(())
230    } else {
231        let stderr = String::from_utf8_lossy(&output.stderr);
232        if stderr.contains("No such container") {
233            Ok(())
234        } else {
235            Err(format!("Failed to stop container: {}", stderr))
236        }
237    }
238}
239
240pub fn remove_container(
241    container_id: &str,
242    force: bool,
243    remove_volumes: bool,
244) -> Result<(), String> {
245    let mut cmd = Command::new("docker");
246
247    cmd.arg("rm");
248
249    if force {
250        cmd.arg("-f");
251    }
252
253    if remove_volumes {
254        cmd.arg("-v");
255    }
256
257    cmd.arg(container_id);
258
259    let output = cmd
260        .output()
261        .map_err(|e| format!("Failed to execute docker rm: {}", e))?;
262
263    if output.status.success() {
264        Ok(())
265    } else {
266        let stderr = String::from_utf8_lossy(&output.stderr);
267        if stderr.contains("No such container") {
268            Ok(())
269        } else {
270            Err(format!("Failed to remove container: {}", stderr))
271        }
272    }
273}
274
275pub fn get_container_host_port(container_id: &str, container_port: u16) -> Result<u16, String> {
276    let output = Command::new("docker")
277        .arg("port")
278        .arg(container_id)
279        .arg(container_port.to_string())
280        .output()
281        .map_err(|e| format!("Failed to get container port: {}", e))?;
282
283    if output.status.success() {
284        let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
285        let port = stdout.split(':').next_back().unwrap_or("");
286        Ok(port.parse().unwrap())
287    } else {
288        let stderr = String::from_utf8_lossy(&output.stderr).to_string();
289        Err(format!("Failed to get container port: {}", stderr))
290    }
291}