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
+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,
};