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 <[email protected]>

* 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 <[email protected]>

* 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 <[email protected]>

* 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 <[email protected]>

* 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 <[email protected]>

* 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 <[email protected]>

* 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 <[email protected]>

* 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 <[email protected]>

* 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 <[email protected]>

* 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 <[email protected]>

* 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 <[email protected]>

* 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 <name>.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* 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 <[email protected]>

* 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 <[email protected]>

* 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<HashMap> 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 <[email protected]>

* 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 <[email protected]>

* 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 <[email protected]>

* 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 <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
This commit is contained in:
Illia Polosukhin
2026-02-18 00:48:43 +00:00
committed by GitHub
co-authored by Claude Opus 4.6
parent bac2d75713
commit cfb579a4bb
41 changed files with 3435 additions and 686 deletions
+1 -1
View File
@@ -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
+10 -4
View File
@@ -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
+1 -1
View File
@@ -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
+26 -7
View File
@@ -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 <name> to switch.");
}
Ok(_) => {
out.push_str(
"\nCould not fetch model list. Use /model <name> to switch.",
);
}
Err(e) => {
out.push_str(&format!(
"\nCould not fetch models: {}. Use /model <name> 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
}
}
+245
View File
@@ -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<IncomingMessage>,
) -> 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::<IncomingMessage>(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::<IncomingMessage>(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::<IncomingMessage>(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::<IncomingMessage>(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"
);
}
}
+1
View File
@@ -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;
+29 -2
View File
@@ -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<RwLock<HashMap<String, Box<dyn Channel>>>>,
inject_tx: mpsc::Sender<IncomingMessage>,
/// Taken once in `start_all()` and merged into the stream.
inject_rx: tokio::sync::Mutex<Option<mpsc::Receiver<IncomingMessage>>>,
}
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<IncomingMessage> {
self.inject_tx.clone()
}
/// Add a channel to the manager.
pub fn add(&mut self, channel: Box<dyn Channel>) {
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<MessageStream, ChannelError> {
let channels = self.channels.read().await;
let mut streams = Vec::new();
let mut streams: Vec<MessageStream> = 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))
+22 -2
View File
@@ -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<crate::orchestrator::auth::CredentialGrant> =
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()))?;
+69 -8
View File
@@ -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<String>,
/// 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<String> {
[
"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<String> {
// 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<Self, ConfigError> {
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<String> {
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 {
+12
View File
@@ -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<StateTransition>,
/// 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<HashMap<String, String>>,
}
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,
}
}
+268 -108
View File
@@ -116,10 +116,19 @@ impl LibSqlBackend {
}
/// Create a new connection to the database.
pub fn connect(&self) -> Result<Connection, DatabaseError> {
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<Connection, DatabaseError> {
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<DateTime<Utc>> {
#[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<Uuid, DatabaseError> {
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<Uuid, DatabaseError> {
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<Vec<ConversationSummary>, 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<Uuid, DatabaseError> {
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<Uuid, DatabaseError> {
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<DateTime<Utc>>,
limit: i64,
) -> Result<(Vec<ConversationMessage>, 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<Option<serde_json::Value>, 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<Vec<ConversationMessage>, 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<bool, DatabaseError> {
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<Option<JobContext>, 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<Vec<Uuid>, 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<Vec<ActionRecord>, 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<Uuid, DatabaseError> {
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<Uuid, DatabaseError> {
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<Decimal>,
) -> 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<Option<SandboxJobRecord>, 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<Vec<SandboxJobRecord>, 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<DateTime<Utc>>,
completed_at: Option<DateTime<Utc>>,
) -> 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<u64, DatabaseError> {
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<SandboxJobSummary, DatabaseError> {
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<Vec<SandboxJobRecord>, 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<SandboxJobSummary, DatabaseError> {
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<bool, DatabaseError> {
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<Option<String>, 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<Vec<JobEventRecord>, DatabaseError> {
let conn = self.connect()?;
let mut rows = conn
.query(
async fn list_job_events(
&self,
job_id: Uuid,
limit: Option<i64>,
) -> Result<Vec<JobEventRecord>, 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<Option<Routine>, 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<Option<Routine>, 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<Vec<Routine>, 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<Vec<Routine>, 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<Vec<Routine>, 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<bool, DatabaseError> {
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<i32>,
) -> 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<Vec<RoutineRun>, 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<i64, DatabaseError> {
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<Vec<BrokenTool>, 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<Option<serde_json::Value>, 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<Option<SettingRow>, 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<bool, DatabaseError> {
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<Vec<SettingRow>, 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<HashMap<String, serde_json::Value>, 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<String, serde_json::Value>,
) -> 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<bool, DatabaseError> {
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,7 +2012,10 @@ impl Database for LibSqlBackend {
agent_id: Option<Uuid>,
path: &str,
) -> Result<MemoryDocument, WorkspaceError> {
let conn = self.connect().map_err(|e| WorkspaceError::SearchFailed {
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());
@@ -2006,7 +2049,10 @@ impl Database for LibSqlBackend {
}
async fn get_document_by_id(&self, id: Uuid) -> Result<MemoryDocument, WorkspaceError> {
let conn = self.connect().map_err(|e| WorkspaceError::SearchFailed {
let conn = self
.connect()
.await
.map_err(|e| WorkspaceError::SearchFailed {
reason: e.to_string(),
})?;
let mut rows = conn
@@ -2051,7 +2097,10 @@ impl Database for LibSqlBackend {
}
// Create
let conn = self.connect().map_err(|e| WorkspaceError::SearchFailed {
let conn = self
.connect()
.await
.map_err(|e| WorkspaceError::SearchFailed {
reason: e.to_string(),
})?;
let id = Uuid::new_v4();
@@ -2073,7 +2122,10 @@ impl Database for LibSqlBackend {
}
async fn update_document(&self, id: Uuid, content: &str) -> Result<(), WorkspaceError> {
let conn = self.connect().map_err(|e| WorkspaceError::SearchFailed {
let conn = self
.connect()
.await
.map_err(|e| WorkspaceError::SearchFailed {
reason: e.to_string(),
})?;
let now = fmt_ts(&Utc::now());
@@ -2097,7 +2149,10 @@ 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 {
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());
@@ -2118,7 +2173,10 @@ impl Database for LibSqlBackend {
agent_id: Option<Uuid>,
directory: &str,
) -> Result<Vec<WorkspaceEntry>, WorkspaceError> {
let conn = self.connect().map_err(|e| WorkspaceError::SearchFailed {
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.
@@ -2223,7 +2281,10 @@ impl Database for LibSqlBackend {
user_id: &str,
agent_id: Option<Uuid>,
) -> Result<Vec<String>, WorkspaceError> {
let conn = self.connect().map_err(|e| WorkspaceError::SearchFailed {
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());
@@ -2255,7 +2316,10 @@ impl Database for LibSqlBackend {
user_id: &str,
agent_id: Option<Uuid>,
) -> Result<Vec<MemoryDocument>, WorkspaceError> {
let conn = self.connect().map_err(|e| WorkspaceError::SearchFailed {
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());
@@ -2291,7 +2355,10 @@ 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 {
let conn = self
.connect()
.await
.map_err(|e| WorkspaceError::ChunkingFailed {
reason: e.to_string(),
})?;
conn.execute(
@@ -2312,7 +2379,10 @@ impl Database for LibSqlBackend {
content: &str,
embedding: Option<&[f32]>,
) -> Result<Uuid, WorkspaceError> {
let conn = self.connect().map_err(|e| WorkspaceError::ChunkingFailed {
let conn = self
.connect()
.await
.map_err(|e| WorkspaceError::ChunkingFailed {
reason: e.to_string(),
})?;
let id = Uuid::new_v4();
@@ -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,7 +2442,10 @@ impl Database for LibSqlBackend {
agent_id: Option<Uuid>,
limit: usize,
) -> Result<Vec<MemoryChunk>, WorkspaceError> {
let conn = self.connect().map_err(|e| WorkspaceError::SearchFailed {
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());
@@ -2422,7 +2496,10 @@ impl Database for LibSqlBackend {
embedding: Option<&[f32]>,
config: &SearchConfig,
) -> Result<Vec<SearchResult>, WorkspaceError> {
let conn = self.connect().map_err(|e| WorkspaceError::SearchFailed {
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());
@@ -2607,3 +2684,86 @@ fn row_to_routine_run_libsql(row: &libsql::Row) -> Result<RoutineRun, DatabaseEr
created_at: get_ts(row, 10),
})
}
#[cfg(test)]
mod tests {
use crate::db::Database;
use crate::db::libsql_backend::LibSqlBackend;
#[tokio::test]
async fn test_wal_mode_after_migrations() {
let backend = LibSqlBackend::new_memory().await.unwrap();
backend.run_migrations().await.unwrap();
let conn = backend.connect().await.unwrap();
let mut rows = conn.query("PRAGMA journal_mode", ()).await.unwrap();
let row = rows.next().await.unwrap().unwrap();
let mode: String = row.get(0).unwrap();
// In-memory databases use "memory" journal mode (WAL doesn't apply to :memory:),
// but the PRAGMA still executes without error. For file-based databases it returns "wal".
assert!(
mode == "wal" || mode == "memory",
"expected wal or memory, got: {}",
mode,
);
}
#[tokio::test]
async fn test_busy_timeout_set_on_connect() {
let backend = LibSqlBackend::new_memory().await.unwrap();
backend.run_migrations().await.unwrap();
let conn = backend.connect().await.unwrap();
let mut rows = conn.query("PRAGMA busy_timeout", ()).await.unwrap();
let row = rows.next().await.unwrap().unwrap();
let timeout: i64 = row.get(0).unwrap();
assert_eq!(timeout, 5000);
}
#[tokio::test]
async fn test_concurrent_writes_succeed() {
// Use a temp file so connections share state (in-memory DBs are connection-local)
let dir = tempfile::tempdir().unwrap();
let db_path = dir.path().join("test_concurrent.db");
let backend = LibSqlBackend::new_local(&db_path).await.unwrap();
backend.run_migrations().await.unwrap();
// Spawn 20 concurrent inserts into the conversations table
let mut handles = Vec::new();
for i in 0..20 {
let conn = backend.connect().await.unwrap();
let handle = tokio::spawn(async move {
let id = uuid::Uuid::new_v4().to_string();
let val = format!("ch_{}", i);
conn.execute(
"INSERT INTO conversations (id, channel, user_id) VALUES (?1, ?2, ?3)",
libsql::params![id, val, "test_user"],
)
.await
});
handles.push(handle);
}
for handle in handles {
let result = handle.await.unwrap();
assert!(
result.is_ok(),
"concurrent write failed: {:?}",
result.err()
);
}
// Verify all 20 rows landed
let conn = backend.connect().await.unwrap();
let mut rows = conn
.query(
"SELECT COUNT(*) FROM conversations WHERE user_id = ?1",
libsql::params!["test_user"],
)
.await
.unwrap();
let row = rows.next().await.unwrap().unwrap();
let count: i64 = row.get(0).unwrap();
assert_eq!(count, 20);
}
}
+6 -2
View File
@@ -309,8 +309,12 @@ pub trait Database: Send + Sync {
data: &serde_json::Value,
) -> Result<(), DatabaseError>;
/// Load all job events.
async fn list_job_events(&self, job_id: Uuid) -> Result<Vec<JobEventRecord>, 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<i64>,
) -> Result<Vec<JobEventRecord>, DatabaseError>;
// ==================== Routines ====================
+6 -2
View File
@@ -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<Vec<JobEventRecord>, DatabaseError> {
self.store.list_job_events(job_id).await
async fn list_job_events(
&self,
job_id: Uuid,
limit: Option<i64>,
) -> Result<Vec<JobEventRecord>, DatabaseError> {
self.store.list_job_events(job_id, limit).await
}
// ==================== Routines ====================
-6
View File
@@ -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).
+38 -8
View File
@@ -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<Utc>,
pub started_at: Option<DateTime<Utc>>,
pub completed_at: Option<DateTime<Utc>>,
/// Serialized JSON of `Vec<CredentialGrant>` 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<i64>,
) -> Result<Vec<JobEventRecord>, 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 {
+53 -1
View File
@@ -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<Vec<String>, 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<String>,
input_cost: Decimal,
output_cost: Decimal,
complete_result: Mutex<Option<Result<CompletionResponse, LlmError>>>,
@@ -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<Vec<String>, 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<MockProvider> = Arc::new(MockProvider::succeeding("model-a", "ok"));
let p2: Arc<MockProvider> = Arc::new(MockProvider::succeeding("model-b", "ok"));
let failover = FailoverProvider::new(vec![
Arc::clone(&p1) as Arc<dyn LlmProvider>,
Arc::clone(&p2) as Arc<dyn LlmProvider>,
])
.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");
}
}
+14 -2
View File
@@ -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<Vec<ModelInfo>, 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<Vec<ModelInfo>, 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(),
+13 -6
View File
@@ -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)
+442 -3
View File
@@ -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<Mutex<HashMap<Uuid, VecDeque<PendingPrompt>>>>,
/// Database handle for persisting job events.
pub store: Option<Arc<dyn Database>>,
/// Encrypted secrets store for credential injection into containers.
pub secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
/// 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<OrchestratorState>,
Path(job_id): Path<Uuid>,
Json(update): Json<StatusUpdate>,
) -> Result<StatusCode, StatusCode> {
@@ -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<OrchestratorState>,
Path(job_id): Path<Uuid>,
) -> Result<(StatusCode, Json<serde_json::Value>), 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<CredentialResponse> = 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::Value> = 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"));
}
}
+136 -7
View File
@@ -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<RwLock<HashMap<Uuid, String>>>,
/// Maps job_id -> granted credentials. Revoked alongside the token.
credential_grants: Arc<RwLock<HashMap<Uuid, Vec<CredentialGrant>>>>,
}
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<CredentialGrant>) {
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<Vec<CredentialGrant>> {
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");
}
}
+223 -48
View File
@@ -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<PathBuf>,
/// Anthropic API key for Claude Code containers (read from ANTHROPIC_API_KEY).
/// Takes priority over OAuth token.
pub claude_code_api_key: Option<String>,
/// 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<String>,
/// 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<Utc>,
pub project_dir: Option<PathBuf>,
pub task_description: String,
/// Last status message reported by the worker (iteration count, progress, etc.).
pub last_worker_status: Option<String>,
/// 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<CompletionResult>,
// NOTE: auth_token is intentionally NOT in this struct.
@@ -121,11 +131,84 @@ pub struct CompletionResult {
pub message: Option<String>,
}
/// 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<PathBuf, OrchestratorError> {
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<RwLock<HashMap<Uuid, ContainerHandle>>>,
pub(crate) containers: Arc<RwLock<HashMap<Uuid, ContainerHandle>>>,
/// Cached Docker connection (created on first use).
docker: Arc<RwLock<Option<bollard::Docker>>>,
}
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<bollard::Docker, OrchestratorError> {
{
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<PathBuf>,
mode: JobMode,
credential_grants: Vec<CredentialGrant>,
) -> Result<String, OrchestratorError> {
// 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<PathBuf>,
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<String>,
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<ContainerHandle> {
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"));
}
}
+4 -2
View File
@@ -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,
};
+6 -48
View File
@@ -159,56 +159,14 @@ pub fn default_allowlist() -> Vec<String> {
]
}
/// 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 <token>
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<CredentialMapping> {
pub fn default_credential_mappings() -> Vec<crate::secrets::CredentialMapping> {
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"),
]
}
+3 -2
View File
@@ -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<Docker> {
}
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(),
})
}
+65 -15
View File
@@ -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<ContainerOutput> for ExecOutput {
pub struct SandboxManager {
config: SandboxConfig,
proxy: Arc<RwLock<Option<HttpProxy>>>,
runner: Arc<RwLock<Option<ContainerRunner>>>,
docker: Arc<RwLock<Option<Docker>>>,
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);
}
}
+2 -4
View File
@@ -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<String> {
}
/// Default credential mappings getter (re-export for convenience).
pub fn default_credential_mappings() -> Vec<CredentialMapping> {
pub fn default_credential_mappings() -> Vec<crate::secrets::CredentialMapping> {
config::default_credential_mappings()
}
+142 -30
View File
@@ -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<dyn NetworkPolicyDecider>,
/// Credential resolver (maps secret names to values).
credential_resolver: Arc<dyn CredentialResolver>,
/// 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<hyper::body::Incoming>,
state: Arc<ProxyState>,
) -> Response<BoxBody<Bytes, Infallible>> {
// 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<BoxBody<Bytes, Infallible>> {
/// Build a response with guaranteed success (valid status + simple body cannot fail).
fn make_response(
status: StatusCode,
body: BoxBody<Bytes, Infallible>,
) -> Response<BoxBody<Bytes, Infallible>> {
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<Bytes, Infallible>,
) -> Response<BoxBody<Bytes, Infallible>> {
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<BoxBody<Bytes, Infallible>> {
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);
}
}
+2 -3
View File
@@ -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 {
+72 -9
View File
@@ -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
self.credential_mappings.iter().find(|m| {
m.host_patterns
.iter()
.find(|m| m.domain.to_lowercase() == host_lower)
.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"));
}
}
+14 -9
View File
@@ -323,10 +323,15 @@ impl LibSqlSecretsStore {
Self { db, crypto }
}
fn connect(&self) -> Result<libsql::Connection, SecretError> {
self.db
async fn connect(&self) -> Result<libsql::Connection, SecretError> {
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<Secret, SecretError> {
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<bool, SecretError> {
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<Vec<SecretRef>, 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<bool, SecretError> {
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#"
+1 -1
View File
@@ -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());
+898 -34
View File
File diff suppressed because it is too large Load Diff
+4 -1
View File
@@ -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::{
+14 -4
View File
@@ -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<String, String>,
) -> 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<u64>,
extra_env: &HashMap<String, String>,
) -> 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<ToolOutput, ToolError> {
let command = require_str(&params, "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();
-2
View File
@@ -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};
+42 -8
View File
@@ -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<ContextManager>,
job_manager: Option<Arc<ContainerJobManager>>,
store: Option<Arc<dyn Database>>,
job_event_tx: Option<
tokio::sync::broadcast::Sender<(uuid::Uuid, crate::channels::web::types::SseEvent)>,
>,
inject_tx: Option<tokio::sync::mpsc::Sender<crate::channels::IncomingMessage>>,
prompt_queue: Option<PromptQueue>,
secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
) {
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).
-136
View File
@@ -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<String>,
/// Allowed filesystem paths (empty = no filesystem).
pub allowed_paths: Vec<String>,
/// 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<u64>,
}
/// 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<SandboxResult, ToolError> {
// 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<SandboxResult, ToolError> {
// 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<SandboxResult, ToolError> {
// 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());
}
}
+17 -10
View File
@@ -581,10 +581,17 @@ impl LibSqlWasmToolStore {
Self { db }
}
fn connect(&self) -> Result<libsql::Connection, WasmStorageError> {
self.db
async fn connect(&self) -> Result<libsql::Connection, WasmStorageError> {
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<StoredWasmTool, WasmStorageError> {
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<StoredWasmToolWithBinary, WasmStorageError> {
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<Option<StoredCapabilities>, 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<Vec<StoredWasmTool>, 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<bool, WasmStorageError> {
let conn = self.connect()?;
let conn = self.connect().await?;
let result = conn
.execute(
"DELETE FROM wasm_tools WHERE user_id = ?1 AND name = ?2",
+65
View File
@@ -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<Vec<CredentialResponse>, 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}"#;
+389 -97
View File
@@ -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<String>,
/// For `assistant` events: the text content blocks.
#[serde(default)]
pub content: Option<Vec<ContentBlock>>,
/// For `result` events: final status info.
#[serde(default)]
pub result: Option<ResultInfo>,
/// For `tool_use`/`tool_result`: the tool name.
#[serde(default)]
pub tool_name: Option<String>,
/// For `tool_use`: the input parameters.
#[serde(default)]
pub input: Option<serde_json::Value>,
/// For `tool_result`: the output content.
#[serde(default)]
pub output: Option<String>,
/// Subtype discriminator (e.g. "text", "tool_use", "tool_result").
#[serde(default)]
pub subtype: Option<String>,
/// For `assistant` and `user` events: the message wrapper containing content blocks.
#[serde(default)]
pub message: Option<MessageWrapper>,
/// For `result` events: the final text output.
#[serde(default)]
pub result: Option<serde_json::Value>,
/// For `result` events: whether the session ended in error.
#[serde(default)]
pub is_error: Option<bool>,
/// For `result` events: total wall-clock duration.
#[serde(default)]
pub duration_ms: Option<u64>,
/// For `result` events: number of agentic turns used.
#[serde(default)]
pub num_turns: Option<u32>,
}
/// 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<String>,
#[serde(default)]
pub content: Option<Vec<ContentBlock>>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContentBlock {
#[serde(rename = "type")]
pub block_type: String,
/// Text block content.
#[serde(default)]
pub text: Option<String>,
/// Tool name (for tool_use blocks).
#[serde(default)]
pub name: Option<String>,
/// Tool use ID (for tool_use and tool_result blocks).
#[serde(default)]
pub id: Option<String>,
/// Tool input params (for tool_use blocks).
#[serde(default)]
pub input: Option<serde_json::Value>,
/// Tool result content (for tool_result blocks), or general content.
#[serde(default)]
pub content: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResultInfo {
pub content: Option<serde_json::Value>,
/// Tool use ID reference (for tool_result blocks).
#[serde(default)]
pub is_error: Option<bool>,
#[serde(default)]
pub duration_ms: Option<u64>,
#[serde(default)]
pub num_turns: Option<u32>,
pub tool_use_id: Option<String>,
}
/// 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<String, String>,
) -> Result<Option<String>, 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<JobEventPayload> {
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<JobEventPayload> {
});
}
"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<JobEventPayload> {
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<JobEventPayload> {
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<usize> {
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,
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,
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);
@@ -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 {
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),
}),
tool_name: None,
input: None,
output: None,
subtype: None,
};
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 {
subtype: Some("error_max_turns".to_string()),
message: None,
result: None,
is_error: Some(true),
duration_ms: None,
num_turns: None,
}),
tool_name: None,
input: None,
output: None,
subtype: 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);
}
}
+29 -2
View File
@@ -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<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
tools: Arc<ToolRegistry>,
/// 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<HashMap<String, String>>,
}
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);