From cfb579a4bb82a1fd9e455b3c5f9baac3c409056d Mon Sep 17 00:00:00 2001 From: Illia Polosukhin Date: Tue, 17 Feb 2026 16:48:43 -0800 Subject: [PATCH] feat: Add PR review tools, job monitor, and channel injection for E2E sandbox workflows (#57) * feat: Add PR review tools, job monitor, and channel injection for E2E sandbox workflows Adds JobEventsTool and JobPromptTool so the main agent can read container event logs and send follow-up prompts to running Claude Code sessions. A background JobMonitor forwards container assistant messages into the agent loop via a new inject channel on ChannelManager. CreateJobTool now accepts a project_dir parameter for mounting existing cloned repos into containers, and spawns the monitor automatically for async jobs. Also: Dockerfile bumped to Rust 1.88 (rig-core needs let chains), GITHUB_TOKEN forwarded into containers for gh CLI auth, and truncate() fixed for multi-byte char boundary panics. Co-Authored-By: Claude Opus 4.6 * fix: Address PR #57 review comments (IDOR, Dockerfile, truncate, logging) - Add ownership checks to JobEventsTool and JobPromptTool via ContextManager to prevent users from accessing other users' jobs (IDOR) - Combine Dockerfile gh CLI install into single apt-get layer - Handle truncate() edge case when max falls inside first multi-byte char - Log actual count of registered job management tools - Document fire-and-forget job monitor lifecycle - Add tests for ownership rejection and schema validation Co-Authored-By: Claude Opus 4.6 * feat: Replace hardcoded GITHUB_TOKEN with on-demand credential delivery Containers now fetch credentials via authenticated GET /worker/{id}/credentials endpoint instead of receiving them baked into env vars at creation time. Secrets are decrypted from SecretsStore on demand, scoped per-job via CredentialGrant, and revoked automatically when the job completes. Co-Authored-By: Claude Opus 4.6 * fix: Address sandbox audit findings (CONNECT tunnel, readonly_rootfs, type consolidation) - Implement real CONNECT tunnel with bidirectional TCP piping via hyper upgrade - Fix readonly_rootfs to apply for both ReadOnly and WorkspaceWrite policies - Consolidate duplicate CredentialMapping/CredentialLocation into secrets::types - Share reqwest::Client across proxy requests instead of per-request allocation - Store Docker connection and reuse across executions - Remove .unwrap() from proxy response builders with safe fallbacks - Add output truncation to direct (non-container) execution (64KB limit) - Delete dead src/tools/sandbox.rs (ToolSandbox never used) - Fix connect_docker error message to list all attempted socket paths - Update proxy credential injection to handle all CredentialLocation variants - Use glob-based host_patterns matching for credential lookup in proxy policy Co-Authored-By: Claude Opus 4.6 * fix: Address PR #57 review comments (IDOR, Dockerfile, truncate, logging) - Dockerfile: install curl+ca-certificates before fetching GitHub CLI GPG key - JobEventsTool/JobPromptTool: reject missing context (prevents IDOR bypass) - parse_credentials: validate env var names against denylist and pattern - resolve_project_dir: require explicit paths to exist before validation - Credential serving: lower log level from info to debug Co-Authored-By: Claude Opus 4.6 * fix: Address orchestrator audit findings (constant-time auth, error handling, tests) - auth: constant-time token comparison via subtle::ConstantTimeEq - auth: replace hand-rolled hex_encode with std::fmt::Write fold - api: report_status now updates ContainerHandle (was a no-op) - api: log complete_job errors instead of silently discarding - job_manager: log Docker cleanup errors in stop_job/complete_job - job_manager: extract validate_bind_mount_path with proper error on missing home_dir and mandatory base dir creation before canonicalize - job_manager: cache Docker connection across operations - error: remove dead OrchestratorError::AuthFailed and ContainerTimeout - Add 13 new tests (prompt queue, credentials, events, status, paths) Co-Authored-By: Claude Opus 4.6 * fix: use floor_char_boundary in sandbox manager truncate to prevent multi-byte panics String::truncate() panics when the index falls mid-way through a multi-byte UTF-8 character. Use the same floor_char_boundary utility already used in worker/runtime.rs and tools/builtin/shell.rs. Co-Authored-By: Claude Opus 4.6 * fix: default base_url to private.near.ai for Responses API mode Session tokens only authenticate against private.near.ai, not cloud-api.near.ai. The default base_url now matches the api_mode: - Responses (session token): https://private.near.ai - ChatCompletions (API key): https://cloud-api.near.ai This broke when the multi-provider merge introduced cloud-api.near.ai as the unconditional default. Co-Authored-By: Claude Opus 4.6 * fix: use private.near.ai as default base URL for all API modes private.near.ai now supports both Responses and ChatCompletions endpoints, so there is no reason to route through cloud-api.near.ai. This also fixes session token auth which only works against private.near.ai. Co-Authored-By: Claude Opus 4.6 * fix: harden libSQL concurrency, fix Claude Code Docker auth and permissions Three fixes for the sandbox/Claude Code pipeline: 1. SQLite "database is locked": set WAL journal mode in migrations and PRAGMA busy_timeout=5000 on every connection across LibSqlBackend, LibSqlSecretsStore, and LibSqlWasmToolStore (~83 async call sites). 2. Claude Code container auth: extract OAuth token from macOS Keychain (or Linux ~/.claude/.credentials.json) at startup and inject via CLAUDE_CODE_OAUTH_TOKEN env var. Removes the broken bind-mount approach that failed on uid mismatch. 3. Claude Code tool permissions: wire CLAUDE_CODE_ALLOWED_TOOLS env var through to the worker binary (was hardcoded to empty vec), and expand defaults to include all standard tools (Read, Write, Edit, Glob, Grep, NotebookEdit, Bash, Task, WebFetch, WebSearch). Also adds --verbose flag to claude CLI (required with stream-json + -p), failover provider model switching, and nearai models endpoint fix. Co-Authored-By: Claude Opus 4.6 * fix: stream event parsing, job ID prefix resolution, session renewal in list_models Three fixes for the Docker/gateway pipeline: 1. Claude Code stream event parsing (claude_bridge.rs): Rewrite ClaudeStreamEvent to match actual NDJSON format where content blocks are nested under message.content[], not at the top level. Add handler for "user" events (tool_result blocks) and emit result text as a "message" event so reviews appear in gateway activity view. 2. Job ID prefix resolution (job.rs): Add resolve_job_id() that accepts short hex prefixes (like git short SHAs) in addition to full UUIDs. The LLM sees truncated IDs in job monitor messages like "[Job f2854dd8]" and can now use them directly with job_status/cancel/events/prompt tools. 3. Session renewal in list_models (nearai.rs): list_models() now retries with OAuth renewal on 401, matching send_request()'s existing behavior. Previously it returned SessionExpired immediately, causing the setup wizard to fall back to defaults instead of prompting re-authentication. Co-Authored-By: Claude Opus 4.6 * fix: /model command now lists available models Previously /model with no args only showed the current model name. Now it fetches and displays all available models from the provider, marking the active one, so users can see what's available before switching with /model . Co-Authored-By: Claude Opus 4.6 * fix: address PR #57 review findings (set_var UB, tunnel timeout, restart creds) - Replace unsafe `std::env::set_var` in worker runtime and Claude bridge with `Command::envs()` injection via a new `extra_env` field on `JobContext`, avoiding undefined behavior in the multi-threaded tokio runtime. - Add 30-minute timeout to CONNECT tunnel `copy_bidirectional` in the sandbox proxy to prevent stuck connections from leaking spawned tasks. - Persist credential grants (as JSON in the description column) on `SandboxJobRecord` so `jobs_restart_handler` can restore them instead of passing `vec![]`, which caused restarted containers to lose access to their original secrets. Co-Authored-By: Claude Opus 4.6 * fix: address second round of PR #57 review comments - Normalize host_patterns to lowercase in proxy policy matching - Push LIMIT into SQL for list_job_events (Database trait + both backends) - Remove unused was_explicit binding in job tool - Return 500 instead of 200 in make_response fallback path - Update copy_auth_from_mount docstring for env-var default - Use entry.file_type() instead of is_dir() to avoid following symlinks Co-Authored-By: Claude Opus 4.6 * fix: address third round of PR #57 review comments - Restore glob patterns in default_claude_code_allowed_tools (Bash -> Bash(*)) - Add tracing::warn for credential grant serialize/deserialize failures - Wrap extra_env in Arc to avoid deep cloning per tool call - Document unsupported credential locations (AuthorizationBasic, UrlPath) - Document TOCTOU window in validate_bind_mount_path - Expand doc comments on JobEventsTool and JobPromptTool Co-Authored-By: Claude Opus 4.6 * fix: address fourth round of PR #57 review comments - Document CONNECT tunnel task lifecycle (timeout is the cleanup mechanism) - Remove secret names from error-level credential logs to prevent leaking - Expand DANGEROUS_ENV_VARS denylist with language runtime hijack vectors Co-Authored-By: Claude Opus 4.6 * fix: address fifth round of PR #57 review comments - Promote job monitor startup log to info level for observability - Require minimum 4-char prefix in resolve_job_id to limit enumeration - Cap credential grants at 20 per job to bound column storage - Clamp job events limit to 1..1000 to prevent memory abuse Co-Authored-By: Claude Opus 4.6 * fix: add missing closing brace for SkillsConfig impl block The merge resolution dropped the closing `}` for `impl SkillsConfig`, causing a compilation error in CI. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 --- .env.example | 2 +- Dockerfile.worker | 14 +- deploy/env.example | 2 +- src/agent/commands.rs | 33 +- src/agent/job_monitor.rs | 245 ++++++++ src/agent/mod.rs | 1 + src/channels/manager.rs | 31 +- src/channels/web/server.rs | 24 +- src/config.rs | 77 ++- src/context/state.rs | 12 + src/db/libsql_backend.rs | 424 +++++++++----- src/db/mod.rs | 8 +- src/db/postgres.rs | 8 +- src/error.rs | 6 - src/history/store.rs | 46 +- src/llm/failover.rs | 54 +- src/llm/nearai.rs | 16 +- src/main.rs | 19 +- src/orchestrator/api.rs | 445 ++++++++++++++- src/orchestrator/auth.rs | 143 ++++- src/orchestrator/job_manager.rs | 271 +++++++-- src/orchestrator/mod.rs | 6 +- src/sandbox/config.rs | 54 +- src/sandbox/container.rs | 5 +- src/sandbox/manager.rs | 80 ++- src/sandbox/mod.rs | 6 +- src/sandbox/proxy/http.rs | 172 +++++- src/sandbox/proxy/mod.rs | 5 +- src/sandbox/proxy/policy.rs | 83 ++- src/secrets/store.rs | 23 +- src/setup/wizard.rs | 2 +- src/tools/builtin/job.rs | 958 ++++++++++++++++++++++++++++++-- src/tools/builtin/mod.rs | 5 +- src/tools/builtin/shell.rs | 18 +- src/tools/mod.rs | 2 - src/tools/registry.rs | 50 +- src/tools/sandbox.rs | 136 ----- src/tools/wasm/storage.rs | 27 +- src/worker/api.rs | 65 +++ src/worker/claude_bridge.rs | 512 +++++++++++++---- src/worker/runtime.rs | 31 +- 41 files changed, 3435 insertions(+), 686 deletions(-) create mode 100644 src/agent/job_monitor.rs delete mode 100644 src/tools/sandbox.rs 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);