Skip to main content

clawbox_server/
state.rs

1//! Shared application state.
2
3use crate::config::ClawboxConfig;
4#[cfg(feature = "docker")]
5use crate::container_proxy::ContainerProxy;
6#[cfg(feature = "docker")]
7use clawbox_containers::{AgentOrchestrator, ContainerBackend, DockerBackend};
8use clawbox_proxy::OutputScanner;
9use clawbox_proxy::{AuditLog, CredentialStore, RateLimiter, parse_master_key};
10use clawbox_sandbox::{SandboxEngine, ToolWatcherHandle, start_watching};
11use clawbox_types::{ToolManifest, ToolMeta};
12use metrics_exporter_prometheus::PrometheusHandle;
13use thiserror::Error;
14
15/// Errors from application state initialization.
16#[derive(Debug, Error)]
17#[non_exhaustive]
18pub enum StateError {
19    /// Sandbox engine initialization failed.
20    #[error("sandbox error: {0}")]
21    Sandbox(#[from] clawbox_sandbox::SandboxError),
22    /// Container backend initialization failed.
23    #[cfg(feature = "docker")]
24    #[error("container error: {0}")]
25    Container(#[from] clawbox_containers::ContainerError),
26    /// I/O error.
27    #[error("I/O error: {0}")]
28    Io(#[from] std::io::Error),
29    /// Watcher initialization failed.
30    #[error("watcher error: {0}")]
31    Watcher(#[from] clawbox_sandbox::WatcherError),
32}
33use std::collections::HashMap;
34use std::sync::Arc;
35use std::sync::Mutex;
36use tokio::sync::RwLock;
37use tracing::{info, warn};
38
39/// Returns the proxy socket directory for a container.
40///
41/// Uses `~/.clawbox/proxies/` instead of `/tmp` to avoid predictable paths.
42/// Creates the base directory with mode 0o700 if it does not exist.
43pub fn proxy_socket_dir(container_id: &str) -> std::path::PathBuf {
44    let base = crate::config::expand_tilde("~/.clawbox/proxies");
45    if !base.exists() {
46        if let Err(e) = std::fs::create_dir_all(&base) {
47            tracing::warn!("Failed to create proxy base dir {}: {e}", base.display());
48        } else {
49            #[cfg(unix)]
50            {
51                use std::os::unix::fs::PermissionsExt;
52                let _ = std::fs::set_permissions(&base, std::fs::Permissions::from_mode(0o700));
53            }
54        }
55    }
56    base.join(container_id)
57}
58
59/// Returns the proxy socket path for a container.
60pub fn proxy_socket_path(container_id: &str) -> std::path::PathBuf {
61    proxy_socket_dir(container_id).join("proxy.sock")
62}
63
64/// Docker/container-specific state, only compiled when `docker` feature is enabled.
65#[cfg(feature = "docker")]
66pub struct DockerState {
67    pub container_manager: Arc<DockerBackend>,
68    pub agent_orchestrator: Arc<AgentOrchestrator>,
69    pub container_proxies: RwLock<HashMap<String, ContainerProxy>>,
70    reaper_shutdown: tokio::sync::watch::Sender<bool>,
71}
72
73/// Shared state for the clawbox HTTP server.
74#[non_exhaustive]
75pub struct AppState {
76    pub sandbox_engine: Arc<SandboxEngine>,
77    pub output_scanner: OutputScanner,
78    pub config: ClawboxConfig,
79    pub tools: RwLock<HashMap<String, ToolManifest>>,
80    pub start_time: std::time::Instant,
81    pub credential_store: Option<CredentialStore>,
82    pub audit_log: AuditLog,
83    pub rate_limiter: Arc<RateLimiter>,
84    pub metrics_handle: PrometheusHandle,
85    /// Handle to the filesystem watcher (if enabled).
86    watcher_handle: Mutex<Option<ToolWatcherHandle>>,
87    /// Docker/container state (only present with `docker` feature).
88    #[cfg(feature = "docker")]
89    pub docker: DockerState,
90}
91
92impl AppState {
93    pub async fn new(config: ClawboxConfig) -> Result<Self, StateError> {
94        let metrics_handle = crate::metrics::init_metrics();
95        let epoch_interval_ms = clawbox_sandbox::resource_limits::EPOCH_INTERVAL_MS;
96        let sandbox_config = clawbox_sandbox::SandboxConfig::new(&config.sandbox.tool_dir)
97            .with_fuel_limit(config.sandbox.default_fuel)
98            .with_epoch_deadline((config.sandbox.default_timeout_ms / epoch_interval_ms).max(1))
99            .with_epoch_interval_ms(epoch_interval_ms)
100            .with_max_host_calls(100)
101            .with_max_memory_bytes(clawbox_sandbox::resource_limits::DEFAULT_MAX_MEMORY_BYTES)
102            .with_max_table_elements(clawbox_sandbox::resource_limits::DEFAULT_MAX_TABLE_ELEMENTS);
103
104        let engine = Arc::new(SandboxEngine::new(sandbox_config)?);
105
106        // Load all .wasm modules from the tool directory
107        let count = engine.load_all_modules()?;
108        info!(count, tool_dir = %config.sandbox.tool_dir, "loaded WASM tools at startup");
109
110        // D4: Auto-register minimal manifests for loaded WASM tools
111        let mut auto_tools = HashMap::new();
112        for name in engine.list_modules() {
113            let meta =
114                ToolMeta::new(&name, "Auto-loaded WASM tool").with_version("0.0.0".to_string());
115            auto_tools.insert(name, ToolManifest::new(meta));
116        }
117
118        // Start filesystem watcher if enabled
119        let watcher_handle = if config.sandbox.watch_tools {
120            let tool_dir = std::path::PathBuf::from(&config.sandbox.tool_dir);
121            if tool_dir.exists() {
122                match start_watching(Arc::clone(&engine), tool_dir) {
123                    Ok(handle) => Some(handle),
124                    Err(e) => {
125                        warn!("Failed to start tool watcher: {e}. Hot-reload disabled.");
126                        None
127                    }
128                }
129            } else {
130                warn!(dir = %config.sandbox.tool_dir, "tool directory does not exist, skipping watcher");
131                None
132            }
133        } else {
134            None
135        };
136
137        // Load credential store if CLAWBOX_MASTER_KEY is set
138        let credential_store = match std::env::var("CLAWBOX_MASTER_KEY") {
139            Ok(hex_key) => match parse_master_key(&hex_key) {
140                Ok(key) => {
141                    let cred_path = crate::config::expand_tilde(&config.credentials.store_path);
142                    match CredentialStore::load(cred_path.to_string_lossy().as_ref(), key) {
143                        Ok(store) => {
144                            info!(
145                                "Credential store loaded from {}",
146                                config.credentials.store_path
147                            );
148                            Some(store)
149                        }
150                        Err(e) => {
151                            warn!(
152                                "Failed to load credential store: {e}. Continuing without credentials."
153                            );
154                            None
155                        }
156                    }
157                }
158                Err(e) => {
159                    warn!("Invalid CLAWBOX_MASTER_KEY: {e}. Continuing without credentials.");
160                    None
161                }
162            },
163            Err(_) => {
164                warn!("CLAWBOX_MASTER_KEY not set. Credential injection disabled.");
165                None
166            }
167        };
168
169        // Initialize audit log
170        let audit_log = AuditLog::new(format!("{}/proxy.jsonl", config.logging.audit_dir));
171
172        // Fix 5: Create shared rate limiter
173        let rate_limiter = Arc::new(RateLimiter::new(50, 10.0)); // 50 burst, 10/sec refill
174
175        // Docker/container initialization
176        #[cfg(feature = "docker")]
177        let docker = {
178            let container_manager = Arc::new(DockerBackend::new().await?);
179
180            let workspace_root = crate::config::expand_tilde(&config.containers.workspace_root)
181                .to_string_lossy()
182                .into_owned();
183            let workspace_root = std::path::PathBuf::from(&workspace_root);
184            std::fs::create_dir_all(&workspace_root)?;
185            let agent_orchestrator = Arc::new(AgentOrchestrator::new(
186                Arc::clone(&container_manager) as Arc<dyn clawbox_containers::ContainerBackend>,
187                workspace_root,
188            ));
189
190            let (reaper_shutdown_tx, reaper_shutdown_rx) = tokio::sync::watch::channel(false);
191            container_manager.spawn_reaper(reaper_shutdown_rx);
192
193            DockerState {
194                container_manager,
195                agent_orchestrator,
196                container_proxies: RwLock::new(HashMap::new()),
197                reaper_shutdown: reaper_shutdown_tx,
198            }
199        };
200
201        // Spawn periodic rate limiter cleanup
202        let cleanup_limiter = Arc::clone(&rate_limiter);
203        tokio::spawn(async move {
204            let mut interval = tokio::time::interval(std::time::Duration::from_secs(300));
205            loop {
206                interval.tick().await;
207                cleanup_limiter.cleanup_stale(std::time::Duration::from_secs(600));
208            }
209        });
210
211        Ok(Self {
212            sandbox_engine: engine,
213            output_scanner: OutputScanner::new(),
214            config,
215            tools: RwLock::new(auto_tools),
216            start_time: std::time::Instant::now(),
217            credential_store,
218            audit_log,
219            rate_limiter,
220            metrics_handle,
221            watcher_handle: Mutex::new(watcher_handle),
222            #[cfg(feature = "docker")]
223            docker,
224        })
225    }
226
227    /// Returns `(docker_available, active_containers, active_agents)`.
228    /// When compiled without the `docker` feature this always returns `(false, 0, 0)`.
229    #[cfg(feature = "docker")]
230    pub async fn docker_status(&self) -> (bool, usize, usize) {
231        let available = self.docker.container_manager.is_available().await;
232        let containers = self.docker.container_manager.list().await.len();
233        let agents = self.docker.agent_orchestrator.list_agents().await.len();
234        (available, containers, agents)
235    }
236
237    #[cfg(not(feature = "docker"))]
238    pub async fn docker_status(&self) -> (bool, usize, usize) {
239        (false, 0, 0)
240    }
241
242    /// Gracefully shut down all managed resources.
243    pub async fn shutdown(&self) {
244        // Shut down the tool watcher
245        if let Some(handle) = self
246            .watcher_handle
247            .lock()
248            .unwrap_or_else(|e| e.into_inner())
249            .take()
250        {
251            handle.shutdown();
252        }
253
254        // Shut down Docker resources
255        #[cfg(feature = "docker")]
256        {
257            let _ = self.docker.reaper_shutdown.send(true);
258
259            let mut proxies = self.docker.container_proxies.write().await;
260            for (id, proxy) in proxies.iter_mut() {
261                tracing::info!(container = %id, "Shutting down container proxy");
262                proxy.shutdown();
263            }
264            proxies.clear();
265
266            let containers = self.docker.container_manager.list().await;
267            for info in &containers {
268                if info.status == clawbox_types::ContainerStatus::Running {
269                    tracing::info!(container = %info.container_id, "Killing container on shutdown");
270                    let _ = self.docker.container_manager.kill(&info.container_id).await;
271                }
272            }
273        }
274
275        // Stop the sandbox epoch ticker
276        self.sandbox_engine.shutdown();
277    }
278}
279
280#[cfg(test)]
281mod tests {
282    use super::*;
283
284    #[test]
285    fn test_proxy_socket_dir() {
286        let dir = proxy_socket_dir("clawbox-abc123");
287        let home = std::env::var("HOME").unwrap_or_else(|_| "/home/test".into());
288        assert_eq!(
289            dir,
290            std::path::PathBuf::from(format!("{home}/.clawbox/proxies/clawbox-abc123"))
291        );
292    }
293
294    #[test]
295    fn test_proxy_socket_path() {
296        let path = proxy_socket_path("clawbox-abc123");
297        let home = std::env::var("HOME").unwrap_or_else(|_| "/home/test".into());
298        assert_eq!(
299            path,
300            std::path::PathBuf::from(format!("{home}/.clawbox/proxies/clawbox-abc123/proxy.sock"))
301        );
302    }
303}