diff --git a/.env.example b/.env.example index 8de0c2fc..7ff45444 100644 --- a/.env.example +++ b/.env.example @@ -7,7 +7,7 @@ DATABASE_POOL_SIZE=10 # Session token is stored in ~/.ironclaw/session.json and managed automatically. # On first run, the agent will open a browser for OAuth authentication. NEARAI_MODEL=claude-3-5-sonnet-20241022 -NEARAI_BASE_URL=https://cloud-api.near.ai +NEARAI_BASE_URL=https://private.near.ai NEARAI_AUTH_URL=https://private.near.ai # NEARAI_SESSION_PATH=~/.ironclaw/session.json # optional, default shown diff --git a/Dockerfile.worker b/Dockerfile.worker index 6c58a5e5..8f556700 100644 --- a/Dockerfile.worker +++ b/Dockerfile.worker @@ -21,10 +21,15 @@ RUN cargo build --release --bin ironclaw FROM debian:bookworm-slim -# Install common development tools -RUN apt-get update && apt-get install -y --no-install-recommends \ - ca-certificates \ - curl \ +# Install curl first (needed to fetch the GitHub CLI GPG key), then add the +# gh CLI apt repository, then install all remaining dev tools in one layer. +RUN apt-get update \ + && apt-get install -y --no-install-recommends ca-certificates curl \ + && curl -fsSL https://cli.github.com/packages/githubcli-archive-keyring.gpg \ + | dd of=/usr/share/keyrings/githubcli-archive-keyring.gpg \ + && echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/githubcli-archive-keyring.gpg] https://cli.github.com/packages stable main" \ + > /etc/apt/sources.list.d/github-cli.list \ + && apt-get update && apt-get install -y --no-install-recommends \ git \ build-essential \ pkg-config \ @@ -34,6 +39,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ python3 \ python3-pip \ python3-venv \ + gh \ && rm -rf /var/lib/apt/lists/* # Install Rust toolchain for the sandbox user diff --git a/deploy/env.example b/deploy/env.example index 1c7dac9e..046d5b0a 100644 --- a/deploy/env.example +++ b/deploy/env.example @@ -5,7 +5,7 @@ DATABASE_URL=postgres://ironclaw:CHANGE_ME@localhost:5432/ironclaw # NEAR AI NEARAI_SESSION_TOKEN=CHANGE_ME NEARAI_MODEL=claude-3-5-sonnet-20241022 -NEARAI_BASE_URL=https://cloud-api.near.ai +NEARAI_BASE_URL=https://private.near.ai NEARAI_AUTH_URL=https://private.near.ai NEARAI_API_MODE=chat_completions diff --git a/src/agent/commands.rs b/src/agent/commands.rs index 5525a528..365e0ad5 100644 --- a/src/agent/commands.rs +++ b/src/agent/commands.rs @@ -415,13 +415,33 @@ impl Agent { } "model" => { + let current = self.llm().active_model_name(); + if args.is_empty() { - // Show current model - let name = self.llm().active_model_name(); - Ok(SubmissionResult::response(format!( - "Active model: {}", - name - ))) + // Show current model and list available models + let mut out = format!("Active model: {}\n", current); + match self.llm().list_models().await { + Ok(models) if !models.is_empty() => { + out.push_str("\nAvailable models:\n"); + for m in &models { + let marker = if *m == current { " (active)" } else { "" }; + out.push_str(&format!(" {}{}\n", m, marker)); + } + out.push_str("\nUse /model to switch."); + } + Ok(_) => { + out.push_str( + "\nCould not fetch model list. Use /model to switch.", + ); + } + Err(e) => { + out.push_str(&format!( + "\nCould not fetch models: {}. Use /model to switch.", + e + )); + } + } + Ok(SubmissionResult::response(out)) } else { let requested = &args[0]; @@ -441,7 +461,6 @@ impl Agent { } Err(e) => { tracing::warn!("Could not fetch model list for validation: {}", e); - // Proceed anyway, the provider will error on the next call if invalid } } diff --git a/src/agent/job_monitor.rs b/src/agent/job_monitor.rs new file mode 100644 index 00000000..b2db8852 --- /dev/null +++ b/src/agent/job_monitor.rs @@ -0,0 +1,245 @@ +//! Background job monitor that forwards Claude Code output to the main agent loop. +//! +//! When the main agent kicks off a sandbox job (especially Claude Code), this +//! monitor subscribes to the broadcast event channel and injects relevant +//! assistant messages back into the channel manager's stream. This lets the +//! main agent see what the sub-agent is producing and surface it to the user. +//! +//! ```text +//! Container ──NDJSON──► Orchestrator ──broadcast──► JobMonitor +//! │ +//! inject_tx (mpsc) +//! │ +//! ▼ +//! Agent Loop +//! ``` + +use tokio::sync::{broadcast, mpsc}; +use tokio::task::JoinHandle; +use uuid::Uuid; + +use crate::channels::IncomingMessage; +use crate::channels::web::types::SseEvent; + +/// Spawn a background task that watches for events from a specific job and +/// injects assistant messages into the agent loop. +/// +/// The monitor forwards: +/// - `SseEvent::JobMessage` (assistant role): injected as incoming messages so +/// the main agent can read and relay to the user. +/// - `SseEvent::JobResult`: injected as a completion notice, then the task exits. +/// +/// Tool use/result and status events are intentionally skipped (too noisy for +/// the main agent's context window). +pub fn spawn_job_monitor( + job_id: Uuid, + mut event_rx: broadcast::Receiver<(Uuid, SseEvent)>, + inject_tx: mpsc::Sender, +) -> JoinHandle<()> { + let short_id = job_id.to_string()[..8].to_string(); + + tokio::spawn(async move { + tracing::info!(job_id = %short_id, "Job monitor started successfully"); + + loop { + match event_rx.recv().await { + Ok((ev_job_id, event)) => { + if ev_job_id != job_id { + continue; + } + + match event { + SseEvent::JobMessage { role, content, .. } if role == "assistant" => { + let msg = IncomingMessage::new( + "job_monitor", + "system", + format!("[Job {}] Claude Code: {}", short_id, content), + ); + if inject_tx.send(msg).await.is_err() { + tracing::debug!( + job_id = %short_id, + "Inject channel closed, stopping monitor" + ); + break; + } + } + SseEvent::JobResult { status, .. } => { + let msg = IncomingMessage::new( + "job_monitor", + "system", + format!( + "[Job {}] Container finished (status: {})", + short_id, status + ), + ); + let _ = inject_tx.send(msg).await; + tracing::debug!( + job_id = %short_id, + status = %status, + "Job monitor exiting (job finished)" + ); + break; + } + _ => { + // Skip tool_use, tool_result, status events + } + } + } + Err(broadcast::error::RecvError::Lagged(n)) => { + tracing::warn!( + job_id = %short_id, + skipped = n, + "Job monitor lagged, some events were dropped" + ); + } + Err(broadcast::error::RecvError::Closed) => { + tracing::debug!( + job_id = %short_id, + "Broadcast channel closed, stopping monitor" + ); + break; + } + } + } + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_monitor_forwards_assistant_messages() { + let (event_tx, _) = broadcast::channel::<(Uuid, SseEvent)>(16); + let (inject_tx, mut inject_rx) = mpsc::channel::(16); + + let job_id = Uuid::new_v4(); + let _handle = spawn_job_monitor(job_id, event_tx.subscribe(), inject_tx); + + // Send an assistant message + event_tx + .send(( + job_id, + SseEvent::JobMessage { + job_id: job_id.to_string(), + role: "assistant".to_string(), + content: "I found a bug".to_string(), + }, + )) + .unwrap(); + + let msg = tokio::time::timeout(std::time::Duration::from_secs(1), inject_rx.recv()) + .await + .unwrap() + .unwrap(); + + assert_eq!(msg.channel, "job_monitor"); + assert_eq!(msg.user_id, "system"); + assert!(msg.content.contains("I found a bug")); + } + + #[tokio::test] + async fn test_monitor_ignores_other_jobs() { + let (event_tx, _) = broadcast::channel::<(Uuid, SseEvent)>(16); + let (inject_tx, mut inject_rx) = mpsc::channel::(16); + + let job_id = Uuid::new_v4(); + let other_job_id = Uuid::new_v4(); + let _handle = spawn_job_monitor(job_id, event_tx.subscribe(), inject_tx); + + // Send a message for a different job + event_tx + .send(( + other_job_id, + SseEvent::JobMessage { + job_id: other_job_id.to_string(), + role: "assistant".to_string(), + content: "wrong job".to_string(), + }, + )) + .unwrap(); + + // Should not receive anything + let result = + tokio::time::timeout(std::time::Duration::from_millis(100), inject_rx.recv()).await; + assert!( + result.is_err(), + "should have timed out, no message expected" + ); + } + + #[tokio::test] + async fn test_monitor_exits_on_job_result() { + let (event_tx, _) = broadcast::channel::<(Uuid, SseEvent)>(16); + let (inject_tx, mut inject_rx) = mpsc::channel::(16); + + let job_id = Uuid::new_v4(); + let handle = spawn_job_monitor(job_id, event_tx.subscribe(), inject_tx); + + // Send a completion event + event_tx + .send(( + job_id, + SseEvent::JobResult { + job_id: job_id.to_string(), + status: "completed".to_string(), + session_id: None, + }, + )) + .unwrap(); + + // Should receive the completion message + let msg = tokio::time::timeout(std::time::Duration::from_secs(1), inject_rx.recv()) + .await + .unwrap() + .unwrap(); + assert!(msg.content.contains("finished")); + + // The monitor task should exit + tokio::time::timeout(std::time::Duration::from_secs(1), handle) + .await + .expect("monitor should have exited") + .expect("monitor task should not panic"); + } + + #[tokio::test] + async fn test_monitor_skips_tool_events() { + let (event_tx, _) = broadcast::channel::<(Uuid, SseEvent)>(16); + let (inject_tx, mut inject_rx) = mpsc::channel::(16); + + let job_id = Uuid::new_v4(); + let _handle = spawn_job_monitor(job_id, event_tx.subscribe(), inject_tx); + + // Send tool use event (should be skipped) + event_tx + .send(( + job_id, + SseEvent::JobToolUse { + job_id: job_id.to_string(), + tool_name: "shell".to_string(), + input: serde_json::json!({"command": "ls"}), + }, + )) + .unwrap(); + + // Send user message (should be skipped) + event_tx + .send(( + job_id, + SseEvent::JobMessage { + job_id: job_id.to_string(), + role: "user".to_string(), + content: "user prompt".to_string(), + }, + )) + .unwrap(); + + // Should not receive anything for tool events or user messages + let result = + tokio::time::timeout(std::time::Duration::from_millis(100), inject_rx.recv()).await; + assert!( + result.is_err(), + "should have timed out, no message expected" + ); + } +} diff --git a/src/agent/mod.rs b/src/agent/mod.rs index 5e1bf64c..d0c96bc1 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -17,6 +17,7 @@ pub mod context_monitor; pub mod cost_guard; mod dispatcher; mod heartbeat; +pub mod job_monitor; mod router; pub mod routine; pub mod routine_engine; diff --git a/src/channels/manager.rs b/src/channels/manager.rs index 6d9b99a6..5cdc2f99 100644 --- a/src/channels/manager.rs +++ b/src/channels/manager.rs @@ -4,24 +4,41 @@ use std::collections::HashMap; use std::sync::Arc; use futures::stream; -use tokio::sync::RwLock; +use tokio::sync::{RwLock, mpsc}; use crate::channels::{Channel, IncomingMessage, MessageStream, OutgoingResponse, StatusUpdate}; use crate::error::ChannelError; /// Manages multiple input channels and merges their message streams. +/// +/// Includes an injection channel so background tasks (e.g., job monitors) can +/// push messages into the agent loop without being a full `Channel` impl. pub struct ChannelManager { channels: Arc>>>, + inject_tx: mpsc::Sender, + /// Taken once in `start_all()` and merged into the stream. + inject_rx: tokio::sync::Mutex>>, } impl ChannelManager { /// Create a new channel manager. pub fn new() -> Self { + let (inject_tx, inject_rx) = mpsc::channel(64); Self { channels: Arc::new(RwLock::new(HashMap::new())), + inject_tx, + inject_rx: tokio::sync::Mutex::new(Some(inject_rx)), } } + /// Get a clone of the injection sender. + /// + /// Background tasks (like job monitors) use this to push messages into the + /// agent loop without being a full `Channel` implementation. + pub fn inject_sender(&self) -> mpsc::Sender { + self.inject_tx.clone() + } + /// Add a channel to the manager. pub fn add(&mut self, channel: Box) { let name = channel.name().to_string(); @@ -36,9 +53,12 @@ impl ChannelManager { } /// Start all channels and return a merged stream of messages. + /// + /// Also merges the injection channel so background tasks can push messages + /// into the same stream. pub async fn start_all(&self) -> Result { let channels = self.channels.read().await; - let mut streams = Vec::new(); + let mut streams: Vec = Vec::new(); for (name, channel) in channels.iter() { match channel.start().await { @@ -60,6 +80,13 @@ impl ChannelManager { }); } + // Take the injection receiver (can only be taken once) + if let Some(inject_rx) = self.inject_rx.lock().await.take() { + let inject_stream = tokio_stream::wrappers::ReceiverStream::new(inject_rx); + streams.push(Box::pin(inject_stream)); + tracing::debug!("Injection channel merged into message stream"); + } + // Merge all streams into one let merged = stream::select_all(streams); Ok(Box::pin(merged)) diff --git a/src/channels/web/server.rs b/src/channels/web/server.rs index e199d3a6..04068b4f 100644 --- a/src/channels/web/server.rs +++ b/src/channels/web/server.rs @@ -1293,6 +1293,7 @@ async fn jobs_restart_handler( created_at: now, started_at: None, completed_at: None, + credential_grants_json: old_job.credential_grants_json.clone(), }; store .save_sandbox_job(&record) @@ -1305,9 +1306,28 @@ async fn jobs_restart_handler( _ => crate::orchestrator::job_manager::JobMode::Worker, }; + // Restore credential grants from the original job so the restarted container + // has access to the same secrets. + let credential_grants: Vec = + serde_json::from_str(&old_job.credential_grants_json).unwrap_or_else(|e| { + tracing::warn!( + job_id = %old_job.id, + "Failed to deserialize credential grants from stored job: {}. \ + Restarted job will have no credentials.", + e + ); + vec![] + }); + let project_dir = std::path::PathBuf::from(&old_job.project_dir); let _token = jm - .create_job(new_job_id, &old_job.task, Some(project_dir), mode) + .create_job( + new_job_id, + &old_job.task, + Some(project_dir), + mode, + credential_grants, + ) .await .map_err(|e| { ( @@ -1403,7 +1423,7 @@ async fn jobs_events_handler( } let events = store - .list_job_events(job_id) + .list_job_events(job_id, None) .await .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; diff --git a/src/config.rs b/src/config.rs index ca5c9dc6..5ea636d8 100644 --- a/src/config.rs +++ b/src/config.rs @@ -533,7 +533,7 @@ pub struct NearAiConfig { /// Cheap/fast model for lightweight tasks (heartbeat, routing, evaluation). /// Falls back to the main model if not set. pub cheap_model: Option, - /// Base URL for the NEAR AI API (default: https://api.near.ai) + /// Base URL for the NEAR AI API (default: https://private.near.ai). pub base_url: String, /// Base URL for auth/refresh endpoints (default: https://private.near.ai) pub auth_base_url: String, @@ -621,7 +621,7 @@ impl LlmConfig { }), cheap_model: optional_env("NEARAI_CHEAP_MODEL")?, base_url: optional_env("NEARAI_BASE_URL")? - .unwrap_or_else(|| "https://cloud-api.near.ai".to_string()), + .unwrap_or_else(|| "https://private.near.ai".to_string()), auth_base_url: optional_env("NEARAI_AUTH_URL")? .unwrap_or_else(|| "https://private.near.ai".to_string()), session_path: optional_env("NEARAI_SESSION_PATH")? @@ -1463,7 +1463,8 @@ impl SandboxModeConfig { pub struct ClaudeCodeConfig { /// Whether Claude Code sandbox mode is available. pub enabled: bool, - /// Host directory containing Claude auth session (mounted read-only). + /// Host directory containing Claude auth config (not mounted into containers; + /// auth is handled via ANTHROPIC_API_KEY env var instead). pub config_dir: std::path::PathBuf, /// Claude model to use (e.g. "sonnet", "opus"). pub model: String, @@ -1490,13 +1491,19 @@ pub struct ClaudeCodeConfig { /// silently auto-approved. fn default_claude_code_allowed_tools() -> Vec { [ - "Bash(*)", - "Read", + // File system -- glob patterns match Claude Code's settings.json format + "Read(*)", + "Write(*)", "Edit(*)", - "Glob", - "Grep", - "WebFetch(*)", + "Glob(*)", + "Grep(*)", + "NotebookEdit(*)", + // Execution + "Bash(*)", "Task(*)", + // Network + "WebFetch(*)", + "WebSearch(*)", ] .into_iter() .map(String::from) @@ -1531,6 +1538,50 @@ impl ClaudeCodeConfig { } } + /// Extract the OAuth access token from the host's credential store. + /// + /// On macOS: reads from Keychain (`Claude Code-credentials` service). + /// On Linux: reads from `~/.claude/.credentials.json`. + /// + /// Returns the access token if found. The token typically expires in + /// 8-12 hours, which is sufficient for any single container job. + pub fn extract_oauth_token() -> Option { + // macOS: extract from Keychain + if cfg!(target_os = "macos") { + match std::process::Command::new("security") + .args([ + "find-generic-password", + "-s", + "Claude Code-credentials", + "-w", + ]) + .output() + { + Ok(output) if output.status.success() => { + if let Ok(json) = String::from_utf8(output.stdout) { + return parse_oauth_access_token(json.trim()); + } + } + Ok(_) => { + tracing::debug!("No Claude Code credentials in macOS Keychain"); + } + Err(e) => { + tracing::debug!("Failed to query macOS Keychain: {e}"); + } + } + } + + // Linux / fallback: read from ~/.claude/.credentials.json + if let Some(home) = dirs::home_dir() { + let creds_path = home.join(".claude").join(".credentials.json"); + if let Ok(json) = std::fs::read_to_string(&creds_path) { + return parse_oauth_access_token(&json); + } + } + + None + } + fn resolve() -> Result { let defaults = Self::default(); Ok(Self { @@ -1563,6 +1614,16 @@ impl ClaudeCodeConfig { } } +/// Parse the OAuth access token from a Claude Code credentials JSON blob. +/// +/// Expected shape: `{"claudeAiOauth": {"accessToken": "sk-ant-oat01-..."}}` +fn parse_oauth_access_token(json: &str) -> Option { + let creds: serde_json::Value = serde_json::from_str(json).ok()?; + creds["claudeAiOauth"]["accessToken"] + .as_str() + .map(String::from) +} + /// Skills system configuration. #[derive(Debug, Clone)] pub struct SkillsConfig { diff --git a/src/context/state.rs b/src/context/state.rs index 5d008d17..66eaca8d 100644 --- a/src/context/state.rs +++ b/src/context/state.rs @@ -1,5 +1,7 @@ //! Job state machine. +use std::collections::HashMap; +use std::sync::Arc; use std::time::Duration; use chrono::{DateTime, Utc}; @@ -135,6 +137,15 @@ pub struct JobContext { pub transitions: Vec, /// Metadata. pub metadata: serde_json::Value, + /// Extra environment variables to inject into spawned child processes. + /// + /// Used by the worker runtime to pass fetched credentials to tools + /// (e.g., shell commands) without mutating the global process environment + /// via `std::env::set_var`, which is unsafe in multi-threaded programs. + /// + /// Wrapped in `Arc` for cheap cloning on every tool invocation. + #[serde(skip)] + pub extra_env: Arc>, } impl JobContext { @@ -170,6 +181,7 @@ impl JobContext { completed_at: None, repair_attempts: 0, transitions: Vec::new(), + extra_env: Arc::new(HashMap::new()), metadata: serde_json::Value::Null, } } diff --git a/src/db/libsql_backend.rs b/src/db/libsql_backend.rs index 8dde3ad3..31f79e77 100644 --- a/src/db/libsql_backend.rs +++ b/src/db/libsql_backend.rs @@ -116,10 +116,19 @@ impl LibSqlBackend { } /// Create a new connection to the database. - pub fn connect(&self) -> Result { - self.db + /// + /// Sets `PRAGMA busy_timeout = 5000` on every connection so concurrent + /// writers wait up to 5 seconds instead of failing instantly with + /// "database is locked". + pub async fn connect(&self) -> Result { + let conn = self + .db .connect() - .map_err(|e| DatabaseError::Pool(format!("Failed to create connection: {}", e))) + .map_err(|e| DatabaseError::Pool(format!("Failed to create connection: {}", e)))?; + conn.query("PRAGMA busy_timeout = 5000", ()) + .await + .map_err(|e| DatabaseError::Pool(format!("Failed to set busy_timeout: {}", e)))?; + Ok(conn) } } @@ -276,7 +285,12 @@ fn get_opt_ts(row: &libsql::Row, idx: i32) -> Option> { #[async_trait] impl Database for LibSqlBackend { async fn run_migrations(&self) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; + // WAL mode persists in the database file: all future connections benefit. + // Readers no longer block writers and vice versa. + conn.query("PRAGMA journal_mode=WAL", ()) + .await + .map_err(|e| DatabaseError::Migration(format!("Failed to enable WAL mode: {}", e)))?; conn.execute_batch(libsql_migrations::SCHEMA) .await .map_err(|e| DatabaseError::Migration(format!("libSQL migration failed: {}", e)))?; @@ -291,7 +305,7 @@ impl Database for LibSqlBackend { user_id: &str, thread_id: Option<&str>, ) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let id = Uuid::new_v4(); conn.execute( "INSERT INTO conversations (id, channel, user_id, thread_id) VALUES (?1, ?2, ?3, ?4)", @@ -303,7 +317,7 @@ impl Database for LibSqlBackend { } async fn touch_conversation(&self, id: Uuid) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let now = fmt_ts(&Utc::now()); conn.execute( "UPDATE conversations SET last_activity = ?2 WHERE id = ?1", @@ -320,7 +334,7 @@ impl Database for LibSqlBackend { role: &str, content: &str, ) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let id = Uuid::new_v4(); conn.execute( "INSERT INTO conversation_messages (id, conversation_id, role, content) VALUES (?1, ?2, ?3, ?4)", @@ -339,7 +353,7 @@ impl Database for LibSqlBackend { user_id: &str, thread_id: Option<&str>, ) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let now = fmt_ts(&Utc::now()); conn.execute( r#" @@ -360,7 +374,7 @@ impl Database for LibSqlBackend { channel: &str, limit: i64, ) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( r#" @@ -418,7 +432,7 @@ impl Database for LibSqlBackend { user_id: &str, channel: &str, ) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; // Try to find existing let mut rows = conn .query( @@ -462,7 +476,7 @@ impl Database for LibSqlBackend { user_id: &str, metadata: &serde_json::Value, ) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let id = Uuid::new_v4(); conn.execute( "INSERT INTO conversations (id, channel, user_id, metadata) VALUES (?1, ?2, ?3, ?4)", @@ -479,7 +493,7 @@ impl Database for LibSqlBackend { before: Option>, limit: i64, ) -> Result<(Vec, bool), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let fetch_limit = limit + 1; let cid = conversation_id.to_string(); @@ -536,7 +550,7 @@ impl Database for LibSqlBackend { key: &str, value: &serde_json::Value, ) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; // SQLite: use json_patch to merge the key let patch = serde_json::json!({ key: value }); conn.execute( @@ -552,7 +566,7 @@ impl Database for LibSqlBackend { &self, id: Uuid, ) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( "SELECT metadata FROM conversations WHERE id = ?1", @@ -575,7 +589,7 @@ impl Database for LibSqlBackend { &self, conversation_id: Uuid, ) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( r#" @@ -610,7 +624,7 @@ impl Database for LibSqlBackend { conversation_id: Uuid, user_id: &str, ) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( "SELECT 1 FROM conversations WHERE id = ?1 AND user_id = ?2", @@ -628,7 +642,7 @@ impl Database for LibSqlBackend { // ==================== Jobs ==================== async fn save_job(&self, ctx: &JobContext) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let status = ctx.state.to_string(); let estimated_time_secs = ctx.estimated_duration.map(|d| d.as_secs() as i64); @@ -678,7 +692,7 @@ impl Database for LibSqlBackend { } async fn get_job(&self, id: Uuid) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( r#" @@ -725,6 +739,7 @@ impl Database for LibSqlBackend { completed_at: get_opt_ts(&row, 16), transitions: Vec::new(), metadata: serde_json::Value::Null, + extra_env: std::sync::Arc::new(std::collections::HashMap::new()), })) } None => Ok(None), @@ -737,7 +752,7 @@ impl Database for LibSqlBackend { status: JobState, failure_reason: Option<&str>, ) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; conn.execute( "UPDATE agent_jobs SET status = ?2, failure_reason = ?3 WHERE id = ?1", params![id.to_string(), status.to_string(), opt_text(failure_reason)], @@ -748,7 +763,7 @@ impl Database for LibSqlBackend { } async fn mark_job_stuck(&self, id: Uuid) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let now = fmt_ts(&Utc::now()); conn.execute( "UPDATE agent_jobs SET status = 'stuck', stuck_since = ?2 WHERE id = ?1", @@ -760,7 +775,7 @@ impl Database for LibSqlBackend { } async fn get_stuck_jobs(&self) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query("SELECT id FROM agent_jobs WHERE status = 'stuck'", ()) .await @@ -784,7 +799,7 @@ impl Database for LibSqlBackend { // ==================== Actions ==================== async fn save_action(&self, job_id: Uuid, action: &ActionRecord) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let duration_ms = action.duration.as_millis() as i64; let warnings_json = serde_json::to_string(&action.sanitization_warnings) .map_err(|e| DatabaseError::Serialization(e.to_string()))?; @@ -818,7 +833,7 @@ impl Database for LibSqlBackend { } async fn get_job_actions(&self, job_id: Uuid) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( r#" @@ -860,7 +875,7 @@ impl Database for LibSqlBackend { // ==================== LLM Calls ==================== async fn record_llm_call(&self, record: &LlmCallRecord<'_>) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let id = Uuid::new_v4(); conn.execute( r#" @@ -895,7 +910,7 @@ impl Database for LibSqlBackend { estimated_time_secs: i32, estimated_value: Decimal, ) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let id = Uuid::new_v4(); let tools_json = serde_json::to_string(tool_names) .map_err(|e| DatabaseError::Serialization(e.to_string()))?; @@ -927,7 +942,7 @@ impl Database for LibSqlBackend { actual_time_secs: i32, actual_value: Option, ) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; conn.execute( "UPDATE estimation_snapshots SET actual_cost = ?2, actual_time_secs = ?3, actual_value = ?4 WHERE id = ?1", params![ @@ -945,13 +960,13 @@ impl Database for LibSqlBackend { // ==================== Sandbox Jobs ==================== async fn save_sandbox_job(&self, job: &SandboxJobRecord) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; conn.execute( r#" INSERT INTO agent_jobs ( id, title, description, status, source, user_id, project_dir, success, failure_reason, created_at, started_at, completed_at - ) VALUES (?1, ?2, '', ?3, 'sandbox', ?4, ?5, ?6, ?7, ?8, ?9, ?10) + ) VALUES (?1, ?2, ?3, ?4, 'sandbox', ?5, ?6, ?7, ?8, ?9, ?10, ?11) ON CONFLICT (id) DO UPDATE SET status = excluded.status, success = excluded.success, @@ -962,6 +977,7 @@ impl Database for LibSqlBackend { params![ job.id.to_string(), job.task.as_str(), + job.credential_grants_json.as_str(), job.status.as_str(), job.user_id.as_str(), job.project_dir.as_str(), @@ -978,11 +994,11 @@ impl Database for LibSqlBackend { } async fn get_sandbox_job(&self, id: Uuid) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( r#" - SELECT id, title, status, user_id, project_dir, + SELECT id, title, description, status, user_id, project_dir, success, failure_reason, created_at, started_at, completed_at FROM agent_jobs WHERE id = ?1 AND source = 'sandbox' "#, @@ -999,25 +1015,26 @@ impl Database for LibSqlBackend { Some(row) => Ok(Some(SandboxJobRecord { id: get_text(&row, 0).parse().unwrap_or_default(), task: get_text(&row, 1), - status: get_text(&row, 2), - user_id: get_text(&row, 3), - project_dir: get_text(&row, 4), - success: get_opt_bool(&row, 5), - failure_reason: get_opt_text(&row, 6), - created_at: get_ts(&row, 7), - started_at: get_opt_ts(&row, 8), - completed_at: get_opt_ts(&row, 9), + credential_grants_json: get_text(&row, 2), + status: get_text(&row, 3), + user_id: get_text(&row, 4), + project_dir: get_text(&row, 5), + success: get_opt_bool(&row, 6), + failure_reason: get_opt_text(&row, 7), + created_at: get_ts(&row, 8), + started_at: get_opt_ts(&row, 9), + completed_at: get_opt_ts(&row, 10), })), None => Ok(None), } } async fn list_sandbox_jobs(&self) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( r#" - SELECT id, title, status, user_id, project_dir, + SELECT id, title, description, status, user_id, project_dir, success, failure_reason, created_at, started_at, completed_at FROM agent_jobs WHERE source = 'sandbox' ORDER BY created_at DESC @@ -1036,14 +1053,15 @@ impl Database for LibSqlBackend { jobs.push(SandboxJobRecord { id: get_text(&row, 0).parse().unwrap_or_default(), task: get_text(&row, 1), - status: get_text(&row, 2), - user_id: get_text(&row, 3), - project_dir: get_text(&row, 4), - success: get_opt_bool(&row, 5), - failure_reason: get_opt_text(&row, 6), - created_at: get_ts(&row, 7), - started_at: get_opt_ts(&row, 8), - completed_at: get_opt_ts(&row, 9), + credential_grants_json: get_text(&row, 2), + status: get_text(&row, 3), + user_id: get_text(&row, 4), + project_dir: get_text(&row, 5), + success: get_opt_bool(&row, 6), + failure_reason: get_opt_text(&row, 7), + created_at: get_ts(&row, 8), + started_at: get_opt_ts(&row, 9), + completed_at: get_opt_ts(&row, 10), }); } Ok(jobs) @@ -1058,7 +1076,7 @@ impl Database for LibSqlBackend { started_at: Option>, completed_at: Option>, ) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; conn.execute( r#" UPDATE agent_jobs SET @@ -1084,7 +1102,7 @@ impl Database for LibSqlBackend { } async fn cleanup_stale_sandbox_jobs(&self) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let now = fmt_ts(&Utc::now()); let count = conn .execute( @@ -1106,7 +1124,7 @@ impl Database for LibSqlBackend { } async fn sandbox_job_summary(&self) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( "SELECT status, COUNT(*) as cnt FROM agent_jobs WHERE source = 'sandbox' GROUP BY status", @@ -1140,11 +1158,11 @@ impl Database for LibSqlBackend { &self, user_id: &str, ) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( r#" - SELECT id, title, status, user_id, project_dir, + SELECT id, title, description, status, user_id, project_dir, success, failure_reason, created_at, started_at, completed_at FROM agent_jobs WHERE source = 'sandbox' AND user_id = ?1 ORDER BY created_at DESC @@ -1163,14 +1181,15 @@ impl Database for LibSqlBackend { jobs.push(SandboxJobRecord { id: get_text(&row, 0).parse().unwrap_or_default(), task: get_text(&row, 1), - status: get_text(&row, 2), - user_id: get_text(&row, 3), - project_dir: get_text(&row, 4), - success: get_opt_bool(&row, 5), - failure_reason: get_opt_text(&row, 6), - created_at: get_ts(&row, 7), - started_at: get_opt_ts(&row, 8), - completed_at: get_opt_ts(&row, 9), + credential_grants_json: get_text(&row, 2), + status: get_text(&row, 3), + user_id: get_text(&row, 4), + project_dir: get_text(&row, 5), + success: get_opt_bool(&row, 6), + failure_reason: get_opt_text(&row, 7), + created_at: get_ts(&row, 8), + started_at: get_opt_ts(&row, 9), + completed_at: get_opt_ts(&row, 10), }); } Ok(jobs) @@ -1180,7 +1199,7 @@ impl Database for LibSqlBackend { &self, user_id: &str, ) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( "SELECT status, COUNT(*) as cnt FROM agent_jobs WHERE source = 'sandbox' AND user_id = ?1 GROUP BY status", @@ -1215,7 +1234,7 @@ impl Database for LibSqlBackend { job_id: Uuid, user_id: &str, ) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( "SELECT 1 FROM agent_jobs WHERE id = ?1 AND user_id = ?2 AND source = 'sandbox'", @@ -1231,7 +1250,7 @@ impl Database for LibSqlBackend { } async fn update_sandbox_job_mode(&self, id: Uuid, mode: &str) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; conn.execute( "UPDATE agent_jobs SET job_mode = ?2 WHERE id = ?1", params![id.to_string(), mode], @@ -1242,7 +1261,7 @@ impl Database for LibSqlBackend { } async fn get_sandbox_job_mode(&self, id: Uuid) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( "SELECT job_mode FROM agent_jobs WHERE id = ?1", @@ -1269,7 +1288,7 @@ impl Database for LibSqlBackend { event_type: &str, data: &serde_json::Value, ) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; conn.execute( "INSERT INTO job_events (job_id, event_type, data) VALUES (?1, ?2, ?3)", params![job_id.to_string(), event_type, data.to_string()], @@ -1279,10 +1298,30 @@ impl Database for LibSqlBackend { Ok(()) } - async fn list_job_events(&self, job_id: Uuid) -> Result, DatabaseError> { - let conn = self.connect()?; - let mut rows = conn - .query( + async fn list_job_events( + &self, + job_id: Uuid, + limit: Option, + ) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = if let Some(n) = limit { + conn.query( + r#" + SELECT id, job_id, event_type, data, created_at + FROM ( + SELECT id, job_id, event_type, data, created_at + FROM job_events WHERE job_id = ?1 + ORDER BY id DESC + LIMIT ?2 + ) + ORDER BY id ASC + "#, + params![job_id.to_string(), n], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + } else { + conn.query( r#" SELECT id, job_id, event_type, data, created_at FROM job_events WHERE job_id = ?1 ORDER BY id ASC @@ -1290,7 +1329,8 @@ impl Database for LibSqlBackend { params![job_id.to_string()], ) .await - .map_err(|e| DatabaseError::Query(e.to_string()))?; + .map_err(|e| DatabaseError::Query(e.to_string()))? + }; let mut events = Vec::new(); while let Some(row) = rows @@ -1312,7 +1352,7 @@ impl Database for LibSqlBackend { // ==================== Routines ==================== async fn create_routine(&self, routine: &Routine) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let trigger_type = routine.trigger.type_tag(); let trigger_config = routine.trigger.to_config_json(); let action_type = routine.action.type_tag(); @@ -1367,7 +1407,7 @@ impl Database for LibSqlBackend { } async fn get_routine(&self, id: Uuid) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( &format!("SELECT {} FROM routines WHERE id = ?1", ROUTINE_COLUMNS), @@ -1391,7 +1431,7 @@ impl Database for LibSqlBackend { user_id: &str, name: &str, ) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( &format!( @@ -1414,7 +1454,7 @@ impl Database for LibSqlBackend { } async fn list_routines(&self, user_id: &str) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( &format!( @@ -1438,7 +1478,7 @@ impl Database for LibSqlBackend { } async fn list_event_routines(&self) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( &format!( @@ -1462,7 +1502,7 @@ impl Database for LibSqlBackend { } async fn list_due_cron_routines(&self) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let now = fmt_ts(&Utc::now()); let mut rows = conn .query( @@ -1487,7 +1527,7 @@ impl Database for LibSqlBackend { } async fn update_routine(&self, routine: &Routine) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let trigger_type = routine.trigger.type_tag(); let trigger_config = routine.trigger.to_config_json(); let action_type = routine.action.type_tag(); @@ -1546,7 +1586,7 @@ impl Database for LibSqlBackend { consecutive_failures: u32, state: &serde_json::Value, ) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let now = fmt_ts(&Utc::now()); conn.execute( r#" @@ -1572,7 +1612,7 @@ impl Database for LibSqlBackend { } async fn delete_routine(&self, id: Uuid) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let count = conn .execute( "DELETE FROM routines WHERE id = ?1", @@ -1586,7 +1626,7 @@ impl Database for LibSqlBackend { // ==================== Routine Runs ==================== async fn create_routine_run(&self, run: &RoutineRun) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; conn.execute( r#" INSERT INTO routine_runs ( @@ -1616,7 +1656,7 @@ impl Database for LibSqlBackend { result_summary: Option<&str>, tokens_used: Option, ) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let now = fmt_ts(&Utc::now()); conn.execute( r#" @@ -1643,7 +1683,7 @@ impl Database for LibSqlBackend { routine_id: Uuid, limit: i64, ) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( &format!( @@ -1667,7 +1707,7 @@ impl Database for LibSqlBackend { } async fn count_running_routine_runs(&self, routine_id: Uuid) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( "SELECT COUNT(*) as cnt FROM routine_runs WHERE routine_id = ?1 AND status = 'running'", @@ -1693,7 +1733,7 @@ impl Database for LibSqlBackend { tool_name: &str, error_message: &str, ) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let now = fmt_ts(&Utc::now()); conn.execute( r#" @@ -1712,7 +1752,7 @@ impl Database for LibSqlBackend { } async fn get_broken_tools(&self, threshold: i32) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( r#" @@ -1748,7 +1788,7 @@ impl Database for LibSqlBackend { } async fn mark_tool_repaired(&self, tool_name: &str) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let now = fmt_ts(&Utc::now()); conn.execute( "UPDATE tool_failures SET repaired_at = ?2, error_count = 0 WHERE tool_name = ?1", @@ -1760,7 +1800,7 @@ impl Database for LibSqlBackend { } async fn increment_repair_attempts(&self, tool_name: &str) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; conn.execute( "UPDATE tool_failures SET repair_attempts = repair_attempts + 1 WHERE tool_name = ?1", params![tool_name], @@ -1777,7 +1817,7 @@ impl Database for LibSqlBackend { user_id: &str, key: &str, ) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( "SELECT value FROM settings WHERE user_id = ?1 AND key = ?2", @@ -1801,7 +1841,7 @@ impl Database for LibSqlBackend { user_id: &str, key: &str, ) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( "SELECT key, value, updated_at FROM settings WHERE user_id = ?1 AND key = ?2", @@ -1830,7 +1870,7 @@ impl Database for LibSqlBackend { key: &str, value: &serde_json::Value, ) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let now = fmt_ts(&Utc::now()); conn.execute( r#" @@ -1848,7 +1888,7 @@ impl Database for LibSqlBackend { } async fn delete_setting(&self, user_id: &str, key: &str) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let count = conn .execute( "DELETE FROM settings WHERE user_id = ?1 AND key = ?2", @@ -1860,7 +1900,7 @@ impl Database for LibSqlBackend { } async fn list_settings(&self, user_id: &str) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( "SELECT key, value, updated_at FROM settings WHERE user_id = ?1 ORDER BY key", @@ -1888,7 +1928,7 @@ impl Database for LibSqlBackend { &self, user_id: &str, ) -> Result, DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( "SELECT key, value FROM settings WHERE user_id = ?1", @@ -1913,7 +1953,7 @@ impl Database for LibSqlBackend { user_id: &str, settings: &HashMap, ) -> Result<(), DatabaseError> { - let conn = self.connect()?; + let conn = self.connect().await?; let now = fmt_ts(&Utc::now()); conn.execute("BEGIN", ()) .await @@ -1945,7 +1985,7 @@ impl Database for LibSqlBackend { } async fn has_settings(&self, user_id: &str) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( "SELECT COUNT(*) as cnt FROM settings WHERE user_id = ?1", @@ -1972,9 +2012,12 @@ impl Database for LibSqlBackend { agent_id: Option, path: &str, ) -> Result { - let conn = self.connect().map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; let agent_id_str = agent_id.map(|id| id.to_string()); let mut rows = conn .query( @@ -2006,9 +2049,12 @@ impl Database for LibSqlBackend { } async fn get_document_by_id(&self, id: Uuid) -> Result { - let conn = self.connect().map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; let mut rows = conn .query( r#" @@ -2051,9 +2097,12 @@ impl Database for LibSqlBackend { } // Create - let conn = self.connect().map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; let id = Uuid::new_v4(); let agent_id_str = agent_id.map(|id| id.to_string()); conn.execute( @@ -2073,9 +2122,12 @@ impl Database for LibSqlBackend { } async fn update_document(&self, id: Uuid, content: &str) -> Result<(), WorkspaceError> { - let conn = self.connect().map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; let now = fmt_ts(&Utc::now()); conn.execute( "UPDATE memory_documents SET content = ?2, updated_at = ?3 WHERE id = ?1", @@ -2097,9 +2149,12 @@ impl Database for LibSqlBackend { let doc = self.get_document_by_path(user_id, agent_id, path).await?; self.delete_chunks(doc.id).await?; - let conn = self.connect().map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; let agent_id_str = agent_id.map(|id| id.to_string()); conn.execute( "DELETE FROM memory_documents WHERE user_id = ?1 AND agent_id IS ?2 AND path = ?3", @@ -2118,9 +2173,12 @@ impl Database for LibSqlBackend { agent_id: Option, directory: &str, ) -> Result, WorkspaceError> { - let conn = self.connect().map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; // Implement the list_workspace_files logic in Rust instead of PL/pgSQL. let dir = if !directory.is_empty() && !directory.ends_with('/') { format!("{}/", directory) @@ -2223,9 +2281,12 @@ impl Database for LibSqlBackend { user_id: &str, agent_id: Option, ) -> Result, WorkspaceError> { - let conn = self.connect().map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; let agent_id_str = agent_id.map(|id| id.to_string()); let mut rows = conn .query( @@ -2255,9 +2316,12 @@ impl Database for LibSqlBackend { user_id: &str, agent_id: Option, ) -> Result, WorkspaceError> { - let conn = self.connect().map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; let agent_id_str = agent_id.map(|id| id.to_string()); let mut rows = conn .query( @@ -2291,9 +2355,12 @@ impl Database for LibSqlBackend { // ==================== Workspace: Chunks ==================== async fn delete_chunks(&self, document_id: Uuid) -> Result<(), WorkspaceError> { - let conn = self.connect().map_err(|e| WorkspaceError::ChunkingFailed { - reason: e.to_string(), - })?; + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::ChunkingFailed { + reason: e.to_string(), + })?; conn.execute( "DELETE FROM memory_chunks WHERE document_id = ?1", params![document_id.to_string()], @@ -2312,9 +2379,12 @@ impl Database for LibSqlBackend { content: &str, embedding: Option<&[f32]>, ) -> Result { - let conn = self.connect().map_err(|e| WorkspaceError::ChunkingFailed { - reason: e.to_string(), - })?; + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::ChunkingFailed { + reason: e.to_string(), + })?; let id = Uuid::new_v4(); let embedding_blob = embedding.map(|e| { // Convert f32 slice to bytes for F32_BLOB @@ -2349,6 +2419,7 @@ impl Database for LibSqlBackend { ) -> Result<(), WorkspaceError> { let conn = self .connect() + .await .map_err(|e| WorkspaceError::EmbeddingFailed { reason: e.to_string(), })?; @@ -2371,9 +2442,12 @@ impl Database for LibSqlBackend { agent_id: Option, limit: usize, ) -> Result, WorkspaceError> { - let conn = self.connect().map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; let agent_id_str = agent_id.map(|id| id.to_string()); let mut rows = conn .query( @@ -2422,9 +2496,12 @@ impl Database for LibSqlBackend { embedding: Option<&[f32]>, config: &SearchConfig, ) -> Result, WorkspaceError> { - let conn = self.connect().map_err(|e| WorkspaceError::SearchFailed { - reason: e.to_string(), - })?; + let conn = self + .connect() + .await + .map_err(|e| WorkspaceError::SearchFailed { + reason: e.to_string(), + })?; let agent_id_str = agent_id.map(|id| id.to_string()); let pre_limit = config.pre_fusion_limit as i64; @@ -2607,3 +2684,86 @@ fn row_to_routine_run_libsql(row: &libsql::Row) -> Result Result<(), DatabaseError>; - /// Load all job events. - async fn list_job_events(&self, job_id: Uuid) -> Result, DatabaseError>; + /// Load job events, returning the most recent `limit` entries (or all if `None`). + async fn list_job_events( + &self, + job_id: Uuid, + limit: Option, + ) -> Result, DatabaseError>; // ==================== Routines ==================== diff --git a/src/db/postgres.rs b/src/db/postgres.rs index 9676144c..096f3a95 100644 --- a/src/db/postgres.rs +++ b/src/db/postgres.rs @@ -334,8 +334,12 @@ impl Database for PgBackend { self.store.save_job_event(job_id, event_type, data).await } - async fn list_job_events(&self, job_id: Uuid) -> Result, DatabaseError> { - self.store.list_job_events(job_id).await + async fn list_job_events( + &self, + job_id: Uuid, + limit: Option, + ) -> Result, DatabaseError> { + self.store.list_job_events(job_id, limit).await } // ==================== Routines ==================== diff --git a/src/error.rs b/src/error.rs index 5aa7d461..ba3b1ccf 100644 --- a/src/error.rs +++ b/src/error.rs @@ -336,17 +336,11 @@ pub enum OrchestratorError { #[error("Container for job {job_id} is in unexpected state: {state}")] InvalidContainerState { job_id: Uuid, state: String }, - #[error("Worker authentication failed: {reason}")] - AuthFailed { reason: String }, - #[error("Internal API error: {reason}")] ApiError { reason: String }, #[error("Docker error: {reason}")] Docker { reason: String }, - - #[error("Job {job_id} timed out in container")] - ContainerTimeout { job_id: Uuid }, } /// Worker errors (container-side execution). diff --git a/src/history/store.rs b/src/history/store.rs index 94812908..e97ec015 100644 --- a/src/history/store.rs +++ b/src/history/store.rs @@ -239,6 +239,7 @@ impl Store { metadata: serde_json::Value::Null, total_tokens_used: 0, max_tokens: 0, + extra_env: std::sync::Arc::new(std::collections::HashMap::new()), })) } None => Ok(None), @@ -470,6 +471,9 @@ pub struct SandboxJobRecord { pub created_at: DateTime, pub started_at: Option>, pub completed_at: Option>, + /// Serialized JSON of `Vec` for restart support. + /// Stored in the `description` column of `agent_jobs` (unused for sandbox jobs). + pub credential_grants_json: String, } /// Summary of sandbox job counts grouped by status. @@ -493,7 +497,7 @@ impl Store { INSERT INTO agent_jobs ( id, title, description, status, source, user_id, project_dir, success, failure_reason, created_at, started_at, completed_at - ) VALUES ($1, $2, '', $3, 'sandbox', $4, $5, $6, $7, $8, $9, $10) + ) VALUES ($1, $2, $3, $4, 'sandbox', $5, $6, $7, $8, $9, $10, $11) ON CONFLICT (id) DO UPDATE SET status = EXCLUDED.status, success = EXCLUDED.success, @@ -504,6 +508,7 @@ impl Store { &[ &job.id, &job.task, + &job.credential_grants_json, &job.status, &job.user_id, &job.project_dir, @@ -527,7 +532,7 @@ impl Store { let row = conn .query_opt( r#" - SELECT id, title, status, user_id, project_dir, + SELECT id, title, description, status, user_id, project_dir, success, failure_reason, created_at, started_at, completed_at FROM agent_jobs WHERE id = $1 AND source = 'sandbox' "#, @@ -548,6 +553,7 @@ impl Store { created_at: r.get("created_at"), started_at: r.get("started_at"), completed_at: r.get("completed_at"), + credential_grants_json: r.get::<_, String>("description"), })) } @@ -557,7 +563,7 @@ impl Store { let rows = conn .query( r#" - SELECT id, title, status, user_id, project_dir, + SELECT id, title, description, status, user_id, project_dir, success, failure_reason, created_at, started_at, completed_at FROM agent_jobs WHERE source = 'sandbox' ORDER BY created_at DESC @@ -581,6 +587,7 @@ impl Store { created_at: r.get("created_at"), started_at: r.get("started_at"), completed_at: r.get("completed_at"), + credential_grants_json: r.get::<_, String>("description"), }) .collect()) } @@ -594,7 +601,7 @@ impl Store { let rows = conn .query( r#" - SELECT id, title, status, user_id, project_dir, + SELECT id, title, description, status, user_id, project_dir, success, failure_reason, created_at, started_at, completed_at FROM agent_jobs WHERE source = 'sandbox' AND user_id = $1 ORDER BY created_at DESC @@ -618,6 +625,7 @@ impl Store { created_at: r.get("created_at"), started_at: r.get("started_at"), completed_at: r.get("completed_at"), + credential_grants_json: r.get::<_, String>("description"), }) .collect()) } @@ -781,14 +789,35 @@ impl Store { Ok(()) } - /// Load all job events for a job, ordered by id. + /// Load job events for a job, ordered by id. + /// + /// When `limit` is `Some(n)`, returns the **most recent** `n` events + /// (ordered ascending by id). When `None`, returns all events. pub async fn list_job_events( &self, job_id: Uuid, + limit: Option, ) -> Result, DatabaseError> { let conn = self.conn().await?; - let rows = conn - .query( + let rows = if let Some(n) = limit { + // Sub-select the last N rows by id DESC, then re-sort ASC. + conn.query( + r#" + SELECT id, job_id, event_type, data, created_at + FROM ( + SELECT id, job_id, event_type, data, created_at + FROM job_events + WHERE job_id = $1 + ORDER BY id DESC + LIMIT $2 + ) sub + ORDER BY id ASC + "#, + &[&job_id, &n], + ) + .await? + } else { + conn.query( r#" SELECT id, job_id, event_type, data, created_at FROM job_events @@ -797,7 +826,8 @@ impl Store { "#, &[&job_id], ) - .await?; + .await? + }; Ok(rows .iter() .map(|r| JobEventRecord { diff --git a/src/llm/failover.rs b/src/llm/failover.rs index 8a49dcf1..c53d725f 100644 --- a/src/llm/failover.rs +++ b/src/llm/failover.rs @@ -300,6 +300,17 @@ impl LlmProvider for FailoverProvider { .await } + fn active_model_name(&self) -> String { + self.providers[self.last_used.load(Ordering::Relaxed)].active_model_name() + } + + fn set_model(&self, model: &str) -> Result<(), LlmError> { + for provider in &self.providers { + provider.set_model(model)?; + } + Ok(()) + } + async fn list_models(&self) -> Result, LlmError> { let mut all_models = Vec::new(); @@ -326,13 +337,15 @@ impl LlmProvider for FailoverProvider { mod tests { use super::*; - use std::sync::Mutex; + use std::sync::{Mutex, RwLock}; + use std::time::Duration; use crate::llm::provider::{CompletionResponse, FinishReason, ToolCompletionResponse}; /// A mock LLM provider that returns a predetermined result. struct MockProvider { name: String, + active_model: RwLock, input_cost: Decimal, output_cost: Decimal, complete_result: Mutex>>, @@ -343,6 +356,7 @@ mod tests { fn succeeding(name: &str, content: &str) -> Self { Self { name: name.to_string(), + active_model: RwLock::new(name.to_string()), input_cost: Decimal::ZERO, output_cost: Decimal::ZERO, complete_result: Mutex::new(Some(Ok(CompletionResponse { @@ -379,6 +393,7 @@ mod tests { fn failing_retryable(name: &str) -> Self { Self { name: name.to_string(), + active_model: RwLock::new(name.to_string()), input_cost: Decimal::ZERO, output_cost: Decimal::ZERO, complete_result: Mutex::new(Some(Err(LlmError::RequestFailed { @@ -395,6 +410,7 @@ mod tests { fn failing_non_retryable(name: &str) -> Self { Self { name: name.to_string(), + active_model: RwLock::new(name.to_string()), input_cost: Decimal::ZERO, output_cost: Decimal::ZERO, complete_result: Mutex::new(Some(Err(LlmError::AuthFailed { @@ -409,6 +425,7 @@ mod tests { fn failing_rate_limited(name: &str) -> Self { Self { name: name.to_string(), + active_model: RwLock::new(name.to_string()), input_cost: Decimal::ZERO, output_cost: Decimal::ZERO, complete_result: Mutex::new(Some(Err(LlmError::RateLimited { @@ -458,6 +475,15 @@ mod tests { async fn list_models(&self) -> Result, LlmError> { Ok(vec![self.name.clone()]) } + + fn active_model_name(&self) -> String { + self.active_model.read().unwrap().clone() + } + + fn set_model(&self, model: &str) -> Result<(), LlmError> { + *self.active_model.write().unwrap() = model.to_string(); + Ok(()) + } } fn make_request() -> CompletionRequest { @@ -1014,4 +1040,30 @@ mod tests { let result = FailoverProvider::new(vec![]); assert!(result.is_err()); } + + // Test: set_model propagates to all providers and active_model_name reflects change. + #[test] + fn set_model_propagates_to_all_providers() { + let p1: Arc = Arc::new(MockProvider::succeeding("model-a", "ok")); + let p2: Arc = Arc::new(MockProvider::succeeding("model-b", "ok")); + + let failover = FailoverProvider::new(vec![ + Arc::clone(&p1) as Arc, + Arc::clone(&p2) as Arc, + ]) + .unwrap(); + + // Before: active_model_name delegates to last_used (index 0 = p1). + assert_eq!(failover.active_model_name(), "model-a"); + + // Switch model. + failover.set_model("new-model").unwrap(); + + // Both inner providers should reflect the change. + assert_eq!(p1.active_model_name(), "new-model"); + assert_eq!(p2.active_model_name(), "new-model"); + + // FailoverProvider itself should report the new model. + assert_eq!(failover.active_model_name(), "new-model"); + } } diff --git a/src/llm/nearai.rs b/src/llm/nearai.rs index 05db7c36..2a995805 100644 --- a/src/llm/nearai.rs +++ b/src/llm/nearai.rs @@ -127,11 +127,24 @@ impl NearAiProvider { } /// Fetch available models from the NEAR AI API. + /// + /// Handles session renewal on 401 (same pattern as `send_request`). pub async fn list_models(&self) -> Result, LlmError> { + match self.list_models_inner().await { + Ok(models) => Ok(models), + Err(LlmError::SessionExpired { .. }) => { + self.session.handle_auth_failure().await?; + self.list_models_inner().await + } + Err(e) => Err(e), + } + } + + async fn list_models_inner(&self) -> Result, LlmError> { use secrecy::ExposeSecret; let token = self.session.get_token().await?; - let url = self.api_url("model/list"); + let url = self.api_url("models"); tracing::debug!("Fetching models from: {}", url); @@ -150,7 +163,6 @@ impl NearAiProvider { let response_text = response.text().await.unwrap_or_default(); if !status.is_success() { - // Check for session expiration if status.as_u16() == 401 { return Err(LlmError::SessionExpired { provider: "nearai".to_string(), diff --git a/src/main.rs b/src/main.rs index b5086c7c..26d37c49 100644 --- a/src/main.rs +++ b/src/main.rs @@ -247,7 +247,7 @@ async fn main() -> anyhow::Result<()> { max_turns: *max_turns, model: model.clone(), timeout: std::time::Duration::from_secs(1800), - allowed_tools: Vec::new(), + allowed_tools: ironclaw::config::ClaudeCodeConfig::from_env().allowed_tools, }; let runtime = ironclaw::worker::ClaudeBridgeRuntime::new(config) @@ -964,11 +964,8 @@ async fn main() -> anyhow::Result<()> { memory_limit_mb: config.sandbox.memory_limit_mb, cpu_shares: config.sandbox.cpu_shares, orchestrator_port: 50051, - claude_config_dir: if config.claude_code.enabled { - Some(config.claude_code.config_dir.clone()) - } else { - None - }, + claude_code_api_key: std::env::var("ANTHROPIC_API_KEY").ok(), + claude_code_oauth_token: ironclaw::config::ClaudeCodeConfig::extract_oauth_token(), claude_code_model: config.claude_code.model.clone(), claude_code_max_turns: config.claude_code.max_turns, claude_code_memory_limit_mb: config.claude_code.memory_limit_mb, @@ -984,6 +981,8 @@ async fn main() -> anyhow::Result<()> { job_event_tx: job_event_tx.clone(), prompt_queue: Arc::clone(&prompt_queue), store: db.clone(), + secrets_store: secrets_store.clone(), + user_id: "default".to_string(), }; tokio::spawn(async move { @@ -1274,6 +1273,14 @@ async fn main() -> anyhow::Result<()> { Arc::clone(&context_manager), container_job_manager.clone(), db.clone(), + job_event_tx.clone(), + Some(channels.inject_sender()), + if config.sandbox.enabled { + Some(Arc::clone(&prompt_queue)) + } else { + None + }, + secrets_store.clone(), ); // Initialize skills system (before gateway so we can wire into GatewayState) diff --git a/src/orchestrator/api.rs b/src/orchestrator/api.rs index afa5ce2a..6ffd411f 100644 --- a/src/orchestrator/api.rs +++ b/src/orchestrator/api.rs @@ -19,10 +19,11 @@ use crate::db::Database; use crate::llm::{CompletionRequest, LlmProvider, ToolCompletionRequest}; use crate::orchestrator::auth::{TokenStore, worker_auth_middleware}; use crate::orchestrator::job_manager::ContainerJobManager; +use crate::secrets::SecretsStore; use crate::worker::api::JobEventPayload; use crate::worker::api::{ - CompletionReport, JobDescription, ProxyCompletionRequest, ProxyCompletionResponse, - ProxyToolCompletionRequest, ProxyToolCompletionResponse, StatusUpdate, + CompletionReport, CredentialResponse, JobDescription, ProxyCompletionRequest, + ProxyCompletionResponse, ProxyToolCompletionRequest, ProxyToolCompletionResponse, StatusUpdate, }; /// A follow-up prompt queued for a Claude Code bridge. @@ -44,6 +45,10 @@ pub struct OrchestratorState { pub prompt_queue: Arc>>>, /// Database handle for persisting job events. pub store: Option>, + /// Encrypted secrets store for credential injection into containers. + pub secrets_store: Option>, + /// User ID for secret lookups (single-tenant, typically "default"). + pub user_id: String, } /// The orchestrator's internal API server. @@ -64,6 +69,7 @@ impl OrchestratorApi { .route("/worker/{job_id}/complete", post(report_complete)) .route("/worker/{job_id}/event", post(job_event_handler)) .route("/worker/{job_id}/prompt", get(get_prompt_handler)) + .route("/worker/{job_id}/credentials", get(get_credentials_handler)) .route_layer(axum::middleware::from_fn_with_state( state.token_store.clone(), worker_auth_middleware, @@ -185,6 +191,7 @@ async fn llm_complete_with_tools( } async fn report_status( + State(state): State, Path(job_id): Path, Json(update): Json, ) -> Result { @@ -195,6 +202,11 @@ async fn report_status( "Worker status update" ); + state + .job_manager + .update_worker_status(job_id, update.message, update.iteration) + .await; + Ok(StatusCode::OK) } @@ -221,7 +233,9 @@ async fn report_complete( success: report.success, message: report.message.clone(), }; - let _ = state.job_manager.complete_job(job_id, result).await; + if let Err(e) = state.job_manager.complete_job(job_id, result).await { + tracing::error!(job_id = %job_id, "Failed to complete job cleanup: {}", e); + } Ok(Json(serde_json::json!({"status": "ok"}))) } @@ -356,6 +370,66 @@ async fn get_prompt_handler( Ok((StatusCode::NO_CONTENT, Json(serde_json::Value::Null))) } +/// Serve decrypted credentials for a job's granted secrets. +/// +/// Returns 204 if no grants exist, 503 if no secrets store is configured, +/// or a JSON array of `{ env_var, value }` pairs. +async fn get_credentials_handler( + State(state): State, + Path(job_id): Path, +) -> Result<(StatusCode, Json), StatusCode> { + let grants = match state.token_store.get_grants(job_id).await { + Some(g) if !g.is_empty() => g, + _ => return Ok((StatusCode::NO_CONTENT, Json(serde_json::Value::Null))), + }; + + let secrets = state.secrets_store.as_ref().ok_or_else(|| { + tracing::error!("Credentials requested but no secrets store configured"); + StatusCode::SERVICE_UNAVAILABLE + })?; + + let mut credentials: Vec = Vec::with_capacity(grants.len()); + + for grant in &grants { + let decrypted = secrets + .get_decrypted(&state.user_id, &grant.secret_name) + .await + .map_err(|e| { + tracing::error!( + job_id = %job_id, + "Failed to decrypt secret for credential grant: {}", e + ); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + // Record usage for audit trail + if let Ok(secret) = secrets.get(&state.user_id, &grant.secret_name).await + && let Err(e) = secrets.record_usage(secret.id).await + { + tracing::warn!( + job_id = %job_id, + "Failed to record credential usage: {}", e + ); + } + + tracing::debug!( + job_id = %job_id, + env_var = %grant.env_var, + "Serving credential to container" + ); + + credentials.push(CredentialResponse { + env_var: grant.env_var.clone(), + value: decrypted.expose().to_string(), + }); + } + + Ok(( + StatusCode::OK, + Json(serde_json::to_value(&credentials).unwrap_or(serde_json::Value::Null)), + )) +} + fn format_finish_reason(reason: crate::llm::FinishReason) -> String { match reason { crate::llm::FinishReason::Stop => "stop".to_string(), @@ -422,6 +496,8 @@ mod tests { job_event_tx: None, prompt_queue: Arc::new(Mutex::new(HashMap::new())), store: None, + secrets_store: None, + user_id: "default".to_string(), } } @@ -508,4 +584,367 @@ mod tests { let resp = router.oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); } + + // -- Prompt queue tests -- + + #[tokio::test] + async fn prompt_returns_204_when_queue_empty() { + let state = test_state(); + let job_id = Uuid::new_v4(); + let token = state.token_store.create_token(job_id).await; + let router = OrchestratorApi::router(state); + + let req = Request::builder() + .uri(format!("/worker/{}/prompt", job_id)) + .header("Authorization", format!("Bearer {}", token)) + .body(Body::empty()) + .unwrap(); + + let resp = router.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::NO_CONTENT); + } + + #[tokio::test] + async fn prompt_returns_queued_prompt() { + let state = test_state(); + let job_id = Uuid::new_v4(); + let token = state.token_store.create_token(job_id).await; + + // Queue a prompt + { + let mut q = state.prompt_queue.lock().await; + q.entry(job_id).or_default().push_back(PendingPrompt { + content: "What is the status?".to_string(), + done: false, + }); + } + + let router = OrchestratorApi::router(state); + let req = Request::builder() + .uri(format!("/worker/{}/prompt", job_id)) + .header("Authorization", format!("Bearer {}", token)) + .body(Body::empty()) + .unwrap(); + + let resp = router.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), 4096).await.unwrap(); + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(json["content"], "What is the status?"); + assert_eq!(json["done"], false); + } + + // -- Credentials handler tests -- + + #[tokio::test] + async fn credentials_returns_204_when_no_grants() { + let state = test_state(); + let job_id = Uuid::new_v4(); + let token = state.token_store.create_token(job_id).await; + let router = OrchestratorApi::router(state); + + let req = Request::builder() + .uri(format!("/worker/{}/credentials", job_id)) + .header("Authorization", format!("Bearer {}", token)) + .body(Body::empty()) + .unwrap(); + + let resp = router.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::NO_CONTENT); + } + + #[tokio::test] + async fn credentials_returns_503_when_no_secrets_store() { + let state = test_state(); + let job_id = Uuid::new_v4(); + let token = state.token_store.create_token(job_id).await; + + // Store grants so we get past the 204 check + state + .token_store + .store_grants( + job_id, + vec![crate::orchestrator::auth::CredentialGrant { + secret_name: "test_secret".to_string(), + env_var: "TEST_SECRET".to_string(), + }], + ) + .await; + + let router = OrchestratorApi::router(state); + let req = Request::builder() + .uri(format!("/worker/{}/credentials", job_id)) + .header("Authorization", format!("Bearer {}", token)) + .body(Body::empty()) + .unwrap(); + + let resp = router.oneshot(req).await.unwrap(); + // No secrets_store configured → 503 + assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE); + } + + #[tokio::test] + async fn credentials_returns_secrets_when_store_configured() { + use secrecy::SecretString; + let key = "0123456789abcdef0123456789abcdef"; + let crypto = Arc::new( + crate::secrets::SecretsCrypto::new(SecretString::from(key.to_string())).unwrap(), + ); + let secrets_store = Arc::new(crate::secrets::InMemorySecretsStore::new(crypto)); + + // Create a secret + secrets_store + .create( + "default", + crate::secrets::CreateSecretParams { + name: "test_secret".to_string(), + value: SecretString::from("supersecretvalue".to_string()), + provider: None, + expires_at: None, + }, + ) + .await + .unwrap(); + + let token_store = TokenStore::new(); + let jm = ContainerJobManager::new(ContainerJobConfig::default(), token_store.clone()); + let job_id = Uuid::new_v4(); + let token = token_store.create_token(job_id).await; + token_store + .store_grants( + job_id, + vec![crate::orchestrator::auth::CredentialGrant { + secret_name: "test_secret".to_string(), + env_var: "MY_SECRET".to_string(), + }], + ) + .await; + + let state = OrchestratorState { + llm: Arc::new(StubLlm), + job_manager: Arc::new(jm), + token_store, + job_event_tx: None, + prompt_queue: Arc::new(Mutex::new(HashMap::new())), + store: None, + secrets_store: Some(secrets_store), + user_id: "default".to_string(), + }; + + let router = OrchestratorApi::router(state); + let req = Request::builder() + .uri(format!("/worker/{}/credentials", job_id)) + .header("Authorization", format!("Bearer {}", token)) + .body(Body::empty()) + .unwrap(); + + let resp = router.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), 4096).await.unwrap(); + let json: Vec = serde_json::from_slice(&body).unwrap(); + assert_eq!(json.len(), 1); + assert_eq!(json[0]["env_var"], "MY_SECRET"); + assert_eq!(json[0]["value"], "supersecretvalue"); + } + + // -- Job event handler tests -- + + #[tokio::test] + async fn job_event_broadcasts_message() { + let (tx, mut rx) = broadcast::channel(16); + let token_store = TokenStore::new(); + let jm = ContainerJobManager::new(ContainerJobConfig::default(), token_store.clone()); + let state = OrchestratorState { + llm: Arc::new(StubLlm), + job_manager: Arc::new(jm), + token_store: token_store.clone(), + job_event_tx: Some(tx), + prompt_queue: Arc::new(Mutex::new(HashMap::new())), + store: None, + secrets_store: None, + user_id: "default".to_string(), + }; + + let job_id = Uuid::new_v4(); + let token = token_store.create_token(job_id).await; + let router = OrchestratorApi::router(state); + + let payload = serde_json::json!({ + "event_type": "message", + "data": { + "role": "assistant", + "content": "Hello from worker" + } + }); + + let req = Request::builder() + .method("POST") + .uri(format!("/worker/{}/event", job_id)) + .header("Authorization", format!("Bearer {}", token)) + .header("Content-Type", "application/json") + .body(Body::from(serde_json::to_vec(&payload).unwrap())) + .unwrap(); + + let resp = router.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let (recv_id, event) = rx.recv().await.unwrap(); + assert_eq!(recv_id, job_id); + match event { + SseEvent::JobMessage { + job_id: jid, + role, + content, + } => { + assert_eq!(jid, job_id.to_string()); + assert_eq!(role, "assistant"); + assert_eq!(content, "Hello from worker"); + } + other => panic!("Expected JobMessage, got {:?}", other), + } + } + + #[tokio::test] + async fn job_event_handles_tool_use() { + let (tx, mut rx) = broadcast::channel(16); + let token_store = TokenStore::new(); + let jm = ContainerJobManager::new(ContainerJobConfig::default(), token_store.clone()); + let state = OrchestratorState { + llm: Arc::new(StubLlm), + job_manager: Arc::new(jm), + token_store: token_store.clone(), + job_event_tx: Some(tx), + prompt_queue: Arc::new(Mutex::new(HashMap::new())), + store: None, + secrets_store: None, + user_id: "default".to_string(), + }; + + let job_id = Uuid::new_v4(); + let token = token_store.create_token(job_id).await; + let router = OrchestratorApi::router(state); + + let payload = serde_json::json!({ + "event_type": "tool_use", + "data": { + "tool_name": "shell", + "input": {"command": "ls"} + } + }); + + let req = Request::builder() + .method("POST") + .uri(format!("/worker/{}/event", job_id)) + .header("Authorization", format!("Bearer {}", token)) + .header("Content-Type", "application/json") + .body(Body::from(serde_json::to_vec(&payload).unwrap())) + .unwrap(); + + let resp = router.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let (_recv_id, event) = rx.recv().await.unwrap(); + match event { + SseEvent::JobToolUse { tool_name, .. } => { + assert_eq!(tool_name, "shell"); + } + other => panic!("Expected JobToolUse, got {:?}", other), + } + } + + #[tokio::test] + async fn job_event_handles_unknown_type() { + let (tx, mut rx) = broadcast::channel(16); + let token_store = TokenStore::new(); + let jm = ContainerJobManager::new(ContainerJobConfig::default(), token_store.clone()); + let state = OrchestratorState { + llm: Arc::new(StubLlm), + job_manager: Arc::new(jm), + token_store: token_store.clone(), + job_event_tx: Some(tx), + prompt_queue: Arc::new(Mutex::new(HashMap::new())), + store: None, + secrets_store: None, + user_id: "default".to_string(), + }; + + let job_id = Uuid::new_v4(); + let token = token_store.create_token(job_id).await; + let router = OrchestratorApi::router(state); + + let payload = serde_json::json!({ + "event_type": "custom_thing", + "data": { "message": "something custom" } + }); + + let req = Request::builder() + .method("POST") + .uri(format!("/worker/{}/event", job_id)) + .header("Authorization", format!("Bearer {}", token)) + .header("Content-Type", "application/json") + .body(Body::from(serde_json::to_vec(&payload).unwrap())) + .unwrap(); + + let resp = router.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let (_recv_id, event) = rx.recv().await.unwrap(); + // Unknown event types fall through to JobStatus + assert!(matches!(event, SseEvent::JobStatus { .. })); + } + + // -- Status update test -- + + #[tokio::test] + async fn report_status_updates_handle() { + let state = test_state(); + let job_id = Uuid::new_v4(); + let token = state.token_store.create_token(job_id).await; + + // Insert a handle so update_worker_status has something to update + { + let mut containers = state.job_manager.containers.write().await; + containers.insert( + job_id, + crate::orchestrator::job_manager::ContainerHandle { + job_id, + container_id: "test-container".to_string(), + state: crate::orchestrator::job_manager::ContainerState::Running, + mode: crate::orchestrator::job_manager::JobMode::Worker, + created_at: chrono::Utc::now(), + project_dir: None, + task_description: "test".to_string(), + last_worker_status: None, + worker_iteration: 0, + completion_result: None, + }, + ); + } + + let jm = Arc::clone(&state.job_manager); + let router = OrchestratorApi::router(state); + + let update = serde_json::json!({ + "state": "in_progress", + "message": "Iteration 5", + "iteration": 5 + }); + + let req = Request::builder() + .method("POST") + .uri(format!("/worker/{}/status", job_id)) + .header("Authorization", format!("Bearer {}", token)) + .header("Content-Type", "application/json") + .body(Body::from(serde_json::to_vec(&update).unwrap())) + .unwrap(); + + let resp = router.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let handle = jm.get_handle(job_id).await.unwrap(); + assert_eq!(handle.worker_iteration, 5); + assert_eq!(handle.last_worker_status.as_deref(), Some("Iteration 5")); + } } diff --git a/src/orchestrator/auth.rs b/src/orchestrator/auth.rs index 894b0ffb..cf1819d2 100644 --- a/src/orchestrator/auth.rs +++ b/src/orchestrator/auth.rs @@ -5,6 +5,7 @@ //! - Tokens are scoped to a specific job_id //! - Tokens are ephemeral (in-memory only, never persisted) //! - A token for Job A cannot access endpoints for Job B +//! - Credential grants are per-job: only secrets explicitly granted are accessible use std::collections::HashMap; use std::sync::Arc; @@ -14,21 +15,37 @@ use axum::http::StatusCode; use axum::middleware::Next; use axum::response::Response; use rand::Rng; +use serde::{Deserialize, Serialize}; use subtle::ConstantTimeEq; use tokio::sync::RwLock; use uuid::Uuid; -/// In-memory store for per-job authentication tokens. +/// A credential grant that maps a secret (stored in SecretsStore) to an +/// environment variable name the container worker expects. +/// +/// For example: `{ secret_name: "github_token", env_var: "GITHUB_TOKEN" }` +/// means "decrypt the secret named `github_token` and provide it as the +/// env var `GITHUB_TOKEN` to the container". +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CredentialGrant { + pub secret_name: String, + pub env_var: String, +} + +/// In-memory store for per-job authentication tokens and credential grants. #[derive(Clone)] pub struct TokenStore { /// Maps job_id -> bearer token. Never logged or persisted. tokens: Arc>>, + /// Maps job_id -> granted credentials. Revoked alongside the token. + credential_grants: Arc>>>, } impl TokenStore { pub fn new() -> Self { Self { tokens: Arc::new(RwLock::new(HashMap::new())), + credential_grants: Arc::new(RwLock::new(HashMap::new())), } } @@ -49,15 +66,28 @@ impl TokenStore { .unwrap_or(false) } - /// Remove a token (on container cleanup). + /// Remove a token and its credential grants (on container cleanup). pub async fn revoke(&self, job_id: Uuid) { self.tokens.write().await.remove(&job_id); + self.credential_grants.write().await.remove(&job_id); } /// Get the number of active tokens (for diagnostics). pub async fn active_count(&self) -> usize { self.tokens.read().await.len() } + + /// Store credential grants for a job. Call right after `create_token()`. + pub async fn store_grants(&self, job_id: Uuid, grants: Vec) { + if !grants.is_empty() { + self.credential_grants.write().await.insert(job_id, grants); + } + } + + /// Retrieve credential grants for a job. + pub async fn get_grants(&self, job_id: Uuid) -> Option> { + self.credential_grants.read().await.get(&job_id).cloned() + } } impl Default for TokenStore { @@ -70,11 +100,12 @@ impl Default for TokenStore { fn generate_token() -> String { let mut bytes = [0u8; 32]; rand::thread_rng().fill(&mut bytes); - hex_encode(&bytes) -} - -fn hex_encode(bytes: &[u8]) -> String { - bytes.iter().map(|b| format!("{:02x}", b)).collect() + // Hex-encode without pulling in a crate: fixed-size array, no allocation concern. + bytes.iter().fold(String::with_capacity(64), |mut s, b| { + use std::fmt::Write; + let _ = write!(s, "{b:02x}"); + s + }) } /// Axum middleware that validates worker bearer tokens. @@ -160,4 +191,102 @@ mod tests { let t2 = generate_token(); assert_ne!(t1, t2); } + + #[tokio::test] + async fn test_store_and_get_grants() { + let store = TokenStore::new(); + let job_id = Uuid::new_v4(); + + // No grants initially + assert!(store.get_grants(job_id).await.is_none()); + + let grants = vec![ + CredentialGrant { + secret_name: "github_token".to_string(), + env_var: "GITHUB_TOKEN".to_string(), + }, + CredentialGrant { + secret_name: "npm_token".to_string(), + env_var: "NPM_TOKEN".to_string(), + }, + ]; + + store.store_grants(job_id, grants).await; + + let retrieved = store.get_grants(job_id).await.unwrap(); + assert_eq!(retrieved.len(), 2); + assert_eq!(retrieved[0].secret_name, "github_token"); + assert_eq!(retrieved[0].env_var, "GITHUB_TOKEN"); + assert_eq!(retrieved[1].secret_name, "npm_token"); + } + + #[tokio::test] + async fn test_revoke_clears_grants() { + let store = TokenStore::new(); + let job_id = Uuid::new_v4(); + + let _token = store.create_token(job_id).await; + store + .store_grants( + job_id, + vec![CredentialGrant { + secret_name: "my_secret".to_string(), + env_var: "MY_SECRET".to_string(), + }], + ) + .await; + + assert!(store.get_grants(job_id).await.is_some()); + + store.revoke(job_id).await; + + assert!(!store.validate(job_id, "anything").await); + assert!(store.get_grants(job_id).await.is_none()); + } + + #[tokio::test] + async fn test_empty_grants_not_stored() { + let store = TokenStore::new(); + let job_id = Uuid::new_v4(); + + store.store_grants(job_id, vec![]).await; + + // Empty vec should not create an entry + assert!(store.get_grants(job_id).await.is_none()); + } + + #[tokio::test] + async fn test_grants_isolated_per_job() { + let store = TokenStore::new(); + let job_a = Uuid::new_v4(); + let job_b = Uuid::new_v4(); + + store + .store_grants( + job_a, + vec![CredentialGrant { + secret_name: "secret_a".to_string(), + env_var: "SECRET_A".to_string(), + }], + ) + .await; + + store + .store_grants( + job_b, + vec![CredentialGrant { + secret_name: "secret_b".to_string(), + env_var: "SECRET_B".to_string(), + }], + ) + .await; + + let grants_a = store.get_grants(job_a).await.unwrap(); + assert_eq!(grants_a.len(), 1); + assert_eq!(grants_a[0].secret_name, "secret_a"); + + let grants_b = store.get_grants(job_b).await.unwrap(); + assert_eq!(grants_b.len(), 1); + assert_eq!(grants_b[0].secret_name, "secret_b"); + } } diff --git a/src/orchestrator/job_manager.rs b/src/orchestrator/job_manager.rs index b1d20306..9bcc969e 100644 --- a/src/orchestrator/job_manager.rs +++ b/src/orchestrator/job_manager.rs @@ -12,7 +12,7 @@ use tokio::sync::RwLock; use uuid::Uuid; use crate::error::OrchestratorError; -use crate::orchestrator::auth::TokenStore; +use crate::orchestrator::auth::{CredentialGrant, TokenStore}; use crate::sandbox::connect_docker; /// Which mode a sandbox container runs in. @@ -50,8 +50,13 @@ pub struct ContainerJobConfig { pub cpu_shares: u32, /// Port the orchestrator internal API listens on. pub orchestrator_port: u16, - /// Host directory containing Claude auth config (mounted read-only for ClaudeCode mode). - pub claude_config_dir: Option, + /// Anthropic API key for Claude Code containers (read from ANTHROPIC_API_KEY). + /// Takes priority over OAuth token. + pub claude_code_api_key: Option, + /// OAuth access token extracted from the host's `claude login` session. + /// Passed as CLAUDE_CODE_OAUTH_TOKEN to containers. Falls back to this + /// when no ANTHROPIC_API_KEY is available. + pub claude_code_oauth_token: Option, /// Claude model to use in ClaudeCode mode. pub claude_code_model: String, /// Maximum turns for Claude Code. @@ -69,7 +74,8 @@ impl Default for ContainerJobConfig { memory_limit_mb: 2048, cpu_shares: 1024, orchestrator_port: 50051, - claude_config_dir: None, + claude_code_api_key: None, + claude_code_oauth_token: None, claude_code_model: "sonnet".to_string(), claude_code_max_turns: 50, claude_code_memory_limit_mb: 4096, @@ -108,6 +114,10 @@ pub struct ContainerHandle { pub created_at: DateTime, pub project_dir: Option, pub task_description: String, + /// Last status message reported by the worker (iteration count, progress, etc.). + pub last_worker_status: Option, + /// Which iteration the worker is on (updated via status reports). + pub worker_iteration: u32, /// Completion result from the worker (set when the worker reports done). pub completion_result: Option, // NOTE: auth_token is intentionally NOT in this struct. @@ -121,11 +131,84 @@ pub struct CompletionResult { pub message: Option, } +/// Validate that a project directory is under `~/.ironclaw/projects/`. +/// +/// Returns the canonicalized path if valid. Creates the base directory if +/// it doesn't exist (so the prefix check always runs). +/// +/// # TOCTOU note +/// +/// There is a time-of-check/time-of-use gap between `canonicalize()` here +/// and the actual Docker `binds.push()` in the caller. In a multi-tenant +/// system a malicious actor could swap a symlink after validation. This is +/// acceptable in IronClaw's single-tenant design where the user controls +/// the filesystem. +fn validate_bind_mount_path( + dir: &std::path::Path, + job_id: Uuid, +) -> Result { + let canonical = dir + .canonicalize() + .map_err(|e| OrchestratorError::ContainerCreationFailed { + job_id, + reason: format!( + "failed to canonicalize project dir {}: {}", + dir.display(), + e + ), + })?; + + let home = dirs::home_dir().ok_or_else(|| OrchestratorError::ContainerCreationFailed { + job_id, + reason: "could not determine home directory for path validation".to_string(), + })?; + let projects_base = home.join(".ironclaw").join("projects"); + + // Ensure the base exists so canonicalize always succeeds. + std::fs::create_dir_all(&projects_base).map_err(|e| { + OrchestratorError::ContainerCreationFailed { + job_id, + reason: format!( + "failed to create projects base {}: {}", + projects_base.display(), + e + ), + } + })?; + + let canonical_base = + projects_base + .canonicalize() + .map_err(|e| OrchestratorError::ContainerCreationFailed { + job_id, + reason: format!( + "failed to canonicalize projects base {}: {}", + projects_base.display(), + e + ), + })?; + + if !canonical.starts_with(&canonical_base) { + return Err(OrchestratorError::ContainerCreationFailed { + job_id, + reason: format!( + "project directory {} is outside allowed base {}", + canonical.display(), + canonical_base.display() + ), + }); + } + + Ok(canonical) +} + /// Manages the lifecycle of Docker containers for sandboxed job execution. pub struct ContainerJobManager { config: ContainerJobConfig, token_store: TokenStore, - containers: Arc>>, + pub(crate) containers: Arc>>, + /// Cached Docker connection (created on first use). + docker: Arc>>, } impl ContainerJobManager { @@ -134,23 +217,49 @@ impl ContainerJobManager { config, token_store, containers: Arc::new(RwLock::new(HashMap::new())), + docker: Arc::new(RwLock::new(None)), } } + /// Get or create a Docker connection. + async fn docker(&self) -> Result { + { + let guard = self.docker.read().await; + if let Some(ref d) = *guard { + return Ok(d.clone()); + } + } + let docker = connect_docker() + .await + .map_err(|e| OrchestratorError::Docker { + reason: e.to_string(), + })?; + *self.docker.write().await = Some(docker.clone()); + Ok(docker) + } + /// Create and start a new container for a job. /// /// The caller provides the `job_id` so it can be persisted to the database - /// before the container is created. Returns the auth token for the worker. + /// before the container is created. Credential grants are stored in the + /// TokenStore and served on-demand via the `/credentials` endpoint. + /// Returns the auth token for the worker. pub async fn create_job( &self, job_id: Uuid, task: &str, project_dir: Option, mode: JobMode, + credential_grants: Vec, ) -> Result { // Generate auth token (stored in TokenStore, never logged) let token = self.token_store.create_token(job_id).await; + // Store credential grants (revoked automatically when the token is revoked) + self.token_store + .store_grants(job_id, credential_grants) + .await; + // Record the handle let handle = ContainerHandle { job_id, @@ -160,6 +269,8 @@ impl ContainerJobManager { created_at: Utc::now(), project_dir: project_dir.clone(), task_description: task.to_string(), + last_worker_status: None, + worker_iteration: 0, completion_result: None, }; self.containers.write().await.insert(job_id, handle); @@ -187,12 +298,8 @@ impl ContainerJobManager { project_dir: Option, mode: JobMode, ) -> Result<(), OrchestratorError> { - // Connect to Docker - let docker = connect_docker() - .await - .map_err(|e| OrchestratorError::Docker { - reason: e.to_string(), - })?; + // Connect to Docker (reuses cached connection) + let docker = self.docker().await?; // Build container configuration let orchestrator_host = if cfg!(target_os = "linux") { @@ -215,41 +322,22 @@ impl ContainerJobManager { // Build volume mounts (validate project_dir stays within ~/.ironclaw/projects/) let mut binds = Vec::new(); if let Some(ref dir) = project_dir { - let canonical = - dir.canonicalize() - .map_err(|e| OrchestratorError::ContainerCreationFailed { - job_id, - reason: format!( - "failed to canonicalize project dir {}: {}", - dir.display(), - e - ), - })?; - let projects_base = dirs::home_dir() - .unwrap_or_else(|| PathBuf::from(".")) - .join(".ironclaw") - .join("projects"); - if let Ok(canonical_base) = projects_base.canonicalize() - && !canonical.starts_with(&canonical_base) - { - return Err(OrchestratorError::ContainerCreationFailed { - job_id, - reason: format!( - "project directory {} is outside allowed base {}", - canonical.display(), - canonical_base.display() - ), - }); - } + let canonical = validate_bind_mount_path(dir, job_id)?; binds.push(format!("{}:/workspace:rw", canonical.display())); env_vec.push("IRONCLAW_WORKSPACE=/workspace".to_string()); } - // Claude Code mode: mount host ~/.claude read-only for auth, - // and pass the tool allowlist so the bridge can write settings.json. + // Claude Code mode: auth + tool allowlist. + // + // Auth strategies (first match wins): + // 1. ANTHROPIC_API_KEY: direct API key (pay-as-you-go billing). + // 2. CLAUDE_CODE_OAUTH_TOKEN: OAuth access token from `claude login` + // session, extracted from the host's credential store. if mode == JobMode::ClaudeCode { - if let Some(ref claude_dir) = self.config.claude_config_dir { - binds.push(format!("{}:/home/sandbox/.claude:ro", claude_dir.display())); + if let Some(ref api_key) = self.config.claude_code_api_key { + env_vec.push(format!("ANTHROPIC_API_KEY={}", api_key)); + } else if let Some(ref oauth_token) = self.config.claude_code_oauth_token { + env_vec.push(format!("CLAUDE_CODE_OAUTH_TOKEN={}", oauth_token)); } if !self.config.claude_code_allowed_tools.is_empty() { env_vec.push(format!( @@ -377,11 +465,7 @@ impl ContainerJobManager { }); } - let docker = connect_docker() - .await - .map_err(|e| OrchestratorError::Docker { - reason: e.to_string(), - })?; + let docker = self.docker().await?; // Stop the container (10 second grace period) if let Err(e) = docker @@ -445,7 +529,7 @@ impl ContainerJobManager { if let Some(cid) = container_id && !cid.is_empty() { - match connect_docker().await { + match self.docker().await { Ok(docker) => { if let Err(e) = docker .stop_container( @@ -485,6 +569,19 @@ impl ContainerJobManager { self.containers.write().await.remove(&job_id); } + /// Update the worker-reported status for a job. + pub async fn update_worker_status( + &self, + job_id: Uuid, + message: Option, + iteration: u32, + ) { + if let Some(handle) = self.containers.write().await.get_mut(&job_id) { + handle.last_worker_status = message; + handle.worker_iteration = iteration; + } + } + /// Get the handle for a job. pub async fn get_handle(&self, job_id: Uuid) -> Option { self.containers.read().await.get(&job_id).cloned() @@ -517,4 +614,82 @@ mod tests { assert_eq!(ContainerState::Running.to_string(), "running"); assert_eq!(ContainerState::Stopped.to_string(), "stopped"); } + + #[test] + fn test_validate_bind_mount_valid_path() { + let base = dirs::home_dir().unwrap().join(".ironclaw").join("projects"); + std::fs::create_dir_all(&base).unwrap(); + + let test_dir = base.join("test_validate_bind"); + std::fs::create_dir_all(&test_dir).unwrap(); + + let result = validate_bind_mount_path(&test_dir, Uuid::new_v4()); + assert!(result.is_ok()); + let canonical = result.unwrap(); + assert!(canonical.starts_with(base.canonicalize().unwrap())); + + let _ = std::fs::remove_dir_all(&test_dir); + } + + #[test] + fn test_validate_bind_mount_rejects_outside_base() { + let tmp = tempfile::tempdir().unwrap(); + let outside = tmp.path().to_path_buf(); + + let result = validate_bind_mount_path(&outside, Uuid::new_v4()); + assert!(result.is_err()); + let err = result.unwrap_err().to_string(); + assert!( + err.contains("outside allowed base"), + "expected 'outside allowed base', got: {}", + err + ); + } + + #[test] + fn test_validate_bind_mount_rejects_nonexistent() { + let nonexistent = PathBuf::from("/no/such/path/at/all"); + let result = validate_bind_mount_path(&nonexistent, Uuid::new_v4()); + assert!(result.is_err()); + let err = result.unwrap_err().to_string(); + assert!( + err.contains("canonicalize"), + "expected canonicalize error, got: {}", + err + ); + } + + #[tokio::test] + async fn test_update_worker_status() { + let store = TokenStore::new(); + let mgr = ContainerJobManager::new(ContainerJobConfig::default(), store); + let job_id = Uuid::new_v4(); + + // Insert a handle + { + let mut containers = mgr.containers.write().await; + containers.insert( + job_id, + ContainerHandle { + job_id, + container_id: "test".to_string(), + state: ContainerState::Running, + mode: JobMode::Worker, + created_at: chrono::Utc::now(), + project_dir: None, + task_description: "test job".to_string(), + last_worker_status: None, + worker_iteration: 0, + completion_result: None, + }, + ); + } + + mgr.update_worker_status(job_id, Some("Iteration 3".to_string()), 3) + .await; + + let handle = mgr.get_handle(job_id).await.unwrap(); + assert_eq!(handle.worker_iteration, 3); + assert_eq!(handle.last_worker_status.as_deref(), Some("Iteration 3")); + } } diff --git a/src/orchestrator/mod.rs b/src/orchestrator/mod.rs index 8462f6aa..921edd93 100644 --- a/src/orchestrator/mod.rs +++ b/src/orchestrator/mod.rs @@ -9,10 +9,11 @@ //! ┌───────────────────────────────────────────────┐ //! │ Orchestrator │ //! │ │ -//! │ Internal API (:50051) │ +//! │ Internal API (default :50051, configurable) │ //! │ POST /worker/{id}/llm/complete │ //! │ POST /worker/{id}/llm/complete_with_tools │ //! │ GET /worker/{id}/job │ +//! │ GET /worker/{id}/credentials │ //! │ POST /worker/{id}/status │ //! │ POST /worker/{id}/complete │ //! │ │ @@ -23,6 +24,7 @@ //! │ │ //! │ TokenStore │ //! │ per-job bearer tokens (in-memory only) │ +//! │ per-job credential grants (in-memory only) │ //! └───────────────────────────────────────────────┘ //! ``` @@ -31,7 +33,7 @@ pub mod auth; pub mod job_manager; pub use api::OrchestratorApi; -pub use auth::TokenStore; +pub use auth::{CredentialGrant, TokenStore}; pub use job_manager::{ CompletionResult, ContainerHandle, ContainerJobConfig, ContainerJobManager, JobMode, }; diff --git a/src/sandbox/config.rs b/src/sandbox/config.rs index fa01ddc0..41fc36fa 100644 --- a/src/sandbox/config.rs +++ b/src/sandbox/config.rs @@ -159,56 +159,14 @@ pub fn default_allowlist() -> Vec { ] } -/// Credential injection configuration. -#[derive(Debug, Clone)] -pub struct CredentialMapping { - /// Domain this credential applies to. - pub domain: String, - /// Name of the secret to inject. - pub secret_name: String, - /// Where to inject the credential. - pub location: CredentialLocation, -} - -/// Where to inject a credential in an HTTP request. -#[derive(Debug, Clone)] -pub enum CredentialLocation { - /// Inject as Authorization: Bearer - AuthorizationBearer, - /// Inject as a custom header. - Header(String), - /// Inject as a query parameter. - QueryParam(String), -} - -impl Default for CredentialMapping { - fn default() -> Self { - Self { - domain: String::new(), - secret_name: String::new(), - location: CredentialLocation::AuthorizationBearer, - } - } -} - /// Default credential mappings for common APIs. -pub fn default_credential_mappings() -> Vec { +pub fn default_credential_mappings() -> Vec { + use crate::secrets::CredentialMapping; + vec![ - CredentialMapping { - domain: "api.openai.com".to_string(), - secret_name: "OPENAI_API_KEY".to_string(), - location: CredentialLocation::AuthorizationBearer, - }, - CredentialMapping { - domain: "api.anthropic.com".to_string(), - secret_name: "ANTHROPIC_API_KEY".to_string(), - location: CredentialLocation::Header("x-api-key".to_string()), - }, - CredentialMapping { - domain: "api.near.ai".to_string(), - secret_name: "NEARAI_API_KEY".to_string(), - location: CredentialLocation::AuthorizationBearer, - }, + CredentialMapping::bearer("OPENAI_API_KEY", "api.openai.com"), + CredentialMapping::header("ANTHROPIC_API_KEY", "x-api-key", "api.anthropic.com"), + CredentialMapping::bearer("NEARAI_API_KEY", "api.near.ai"), ] } diff --git a/src/sandbox/container.rs b/src/sandbox/container.rs index 1ef0fa0f..2160d450 100644 --- a/src/sandbox/container.rs +++ b/src/sandbox/container.rs @@ -283,7 +283,7 @@ impl ContainerRunner { // Prevent privilege escalation security_opt: Some(vec!["no-new-privileges:true".to_string()]), // Read-only root filesystem (workspace is still writable if policy allows) - readonly_rootfs: Some(policy == SandboxPolicy::ReadOnly), + readonly_rootfs: Some(policy != SandboxPolicy::FullAccess), // Tmpfs mounts for /tmp and cargo cache tmpfs: Some( [ @@ -515,7 +515,8 @@ pub async fn connect_docker() -> Result { } Err(SandboxError::DockerNotAvailable { - reason: "Socket not found: /var/run/docker.sock".to_string(), + reason: "Could not connect to Docker. Tried: default socket, ~/.docker/run/docker.sock" + .to_string(), }) } diff --git a/src/sandbox/manager.rs b/src/sandbox/manager.rs index 3b07eafe..9da7a22e 100644 --- a/src/sandbox/manager.rs +++ b/src/sandbox/manager.rs @@ -35,6 +35,8 @@ use std::time::Duration; use tokio::sync::RwLock; +use bollard::Docker; + use crate::sandbox::config::{ResourceLimits, SandboxConfig, SandboxPolicy}; use crate::sandbox::container::{ContainerOutput, ContainerRunner, connect_docker}; use crate::sandbox::error::{Result, SandboxError}; @@ -82,7 +84,7 @@ impl From for ExecOutput { pub struct SandboxManager { config: SandboxConfig, proxy: Arc>>, - runner: Arc>>, + docker: Arc>>, initialized: std::sync::atomic::AtomicBool, } @@ -92,7 +94,7 @@ impl SandboxManager { Self { config, proxy: Arc::new(RwLock::new(None)), - runner: Arc::new(RwLock::new(None)), + docker: Arc::new(RwLock::new(None)), initialized: std::sync::atomic::AtomicBool::new(false), } } @@ -137,14 +139,15 @@ impl SandboxManager { reason: e.to_string(), })?; - // Create container runner - let runner = - ContainerRunner::new(docker, self.config.image.clone(), self.config.proxy_port); - - // Check for / pull image - if !runner.image_exists().await { + // Check for / pull image using a temporary runner + let checker = ContainerRunner::new( + docker.clone(), + self.config.image.clone(), + self.config.proxy_port, + ); + if !checker.image_exists().await { if self.config.auto_pull_image { - runner.pull_image().await?; + checker.pull_image().await?; } else { return Err(SandboxError::ContainerCreationFailed { reason: format!( @@ -155,7 +158,7 @@ impl SandboxManager { } } - *self.runner.write().await = Some(runner); + *self.docker.write().await = Some(docker); // Start the network proxy if we're using a sandboxed policy if self.config.policy.is_sandboxed() { @@ -221,8 +224,15 @@ impl SandboxManager { 0 }; - // Create a runner with the current proxy port - let docker = connect_docker().await?; + // Reuse the stored Docker connection, create a runner with the current proxy port + let docker = + self.docker + .read() + .await + .clone() + .ok_or_else(|| SandboxError::DockerNotAvailable { + reason: "Docker connection not initialized".to_string(), + })?; let runner = ContainerRunner::new(docker, self.config.image.clone(), proxy_port); let limits = ResourceLimits { @@ -268,8 +278,24 @@ impl SandboxManager { reason: e.to_string(), })?; - let stdout = String::from_utf8_lossy(&output.stdout).to_string(); - let stderr = String::from_utf8_lossy(&output.stderr).to_string(); + let max_output: usize = 64 * 1024; // 64 KB, matching container path + let half_max = max_output / 2; + + let mut stdout = String::from_utf8_lossy(&output.stdout).to_string(); + let mut stderr = String::from_utf8_lossy(&output.stderr).to_string(); + let mut truncated = false; + + if stdout.len() > half_max { + let end = crate::util::floor_char_boundary(&stdout, half_max); + stdout.truncate(end); + truncated = true; + } + if stderr.len() > half_max { + let end = crate::util::floor_char_boundary(&stderr, half_max); + stderr.truncate(end); + truncated = true; + } + let combined = if stderr.is_empty() { stdout.clone() } else if stdout.is_empty() { @@ -284,7 +310,7 @@ impl SandboxManager { stderr, output: combined, duration: start.elapsed(), - truncated: false, + truncated, }) } @@ -471,4 +497,28 @@ mod tests { let output = result.unwrap(); assert!(output.stdout.contains("hello")); } + + #[tokio::test] + async fn test_direct_execution_truncates_large_output() { + let manager = SandboxManager::new(SandboxConfig { + enabled: true, + policy: SandboxPolicy::FullAccess, + ..Default::default() + }); + + // Generate output larger than 32KB (half of 64KB limit) + // printf repeats a 100-char line 400 times = 40KB + let result = manager + .execute( + "printf 'A%.0s' $(seq 1 40000)", + Path::new("."), + HashMap::new(), + ) + .await; + + assert!(result.is_ok()); + let output = result.unwrap(); + assert!(output.truncated); + assert!(output.stdout.len() <= 32 * 1024); + } } diff --git a/src/sandbox/mod.rs b/src/sandbox/mod.rs index 5a7aed0f..73f553e7 100644 --- a/src/sandbox/mod.rs +++ b/src/sandbox/mod.rs @@ -91,9 +91,7 @@ pub mod error; pub mod manager; pub mod proxy; -pub use config::{ - CredentialLocation, CredentialMapping, ResourceLimits, SandboxConfig, SandboxPolicy, -}; +pub use config::{ResourceLimits, SandboxConfig, SandboxPolicy}; pub use container::{ContainerOutput, ContainerRunner, connect_docker}; pub use error::{Result, SandboxError}; pub use manager::{ExecOutput, SandboxManager, SandboxManagerBuilder}; @@ -108,6 +106,6 @@ pub fn default_allowlist() -> Vec { } /// Default credential mappings getter (re-export for convenience). -pub fn default_credential_mappings() -> Vec { +pub fn default_credential_mappings() -> Vec { config::default_credential_mappings() } diff --git a/src/sandbox/proxy/http.rs b/src/sandbox/proxy/http.rs index 3205841b..3b0268e7 100644 --- a/src/sandbox/proxy/http.rs +++ b/src/sandbox/proxy/http.rs @@ -21,12 +21,12 @@ use hyper::server::conn::http1; use hyper::service::service_fn; use hyper::{Method, Request, Response, StatusCode}; use hyper_util::rt::TokioIo; -use tokio::net::TcpListener; +use tokio::net::{TcpListener, TcpStream}; use tokio::sync::RwLock; -use crate::sandbox::config::CredentialLocation; use crate::sandbox::error::{Result, SandboxError}; use crate::sandbox::proxy::policy::{NetworkDecision, NetworkPolicyDecider, NetworkRequest}; +use crate::secrets::CredentialLocation; /// State shared across proxy connections. struct ProxyState { @@ -34,6 +34,8 @@ struct ProxyState { decider: Arc, /// Credential resolver (maps secret names to values). credential_resolver: Arc, + /// Shared HTTP client for forwarding requests. + http_client: reqwest::Client, /// Request counter for logging. request_count: std::sync::atomic::AtomicU64, /// Whether the proxy is running. @@ -84,6 +86,7 @@ impl HttpProxy { state: Arc::new(ProxyState { decider, credential_resolver, + http_client: reqwest::Client::new(), request_count: std::sync::atomic::AtomicU64::new(0), running: std::sync::atomic::AtomicBool::new(false), }), @@ -235,20 +238,29 @@ async fn handle_request( } /// Handle CONNECT method for HTTPS tunneling. +/// +/// Establishes a bidirectional TCP tunnel between the client and the target host. +/// Returns 200 OK to signal the client to begin TLS over the upgraded connection. +/// +/// NOTE: Credential injection is not possible through CONNECT tunnels since the proxy +/// cannot inspect or modify TLS-encrypted traffic without MITM. Containers that need +/// authenticated HTTPS should fetch credentials via the orchestrator's +/// `GET /worker/{id}/credentials` endpoint and set them as environment variables. async fn handle_connect( req: Request, state: Arc, ) -> Response> { - // Extract host from CONNECT target - let host = req.uri().authority().map(|a| a.host().to_string()); - - let host = match host { - Some(h) => h, + // Extract host:port from CONNECT target (e.g. "api.github.com:443") + let authority = match req.uri().authority() { + Some(a) => a.clone(), None => { return error_response(StatusCode::BAD_REQUEST, "Missing host".to_string()); } }; + let host = authority.host().to_string(); + let target_addr = authority.as_str().to_string(); + // Check if host is allowed let network_req = NetworkRequest { method: "CONNECT".to_string(), @@ -259,21 +271,56 @@ async fn handle_connect( let decision = state.decider.decide(&network_req).await; - if !decision.is_allowed() - && let NetworkDecision::Deny { reason } = decision - { + if let NetworkDecision::Deny { reason } = decision { tracing::info!("Proxy: blocked CONNECT {} - {}", host, reason); return error_response(StatusCode::FORBIDDEN, reason); } - tracing::debug!("Proxy: allowing CONNECT to {}", host); + tracing::debug!("Proxy: allowing CONNECT to {}", target_addr); - // For CONNECT, we return 200 OK and the client will upgrade to TLS - // The actual TLS connection goes directly to the target, we just act as a tunnel - Response::builder() - .status(StatusCode::OK) - .body(empty_body()) - .unwrap() + // Spawn a fire-and-forget task to establish the tunnel after the upgrade + // completes. The 30-minute timeout guarantees every tunnel task terminates + // even if the remote peer hangs, so no `JoinSet` tracking is needed. + // On process exit these tasks are dropped by the runtime. + let target = target_addr.clone(); + tokio::spawn(async move { + match hyper::upgrade::on(req).await { + Ok(upgraded) => { + let mut client_stream = TokioIo::new(upgraded); + match TcpStream::connect(&target).await { + Ok(mut server_stream) => { + let tunnel_timeout = std::time::Duration::from_secs(30 * 60); + match tokio::time::timeout( + tunnel_timeout, + tokio::io::copy_bidirectional(&mut client_stream, &mut server_stream), + ) + .await + { + Ok(Ok(_)) => {} + Ok(Err(e)) => { + tracing::debug!("Proxy: tunnel to {} closed: {}", target, e); + } + Err(_) => { + tracing::info!( + "Proxy: tunnel to {} timed out after 30m, closing", + target + ); + } + } + } + Err(e) => { + tracing::error!("Proxy: failed to connect to {}: {}", target, e); + } + } + } + Err(e) => { + tracing::error!("Proxy: upgrade failed for {}: {}", target, e); + } + } + }); + + // Return 200 OK so the client begins the TLS handshake over the upgraded connection + make_response(StatusCode::OK, empty_body()) } /// Forward a request to the target server. @@ -286,8 +333,7 @@ async fn forward_request( let uri = req.uri().clone(); // Build the forwarded request - let client = reqwest::Client::new(); - let mut builder = client.request( + let mut builder = state.http_client.request( reqwest::Method::from_bytes(method.as_str().as_bytes()).unwrap_or(reqwest::Method::GET), uri.to_string(), ); @@ -312,9 +358,27 @@ async fn forward_request( CredentialLocation::AuthorizationBearer => { builder.header("Authorization", format!("Bearer {}", credential)) } - CredentialLocation::Header(header_name) => builder.header(header_name, credential), - CredentialLocation::QueryParam(param_name) => { - builder.query(&[(param_name, credential)]) + CredentialLocation::Header { name, prefix } => { + let value = match prefix { + Some(p) => format!("{}{}", p, credential), + None => credential.clone(), + }; + builder.header(name, value) + } + CredentialLocation::QueryParam { name } => builder.query(&[(name, credential)]), + // Known limitation: AuthorizationBasic requires the proxy to + // construct a Base64 username:password pair from a single secret, + // and UrlPath requires rewriting the request URI. Neither is + // implemented yet. Containers needing these auth styles should + // fetch credentials via the orchestrator's GET /worker/{id}/credentials + // endpoint and set them directly. + CredentialLocation::AuthorizationBasic { .. } + | CredentialLocation::UrlPath { .. } => { + tracing::warn!( + "Proxy: credential location {:?} not supported for forward proxy, skipping", + location + ); + builder } }; tracing::debug!("Proxy: injected credential for {}", secret_name); @@ -347,15 +411,15 @@ async fn forward_request( match response.bytes().await { Ok(body) => { - let mut builder = Response::builder().status(status.as_u16()); + let mut resp_builder = Response::builder().status(status.as_u16()); for (name, value) in headers.iter() { if !is_hop_by_hop_header(name.as_str()) { - builder = builder.header(name.as_str(), value.as_bytes()); + resp_builder = resp_builder.header(name.as_str(), value.as_bytes()); } } - Ok(builder.body(full_body(body)).unwrap()) + Ok(make_response_from_builder(resp_builder, full_body(body))) } Err(e) => { tracing::error!("Proxy: failed to read response body: {}", e); @@ -391,13 +455,52 @@ fn is_hop_by_hop_header(name: &str) -> bool { ) } -/// Create an error response. -fn error_response(status: StatusCode, message: String) -> Response> { +/// Build a response with guaranteed success (valid status + simple body cannot fail). +fn make_response( + status: StatusCode, + body: BoxBody, +) -> Response> { Response::builder() .status(status) - .header("Content-Type", "text/plain") - .body(full_body(Bytes::from(message))) - .unwrap() + .body(body) + .unwrap_or_else(|_| { + let mut resp = Response::new( + Full::new(Bytes::from("Internal error")) + .map_err(|_| unreachable!()) + .boxed(), + ); + *resp.status_mut() = StatusCode::INTERNAL_SERVER_ERROR; + resp + }) +} + +/// Finalize a partially-built response, falling back to 500 on builder error. +fn make_response_from_builder( + builder: hyper::http::response::Builder, + body: BoxBody, +) -> Response> { + builder.body(body).unwrap_or_else(|_| { + Response::builder() + .status(StatusCode::INTERNAL_SERVER_ERROR) + .body(full_body(Bytes::from("Response build error"))) + .unwrap_or_else(|_| { + Response::new( + Full::new(Bytes::from("Internal error")) + .map_err(|_| unreachable!()) + .boxed(), + ) + }) + }) +} + +/// Create an error response. +fn error_response(status: StatusCode, message: String) -> Response> { + make_response_from_builder( + Response::builder() + .status(status) + .header("Content-Type", "text/plain"), + full_body(Bytes::from(message)), + ) } /// Create an empty body. @@ -441,4 +544,13 @@ mod tests { assert!(!is_hop_by_hop_header("content-type")); assert!(!is_hop_by_hop_header("authorization")); } + + #[test] + fn test_make_response_does_not_panic() { + let resp = make_response(StatusCode::OK, empty_body()); + assert_eq!(resp.status(), StatusCode::OK); + + let resp = error_response(StatusCode::FORBIDDEN, "denied".to_string()); + assert_eq!(resp.status(), StatusCode::FORBIDDEN); + } } diff --git a/src/sandbox/proxy/mod.rs b/src/sandbox/proxy/mod.rs index 2feb460c..c4c89850 100644 --- a/src/sandbox/proxy/mod.rs +++ b/src/sandbox/proxy/mod.rs @@ -41,10 +41,9 @@ pub use policy::{ use std::sync::Arc; -use crate::sandbox::config::{ - CredentialMapping, SandboxConfig, SandboxPolicy, default_credential_mappings, -}; +use crate::sandbox::config::{SandboxConfig, SandboxPolicy, default_credential_mappings}; use crate::sandbox::error::Result; +use crate::secrets::CredentialMapping; /// Creates a configured network proxy from sandbox config. pub struct NetworkProxyBuilder { diff --git a/src/sandbox/proxy/policy.rs b/src/sandbox/proxy/policy.rs index 2e2a45b9..2406a694 100644 --- a/src/sandbox/proxy/policy.rs +++ b/src/sandbox/proxy/policy.rs @@ -5,8 +5,8 @@ use async_trait::async_trait; -use crate::sandbox::config::{CredentialLocation, CredentialMapping}; use crate::sandbox::proxy::allowlist::DomainAllowlist; +use crate::secrets::{CredentialLocation, CredentialMapping}; /// A network request to be evaluated. #[derive(Debug, Clone)] @@ -95,12 +95,14 @@ impl DefaultPolicyDecider { } } - /// Find credential mapping for a domain. + /// Find credential mapping for a host (supports glob patterns like `*.example.com`). fn find_credential(&self, host: &str) -> Option<&CredentialMapping> { let host_lower = host.to_lowercase(); - self.credential_mappings - .iter() - .find(|m| m.domain.to_lowercase() == host_lower) + self.credential_mappings.iter().find(|m| { + m.host_patterns + .iter() + .any(|pattern| host_matches_pattern(&host_lower, pattern)) + }) } } @@ -128,6 +130,27 @@ impl NetworkPolicyDecider for DefaultPolicyDecider { } } +/// Check if a host matches a pattern (supports `*.example.com` wildcards). +fn host_matches_pattern(host: &str, pattern: &str) -> bool { + let pattern_lower = pattern.to_lowercase(); + if pattern_lower == host { + return true; + } + + // Support wildcard: *.example.com matches sub.example.com + if let Some(suffix) = pattern_lower.strip_prefix("*.") + && host.ends_with(suffix) + && host.len() > suffix.len() + { + let prefix = &host[..host.len() - suffix.len()]; + if prefix.ends_with('.') || prefix.is_empty() { + return true; + } + } + + false +} + /// A policy decider that allows everything (use with FullAccess policy). pub struct AllowAllDecider; @@ -207,11 +230,10 @@ mod tests { #[tokio::test] async fn test_credential_injection() { let allowlist = DomainAllowlist::new(&["api.openai.com".to_string()]); - let credentials = vec![CredentialMapping { - domain: "api.openai.com".to_string(), - secret_name: "OPENAI_API_KEY".to_string(), - location: CredentialLocation::AuthorizationBearer, - }]; + let credentials = vec![CredentialMapping::bearer( + "OPENAI_API_KEY", + "api.openai.com", + )]; let decider = DefaultPolicyDecider::new(allowlist, credentials); let req = @@ -225,4 +247,45 @@ mod tests { _ => panic!("Expected AllowWithCredentials"), } } + + #[tokio::test] + async fn test_credential_injection_with_wildcard_host_pattern() { + let allowlist = + DomainAllowlist::new(&["api.example.com".to_string(), "sub.example.com".to_string()]); + let credentials = vec![CredentialMapping { + secret_name: "EXAMPLE_KEY".to_string(), + location: CredentialLocation::AuthorizationBearer, + host_patterns: vec!["*.example.com".to_string()], + }]; + let decider = DefaultPolicyDecider::new(allowlist, credentials); + + let req = NetworkRequest::from_url("GET", "https://api.example.com/data").unwrap(); + let decision = decider.decide(&req).await; + + match decision { + NetworkDecision::AllowWithCredentials { secret_name, .. } => { + assert_eq!(secret_name, "EXAMPLE_KEY"); + } + _ => panic!("Expected AllowWithCredentials for wildcard match"), + } + + let req2 = NetworkRequest::from_url("GET", "https://sub.example.com/data").unwrap(); + let decision2 = decider.decide(&req2).await; + assert!( + matches!(decision2, NetworkDecision::AllowWithCredentials { .. }), + "Wildcard pattern should match sub.example.com too" + ); + } + + #[test] + fn test_host_matches_pattern_exact() { + assert!(host_matches_pattern("api.openai.com", "api.openai.com")); + assert!(!host_matches_pattern("api.openai.com", "evil.com")); + } + + #[test] + fn test_host_matches_pattern_wildcard() { + assert!(host_matches_pattern("api.example.com", "*.example.com")); + assert!(!host_matches_pattern("example.com", "*.example.com")); + } } diff --git a/src/secrets/store.rs b/src/secrets/store.rs index 123c4f29..d9eb6d79 100644 --- a/src/secrets/store.rs +++ b/src/secrets/store.rs @@ -323,10 +323,15 @@ impl LibSqlSecretsStore { Self { db, crypto } } - fn connect(&self) -> Result { - self.db + async fn connect(&self) -> Result { + let conn = self + .db .connect() - .map_err(|e| SecretError::Database(format!("Connection failed: {}", e))) + .map_err(|e| SecretError::Database(format!("Connection failed: {}", e)))?; + conn.query("PRAGMA busy_timeout = 5000", ()) + .await + .map_err(|e| SecretError::Database(format!("Failed to set busy_timeout: {}", e)))?; + Ok(conn) } } @@ -349,7 +354,7 @@ impl SecretsStore for LibSqlSecretsStore { .map(|dt| dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true)); // Start transaction for atomic upsert + read-back - let conn = self.connect()?; + let conn = self.connect().await?; let tx = conn .transaction() .await @@ -410,7 +415,7 @@ impl SecretsStore for LibSqlSecretsStore { } async fn get(&self, user_id: &str, name: &str) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( r#" @@ -455,7 +460,7 @@ impl SecretsStore for LibSqlSecretsStore { } async fn exists(&self, user_id: &str, name: &str) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( "SELECT 1 FROM secrets WHERE user_id = ?1 AND name = ?2", @@ -472,7 +477,7 @@ impl SecretsStore for LibSqlSecretsStore { } async fn list(&self, user_id: &str) -> Result, SecretError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( "SELECT name, provider FROM secrets WHERE user_id = ?1 ORDER BY name", @@ -496,7 +501,7 @@ impl SecretsStore for LibSqlSecretsStore { } async fn delete(&self, user_id: &str, name: &str) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let affected = conn .execute( "DELETE FROM secrets WHERE user_id = ?1 AND name = ?2", @@ -510,7 +515,7 @@ impl SecretsStore for LibSqlSecretsStore { async fn record_usage(&self, secret_id: Uuid) -> Result<(), SecretError> { let now = Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true); - let conn = self.connect()?; + let conn = self.connect().await?; conn.execute( r#" diff --git a/src/setup/wizard.rs b/src/setup/wizard.rs index 718989a7..11a8d1e4 100644 --- a/src/setup/wizard.rs +++ b/src/setup/wizard.rs @@ -1006,7 +1006,7 @@ impl SetupWizard { use crate::llm::create_llm_provider; let base_url = std::env::var("NEARAI_BASE_URL") - .unwrap_or_else(|_| "https://cloud-api.near.ai".to_string()); + .unwrap_or_else(|_| "https://private.near.ai".to_string()); let auth_base_url = std::env::var("NEARAI_AUTH_URL") .unwrap_or_else(|_| "https://private.near.ai".to_string()); diff --git a/src/tools/builtin/job.rs b/src/tools/builtin/job.rs index 7f0e9246..1748a2cb 100644 --- a/src/tools/builtin/job.rs +++ b/src/tools/builtin/job.rs @@ -14,12 +14,57 @@ use async_trait::async_trait; use chrono::Utc; use uuid::Uuid; +use crate::channels::IncomingMessage; +use crate::channels::web::types::SseEvent; use crate::context::{ContextManager, JobContext, JobState}; use crate::db::Database; use crate::history::SandboxJobRecord; +use crate::orchestrator::auth::CredentialGrant; use crate::orchestrator::job_manager::{ContainerJobManager, JobMode}; +use crate::secrets::SecretsStore; use crate::tools::tool::{Tool, ToolError, ToolOutput, require_str}; +/// Resolve a job ID from a full UUID or a short prefix (like git short SHAs). +/// +/// Tries full UUID parse first. If that fails, treats the input as a hex prefix +/// and searches the context manager for a unique match. +async fn resolve_job_id(input: &str, context_manager: &ContextManager) -> Result { + // Fast path: full UUID + if let Ok(id) = Uuid::parse_str(input) { + return Ok(id); + } + + // Require a minimum prefix length to limit brute-force enumeration. + if input.len() < 4 { + return Err(ToolError::InvalidParameters( + "job ID prefix must be at least 4 hex characters".to_string(), + )); + } + + // Prefix match against known jobs + let input_lower = input.to_lowercase(); + let all_ids = context_manager.all_jobs().await; + let matches: Vec = all_ids + .into_iter() + .filter(|id| { + let hex = id.to_string().replace('-', ""); + hex.starts_with(&input_lower) + }) + .collect(); + + match matches.len() { + 1 => Ok(matches[0]), + 0 => Err(ToolError::InvalidParameters(format!( + "no job found matching prefix '{}'", + input + ))), + n => Err(ToolError::InvalidParameters(format!( + "ambiguous prefix '{}' matches {} jobs, provide more characters", + input, n + ))), + } +} + /// Tool for creating a new job. /// /// When sandbox deps are injected (via `with_sandbox`), the tool automatically @@ -29,6 +74,12 @@ pub struct CreateJobTool { context_manager: Arc, job_manager: Option>, store: Option>, + /// Broadcast sender for job events (used to subscribe a monitor). + event_tx: Option>, + /// Injection channel for pushing messages into the agent loop. + inject_tx: Option>, + /// Encrypted secrets store for validating credential grants. + secrets_store: Option>, } impl CreateJobTool { @@ -37,6 +88,9 @@ impl CreateJobTool { context_manager, job_manager: None, store: None, + event_tx: None, + inject_tx: None, + secrets_store: None, } } @@ -51,10 +105,98 @@ impl CreateJobTool { self } + /// Inject monitor dependencies so fire-and-forget jobs spawn a background + /// monitor that forwards Claude Code output to the main agent loop. + pub fn with_monitor_deps( + mut self, + event_tx: tokio::sync::broadcast::Sender<(Uuid, SseEvent)>, + inject_tx: tokio::sync::mpsc::Sender, + ) -> Self { + self.event_tx = Some(event_tx); + self.inject_tx = Some(inject_tx); + self + } + + /// Inject secrets store for credential validation. + pub fn with_secrets(mut self, secrets: Arc) -> Self { + self.secrets_store = Some(secrets); + self + } + pub fn sandbox_enabled(&self) -> bool { self.job_manager.is_some() } + /// Parse and validate the `credentials` parameter. + /// + /// Each key is a secret name (must exist in SecretsStore), each value is the + /// env var name the container should receive it as. Returns an empty vec if + /// no credentials were requested. + async fn parse_credentials( + &self, + params: &serde_json::Value, + user_id: &str, + ) -> Result, ToolError> { + let creds_obj = match params.get("credentials").and_then(|v| v.as_object()) { + Some(obj) if !obj.is_empty() => obj, + _ => return Ok(vec![]), + }; + + const MAX_CREDENTIAL_GRANTS: usize = 20; + if creds_obj.len() > MAX_CREDENTIAL_GRANTS { + return Err(ToolError::InvalidParameters(format!( + "too many credential grants ({}, max {})", + creds_obj.len(), + MAX_CREDENTIAL_GRANTS + ))); + } + + let secrets = match &self.secrets_store { + Some(s) => s, + None => { + return Err(ToolError::ExecutionFailed( + "credentials requested but no secrets store is configured. \ + Set SECRETS_MASTER_KEY to enable credential management." + .to_string(), + )); + } + }; + + let mut grants = Vec::with_capacity(creds_obj.len()); + for (secret_name, env_var_value) in creds_obj { + let env_var = env_var_value.as_str().ok_or_else(|| { + ToolError::InvalidParameters(format!( + "credential env var for '{}' must be a string", + secret_name + )) + })?; + + validate_env_var_name(env_var)?; + + // Validate the secret actually exists + let exists = secrets.exists(user_id, secret_name).await.map_err(|e| { + ToolError::ExecutionFailed(format!( + "failed to check secret '{}': {}", + secret_name, e + )) + })?; + + if !exists { + return Err(ToolError::ExecutionFailed(format!( + "secret '{}' not found. Store it first via 'ironclaw tool auth' or the web UI.", + secret_name + ))); + } + + grants.push(CredentialGrant { + secret_name: secret_name.clone(), + env_var: env_var.to_string(), + }); + } + + Ok(grants) + } + /// Persist a sandbox job record (fire-and-forget). fn persist_job(&self, record: SandboxJobRecord) { if let Some(store) = self.store.clone() { @@ -134,6 +276,7 @@ impl CreateJobTool { explicit_dir: Option, wait: bool, mode: JobMode, + credential_grants: Vec, ctx: &JobContext, ) -> Result { let start = std::time::Instant::now(); @@ -143,6 +286,20 @@ impl CreateJobTool { let (project_dir, browse_id) = resolve_project_dir(explicit_dir, job_id)?; let project_dir_str = project_dir.display().to_string(); + // Serialize credential grants so restarts can reload them. + let credential_grants_json = match serde_json::to_string(&credential_grants) { + Ok(json) => json, + Err(e) => { + tracing::warn!( + "Failed to serialize credential grants for job {}: {}. \ + Grants will not survive a restart.", + job_id, + e + ); + String::from("[]") + } + }; + // Persist the job to DB before creating the container. self.persist_job(SandboxJobRecord { id: job_id, @@ -155,6 +312,7 @@ impl CreateJobTool { created_at: Utc::now(), started_at: None, completed_at: None, + credential_grants_json, }); // Persist the job mode to DB @@ -174,7 +332,7 @@ impl CreateJobTool { // Create the container job with the pre-determined job_id. let _token = jm - .create_job(job_id, task, Some(project_dir), mode) + .create_job(job_id, task, Some(project_dir), mode, credential_grants) .await .map_err(|e| { self.update_status( @@ -193,10 +351,23 @@ impl CreateJobTool { self.update_status(job_id, "running", None, None, Some(now), None); if !wait { + // Spawn a background monitor that forwards Claude Code output + // into the main agent loop. + // + // This monitor is intentionally fire-and-forget: its lifetime is + // bound to the broadcast channel (etx) and the inject sender (itx). + // When the broadcast sender is dropped during shutdown the + // subscription closes and the monitor exits. Likewise, if the agent + // loop stops consuming from inject_tx the send will fail and the + // monitor terminates. No JoinHandle is retained. + if let (Some(etx), Some(itx)) = (&self.event_tx, &self.inject_tx) { + crate::agent::job_monitor::spawn_job_monitor(job_id, etx.subscribe(), itx.clone()); + } + let result = serde_json::json!({ "job_id": job_id.to_string(), "status": "started", - "message": "Container started. Use job tools to check status.", + "message": "Container started. Use job_events to check status or job_prompt to send follow-up instructions.", "project_dir": project_dir_str, "browse_url": format!("/projects/{}", browse_id), }); @@ -322,6 +493,66 @@ impl CreateJobTool { } /// The base directory where all project directories must live. +/// Env var names that could be abused to hijack process behavior. +const DANGEROUS_ENV_VARS: &[&str] = &[ + // Dynamic linker hijacking + "LD_PRELOAD", + "LD_LIBRARY_PATH", + "LD_AUDIT", + "DYLD_INSERT_LIBRARIES", + "DYLD_LIBRARY_PATH", + // Shell behavior + "BASH_ENV", + "ENV", + "CDPATH", + "IFS", + "PATH", + "HOME", + // Language runtime library path hijacking + "PYTHONPATH", + "NODE_PATH", + "PERL5LIB", + "RUBYLIB", + "CLASSPATH", + // JVM injection + "JAVA_TOOL_OPTIONS", + "MAVEN_OPTS", + "USER", + "SHELL", + "RUST_LOG", +]; + +/// Validate that an env var name is safe for container injection. +fn validate_env_var_name(name: &str) -> Result<(), ToolError> { + if name.is_empty() { + return Err(ToolError::InvalidParameters( + "env var name cannot be empty".into(), + )); + } + + // Must match ^[A-Z_][A-Z0-9_]*$ + let valid = name + .bytes() + .enumerate() + .all(|(i, b)| matches!(b, b'A'..=b'Z' | b'_') || (i > 0 && b.is_ascii_digit())); + + if !valid { + return Err(ToolError::InvalidParameters(format!( + "env var '{}' must match [A-Z_][A-Z0-9_]* (uppercase, underscores, digits)", + name + ))); + } + + if DANGEROUS_ENV_VARS.contains(&name) { + return Err(ToolError::InvalidParameters(format!( + "env var '{}' is on the denylist (could hijack process behavior)", + name + ))); + } + + Ok(()) +} + fn projects_base() -> PathBuf { dirs::home_dir() .unwrap_or_else(|| PathBuf::from(".")) @@ -352,35 +583,45 @@ fn resolve_project_dir( ToolError::ExecutionFailed(format!("failed to canonicalize projects base: {}", e)) })?; - let dir = match explicit { - Some(d) => d, - None => canonical_base.join(project_id.to_string()), + let (canonical_dir, _was_explicit) = match explicit { + Some(d) => { + // Explicit paths: validate BEFORE creating anything. + // The path must already exist (it comes from a previous job run). + let canonical = d.canonicalize().map_err(|e| { + ToolError::InvalidParameters(format!( + "explicit project dir {} does not exist or is inaccessible: {}", + d.display(), + e + )) + })?; + if !canonical.starts_with(&canonical_base) { + return Err(ToolError::InvalidParameters(format!( + "project directory must be under {}", + canonical_base.display() + ))); + } + (canonical, true) + } + None => { + let dir = canonical_base.join(project_id.to_string()); + std::fs::create_dir_all(&dir).map_err(|e| { + ToolError::ExecutionFailed(format!( + "failed to create project dir {}: {}", + dir.display(), + e + )) + })?; + let canonical = dir.canonicalize().map_err(|e| { + ToolError::ExecutionFailed(format!( + "failed to canonicalize project dir {}: {}", + dir.display(), + e + )) + })?; + (canonical, false) + } }; - std::fs::create_dir_all(&dir).map_err(|e| { - ToolError::ExecutionFailed(format!( - "failed to create project dir {}: {}", - dir.display(), - e - )) - })?; - - // Canonicalize resolves symlinks, `..`, etc. so we can do a reliable prefix check. - let canonical_dir = dir.canonicalize().map_err(|e| { - ToolError::ExecutionFailed(format!( - "failed to canonicalize project dir {}: {}", - dir.display(), - e - )) - })?; - - if !canonical_dir.starts_with(&canonical_base) { - return Err(ToolError::InvalidParameters(format!( - "project directory must be under {}", - canonical_base.display() - ))); - } - let browse_id = canonical_dir .file_name() .map(|n| n.to_string_lossy().to_string()) @@ -431,6 +672,18 @@ impl Tool for CreateJobTool { "enum": ["worker", "claude_code"], "description": "Execution mode. 'worker' (default) uses the IronClaw sub-agent. \ 'claude_code' uses Claude Code CLI for full agentic software engineering." + }, + "project_dir": { + "type": "string", + "description": "Path to an existing project directory to mount into the container. \ + Must be under ~/.ironclaw/projects/. If omitted, a fresh directory is created." + }, + "credentials": { + "type": "object", + "description": "Map of secret names to env var names. Each secret must exist in the \ + secrets store (via 'ironclaw tool auth' or web UI). Example: \ + {\"github_token\": \"GITHUB_TOKEN\", \"npm_token\": \"NPM_TOKEN\"}", + "additionalProperties": { "type": "string" } } }, "required": ["title", "description"] @@ -479,9 +732,18 @@ impl Tool for CreateJobTool { _ => JobMode::Worker, }; + let explicit_dir = params + .get("project_dir") + .and_then(|v| v.as_str()) + .map(PathBuf::from); + + // Parse and validate credential grants + let credential_grants = self.parse_credentials(¶ms, &ctx.user_id).await?; + // Combine title and description into the task prompt for the sub-agent. let task = format!("{}\n\n{}", title, description); - self.execute_sandbox(&task, None, wait, mode, ctx).await + self.execute_sandbox(&task, explicit_dir, wait, mode, credential_grants, ctx) + .await } else { self.execute_local(title, description, ctx).await } @@ -612,7 +874,7 @@ impl Tool for JobStatusTool { "properties": { "job_id": { "type": "string", - "description": "The UUID of the job to check" + "description": "The job ID (full UUID or short prefix, e.g. 'f2854dd8')" } }, "required": ["job_id"] @@ -628,10 +890,7 @@ impl Tool for JobStatusTool { let requester_id = ctx.user_id.clone(); let job_id_str = require_str(¶ms, "job_id")?; - - let job_id = Uuid::parse_str(job_id_str).map_err(|_| { - ToolError::InvalidParameters(format!("invalid job ID format: {}", job_id_str)) - })?; + let job_id = resolve_job_id(job_id_str, &self.context_manager).await?; match self.context_manager.get_context(job_id).await { Ok(job_ctx) => { @@ -694,7 +953,7 @@ impl Tool for CancelJobTool { "properties": { "job_id": { "type": "string", - "description": "The UUID of the job to cancel" + "description": "The job ID (full UUID or short prefix, e.g. 'f2854dd8')" } }, "required": ["job_id"] @@ -710,10 +969,7 @@ impl Tool for CancelJobTool { let requester_id = ctx.user_id.clone(); let job_id_str = require_str(¶ms, "job_id")?; - - let job_id = Uuid::parse_str(job_id_str).map_err(|_| { - ToolError::InvalidParameters(format!("invalid job ID format: {}", job_id_str)) - })?; + let job_id = resolve_job_id(job_id_str, &self.context_manager).await?; // Transition to cancelled state match self @@ -758,6 +1014,269 @@ impl Tool for CancelJobTool { } } +/// Tool for reading sandbox job event logs. +/// +/// Lets the main agent inspect what a running (or completed) container job has +/// been doing: messages, tool calls, results, status changes, etc. +/// +/// Events are streamed from the sandbox worker into the database via the +/// orchestrator's event pipeline. This tool queries them with a DB-level +/// `LIMIT` (default 50, configurable via the `limit` parameter) so the +/// agent sees the most recent activity without loading the full history. +pub struct JobEventsTool { + store: Arc, + context_manager: Arc, +} + +impl JobEventsTool { + pub fn new(store: Arc, context_manager: Arc) -> Self { + Self { + store, + context_manager, + } + } +} + +#[async_trait] +impl Tool for JobEventsTool { + fn name(&self) -> &str { + "job_events" + } + + fn description(&self) -> &str { + "Read the event log for a sandbox job. Shows messages, tool calls, results, \ + and status changes from the container. Use this to check what Claude Code \ + or a worker sub-agent has been doing." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "job_id": { + "type": "string", + "description": "The job ID (full UUID or short prefix, e.g. 'f2854dd8')" + }, + "limit": { + "type": "integer", + "description": "Maximum number of events to return (default 50, most recent)" + } + }, + "required": ["job_id"] + }) + } + + async fn execute( + &self, + params: serde_json::Value, + ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let job_id_str = params + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::InvalidParameters("missing 'job_id' parameter".into()))?; + + let job_id = resolve_job_id(job_id_str, &self.context_manager).await?; + + // Verify the caller owns this job. A missing context is treated as + // unauthorized to prevent leaking events after process restarts. + let job_ctx = self + .context_manager + .get_context(job_id) + .await + .map_err(|_| { + ToolError::ExecutionFailed(format!( + "job {} not found or context unavailable", + job_id + )) + })?; + + if job_ctx.user_id != ctx.user_id { + return Err(ToolError::ExecutionFailed(format!( + "job {} does not belong to current user", + job_id + ))); + } + + const MAX_EVENT_LIMIT: i64 = 1000; + let limit = params + .get("limit") + .and_then(|v| v.as_i64()) + .unwrap_or(50) + .clamp(1, MAX_EVENT_LIMIT); + + let events = self + .store + .list_job_events(job_id, Some(limit)) + .await + .map_err(|e| ToolError::ExecutionFailed(format!("failed to load job events: {}", e)))?; + + let recent: Vec = events + .iter() + .map(|ev| { + serde_json::json!({ + "event_type": ev.event_type, + "data": ev.data, + "created_at": ev.created_at.to_rfc3339(), + }) + }) + .collect(); + + let result = serde_json::json!({ + "job_id": job_id.to_string(), + "total_events": events.len(), + "returned": recent.len(), + "events": recent, + }); + + Ok(ToolOutput::success(result, start.elapsed())) + } + + fn requires_sanitization(&self) -> bool { + true + } +} + +/// Tool for sending follow-up prompts to a running Claude Code sandbox job. +/// +/// The prompt is queued in an in-memory `PromptQueue` (a broadcast channel +/// shared with the web gateway). The Claude Code bridge inside the container +/// polls for queued prompts between turns and feeds them into the next +/// `claude --resume` invocation, enabling interactive multi-turn sessions +/// with long-running sandbox jobs. +pub struct JobPromptTool { + prompt_queue: PromptQueue, + context_manager: Arc, +} + +/// Type alias matching `crate::channels::web::server::PromptQueue`. +pub type PromptQueue = Arc< + tokio::sync::Mutex< + std::collections::HashMap< + Uuid, + std::collections::VecDeque, + >, + >, +>; + +impl JobPromptTool { + pub fn new(prompt_queue: PromptQueue, context_manager: Arc) -> Self { + Self { + prompt_queue, + context_manager, + } + } +} + +#[async_trait] +impl Tool for JobPromptTool { + fn name(&self) -> &str { + "job_prompt" + } + + fn description(&self) -> &str { + "Send a follow-up prompt to a running Claude Code sandbox job. The prompt is \ + queued and delivered on the next poll cycle. Use this to give the sub-agent \ + additional instructions, answer its questions, or tell it to wrap up." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "job_id": { + "type": "string", + "description": "The job ID (full UUID or short prefix, e.g. 'f2854dd8')" + }, + "content": { + "type": "string", + "description": "The follow-up prompt text to send" + }, + "done": { + "type": "boolean", + "description": "If true, signals the sub-agent that no more prompts are coming \ + and it should finish up. Default false." + } + }, + "required": ["job_id", "content"] + }) + } + + async fn execute( + &self, + params: serde_json::Value, + ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let job_id_str = params + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::InvalidParameters("missing 'job_id' parameter".into()))?; + + let job_id = resolve_job_id(job_id_str, &self.context_manager).await?; + + // Verify the caller owns this job. A missing context is treated as + // unauthorized to prevent sending prompts to jobs after process restarts. + let job_ctx = self + .context_manager + .get_context(job_id) + .await + .map_err(|_| { + ToolError::ExecutionFailed(format!( + "job {} not found or context unavailable", + job_id + )) + })?; + + if job_ctx.user_id != ctx.user_id { + return Err(ToolError::ExecutionFailed(format!( + "job {} does not belong to current user", + job_id + ))); + } + + let content = params + .get("content") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::InvalidParameters("missing 'content' parameter".into()))?; + + let done = params + .get("done") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + let prompt = crate::orchestrator::api::PendingPrompt { + content: content.to_string(), + done, + }; + + { + let mut queue = self.prompt_queue.lock().await; + queue.entry(job_id).or_default().push_back(prompt); + } + + let result = serde_json::json!({ + "job_id": job_id.to_string(), + "status": "queued", + "message": "Prompt queued", + "done": done, + }); + + Ok(ToolOutput::success(result, start.elapsed())) + } + + fn requires_approval(&self) -> bool { + true + } + + fn requires_sanitization(&self) -> bool { + false + } +} + #[cfg(test)] mod tests { use super::*; @@ -796,10 +1315,6 @@ mod tests { let props = schema.get("properties").unwrap().as_object().unwrap(); assert!(props.contains_key("title")); assert!(props.contains_key("description")); - assert!( - !props.contains_key("project_dir"), - "project_dir must not be exposed to the LLM" - ); assert!(!props.contains_key("wait")); assert!(!props.contains_key("mode")); } @@ -870,6 +1385,8 @@ mod tests { let base = projects_base(); std::fs::create_dir_all(&base).unwrap(); let explicit = base.join("test_explicit_project"); + // Explicit paths must already exist (no auto-create). + std::fs::create_dir_all(&explicit).unwrap(); let project_id = Uuid::new_v4(); let (dir, browse_id) = resolve_project_dir(Some(explicit.clone()), project_id).unwrap(); @@ -886,10 +1403,28 @@ mod tests { fn test_resolve_project_dir_rejects_outside_base() { let tmp = tempfile::tempdir().unwrap(); let escape_attempt = tmp.path().join("evil_project"); + // Don't create it: explicit paths that don't exist are rejected + // before the prefix check even runs. let result = resolve_project_dir(Some(escape_attempt), Uuid::new_v4()); assert!(result.is_err()); let err = result.unwrap_err().to_string(); + assert!( + err.contains("does not exist"), + "expected 'does not exist' error, got: {}", + err + ); + } + + #[test] + fn test_resolve_project_dir_rejects_outside_base_existing() { + // A directory that exists but is outside the projects base. + let tmp = tempfile::tempdir().unwrap(); + let outside = tmp.path().to_path_buf(); + + let result = resolve_project_dir(Some(outside), Uuid::new_v4()); + assert!(result.is_err()); + let err = result.unwrap_err().to_string(); assert!( err.contains("must be under"), "expected 'must be under' error, got: {}", @@ -899,17 +1434,346 @@ mod tests { #[test] fn test_resolve_project_dir_rejects_traversal() { - // Attempt to escape via `..` components + // Non-existent traversal path is rejected because canonicalize fails. let base = projects_base(); let traversal = base.join("legit").join("..").join("..").join(".ssh"); let result = resolve_project_dir(Some(traversal), Uuid::new_v4()); + assert!(result.is_err(), "traversal path should be rejected"); + + // Traversal path that actually resolves gets the prefix check. + // `base/../` resolves to the parent of projects base, which is outside. + let base_parent = projects_base().join("..").join("definitely_not_projects"); + std::fs::create_dir_all(&base_parent).ok(); + if base_parent.exists() { + let result = resolve_project_dir(Some(base_parent.clone()), Uuid::new_v4()); + assert!(result.is_err(), "path outside base should be rejected"); + let _ = std::fs::remove_dir_all(&base_parent); + } + } + + #[test] + fn test_sandbox_schema_includes_project_dir() { + let manager = Arc::new(ContextManager::new(5)); + let jm = Arc::new(ContainerJobManager::new( + crate::orchestrator::job_manager::ContainerJobConfig::default(), + crate::orchestrator::TokenStore::new(), + )); + let tool = CreateJobTool::new(manager).with_sandbox(jm, None); + let schema = tool.parameters_schema(); + let props = schema.get("properties").unwrap().as_object().unwrap(); + assert!( + props.contains_key("project_dir"), + "sandbox schema must expose project_dir" + ); + } + + #[test] + fn test_sandbox_schema_includes_credentials() { + let manager = Arc::new(ContextManager::new(5)); + let jm = Arc::new(ContainerJobManager::new( + crate::orchestrator::job_manager::ContainerJobConfig::default(), + crate::orchestrator::TokenStore::new(), + )); + let tool = CreateJobTool::new(manager).with_sandbox(jm, None); + let schema = tool.parameters_schema(); + let props = schema.get("properties").unwrap().as_object().unwrap(); + assert!( + props.contains_key("credentials"), + "sandbox schema must expose credentials" + ); + } + + #[tokio::test] + async fn test_parse_credentials_empty() { + let manager = Arc::new(ContextManager::new(5)); + let tool = CreateJobTool::new(manager); + + // No credentials parameter + let params = serde_json::json!({"title": "t", "description": "d"}); + let grants = tool.parse_credentials(¶ms, "user1").await.unwrap(); + assert!(grants.is_empty()); + + // Empty credentials object + let params = serde_json::json!({"credentials": {}}); + let grants = tool.parse_credentials(¶ms, "user1").await.unwrap(); + assert!(grants.is_empty()); + } + + #[tokio::test] + async fn test_parse_credentials_no_secrets_store() { + let manager = Arc::new(ContextManager::new(5)); + let tool = CreateJobTool::new(manager); + + let params = serde_json::json!({"credentials": {"my_secret": "MY_SECRET"}}); + let result = tool.parse_credentials(¶ms, "user1").await; assert!(result.is_err()); let err = result.unwrap_err().to_string(); assert!( - err.contains("must be under"), - "expected 'must be under' error, got: {}", + err.contains("no secrets store"), + "expected 'no secrets store' error, got: {}", err ); } + + #[tokio::test] + async fn test_parse_credentials_missing_secret() { + use crate::secrets::{InMemorySecretsStore, SecretsCrypto}; + use secrecy::SecretString; + + let manager = Arc::new(ContextManager::new(5)); + let key = "0123456789abcdef0123456789abcdef"; + let crypto = Arc::new(SecretsCrypto::new(SecretString::from(key.to_string())).unwrap()); + let secrets: Arc = + Arc::new(InMemorySecretsStore::new(crypto)); + + let tool = CreateJobTool::new(manager).with_secrets(Arc::clone(&secrets)); + + let params = serde_json::json!({"credentials": {"nonexistent_secret": "SOME_VAR"}}); + let result = tool.parse_credentials(¶ms, "user1").await; + assert!(result.is_err()); + let err = result.unwrap_err().to_string(); + assert!( + err.contains("not found"), + "expected 'not found' error, got: {}", + err + ); + } + + #[tokio::test] + async fn test_parse_credentials_valid() { + use crate::secrets::{CreateSecretParams, InMemorySecretsStore, SecretsCrypto}; + use secrecy::SecretString; + + let manager = Arc::new(ContextManager::new(5)); + let key = "0123456789abcdef0123456789abcdef"; + let crypto = Arc::new(SecretsCrypto::new(SecretString::from(key.to_string())).unwrap()); + let secrets: Arc = + Arc::new(InMemorySecretsStore::new(Arc::clone(&crypto))); + + // Store a secret + secrets + .create( + "user1", + CreateSecretParams::new("github_token", "ghp_test123"), + ) + .await + .unwrap(); + + let tool = CreateJobTool::new(manager).with_secrets(Arc::clone(&secrets)); + + let params = serde_json::json!({ + "credentials": {"github_token": "GITHUB_TOKEN"} + }); + let grants = tool.parse_credentials(¶ms, "user1").await.unwrap(); + assert_eq!(grants.len(), 1); + assert_eq!(grants[0].secret_name, "github_token"); + assert_eq!(grants[0].env_var, "GITHUB_TOKEN"); + } + + fn test_prompt_tool(queue: PromptQueue) -> JobPromptTool { + let cm = Arc::new(ContextManager::new(5)); + JobPromptTool::new(queue, cm) + } + + #[tokio::test] + async fn test_job_prompt_tool_queues_prompt() { + let cm = Arc::new(ContextManager::new(5)); + let job_id = cm + .create_job_for_user("default", "Test Job", "desc") + .await + .unwrap(); + + let queue: PromptQueue = + Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())); + let tool = JobPromptTool::new(Arc::clone(&queue), cm); + + let params = serde_json::json!({ + "job_id": job_id.to_string(), + "content": "What's the status?", + "done": false, + }); + + let ctx = JobContext::default(); + let result = tool.execute(params, &ctx).await.unwrap(); + + assert_eq!( + result.result.get("status").unwrap().as_str().unwrap(), + "queued" + ); + + let q = queue.lock().await; + let prompts = q.get(&job_id).unwrap(); + assert_eq!(prompts.len(), 1); + assert_eq!(prompts[0].content, "What's the status?"); + assert!(!prompts[0].done); + } + + #[tokio::test] + async fn test_job_prompt_tool_requires_approval() { + let queue: PromptQueue = + Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())); + let tool = test_prompt_tool(queue); + assert!(tool.requires_approval()); + } + + #[tokio::test] + async fn test_job_prompt_tool_rejects_invalid_uuid() { + let queue: PromptQueue = + Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())); + let tool = test_prompt_tool(queue); + + let params = serde_json::json!({ + "job_id": "not-a-uuid", + "content": "hello", + }); + + let ctx = JobContext::default(); + let result = tool.execute(params, &ctx).await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn test_job_prompt_tool_rejects_missing_content() { + let queue: PromptQueue = + Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())); + let tool = test_prompt_tool(queue); + + let params = serde_json::json!({ + "job_id": Uuid::new_v4().to_string(), + }); + + let ctx = JobContext::default(); + let result = tool.execute(params, &ctx).await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn test_job_events_tool_rejects_other_users_job() { + // JobEventsTool needs a Store (PostgreSQL) for the full path, but the + // ownership check happens first via ContextManager, so we can test that + // without a database by using a Store that will never be reached. + // + // We construct the tool by hand: the store field is never touched + // because the ownership check short-circuits before the query. + let cm = Arc::new(ContextManager::new(5)); + let job_id = cm + .create_job_for_user("owner-user", "Secret Job", "classified") + .await + .unwrap(); + + // We need a Store to construct the tool, but creating one requires + // a database URL. Instead, test the ownership logic directly: + // simulate what execute() does. + let attacker_ctx = JobContext { + user_id: "attacker".to_string(), + ..Default::default() + }; + + let job_ctx = cm.get_context(job_id).await.unwrap(); + assert_ne!(job_ctx.user_id, attacker_ctx.user_id); + assert_eq!(job_ctx.user_id, "owner-user"); + } + + #[test] + fn test_job_events_tool_schema() { + // Verify the schema shape is correct (doesn't need a Store instance). + let schema = serde_json::json!({ + "type": "object", + "properties": { + "job_id": { + "type": "string", + "description": "The job ID (full UUID or short prefix, e.g. 'f2854dd8')" + }, + "limit": { + "type": "integer", + "description": "Maximum number of events to return (default 50, most recent)" + } + }, + "required": ["job_id"] + }); + + let props = schema.get("properties").unwrap().as_object().unwrap(); + assert!(props.contains_key("job_id")); + assert!(props.contains_key("limit")); + let required = schema.get("required").unwrap().as_array().unwrap(); + assert_eq!(required.len(), 1); + assert_eq!(required[0].as_str().unwrap(), "job_id"); + } + + #[tokio::test] + async fn test_job_prompt_tool_rejects_other_users_job() { + let cm = Arc::new(ContextManager::new(5)); + let job_id = cm + .create_job_for_user("owner-user", "Test Job", "desc") + .await + .unwrap(); + + let queue: PromptQueue = + Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::new())); + let tool = JobPromptTool::new(queue, cm); + + let params = serde_json::json!({ + "job_id": job_id.to_string(), + "content": "sneaky prompt", + }); + + // Attacker context with a different user_id. + let ctx = JobContext { + user_id: "attacker".to_string(), + ..Default::default() + }; + + let result = tool.execute(params, &ctx).await; + assert!(result.is_err()); + let err = result.unwrap_err().to_string(); + assert!( + err.contains("does not belong to current user"), + "expected ownership error, got: {}", + err + ); + } + + #[tokio::test] + async fn test_resolve_job_id_full_uuid() { + let cm = ContextManager::new(5); + let job_id = cm.create_job("Test", "Desc").await.unwrap(); + + let resolved = resolve_job_id(&job_id.to_string(), &cm).await.unwrap(); + assert_eq!(resolved, job_id); + } + + #[tokio::test] + async fn test_resolve_job_id_short_prefix() { + let cm = ContextManager::new(5); + let job_id = cm.create_job("Test", "Desc").await.unwrap(); + + // Use first 8 hex chars (without dashes) + let hex = job_id.to_string().replace('-', ""); + let prefix = &hex[..8]; + let resolved = resolve_job_id(prefix, &cm).await.unwrap(); + assert_eq!(resolved, job_id); + } + + #[tokio::test] + async fn test_resolve_job_id_no_match() { + let cm = ContextManager::new(5); + cm.create_job("Test", "Desc").await.unwrap(); + + let result = resolve_job_id("00000000", &cm).await; + assert!(result.is_err()); + let err = result.unwrap_err().to_string(); + assert!( + err.contains("no job found"), + "expected 'no job found', got: {}", + err + ); + } + + #[tokio::test] + async fn test_resolve_job_id_invalid_input() { + let cm = ContextManager::new(5); + let result = resolve_job_id("not-hex-at-all!", &cm).await; + assert!(result.is_err()); + } } diff --git a/src/tools/builtin/mod.rs b/src/tools/builtin/mod.rs index af06ebee..2cdba652 100644 --- a/src/tools/builtin/mod.rs +++ b/src/tools/builtin/mod.rs @@ -18,7 +18,10 @@ pub use extension_tools::{ }; pub use file::{ApplyPatchTool, ListDirTool, ReadFileTool, WriteFileTool}; pub use http::HttpTool; -pub use job::{CancelJobTool, CreateJobTool, JobStatusTool, ListJobsTool}; +pub use job::{ + CancelJobTool, CreateJobTool, JobEventsTool, JobPromptTool, JobStatusTool, ListJobsTool, + PromptQueue, +}; pub use json::JsonTool; pub use memory::{MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool}; pub use routine::{ diff --git a/src/tools/builtin/shell.rs b/src/tools/builtin/shell.rs index 147345ee..c5c0190f 100644 --- a/src/tools/builtin/shell.rs +++ b/src/tools/builtin/shell.rs @@ -18,7 +18,7 @@ //! - Commands run directly on host with basic protections //! - Blocked command patterns are still enforced -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::path::{Path, PathBuf}; use std::process::Stdio; use std::sync::{Arc, LazyLock}; @@ -246,6 +246,7 @@ impl ShellTool { cmd: &str, workdir: &PathBuf, timeout: Duration, + extra_env: &HashMap, ) -> Result<(String, i32), ToolError> { // Build command let mut command = if cfg!(target_os = "windows") { @@ -258,6 +259,10 @@ impl ShellTool { c }; + // Inject extra environment variables (e.g., credentials fetched by the + // worker runtime) into the child process without mutating the global env. + command.envs(extra_env); + command .current_dir(workdir) .stdin(Stdio::null()) @@ -322,6 +327,7 @@ impl ShellTool { cmd: &str, workdir: Option<&str>, timeout: Option, + extra_env: &HashMap, ) -> Result<(String, i64), ToolError> { // Check for blocked commands if let Some(reason) = self.is_blocked(cmd) { @@ -352,7 +358,9 @@ impl ShellTool { } // Only execute directly when no sandbox was configured at all. - let (output, code) = self.execute_direct(cmd, &cwd, timeout_duration).await?; + let (output, code) = self + .execute_direct(cmd, &cwd, timeout_duration, extra_env) + .await?; Ok((output, code as i64)) } } @@ -399,7 +407,7 @@ impl Tool for ShellTool { async fn execute( &self, params: serde_json::Value, - _ctx: &JobContext, + ctx: &JobContext, ) -> Result { let command = require_str(¶ms, "command")?; @@ -407,7 +415,9 @@ impl Tool for ShellTool { let timeout = params.get("timeout").and_then(|v| v.as_u64()); let start = std::time::Instant::now(); - let (output, exit_code) = self.execute_command(command, workdir, timeout).await?; + let (output, exit_code) = self + .execute_command(command, workdir, timeout, &ctx.extra_env) + .await?; let duration = start.elapsed(); let sandboxed = self.sandbox.is_some(); diff --git a/src/tools/mod.rs b/src/tools/mod.rs index d26f79d3..a731d46a 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -13,7 +13,6 @@ pub mod mcp; pub mod wasm; mod registry; -mod sandbox; mod tool; pub use builder::{ @@ -22,5 +21,4 @@ pub use builder::{ TestCase, TestHarness, TestResult, TestSuite, ValidationError, ValidationResult, WasmValidator, }; pub use registry::ToolRegistry; -pub use sandbox::ToolSandbox; pub use tool::{Tool, ToolDomain, ToolError, ToolOutput}; diff --git a/src/tools/registry.rs b/src/tools/registry.rs index d348a243..9a66daa0 100644 --- a/src/tools/registry.rs +++ b/src/tools/registry.rs @@ -16,11 +16,11 @@ use crate::skills::catalog::SkillCatalog; use crate::skills::registry::SkillRegistry; use crate::tools::builder::{BuildSoftwareTool, BuilderConfig, LlmSoftwareBuilder}; use crate::tools::builtin::{ - ApplyPatchTool, CancelJobTool, CreateJobTool, EchoTool, HttpTool, JobStatusTool, JsonTool, - ListDirTool, ListJobsTool, MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool, - ReadFileTool, ShellTool, SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool, - TimeTool, ToolActivateTool, ToolAuthTool, ToolInstallTool, ToolListTool, ToolRemoveTool, - ToolSearchTool, WriteFileTool, + ApplyPatchTool, CancelJobTool, CreateJobTool, EchoTool, HttpTool, JobEventsTool, JobPromptTool, + JobStatusTool, JsonTool, ListDirTool, ListJobsTool, MemoryReadTool, MemorySearchTool, + MemoryTreeTool, MemoryWriteTool, PromptQueue, ReadFileTool, ShellTool, SkillInstallTool, + SkillListTool, SkillRemoveTool, SkillSearchTool, TimeTool, ToolActivateTool, ToolAuthTool, + ToolInstallTool, ToolListTool, ToolRemoveTool, ToolSearchTool, WriteFileTool, }; use crate::tools::tool::{Tool, ToolDomain}; use crate::tools::wasm::{ @@ -247,22 +247,56 @@ impl ToolRegistry { /// Job tools allow the LLM to create, list, check status, and cancel jobs. /// When sandbox deps are provided, `create_job` automatically delegates to /// Docker containers. Otherwise it creates in-memory jobs via ContextManager. + #[allow(clippy::too_many_arguments)] pub fn register_job_tools( &self, context_manager: Arc, job_manager: Option>, store: Option>, + job_event_tx: Option< + tokio::sync::broadcast::Sender<(uuid::Uuid, crate::channels::web::types::SseEvent)>, + >, + inject_tx: Option>, + prompt_queue: Option, + secrets_store: Option>, ) { let mut create_tool = CreateJobTool::new(Arc::clone(&context_manager)); if let Some(jm) = job_manager { - create_tool = create_tool.with_sandbox(jm, store); + create_tool = create_tool.with_sandbox(jm, store.clone()); + } + if let (Some(etx), Some(itx)) = (job_event_tx, inject_tx) { + create_tool = create_tool.with_monitor_deps(etx, itx); + } + if let Some(secrets) = secrets_store { + create_tool = create_tool.with_secrets(secrets); } self.register_sync(Arc::new(create_tool)); self.register_sync(Arc::new(ListJobsTool::new(Arc::clone(&context_manager)))); self.register_sync(Arc::new(JobStatusTool::new(Arc::clone(&context_manager)))); - self.register_sync(Arc::new(CancelJobTool::new(context_manager))); + self.register_sync(Arc::new(CancelJobTool::new(Arc::clone(&context_manager)))); - tracing::info!("Registered 4 job management tools"); + // Base tools: create, list, status, cancel + let mut job_tool_count = 4; + + // Register event reader if store is available + if let Some(store) = store { + self.register_sync(Arc::new(JobEventsTool::new( + store, + Arc::clone(&context_manager), + ))); + job_tool_count += 1; + } + + // Register prompt tool if queue is available + if let Some(pq) = prompt_queue { + self.register_sync(Arc::new(JobPromptTool::new( + pq, + Arc::clone(&context_manager), + ))); + job_tool_count += 1; + } + + tracing::info!("Registered {} job management tools", job_tool_count); } /// Register extension management tools (search, install, auth, activate, list, remove). diff --git a/src/tools/sandbox.rs b/src/tools/sandbox.rs deleted file mode 100644 index ed605429..00000000 --- a/src/tools/sandbox.rs +++ /dev/null @@ -1,136 +0,0 @@ -//! Sandboxed tool execution environment. -//! -//! NOTE: For WASM-based sandboxing with full security, use the `wasm` module instead. -//! This module provides a simpler process-based sandbox for scripts. - -use std::time::Duration; - -use crate::tools::tool::ToolError; - -/// Configuration for the sandbox. -#[derive(Debug, Clone)] -pub struct SandboxConfig { - /// Maximum execution time. - pub max_execution_time: Duration, - /// Maximum memory in bytes. - pub max_memory_bytes: u64, - /// Allowed network hosts (empty = no network). - pub allowed_hosts: Vec, - /// Allowed filesystem paths (empty = no filesystem). - pub allowed_paths: Vec, - /// Environment variables to pass. - pub env_vars: Vec<(String, String)>, -} - -impl Default for SandboxConfig { - fn default() -> Self { - Self { - max_execution_time: Duration::from_secs(30), - max_memory_bytes: 128 * 1024 * 1024, // 128 MB - allowed_hosts: vec![], - allowed_paths: vec![], - env_vars: vec![], - } - } -} - -/// Result of a sandboxed execution. -#[derive(Debug)] -pub struct SandboxResult { - /// Standard output. - pub stdout: String, - /// Standard error. - pub stderr: String, - /// Exit code. - pub exit_code: i32, - /// Execution time. - pub duration: Duration, - /// Memory used (if available). - pub memory_used: Option, -} - -/// Sandbox for executing untrusted code. -pub struct ToolSandbox { - #[allow(dead_code)] // Will be used when sandbox execution is implemented - config: SandboxConfig, -} - -impl ToolSandbox { - /// Create a new sandbox with the given configuration. - pub fn new(config: SandboxConfig) -> Self { - Self { config } - } - - /// Execute code in the sandbox. - /// - /// Currently supports: - /// - Python scripts - /// - JavaScript/Node.js scripts - /// - Shell scripts (limited) - /// - /// TODO: Implement WASM-based sandboxing for better isolation. - pub async fn execute( - &self, - code: &str, - language: &str, - input: &str, - ) -> Result { - // TODO: Implement actual sandboxed execution - // Options: - // 1. WASM (wasmtime) - Best isolation but limited language support - // 2. Docker containers - Good isolation but slower startup - // 3. Process isolation with seccomp/AppArmor - Linux-specific - // 4. Firecracker microVMs - Best isolation but complex - - match language { - "python" => self.execute_python(code, input).await, - "javascript" | "js" => self.execute_javascript(code, input).await, - _ => Err(ToolError::Sandbox(format!( - "Unsupported language: {}", - language - ))), - } - } - - async fn execute_python(&self, _code: &str, _input: &str) -> Result { - // TODO: Execute Python in sandbox - Err(ToolError::Sandbox( - "Python sandbox execution not yet implemented".to_string(), - )) - } - - async fn execute_javascript( - &self, - _code: &str, - _input: &str, - ) -> Result { - // TODO: Execute JavaScript in sandbox (could use Deno or isolated V8) - Err(ToolError::Sandbox( - "JavaScript sandbox execution not yet implemented".to_string(), - )) - } - - /// Check if the sandbox is available. - pub fn is_available() -> bool { - // TODO: Check for required runtime components - false - } -} - -impl Default for ToolSandbox { - fn default() -> Self { - Self::new(SandboxConfig::default()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_sandbox_config_default() { - let config = SandboxConfig::default(); - assert_eq!(config.max_execution_time, Duration::from_secs(30)); - assert!(config.allowed_hosts.is_empty()); - } -} diff --git a/src/tools/wasm/storage.rs b/src/tools/wasm/storage.rs index 40d9857f..a223c247 100644 --- a/src/tools/wasm/storage.rs +++ b/src/tools/wasm/storage.rs @@ -581,10 +581,17 @@ impl LibSqlWasmToolStore { Self { db } } - fn connect(&self) -> Result { - self.db + async fn connect(&self) -> Result { + let conn = self + .db .connect() - .map_err(|e| WasmStorageError::Database(format!("Connection failed: {}", e))) + .map_err(|e| WasmStorageError::Database(format!("Connection failed: {}", e)))?; + conn.query("PRAGMA busy_timeout = 5000", ()) + .await + .map_err(|e| { + WasmStorageError::Database(format!("Failed to set busy_timeout: {}", e)) + })?; + Ok(conn) } } @@ -599,7 +606,7 @@ impl WasmToolStore for LibSqlWasmToolStore { .map_err(|e| WasmStorageError::InvalidData(e.to_string()))?; // Wrap INSERT + read-back in a transaction to prevent TOCTOU races - let conn = self.connect()?; + let conn = self.connect().await?; let tx = conn .transaction() .await @@ -671,7 +678,7 @@ impl WasmToolStore for LibSqlWasmToolStore { } async fn get(&self, user_id: &str, name: &str) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( r#" @@ -709,7 +716,7 @@ impl WasmToolStore for LibSqlWasmToolStore { user_id: &str, name: &str, ) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( r#" @@ -768,7 +775,7 @@ impl WasmToolStore for LibSqlWasmToolStore { &self, tool_id: Uuid, ) -> Result, WasmStorageError> { - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( r#" @@ -838,7 +845,7 @@ impl WasmToolStore for LibSqlWasmToolStore { async fn list(&self, user_id: &str) -> Result, WasmStorageError> { // SQLite doesn't have DISTINCT ON, so we use a subquery to get latest version per name - let conn = self.connect()?; + let conn = self.connect().await?; let mut rows = conn .query( r#" @@ -877,7 +884,7 @@ impl WasmToolStore for LibSqlWasmToolStore { status: ToolStatus, ) -> Result<(), WasmStorageError> { let now = Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true); - let conn = self.connect()?; + let conn = self.connect().await?; let result = conn .execute( @@ -895,7 +902,7 @@ impl WasmToolStore for LibSqlWasmToolStore { } async fn delete(&self, user_id: &str, name: &str) -> Result { - let conn = self.connect()?; + let conn = self.connect().await?; let result = conn .execute( "DELETE FROM wasm_tools WHERE user_id = ?1 AND name = ?2", diff --git a/src/worker/api.rs b/src/worker/api.rs index 756d4d55..1765d66b 100644 --- a/src/worker/api.rs +++ b/src/worker/api.rs @@ -94,6 +94,15 @@ pub struct PromptResponse { pub done: bool, } +/// A single credential delivered from the orchestrator to a container worker. +/// +/// Shared between the orchestrator endpoint and the worker client. +#[derive(Debug, Serialize, Deserialize)] +pub struct CredentialResponse { + pub env_var: String, + pub value: String, +} + impl WorkerHttpClient { /// Create a new client from environment. /// @@ -335,6 +344,45 @@ impl WorkerHttpClient { Ok(Some(prompt)) } + /// Fetch credentials granted to this job from the orchestrator. + /// + /// Returns an empty vec if no credentials are granted (204 No Content) + /// or if the endpoint returns 404. The caller should set each credential + /// as an environment variable before starting the execution loop. + pub async fn fetch_credentials(&self) -> Result, WorkerError> { + let resp = self + .client + .get(self.url("credentials")) + .bearer_auth(&self.token) + .send() + .await + .map_err(|e| WorkerError::ConnectionFailed { + url: self.orchestrator_url.clone(), + reason: e.to_string(), + })?; + + // 204 or 404 means no credentials granted, not an error + if resp.status() == reqwest::StatusCode::NO_CONTENT + || resp.status() == reqwest::StatusCode::NOT_FOUND + { + return Ok(vec![]); + } + + if !resp.status().is_success() { + return Err(WorkerError::SecretResolveFailed { + secret_name: "(all)".to_string(), + reason: format!("credentials endpoint returned {}", resp.status()), + }); + } + + resp.json() + .await + .map_err(|e| WorkerError::SecretResolveFailed { + secret_name: "(all)".to_string(), + reason: format!("failed to parse credentials response: {}", e), + }) + } + /// Signal job completion to the orchestrator. pub async fn report_complete(&self, report: &CompletionReport) -> Result<(), WorkerError> { let _: serde_json::Value = self @@ -382,6 +430,23 @@ mod tests { assert_eq!(parse_finish_reason("unknown"), FinishReason::Unknown); } + #[test] + fn test_credentials_url_construction() { + let client = WorkerHttpClient::new( + "http://host.docker.internal:50051".to_string(), + Uuid::nil(), + "test-token".to_string(), + ); + + assert_eq!( + client.url("credentials"), + format!( + "http://host.docker.internal:50051/worker/{}/credentials", + Uuid::nil() + ) + ); + } + #[test] fn test_job_description_deserialization() { let json = r#"{"title":"Test","description":"desc","project_dir":null}"#; diff --git a/src/worker/claude_bridge.rs b/src/worker/claude_bridge.rs index 1fbc1410..b2f674cf 100644 --- a/src/worker/claude_bridge.rs +++ b/src/worker/claude_bridge.rs @@ -50,64 +50,81 @@ pub struct ClaudeBridgeConfig { /// A Claude Code streaming event (NDJSON line from `--output-format stream-json`). /// -/// Claude Code emits one JSON object per line. We capture the key fields -/// we need and forward the rest as opaque data. +/// Claude Code emits one JSON object per line with these top-level types: +/// +/// system -> session init (session_id, tools, model) +/// assistant -> LLM response, nested under message.content[] as text/tool_use blocks +/// user -> tool results, nested under message.content[] as tool_result blocks +/// result -> final summary (is_error, duration_ms, num_turns, result text) +/// +/// Content blocks live under `message.content`, NOT at the top level. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ClaudeStreamEvent { #[serde(rename = "type")] pub event_type: String, - /// For `system` events: the session ID. #[serde(default)] pub session_id: Option, - /// For `assistant` events: the text content blocks. - #[serde(default)] - pub content: Option>, - - /// For `result` events: final status info. - #[serde(default)] - pub result: Option, - - /// For `tool_use`/`tool_result`: the tool name. - #[serde(default)] - pub tool_name: Option, - - /// For `tool_use`: the input parameters. - #[serde(default)] - pub input: Option, - - /// For `tool_result`: the output content. - #[serde(default)] - pub output: Option, - - /// Subtype discriminator (e.g. "text", "tool_use", "tool_result"). #[serde(default)] pub subtype: Option, + + /// For `assistant` and `user` events: the message wrapper containing content blocks. + #[serde(default)] + pub message: Option, + + /// For `result` events: the final text output. + #[serde(default)] + pub result: Option, + + /// For `result` events: whether the session ended in error. + #[serde(default)] + pub is_error: Option, + + /// For `result` events: total wall-clock duration. + #[serde(default)] + pub duration_ms: Option, + + /// For `result` events: number of agentic turns used. + #[serde(default)] + pub num_turns: Option, +} + +/// Wrapper around the `message` field in assistant/user events. +/// +/// ```text +/// { "type": "assistant", "message": { "content": [ { "type": "text", ... } ] } } +/// ``` +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MessageWrapper { + #[serde(default)] + pub role: Option, + #[serde(default)] + pub content: Option>, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ContentBlock { #[serde(rename = "type")] pub block_type: String, + /// Text block content. #[serde(default)] pub text: Option, + /// Tool name (for tool_use blocks). #[serde(default)] pub name: Option, + /// Tool use ID (for tool_use and tool_result blocks). + #[serde(default)] + pub id: Option, + /// Tool input params (for tool_use blocks). #[serde(default)] pub input: Option, + /// Tool result content (for tool_result blocks), or general content. #[serde(default)] - pub content: Option, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ResultInfo { + pub content: Option, + /// Tool use ID reference (for tool_result blocks). #[serde(default)] - pub is_error: Option, - #[serde(default)] - pub duration_ms: Option, - #[serde(default)] - pub num_turns: Option, + pub tool_use_id: Option, } /// The Claude Code bridge runtime. @@ -153,8 +170,46 @@ impl ClaudeBridgeRuntime { Ok(()) } + /// Copy auth files from a read-only source into the writable home dir. + /// + /// If the orchestrator bind-mounts the host's `~/.claude` at + /// `/home/sandbox/.claude-host:ro`, this copies everything into the + /// container's own `/home/sandbox/.claude` so Claude Code can read auth + /// credentials AND write its state (todos, debug files, etc.) without + /// touching the host filesystem. + /// + /// When no host mount is present (the default orchestrator injects + /// credentials via environment variables), this is a no-op. + fn copy_auth_from_mount(&self) -> Result<(), WorkerError> { + let mount = std::path::Path::new("/home/sandbox/.claude-host"); + if !mount.exists() { + return Ok(()); + } + + let target = std::path::Path::new("/home/sandbox/.claude"); + std::fs::create_dir_all(target).map_err(|e| WorkerError::ExecutionFailed { + reason: format!("failed to create ~/.claude: {e}"), + })?; + + let copied = + copy_dir_recursive(mount, target).map_err(|e| WorkerError::ExecutionFailed { + reason: format!("failed to copy auth from host mount: {e}"), + })?; + + tracing::info!( + job_id = %self.config.job_id, + files_copied = copied, + "Copied auth config from host mount into container" + ); + Ok(()) + } + /// Run the bridge: fetch job, spawn claude, stream events, handle follow-ups. pub async fn run(&self) -> Result<(), WorkerError> { + // Copy auth files from read-only host mount (if present) into the + // writable home directory before Claude Code needs them. + self.copy_auth_from_mount()?; + // Write project-level settings with explicit tool allowlist. // This replaces --dangerously-skip-permissions with defense-in-depth: // only the listed tools are auto-approved, unknown tools fail safely. @@ -169,6 +224,34 @@ impl ClaudeBridgeRuntime { truncate(&job.description, 100) ); + // Fetch credentials for injection into the spawned Command via .envs() + // (avoids unsafe std::env::set_var in multi-threaded runtime). + let credentials = self.client.fetch_credentials().await?; + let mut extra_env = std::collections::HashMap::new(); + for cred in &credentials { + extra_env.insert(cred.env_var.clone(), cred.value.clone()); + } + if !extra_env.is_empty() { + tracing::info!( + job_id = %self.config.job_id, + "Fetched {} credential(s) for child process injection", + extra_env.len() + ); + } + + // Warn if no auth method is available (check both process env and fetched credentials). + let has_api_key = extra_env.contains_key("ANTHROPIC_API_KEY") + || std::env::var("ANTHROPIC_API_KEY").is_ok(); + let has_oauth = extra_env.contains_key("CLAUDE_CODE_OAUTH_TOKEN") + || std::env::var("CLAUDE_CODE_OAUTH_TOKEN").is_ok(); + if !has_api_key && !has_oauth { + tracing::warn!( + job_id = %self.config.job_id, + "No Claude Code auth available. Set ANTHROPIC_API_KEY or run \ + `claude login` on the host to authenticate." + ); + } + // Report that we're running self.client .report_status(&crate::worker::api::StatusUpdate { @@ -179,7 +262,10 @@ impl ClaudeBridgeRuntime { .await?; // Run the initial Claude session - let session_id = match self.run_claude_session(&job.description, None).await { + let session_id = match self + .run_claude_session(&job.description, None, &extra_env) + .await + { Ok(sid) => sid, Err(e) => { tracing::error!(job_id = %self.config.job_id, "Claude session failed: {}", e); @@ -210,7 +296,7 @@ impl ClaudeBridgeRuntime { "Got follow-up prompt, resuming session" ); if let Err(e) = self - .run_claude_session(&prompt.content, session_id.as_deref()) + .run_claude_session(&prompt.content, session_id.as_deref(), &extra_env) .await { tracing::error!( @@ -259,12 +345,14 @@ impl ClaudeBridgeRuntime { &self, prompt: &str, resume_session_id: Option<&str>, + extra_env: &std::collections::HashMap, ) -> Result, WorkerError> { let mut cmd = Command::new("claude"); cmd.arg("-p") .arg(prompt) .arg("--output-format") .arg("stream-json") + .arg("--verbose") .arg("--max-turns") .arg(self.config.max_turns.to_string()) .arg("--model") @@ -274,6 +362,10 @@ impl ClaudeBridgeRuntime { cmd.arg("--resume").arg(sid); } + // Inject credentials into the child process environment without + // mutating the global process env (which is unsafe in multi-threaded programs). + cmd.envs(extra_env); + cmd.current_dir("/workspace") .stdout(std::process::Stdio::piped()) .stderr(std::process::Stdio::piped()); @@ -435,6 +527,9 @@ fn build_permission_settings(allowed_tools: &[String]) -> String { fn stream_event_to_payloads(event: &ClaudeStreamEvent) -> Vec { let mut payloads = Vec::new(); + // Helper: extract content blocks from message wrapper. + let blocks = event.message.as_ref().and_then(|m| m.content.as_ref()); + match event.event_type.as_str() { "system" => { payloads.push(JobEventPayload { @@ -446,12 +541,13 @@ fn stream_event_to_payloads(event: &ClaudeStreamEvent) -> Vec { }); } "assistant" => { - // Extract text content and tool_use blocks - if let Some(ref blocks) = event.content { + // Content blocks are nested under message.content[]. + if let Some(blocks) = blocks { for block in blocks { match block.block_type.as_str() { "text" => { - if let Some(ref text) = block.text { + if let Some(ref text) = block.text.as_deref().filter(|t| !t.is_empty()) + { payloads.push(JobEventPayload { event_type: "message".to_string(), data: serde_json::json!({ @@ -466,37 +562,58 @@ fn stream_event_to_payloads(event: &ClaudeStreamEvent) -> Vec { event_type: "tool_use".to_string(), data: serde_json::json!({ "tool_name": block.name, + "tool_use_id": block.id, "input": block.input, }), }); } - "tool_result" => { - payloads.push(JobEventPayload { - event_type: "tool_result".to_string(), - data: serde_json::json!({ - "tool_name": block.name.as_deref().unwrap_or("unknown"), - "output": block.content.as_deref().unwrap_or(""), - }), - }); - } _ => {} } } } } + "user" => { + // User events carry tool_result blocks under message.content[]. + if let Some(blocks) = blocks { + for block in blocks { + if block.block_type == "tool_result" { + payloads.push(JobEventPayload { + event_type: "tool_result".to_string(), + data: serde_json::json!({ + "tool_use_id": block.tool_use_id, + "output": block.content, + }), + }); + } + } + } + } "result" => { - let is_error = event + let is_error = event.is_error.unwrap_or(false); + + // Emit the final review text as a message so it appears in activity. + if let Some(text) = event .result .as_ref() - .and_then(|r| r.is_error) - .unwrap_or(false); + .and_then(|v| v.as_str()) + .filter(|t| !t.is_empty()) + { + payloads.push(JobEventPayload { + event_type: "message".to_string(), + data: serde_json::json!({ + "role": "assistant", + "content": text, + }), + }); + } + payloads.push(JobEventPayload { event_type: "result".to_string(), data: serde_json::json!({ "status": if is_error { "error" } else { "completed" }, "session_id": event.session_id, - "duration_ms": event.result.as_ref().and_then(|r| r.duration_ms), - "num_turns": event.result.as_ref().and_then(|r| r.num_turns), + "duration_ms": event.duration_ms, + "num_turns": event.num_turns, }), }); } @@ -515,6 +632,66 @@ fn stream_event_to_payloads(event: &ClaudeStreamEvent) -> Vec { payloads } +/// Recursively copy files and directories from `src` to `dst`, skipping +/// entries that can't be read (e.g. permission-restricted files owned by a +/// different uid on a read-only bind mount). Returns the number of files +/// successfully copied. +fn copy_dir_recursive(src: &std::path::Path, dst: &std::path::Path) -> std::io::Result { + let entries = match std::fs::read_dir(src) { + Ok(e) => e, + Err(e) => { + tracing::debug!("Skipping unreadable directory {}: {}", src.display(), e); + return Ok(0); + } + }; + + let mut copied = 0; + for entry in entries { + let entry = match entry { + Ok(e) => e, + Err(e) => { + tracing::debug!("Skipping unreadable entry in {}: {}", src.display(), e); + continue; + } + }; + + let src_path = entry.path(); + let dst_path = dst.join(entry.file_name()); + + let file_type = match entry.file_type() { + Ok(ft) => ft, + Err(e) => { + tracing::debug!( + "Skipping entry with unreadable type {}: {}", + src_path.display(), + e + ); + continue; + } + }; + + // Skip symlinks to avoid following links outside the mount. + if file_type.is_symlink() { + tracing::debug!("Skipping symlink {}", src_path.display()); + continue; + } + + if file_type.is_dir() { + if std::fs::create_dir_all(&dst_path).is_ok() { + copied += copy_dir_recursive(&src_path, &dst_path)?; + } + } else { + match std::fs::copy(&src_path, &dst_path) { + Ok(_) => copied += 1, + Err(e) => { + tracing::debug!("Skipping unreadable file {}: {}", src_path.display(), e); + } + } + } + } + Ok(copied) +} + fn truncate(s: &str, max_len: usize) -> &str { if s.len() <= max_len { s @@ -542,10 +719,11 @@ mod tests { #[test] fn test_parse_assistant_text_event() { - let json = r#"{"type":"assistant","content":[{"type":"text","text":"Hello world"}]}"#; + // Real Claude Code format: content blocks are under message.content + let json = r#"{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"Hello world"}]}}"#; let event: ClaudeStreamEvent = serde_json::from_str(json).unwrap(); assert_eq!(event.event_type, "assistant"); - let blocks = event.content.unwrap(); + let blocks = event.message.unwrap().content.unwrap(); assert_eq!(blocks.len(), 1); assert_eq!(blocks[0].block_type, "text"); assert_eq!(blocks[0].text.as_deref(), Some("Hello world")); @@ -553,32 +731,42 @@ mod tests { #[test] fn test_parse_assistant_tool_use_event() { - let json = r#"{"type":"assistant","content":[{"type":"tool_use","name":"Bash","input":{"command":"ls"}}]}"#; + let json = r#"{"type":"assistant","message":{"content":[{"type":"tool_use","id":"toolu_01abc","name":"Bash","input":{"command":"ls"}}]}}"#; let event: ClaudeStreamEvent = serde_json::from_str(json).unwrap(); - let blocks = event.content.unwrap(); + let blocks = event.message.unwrap().content.unwrap(); assert_eq!(blocks[0].block_type, "tool_use"); assert_eq!(blocks[0].name.as_deref(), Some("Bash")); + assert_eq!(blocks[0].id.as_deref(), Some("toolu_01abc")); assert!(blocks[0].input.is_some()); } + #[test] + fn test_parse_user_tool_result_event() { + let json = r#"{"type":"user","message":{"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_01abc","content":"/workspace"}]}}"#; + let event: ClaudeStreamEvent = serde_json::from_str(json).unwrap(); + assert_eq!(event.event_type, "user"); + let blocks = event.message.unwrap().content.unwrap(); + assert_eq!(blocks[0].block_type, "tool_result"); + assert_eq!(blocks[0].tool_use_id.as_deref(), Some("toolu_01abc")); + } + #[test] fn test_parse_result_event() { - let json = - r#"{"type":"result","result":{"is_error":false,"duration_ms":5000,"num_turns":3}}"#; + let json = r#"{"type":"result","subtype":"success","is_error":false,"duration_ms":5000,"num_turns":3,"result":"Done.","session_id":"sid-1"}"#; let event: ClaudeStreamEvent = serde_json::from_str(json).unwrap(); assert_eq!(event.event_type, "result"); - let result = event.result.unwrap(); - assert_eq!(result.is_error, Some(false)); - assert_eq!(result.duration_ms, Some(5000)); - assert_eq!(result.num_turns, Some(3)); + assert_eq!(event.is_error, Some(false)); + assert_eq!(event.duration_ms, Some(5000)); + assert_eq!(event.num_turns, Some(3)); + assert_eq!(event.result.unwrap().as_str().unwrap(), "Done."); } #[test] fn test_parse_result_error_event() { - let json = r#"{"type":"result","result":{"is_error":true}}"#; + let json = r#"{"type":"result","subtype":"error_max_turns","is_error":true,"duration_ms":60000,"num_turns":50}"#; let event: ClaudeStreamEvent = serde_json::from_str(json).unwrap(); - let result = event.result.unwrap(); - assert_eq!(result.is_error, Some(true)); + assert_eq!(event.is_error, Some(true)); + assert_eq!(event.subtype.as_deref(), Some("error_max_turns")); } #[test] @@ -586,12 +774,12 @@ mod tests { let event = ClaudeStreamEvent { event_type: "system".to_string(), session_id: Some("sid-123".to_string()), - content: None, + subtype: Some("init".to_string()), + message: None, result: None, - tool_name: None, - input: None, - output: None, - subtype: None, + is_error: None, + duration_ms: None, + num_turns: None, }; let payloads = stream_event_to_payloads(&event); assert_eq!(payloads.len(), 1); @@ -604,18 +792,23 @@ mod tests { let event = ClaudeStreamEvent { event_type: "assistant".to_string(), session_id: None, - content: Some(vec![ContentBlock { - block_type: "text".to_string(), - text: Some("Here's the answer".to_string()), - name: None, - input: None, - content: None, - }]), - result: None, - tool_name: None, - input: None, - output: None, subtype: None, + message: Some(MessageWrapper { + role: Some("assistant".to_string()), + content: Some(vec![ContentBlock { + block_type: "text".to_string(), + text: Some("Here's the answer".to_string()), + name: None, + id: None, + input: None, + content: None, + tool_use_id: None, + }]), + }), + result: None, + is_error: None, + duration_ms: None, + num_turns: None, }; let payloads = stream_event_to_payloads(&event); assert_eq!(payloads.len(), 1); @@ -624,26 +817,85 @@ mod tests { assert_eq!(payloads[0].data["content"], "Here's the answer"); } + #[test] + fn test_stream_event_to_payloads_assistant_tool_use() { + let event = ClaudeStreamEvent { + event_type: "assistant".to_string(), + session_id: None, + subtype: None, + message: Some(MessageWrapper { + role: Some("assistant".to_string()), + content: Some(vec![ContentBlock { + block_type: "tool_use".to_string(), + text: None, + name: Some("Bash".to_string()), + id: Some("toolu_01abc".to_string()), + input: Some(serde_json::json!({"command": "ls"})), + content: None, + tool_use_id: None, + }]), + }), + result: None, + is_error: None, + duration_ms: None, + num_turns: None, + }; + let payloads = stream_event_to_payloads(&event); + assert_eq!(payloads.len(), 1); + assert_eq!(payloads[0].event_type, "tool_use"); + assert_eq!(payloads[0].data["tool_name"], "Bash"); + assert_eq!(payloads[0].data["tool_use_id"], "toolu_01abc"); + } + + #[test] + fn test_stream_event_to_payloads_user_tool_result() { + let event = ClaudeStreamEvent { + event_type: "user".to_string(), + session_id: None, + subtype: None, + message: Some(MessageWrapper { + role: Some("user".to_string()), + content: Some(vec![ContentBlock { + block_type: "tool_result".to_string(), + text: None, + name: None, + id: None, + input: None, + content: Some(serde_json::json!("/workspace")), + tool_use_id: Some("toolu_01abc".to_string()), + }]), + }), + result: None, + is_error: None, + duration_ms: None, + num_turns: None, + }; + let payloads = stream_event_to_payloads(&event); + assert_eq!(payloads.len(), 1); + assert_eq!(payloads[0].event_type, "tool_result"); + assert_eq!(payloads[0].data["tool_use_id"], "toolu_01abc"); + assert_eq!(payloads[0].data["output"], "/workspace"); + } + #[test] fn test_stream_event_to_payloads_result_success() { let event = ClaudeStreamEvent { event_type: "result".to_string(), session_id: Some("s1".to_string()), - content: None, - result: Some(ResultInfo { - is_error: Some(false), - duration_ms: Some(12000), - num_turns: Some(5), - }), - tool_name: None, - input: None, - output: None, - subtype: None, + subtype: Some("success".to_string()), + message: None, + result: Some(serde_json::json!("The review is complete.")), + is_error: Some(false), + duration_ms: Some(12000), + num_turns: Some(5), }; let payloads = stream_event_to_payloads(&event); - assert_eq!(payloads.len(), 1); - assert_eq!(payloads[0].event_type, "result"); - assert_eq!(payloads[0].data["status"], "completed"); + // Should emit a message (the result text) + a result event + assert_eq!(payloads.len(), 2); + assert_eq!(payloads[0].event_type, "message"); + assert_eq!(payloads[0].data["content"], "The review is complete."); + assert_eq!(payloads[1].event_type, "result"); + assert_eq!(payloads[1].data["status"], "completed"); } #[test] @@ -651,18 +903,15 @@ mod tests { let event = ClaudeStreamEvent { event_type: "result".to_string(), session_id: None, - content: None, - result: Some(ResultInfo { - is_error: Some(true), - duration_ms: None, - num_turns: None, - }), - tool_name: None, - input: None, - output: None, - subtype: None, + subtype: Some("error_max_turns".to_string()), + message: None, + result: None, + is_error: Some(true), + duration_ms: None, + num_turns: None, }; let payloads = stream_event_to_payloads(&event); + assert_eq!(payloads.len(), 1); assert_eq!(payloads[0].data["status"], "error"); } @@ -671,12 +920,12 @@ mod tests { let event = ClaudeStreamEvent { event_type: "fancy_new_thing".to_string(), session_id: None, - content: None, - result: None, - tool_name: None, - input: None, - output: None, subtype: None, + message: None, + result: None, + is_error: None, + duration_ms: None, + num_turns: None, }; let payloads = stream_event_to_payloads(&event); assert_eq!(payloads.len(), 1); @@ -735,4 +984,47 @@ mod tests { assert!(parsed["permissions"].is_object()); assert!(parsed["permissions"]["allow"].is_array()); } + + #[test] + fn test_copy_dir_recursive() { + let src = tempfile::tempdir().unwrap(); + let dst = tempfile::tempdir().unwrap(); + + // Create a nested structure in src + std::fs::write(src.path().join("auth.json"), r#"{"token":"abc"}"#).unwrap(); + std::fs::create_dir_all(src.path().join("subdir")).unwrap(); + std::fs::write(src.path().join("subdir").join("nested.txt"), "nested").unwrap(); + + let copied = copy_dir_recursive(src.path(), dst.path()).unwrap(); + assert_eq!(copied, 2); + + // Verify files were copied + assert_eq!( + std::fs::read_to_string(dst.path().join("auth.json")).unwrap(), + r#"{"token":"abc"}"# + ); + assert_eq!( + std::fs::read_to_string(dst.path().join("subdir").join("nested.txt")).unwrap(), + "nested" + ); + } + + #[test] + fn test_copy_dir_recursive_empty_source() { + let src = tempfile::tempdir().unwrap(); + let dst = tempfile::tempdir().unwrap(); + + let copied = copy_dir_recursive(src.path(), dst.path()).unwrap(); + assert_eq!(copied, 0); + } + + #[test] + fn test_copy_dir_recursive_skips_nonexistent_source() { + let dst = tempfile::tempdir().unwrap(); + let nonexistent = std::path::Path::new("/no/such/path"); + + // Should gracefully return 0 instead of failing + let copied = copy_dir_recursive(nonexistent, dst.path()).unwrap(); + assert_eq!(copied, 0); + } } diff --git a/src/worker/runtime.rs b/src/worker/runtime.rs index a1da32f2..7de30d49 100644 --- a/src/worker/runtime.rs +++ b/src/worker/runtime.rs @@ -5,6 +5,7 @@ //! Streams real-time events (message, tool_use, tool_result, result) through //! the orchestrator's job event pipeline for UI visibility. +use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; @@ -51,6 +52,11 @@ pub struct WorkerRuntime { llm: Arc, safety: Arc, tools: Arc, + /// Credentials fetched from the orchestrator, injected into child processes + /// via `Command::envs()` rather than mutating the global process environment. + /// + /// Wrapped in `Arc` to avoid deep-cloning the map on every tool invocation. + extra_env: Arc>, } impl WorkerRuntime { @@ -83,11 +89,12 @@ impl WorkerRuntime { llm, safety, tools, + extra_env: Arc::new(HashMap::new()), }) } /// Run the worker until the job is complete or an error occurs. - pub async fn run(self) -> Result<(), WorkerError> { + pub async fn run(mut self) -> Result<(), WorkerError> { tracing::info!("Worker starting for job {}", self.config.job_id); // Fetch job description from orchestrator @@ -99,6 +106,23 @@ impl WorkerRuntime { truncate(&job.description, 100) ); + // Fetch credentials and store them for injection into child processes + // via Command::envs() (avoids unsafe std::env::set_var in multi-threaded runtime). + let credentials = self.client.fetch_credentials().await?; + { + let mut env_map = HashMap::new(); + for cred in &credentials { + env_map.insert(cred.env_var.clone(), cred.value.clone()); + } + self.extra_env = Arc::new(env_map); + } + if !credentials.is_empty() { + tracing::info!( + "Fetched {} credential(s) for child process injection", + credentials.len() + ); + } + // Report that we're starting self.client .report_status(&StatusUpdate { @@ -378,7 +402,10 @@ Work independently to complete this job. Report when done."#, None => return Err(format!("tool '{}' not found", tool_name)), }; - let ctx = JobContext::default(); + let ctx = JobContext { + extra_env: self.extra_env.clone(), + ..Default::default() + }; // Validate params let validation = self.safety.validator().validate_tool_params(params);