mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-27 08:00:17 +00:00
feat: Sandbox jobs (#4)
* Orchestrating jobs and running them in sandboxes * Fix heartbeat: dynamic max_tokens, empty content guard, notification fallback - Query /v1/models API for context_length and set max_tokens to half (floor 4096) instead of hardcoded 1024; reasoning models like GLM-4.7 need much larger budgets - Guard against empty LLM content (reasoning models can burn all tokens on chain-of-thought and return content: null) - Simplify notification routing: try configured channel first, fall back to broadcast_all so heartbeat alerts always reach someone - Add ModelMetadata struct and model_metadata() to LlmProvider trait - Refactor NearAiChatProvider::list_models into shared fetch_models() - Add standalone test_heartbeat example for isolated debugging Co-Authored-By: Claude Opus 4.6 <[email protected]> * Add job detail view with drill-down from jobs list Click a job row to see full details across four sub-tabs: Overview (metadata grid, description, state transitions timeline), Actions (expandable tool call cards with input/output JSON), Thinking (conversation messages styled by role), and Files (embedded workspace tree browser). Co-Authored-By: Claude Opus 4.6 <[email protected]> * Strip model-internal XML tags from LLM responses, fix Telegram parse_mode 400 Some models (GLM-4.7, etc.) emit <tool_call>tool_list</tool_call> in the content field instead of using the OpenAI tool_calls array. This XML leaks through to channels as text, and Telegram's Markdown parser chokes on the underscores, returning 400 "can't parse entities". Two fixes: - Generalize clean_response() to strip <tool_call>, <function_call>, <tool_calls>, and pipe-delimited variants (<|tool_call|>) alongside the existing <thinking> tag stripping - Add Telegram send_message helper with parse_mode fallback: try Markdown first, retry as plain text on "can't parse entities" 400 errors Co-Authored-By: Claude Opus 4.6 <[email protected]> * Add SystemCommand submission type for thread-state-independent commands System commands (/help, /model, /version, /tools, /ping, /debug) now bypass thread-state checks and safety validation via a dedicated Submission::SystemCommand variant. Previously these flowed through process_user_input() which blocked them during Processing/AwaitingApproval /Completed states. - Add /model [name] for runtime model switching with provider validation - Add active_model_name()/set_model() to LlmProvider trait with RwLock hot-swap in both NEAR AI providers - Rewrite /help with aligned columns grouped by category - Expand REPL tab-completion from 10 to 23 slash commands - Remove REPL-local /help interception (now handled by agent) Co-Authored-By: Claude Opus 4.6 <[email protected]> * Add per-tool execution timeouts, auto-create sandbox project dirs, serve built files The sandbox e2e pipeline (agent -> container -> built website -> browsable URL) was broken by three gaps: hardcoded 60s timeouts killed sandbox jobs that need minutes, no auto-created project directory meant container output vanished, and no HTTP route to browse the built files. - Add `execution_timeout()` to the `Tool` trait (default 60s), replace all four hardcoded `Duration::from_secs(60)` call sites (agent_loop, worker, scheduler, worker/runtime) with the per-tool value - Override to 660s in `RunInSandboxTool` (10 min polling + 60s buffer) - Auto-create `~/.ironclaw/projects/{uuid}/` when no `project_dir` is specified, so every sandbox job gets a persistent bind mount - Include `project_dir` and `browse_url` in sandbox tool output JSON - Add `/projects/{id}` and `/projects/{id}/{path}` static file serving routes to the web gateway with path traversal protection and MIME type detection - Add `mime_guess` dependency for content-type detection Co-Authored-By: Claude Opus 4.6 <[email protected]> * Apply cargo fmt to wizard.rs after merge Co-Authored-By: Claude Opus 4.6 <[email protected]> * Persist sandbox jobs in DB, fix web UI, unify job model Sandbox container jobs were invisible to the web UI because they lived only in ContainerJobManager's in-memory HashMap while the API queried ContextManager. This persists them to the agent_jobs table and fixes all six front-end bugs (empty job list, broken back button, empty actions/thinking tabs, wrong files tab, stuck status, no persistence). Key changes: - V4 migration adds project_dir and user_id columns to agent_jobs - Embedded migrations via refinery (no external CLI needed) - SandboxJobRecord CRUD in Store with fire-and-forget DB writes - Unified job_id: sandbox tool generates UUID, passes to ContainerJobManager - Web API queries DB for sandbox jobs, merges with ContextManager direct jobs - New endpoints: restart, project file list/read with path traversal protection - Front-end: rebuild DOM on back navigation, sandbox-aware tabs, job cards in chat stream, source badges, restart button for failed/interrupted jobs - Gateway defaults to enabled, prints Web UI URL on startup - Stale jobs marked "interrupted" on restart for visibility and restartability Co-Authored-By: Claude Opus 4.6 <[email protected]> * Secure in-chat auth: tokens never touch the LLM or chat history Remove the token parameter from tool_auth so the LLM cannot pass raw API keys. Add dedicated REST (POST /api/chat/auth-token) and WebSocket (auth_token) endpoints that route tokens directly to ext_mgr.auth(), completely bypassing the message pipeline, turns, history, and compaction. Web UI shows an auth card (password input + OAuth button) when the agent enters auth mode, submitted via the dedicated endpoint. CLI auth mode interception is unchanged (already secure). New StatusUpdate::AuthRequired/AuthCompleted variants propagate through all channels (SSE, WebSocket, REPL, WASM). Co-Authored-By: Claude Opus 4.6 <[email protected]> * feat: Add Claude Code mode for sandbox jobs Run Claude Code CLI inside Docker containers as an alternative to the standard worker mode. The bridge spawns `claude -p` with stream-json output, posts events to the orchestrator, and supports follow-up prompts via `--resume`. Key additions: - `claude-bridge` CLI subcommand and ClaudeBridgeRuntime - JobMode enum (Worker vs ClaudeCode) with per-mode container config - Orchestrator endpoints for Claude events and prompt polling - SSE event variants for real-time Claude Code streaming to frontend - Claude Code sub-tab in web UI with terminal-style output and input bar - Database migration for job_mode column and claude_code_events table - ClaudeCodeConfig with env var support (CLAUDE_CODE_ENABLED, etc.) - Mode parameter on run_in_sandbox tool schema Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: Skip create_job tool when sandbox is enabled to prevent duplicate jobs When sandbox mode is on, the LLM would call create_job (creating a pending "direct" entry) then run_in_sandbox (creating a second "sandbox" entry), producing two jobs in the list for a single user request. Now register_job_tools() skips create_job when sandbox is enabled since run_in_sandbox already creates tracked jobs. Also improved the run_in_sandbox description to guide the LLM to use it directly and to mention wait=false for async execution. Co-Authored-By: Claude Opus 4.6 <[email protected]> * feat: Web gateway UI quality-of-life improvements Phase 1: Send button disabled state to prevent double-sends, copy button on code blocks, confirm() guards on destructive actions, SSE-driven job list auto-refresh, log filters re-applied on tab switch, jobEvents memory leak fix (cap at 500, cleanup after 60s). Phase 2: Toast notification system replacing chat-based system messages, memory search highlighting with centered snippets, keyboard shortcuts (Ctrl+1-5 tabs, Ctrl+K focus, Ctrl+N new thread, Escape close/blur), activity tab toolbar with event type filter and auto-scroll toggle. Phase 3: Thread sidebar with load/switch/create, thread_id passed with messages, collapsible to hamburger. Memory inline editing with textarea, Save/Cancel, POST to /api/memory/write. Phase 4: Gateway status popover on hover (polls every 30s), extension install form (name/URL/kind), markdown rendering in memory viewer for .md files, mobile responsive layout at 768px breakpoint. Co-Authored-By: Claude Opus 4.6 <[email protected]> * feat: Add routines system, remove non-sandbox job mode from web UI Routines: scheduled & reactive job system with cron and event triggers, lightweight (single LLM call) and full-job execution modes, guardrails (cooldown, max concurrent, dedup), and LLM-facing tools for CRUD. Web UI: remove ContextManager-backed "direct" job mode entirely. Jobs are now exclusively sandbox-backed (DB + container). Simplify job detail response, drop dead types (ActionInfo, MessageInfo, MessageToolCallInfo), fix Browse Files CSS loading (trailing-slash redirect), fix Activity tab event rendering. Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: Re-enable chat input on agent completion, auto-auth on tool_activate, recover tool calls from content XML Three fixes: 1. Chat input stays disabled after agent finishes: the "Done" status SSE event now calls enableChatInput() as a safety net when the response event is empty or lost. Same for auth_completed and cancelAuth(). 2. tool_activate never triggers auth: when activation fails due to missing authentication, it now auto-initiates the auth flow (same pattern as the web API handler). detect_auth_awaiting() also matches tool_activate results now. 3. Models like GLM-4.7 emit tool calls as XML tags in content (<tool_call>tool_list</tool_call>) instead of using the structured tool_calls array. recover_tool_calls_from_content() extracts and validates these before falling back to plain text. Co-Authored-By: Claude Opus 4.6 <[email protected]> * feat: Add routines web UI tab, update docs for sandbox-jobs branch Add full routines management to the web gateway (list, detail, trigger, toggle, delete) with 7 new API endpoints, response types, and frontend (HTML, JS, CSS). Update FEATURE_PARITY.md (~23 rows), CLAUDE.md (new subsystems, config, TODOs), and README.md (architecture diagram, features, components, fix onboard command). Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: Bind Telegram bot to owner account during setup Without owner binding, anyone who discovers the bot can send it messages. The setup wizard now prompts the user to message their bot, captures their Telegram user ID via getUpdates, and persists it as telegram_owner_id in settings. On startup, the owner_id is injected into the WASM channel config so the existing owner restriction logic drops messages from non-owners. Co-Authored-By: Claude Opus 4.6 <[email protected]> * feat: Move settings from disk to PostgreSQL database Settings previously lived in three JSON files on disk (settings.json, mcp-servers.json, session.json). This made them inaccessible from the web UI and caused redundant disk reads (Settings::load() called 8+ times during startup). Now all settings live in a `settings` table (user_id + key -> JSONB) with only 4 bootstrap fields remaining on disk (database_url, pool size, secrets key source, onboard_completed) since they're needed before the DB connection exists. - Add V8 migration for settings table - Add BootstrapConfig (thin disk file) and Settings DB round-trip - Add Store CRUD methods for settings (get/set/delete/list/bulk) - Refactor Config to load from DB (env > DB > default cascade) - Add SessionManager DB persistence for session tokens - Add DB-backed MCP server config load/save functions - Add 6 settings web API endpoints (list/get/set/delete/export/import) - Add one-time disk-to-DB migration on first boot - Make CLI config commands async with DB access (disk fallback) Co-Authored-By: Claude Opus 4.6 <[email protected]> * feat: Seed workspace on boot, fix gateway duplicate logs and URL auto-auth - Add Workspace::seed_if_empty() to create core identity files (README, MEMORY, IDENTITY, SOUL, AGENTS, USER, HEARTBEAT) when missing, called on every boot without overwriting existing user edits - Remove duplicate gateway log lines from web/mod.rs (main.rs has the useful clickable ?token= URL) - Auto-authenticate from ?token= URL parameter in the web UI and strip the token from the address bar after successful auth Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: Harden sandbox security (path traversal + orchestrator auth) Two vulnerabilities fixed: 1. project_dir path traversal: The create_job tool let the LLM specify arbitrary host paths for Docker bind mounts. Removed project_dir from the tool schema entirely, and added canonicalization + prefix validation at both resolve_project_dir() and the job_manager bind mount point. 2. Orchestrator API auth bypass: worker_auth_middleware was defined but never applied. Each handler manually called validate_token(), so any new endpoint that forgot would be publicly accessible. Applied the middleware as route_layer on all /worker/ routes, removed manual auth from all 7 handlers. Bind to 127.0.0.1 on macOS/Windows (Linux keeps 0.0.0.0 since containers reach host via docker bridge, not loopback). Co-Authored-By: Claude Opus 4.6 <[email protected]> * feat: Rework gateway chat with pinned assistant, pagination, and NEAR AI response chaining Implements the 4-phase plan for overhauling the web gateway chat: - Phase 1: Pinned "Assistant" thread at top of sidebar, regular threads below - Phase 2: Cursor-based history pagination with infinite scroll - Phase 3: NEAR AI previous_response_id chaining (delta-only messages), with fallback to full history on chain errors, and DB persistence of chain state across restarts - Phase 4: SSE thread isolation (events filtered by thread_id) Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: Add per-request HTTP timeout to WASM host, redact credentials in errors Three fixes for WASM channel reliability: 1. Per-request timeout: Add optional timeout-ms parameter to http-request in both channel and tool WIT interfaces. Telegram long-poll now specifies 35s (outliving the 30s server-side hold), while regular API calls use the 30s default. Fixes the triple-30s timeout race that caused polling failures. 2. Credential redaction: reqwest::Error includes the full URL (with injected bot tokens) in its Display output. Scrub credential values from error messages before logging or returning to WASM. 3. Webhook route registration: Remove tunnel URL gate so webhook routes are always available when webhook channels exist, not only when TUNNEL_URL is configured. Co-Authored-By: Claude Opus 4.6 <[email protected]> * chore: Fix clippy warnings in WASM tools and channels - slack channel: allow dead_code on signing_secret_name (forward compat field) - gmail tool: use div_ceil() instead of manual (n+2)/3 - google-calendar tool: extract CreateEventParams/UpdateEventParams structs to fix too-many-arguments warnings Co-Authored-By: Claude Opus 4.6 <[email protected]> * Fix approval flow * fix: Rebuild bundled telegram.wasm with updated WIT interface The bundled WASM binary must match the host's WIT definition. Previous binary was compiled against the old 4-arg http-request; this rebuild includes the new timeout-ms parameter. Co-Authored-By: Claude Opus 4.6 <[email protected]> * refactor: Load WASM channels from disk instead of bundling in binary Remove include_bytes! embedding of telegram.wasm. Channels are now loaded from their build output directories (channels-src/<name>/target/) during onboarding, then from ~/.ironclaw/channels/ at runtime. - bundled.rs: locate_channel_artifacts() finds WASM + capabilities from build output; IRONCLAW_CHANNELS_SRC env var overrides the default path - available_channel_names(): only lists channels with build artifacts - bundled_channel_names(): lists all known channels (manifest) - Setup wizard uses available_channel_names() to offer installable channels - Add *.wasm to .gitignore, remove tracked telegram.wasm Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: Persist gateway auth token, fix thread hydration race, polish auth screen Three web gateway UX fixes: 1. Token persistence: Store auth token in sessionStorage so refreshing the page doesn't force re-authentication. Hide the auth screen immediately when a saved token exists to prevent flash. 2. Thread hydration: Remove the !msgs.is_empty() bail-out in maybe_hydrate_thread so that even brand-new (empty) assistant threads get hydrated with their correct DB UUID. Previously resolve_thread would mint a fresh UUID, causing messages to land in the wrong conversation and duplicate threads to appear. 3. Auth screen: Redesign as a centered card with brand, tagline, labeled input, and hint text. Also adds 34 new tests covering session/thread lifecycle, thread resolution isolation (user, channel, external ID), hydration edge cases, serialization round-trips, approval flows, and stale mapping recovery. Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: Use bindgen! for WASM tool wrapper, add dev tool loading Three changes: 1. Rewrite src/tools/wasm/wrapper.rs to use wasmtime::component::bindgen! instead of manual linker.root().func_wrap(). This fixes the "component imports instance 'near:agent/host', but a matching implementation was not found in the linker" error. All 6 host functions (log, now-millis, workspace-read, http-request, secret-exists, tool-invoke) are now properly registered under the near:agent/host namespace. Also adds WASI support, credential injection, and leak detection for HTTP requests made by WASM tools. 2. Add dev tool loading to src/tools/wasm/loader.rs. During startup, the loader now also scans tools-src/*/target/wasm32-wasip2/release/ for build artifacts that are newer than installed copies. This means during development you just rebuild the WASM and restart the host; no manual copy step needed. Set IRONCLAW_TOOLS_SRC to override the source dir. 3. Wire up load_dev_tools() in main.rs alongside the existing load_from_dir() call. Co-Authored-By: Claude Opus 4.6 <[email protected]> * feat: Wire main startup and CLI to use DB-backed settings main.rs now reloads Config from the database after connecting, attaches the store to the session manager for dual-write tokens, and loads MCP servers from DB instead of disk. ExtensionManager and MCP CLI commands use DB when available with disk fallback. 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
co-authored by
Claude Opus 4.6
Claude Opus 4.6
parent
202665a55c
commit
ced83d5b4d
@@ -0,0 +1,400 @@
|
||||
//! HTTP client for worker-to-orchestrator communication.
|
||||
//!
|
||||
//! Every request includes a bearer token from `IRONCLAW_WORKER_TOKEN` env var.
|
||||
//! The orchestrator validates this token is scoped to the correct job.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::error::WorkerError;
|
||||
use crate::llm::{
|
||||
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, ToolCall,
|
||||
ToolCompletionRequest, ToolCompletionResponse, ToolDefinition,
|
||||
};
|
||||
|
||||
/// HTTP client that a container worker uses to talk to the orchestrator.
|
||||
pub struct WorkerHttpClient {
|
||||
client: reqwest::Client,
|
||||
orchestrator_url: String,
|
||||
job_id: Uuid,
|
||||
token: String,
|
||||
}
|
||||
|
||||
/// Status update sent from worker to orchestrator.
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct StatusUpdate {
|
||||
pub state: String,
|
||||
pub message: Option<String>,
|
||||
pub iteration: u32,
|
||||
}
|
||||
|
||||
/// Job description fetched from orchestrator.
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct JobDescription {
|
||||
pub title: String,
|
||||
pub description: String,
|
||||
pub project_dir: Option<String>,
|
||||
}
|
||||
|
||||
/// Completion result from the orchestrator (proxied from the real LLM).
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct ProxyCompletionRequest {
|
||||
pub messages: Vec<ChatMessage>,
|
||||
pub max_tokens: Option<u32>,
|
||||
pub temperature: Option<f32>,
|
||||
pub stop_sequences: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct ProxyCompletionResponse {
|
||||
pub content: String,
|
||||
pub input_tokens: u32,
|
||||
pub output_tokens: u32,
|
||||
pub finish_reason: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct ProxyToolCompletionRequest {
|
||||
pub messages: Vec<ChatMessage>,
|
||||
pub tools: Vec<ToolDefinition>,
|
||||
pub max_tokens: Option<u32>,
|
||||
pub temperature: Option<f32>,
|
||||
pub tool_choice: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct ProxyToolCompletionResponse {
|
||||
pub content: Option<String>,
|
||||
pub tool_calls: Vec<ToolCall>,
|
||||
pub input_tokens: u32,
|
||||
pub output_tokens: u32,
|
||||
pub finish_reason: String,
|
||||
}
|
||||
|
||||
/// Completion result for the worker to report when done.
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct CompletionReport {
|
||||
pub success: bool,
|
||||
pub message: Option<String>,
|
||||
pub iterations: u32,
|
||||
}
|
||||
|
||||
/// Payload sent to the orchestrator for each job event (shared by worker and Claude Code bridge).
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct JobEventPayload {
|
||||
pub event_type: String,
|
||||
pub data: serde_json::Value,
|
||||
}
|
||||
|
||||
/// Response from the prompt polling endpoint.
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct PromptResponse {
|
||||
pub content: String,
|
||||
#[serde(default)]
|
||||
pub done: bool,
|
||||
}
|
||||
|
||||
impl WorkerHttpClient {
|
||||
/// Create a new client from environment.
|
||||
///
|
||||
/// Reads `IRONCLAW_WORKER_TOKEN` from the environment.
|
||||
pub fn from_env(orchestrator_url: String, job_id: Uuid) -> Result<Self, WorkerError> {
|
||||
let token =
|
||||
std::env::var("IRONCLAW_WORKER_TOKEN").map_err(|_| WorkerError::MissingToken)?;
|
||||
|
||||
Ok(Self {
|
||||
client: reqwest::Client::new(),
|
||||
orchestrator_url: orchestrator_url.trim_end_matches('/').to_string(),
|
||||
job_id,
|
||||
token,
|
||||
})
|
||||
}
|
||||
|
||||
/// Create with an explicit token (for testing).
|
||||
pub fn new(orchestrator_url: String, job_id: Uuid, token: String) -> Self {
|
||||
Self {
|
||||
client: reqwest::Client::new(),
|
||||
orchestrator_url: orchestrator_url.trim_end_matches('/').to_string(),
|
||||
job_id,
|
||||
token,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the base orchestrator URL.
|
||||
pub fn orchestrator_url(&self) -> &str {
|
||||
&self.orchestrator_url
|
||||
}
|
||||
|
||||
fn url(&self, path: &str) -> String {
|
||||
format!("{}/worker/{}/{}", self.orchestrator_url, self.job_id, path)
|
||||
}
|
||||
|
||||
/// Fetch the job description from the orchestrator.
|
||||
pub async fn get_job(&self) -> Result<JobDescription, WorkerError> {
|
||||
let resp = self
|
||||
.client
|
||||
.get(self.url("job"))
|
||||
.bearer_auth(&self.token)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| WorkerError::ConnectionFailed {
|
||||
url: self.orchestrator_url.clone(),
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
return Err(WorkerError::OrchestratorRejected {
|
||||
job_id: self.job_id,
|
||||
reason: format!("GET /job returned {}", resp.status()),
|
||||
});
|
||||
}
|
||||
|
||||
resp.json().await.map_err(|e| WorkerError::LlmProxyFailed {
|
||||
reason: format!("failed to parse job description: {}", e),
|
||||
})
|
||||
}
|
||||
|
||||
/// Proxy an LLM completion request through the orchestrator.
|
||||
pub async fn llm_complete(
|
||||
&self,
|
||||
request: &CompletionRequest,
|
||||
) -> Result<CompletionResponse, WorkerError> {
|
||||
let proxy_req = ProxyCompletionRequest {
|
||||
messages: request.messages.clone(),
|
||||
max_tokens: request.max_tokens,
|
||||
temperature: request.temperature,
|
||||
stop_sequences: request.stop_sequences.clone(),
|
||||
};
|
||||
|
||||
let resp = self
|
||||
.client
|
||||
.post(self.url("llm/complete"))
|
||||
.bearer_auth(&self.token)
|
||||
.json(&proxy_req)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| WorkerError::LlmProxyFailed {
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return Err(WorkerError::LlmProxyFailed {
|
||||
reason: format!("orchestrator returned {}: {}", status, body),
|
||||
});
|
||||
}
|
||||
|
||||
let proxy_resp: ProxyCompletionResponse =
|
||||
resp.json().await.map_err(|e| WorkerError::LlmProxyFailed {
|
||||
reason: format!("failed to parse LLM response: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(CompletionResponse {
|
||||
content: proxy_resp.content,
|
||||
input_tokens: proxy_resp.input_tokens,
|
||||
output_tokens: proxy_resp.output_tokens,
|
||||
finish_reason: parse_finish_reason(&proxy_resp.finish_reason),
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// Proxy an LLM tool completion request through the orchestrator.
|
||||
pub async fn llm_complete_with_tools(
|
||||
&self,
|
||||
request: &ToolCompletionRequest,
|
||||
) -> Result<ToolCompletionResponse, WorkerError> {
|
||||
let proxy_req = ProxyToolCompletionRequest {
|
||||
messages: request.messages.clone(),
|
||||
tools: request.tools.clone(),
|
||||
max_tokens: request.max_tokens,
|
||||
temperature: request.temperature,
|
||||
tool_choice: request.tool_choice.clone(),
|
||||
};
|
||||
|
||||
let resp = self
|
||||
.client
|
||||
.post(self.url("llm/complete_with_tools"))
|
||||
.bearer_auth(&self.token)
|
||||
.json(&proxy_req)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| WorkerError::LlmProxyFailed {
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return Err(WorkerError::LlmProxyFailed {
|
||||
reason: format!("orchestrator returned {}: {}", status, body),
|
||||
});
|
||||
}
|
||||
|
||||
let proxy_resp: ProxyToolCompletionResponse =
|
||||
resp.json().await.map_err(|e| WorkerError::LlmProxyFailed {
|
||||
reason: format!("failed to parse tool completion response: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(ToolCompletionResponse {
|
||||
content: proxy_resp.content,
|
||||
tool_calls: proxy_resp.tool_calls,
|
||||
input_tokens: proxy_resp.input_tokens,
|
||||
output_tokens: proxy_resp.output_tokens,
|
||||
finish_reason: parse_finish_reason(&proxy_resp.finish_reason),
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// Report status to the orchestrator.
|
||||
pub async fn report_status(&self, update: &StatusUpdate) -> Result<(), WorkerError> {
|
||||
let resp = self
|
||||
.client
|
||||
.post(self.url("status"))
|
||||
.bearer_auth(&self.token)
|
||||
.json(update)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| WorkerError::ConnectionFailed {
|
||||
url: self.orchestrator_url.clone(),
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
tracing::warn!(
|
||||
"Status report failed with {}: {}",
|
||||
resp.status(),
|
||||
resp.text().await.unwrap_or_default()
|
||||
);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Post a job event to the orchestrator (fire-and-forget style, logs on failure).
|
||||
pub async fn post_event(&self, payload: &JobEventPayload) {
|
||||
let resp = self
|
||||
.client
|
||||
.post(self.url("event"))
|
||||
.bearer_auth(&self.token)
|
||||
.json(payload)
|
||||
.send()
|
||||
.await;
|
||||
|
||||
match resp {
|
||||
Ok(r) if !r.status().is_success() => {
|
||||
tracing::debug!(
|
||||
job_id = %self.job_id,
|
||||
event_type = %payload.event_type,
|
||||
status = %r.status(),
|
||||
"Job event POST rejected"
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::debug!(
|
||||
job_id = %self.job_id,
|
||||
event_type = %payload.event_type,
|
||||
"Job event POST failed: {}", e
|
||||
);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
/// Poll the orchestrator for a follow-up prompt.
|
||||
///
|
||||
/// Returns `None` if no prompt is available (204 No Content).
|
||||
pub async fn poll_prompt(&self) -> Result<Option<PromptResponse>, WorkerError> {
|
||||
let resp = self
|
||||
.client
|
||||
.get(self.url("prompt"))
|
||||
.bearer_auth(&self.token)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| WorkerError::ConnectionFailed {
|
||||
url: self.orchestrator_url.clone(),
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
if resp.status() == reqwest::StatusCode::NO_CONTENT {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if !resp.status().is_success() {
|
||||
return Err(WorkerError::OrchestratorRejected {
|
||||
job_id: self.job_id,
|
||||
reason: format!("prompt endpoint returned {}", resp.status()),
|
||||
});
|
||||
}
|
||||
|
||||
let prompt: PromptResponse =
|
||||
resp.json().await.map_err(|e| WorkerError::LlmProxyFailed {
|
||||
reason: format!("failed to parse prompt response: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(Some(prompt))
|
||||
}
|
||||
|
||||
/// Signal job completion to the orchestrator.
|
||||
pub async fn report_complete(&self, report: &CompletionReport) -> Result<(), WorkerError> {
|
||||
let resp = self
|
||||
.client
|
||||
.post(self.url("complete"))
|
||||
.bearer_auth(&self.token)
|
||||
.json(report)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| WorkerError::ConnectionFailed {
|
||||
url: self.orchestrator_url.clone(),
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
return Err(WorkerError::OrchestratorRejected {
|
||||
job_id: self.job_id,
|
||||
reason: format!("completion report rejected: {}", resp.status()),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_finish_reason(s: &str) -> FinishReason {
|
||||
match s {
|
||||
"stop" => FinishReason::Stop,
|
||||
"length" => FinishReason::Length,
|
||||
"tool_use" | "tool_calls" => FinishReason::ToolUse,
|
||||
"content_filter" => FinishReason::ContentFilter,
|
||||
_ => FinishReason::Unknown,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_url_construction() {
|
||||
let client = WorkerHttpClient::new(
|
||||
"http://host.docker.internal:50051".to_string(),
|
||||
Uuid::nil(),
|
||||
"test-token".to_string(),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
client.url("llm/complete"),
|
||||
format!(
|
||||
"http://host.docker.internal:50051/worker/{}/llm/complete",
|
||||
Uuid::nil()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_finish_reason() {
|
||||
assert_eq!(parse_finish_reason("stop"), FinishReason::Stop);
|
||||
assert_eq!(parse_finish_reason("tool_use"), FinishReason::ToolUse);
|
||||
assert_eq!(parse_finish_reason("unknown"), FinishReason::Unknown);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,644 @@
|
||||
//! Claude Code bridge for sandboxed execution.
|
||||
//!
|
||||
//! Spawns the `claude` CLI inside a Docker container and streams its NDJSON
|
||||
//! output back to the orchestrator via HTTP. Supports follow-up prompts via
|
||||
//! `--resume`.
|
||||
//!
|
||||
//! ```text
|
||||
//! ┌─────────────────────────────────────────────┐
|
||||
//! │ Docker Container │
|
||||
//! │ │
|
||||
//! │ ironclaw claude-bridge --job-id <uuid> │
|
||||
//! │ └─ claude -p "task" --output-format │
|
||||
//! │ stream-json --dangerously-skip-perms │
|
||||
//! │ └─ reads stdout line-by-line │
|
||||
//! │ └─ POSTs events to orchestrator │
|
||||
//! │ └─ polls for follow-up prompts │
|
||||
//! │ └─ on follow-up: claude --resume │
|
||||
//! └─────────────────────────────────────────────┘
|
||||
//! ```
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::io::{AsyncBufReadExt, BufReader};
|
||||
use tokio::process::Command;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::error::WorkerError;
|
||||
use crate::worker::api::{CompletionReport, JobEventPayload, PromptResponse, WorkerHttpClient};
|
||||
|
||||
/// Configuration for the Claude bridge runtime.
|
||||
pub struct ClaudeBridgeConfig {
|
||||
pub job_id: Uuid,
|
||||
pub orchestrator_url: String,
|
||||
pub max_turns: u32,
|
||||
pub model: String,
|
||||
pub timeout: Duration,
|
||||
}
|
||||
|
||||
/// 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.
|
||||
#[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>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ContentBlock {
|
||||
#[serde(rename = "type")]
|
||||
pub block_type: String,
|
||||
#[serde(default)]
|
||||
pub text: Option<String>,
|
||||
#[serde(default)]
|
||||
pub name: Option<String>,
|
||||
#[serde(default)]
|
||||
pub input: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
pub content: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ResultInfo {
|
||||
#[serde(default)]
|
||||
pub is_error: Option<bool>,
|
||||
#[serde(default)]
|
||||
pub duration_ms: Option<u64>,
|
||||
#[serde(default)]
|
||||
pub num_turns: Option<u32>,
|
||||
}
|
||||
|
||||
/// The Claude Code bridge runtime.
|
||||
pub struct ClaudeBridgeRuntime {
|
||||
config: ClaudeBridgeConfig,
|
||||
client: Arc<WorkerHttpClient>,
|
||||
}
|
||||
|
||||
impl ClaudeBridgeRuntime {
|
||||
/// Create a new bridge runtime.
|
||||
///
|
||||
/// Reads `IRONCLAW_WORKER_TOKEN` from the environment for auth.
|
||||
pub fn new(config: ClaudeBridgeConfig) -> Result<Self, WorkerError> {
|
||||
let client = Arc::new(WorkerHttpClient::from_env(
|
||||
config.orchestrator_url.clone(),
|
||||
config.job_id,
|
||||
)?);
|
||||
|
||||
Ok(Self { config, client })
|
||||
}
|
||||
|
||||
/// Run the bridge: fetch job, spawn claude, stream events, handle follow-ups.
|
||||
pub async fn run(&self) -> Result<(), WorkerError> {
|
||||
// Fetch the job description from the orchestrator
|
||||
let job = self.client.get_job().await?;
|
||||
|
||||
tracing::info!(
|
||||
job_id = %self.config.job_id,
|
||||
"Starting Claude Code bridge for: {}",
|
||||
truncate(&job.description, 100)
|
||||
);
|
||||
|
||||
// Report that we're running
|
||||
self.client
|
||||
.report_status(&crate::worker::api::StatusUpdate {
|
||||
state: "running".to_string(),
|
||||
message: Some("Spawning Claude Code".to_string()),
|
||||
iteration: 0,
|
||||
})
|
||||
.await?;
|
||||
|
||||
// Run the initial Claude session
|
||||
let session_id = match self.run_claude_session(&job.description, None).await {
|
||||
Ok(sid) => sid,
|
||||
Err(e) => {
|
||||
tracing::error!(job_id = %self.config.job_id, "Claude session failed: {}", e);
|
||||
self.client
|
||||
.report_complete(&CompletionReport {
|
||||
success: false,
|
||||
message: Some(format!("Claude Code failed: {}", e)),
|
||||
iterations: 1,
|
||||
})
|
||||
.await?;
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
|
||||
// Follow-up loop: poll for prompts, resume Claude sessions
|
||||
let mut iteration = 1u32;
|
||||
loop {
|
||||
// Poll for a follow-up prompt (2 second intervals)
|
||||
match self.poll_for_prompt().await {
|
||||
Ok(Some(prompt)) => {
|
||||
if prompt.done {
|
||||
tracing::info!(job_id = %self.config.job_id, "Orchestrator signaled done");
|
||||
break;
|
||||
}
|
||||
iteration += 1;
|
||||
tracing::info!(
|
||||
job_id = %self.config.job_id,
|
||||
"Got follow-up prompt, resuming session"
|
||||
);
|
||||
if let Err(e) = self
|
||||
.run_claude_session(&prompt.content, session_id.as_deref())
|
||||
.await
|
||||
{
|
||||
tracing::error!(
|
||||
job_id = %self.config.job_id,
|
||||
"Follow-up Claude session failed: {}", e
|
||||
);
|
||||
// Don't fail the whole job on a follow-up error, just report it
|
||||
self.report_event(
|
||||
"status",
|
||||
&serde_json::json!({
|
||||
"message": format!("Follow-up session failed: {}", e),
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
Ok(None) => {
|
||||
// No prompt available, wait and poll again
|
||||
tokio::time::sleep(Duration::from_secs(2)).await;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
job_id = %self.config.job_id,
|
||||
"Prompt polling error: {}", e
|
||||
);
|
||||
tokio::time::sleep(Duration::from_secs(5)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self.client
|
||||
.report_complete(&CompletionReport {
|
||||
success: true,
|
||||
message: Some("Claude Code session completed".to_string()),
|
||||
iterations: iteration,
|
||||
})
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Spawn a `claude` CLI process and stream its output.
|
||||
///
|
||||
/// Returns the session_id if captured from the `system` init message.
|
||||
async fn run_claude_session(
|
||||
&self,
|
||||
prompt: &str,
|
||||
resume_session_id: Option<&str>,
|
||||
) -> Result<Option<String>, WorkerError> {
|
||||
let mut cmd = Command::new("claude");
|
||||
cmd.arg("-p")
|
||||
.arg(prompt)
|
||||
.arg("--output-format")
|
||||
.arg("stream-json")
|
||||
.arg("--dangerously-skip-permissions")
|
||||
.arg("--max-turns")
|
||||
.arg(self.config.max_turns.to_string())
|
||||
.arg("--model")
|
||||
.arg(&self.config.model);
|
||||
|
||||
if let Some(sid) = resume_session_id {
|
||||
cmd.arg("--resume").arg(sid);
|
||||
}
|
||||
|
||||
cmd.current_dir("/workspace")
|
||||
.stdout(std::process::Stdio::piped())
|
||||
.stderr(std::process::Stdio::piped());
|
||||
|
||||
let mut child = cmd.spawn().map_err(|e| WorkerError::ExecutionFailed {
|
||||
reason: format!("failed to spawn claude: {}", e),
|
||||
})?;
|
||||
|
||||
let stdout = child
|
||||
.stdout
|
||||
.take()
|
||||
.ok_or_else(|| WorkerError::ExecutionFailed {
|
||||
reason: "failed to capture claude stdout".to_string(),
|
||||
})?;
|
||||
|
||||
let stderr = child
|
||||
.stderr
|
||||
.take()
|
||||
.ok_or_else(|| WorkerError::ExecutionFailed {
|
||||
reason: "failed to capture claude stderr".to_string(),
|
||||
})?;
|
||||
|
||||
// Spawn stderr reader that forwards lines as log events
|
||||
let client_for_stderr = Arc::clone(&self.client);
|
||||
let job_id = self.config.job_id;
|
||||
let stderr_handle = tokio::spawn(async move {
|
||||
let reader = BufReader::new(stderr);
|
||||
let mut lines = reader.lines();
|
||||
while let Ok(Some(line)) = lines.next_line().await {
|
||||
tracing::debug!(job_id = %job_id, "claude stderr: {}", line);
|
||||
let payload = JobEventPayload {
|
||||
event_type: "status".to_string(),
|
||||
data: serde_json::json!({ "message": line }),
|
||||
};
|
||||
client_for_stderr.post_event(&payload).await;
|
||||
}
|
||||
});
|
||||
|
||||
// Read stdout NDJSON line by line
|
||||
let reader = BufReader::new(stdout);
|
||||
let mut lines = reader.lines();
|
||||
let mut session_id: Option<String> = None;
|
||||
|
||||
while let Ok(Some(line)) = lines.next_line().await {
|
||||
let line = line.trim().to_string();
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
match serde_json::from_str::<ClaudeStreamEvent>(&line) {
|
||||
Ok(event) => {
|
||||
// Capture session_id from system init
|
||||
if event.event_type == "system" {
|
||||
if let Some(ref sid) = event.session_id {
|
||||
session_id = Some(sid.clone());
|
||||
tracing::info!(
|
||||
job_id = %self.config.job_id,
|
||||
session_id = %sid,
|
||||
"Captured Claude session ID"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Convert to our event payload and forward
|
||||
let payloads = stream_event_to_payloads(&event);
|
||||
for payload in payloads {
|
||||
self.report_event(&payload.event_type, &payload.data).await;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
// Not valid JSON, forward as a status message
|
||||
tracing::debug!(
|
||||
job_id = %self.config.job_id,
|
||||
"Non-JSON claude output: {} (parse error: {})", line, e
|
||||
);
|
||||
self.report_event("status", &serde_json::json!({ "message": line }))
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Wait for the process to exit
|
||||
let status = child
|
||||
.wait()
|
||||
.await
|
||||
.map_err(|e| WorkerError::ExecutionFailed {
|
||||
reason: format!("failed waiting for claude: {}", e),
|
||||
})?;
|
||||
|
||||
// Wait for stderr reader to finish
|
||||
let _ = stderr_handle.await;
|
||||
|
||||
if !status.success() {
|
||||
let code = status.code().unwrap_or(-1);
|
||||
tracing::warn!(
|
||||
job_id = %self.config.job_id,
|
||||
exit_code = code,
|
||||
"Claude process exited with non-zero status"
|
||||
);
|
||||
|
||||
// Report result event
|
||||
self.report_event(
|
||||
"result",
|
||||
&serde_json::json!({
|
||||
"status": "error",
|
||||
"exit_code": code,
|
||||
"session_id": session_id,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
return Err(WorkerError::ExecutionFailed {
|
||||
reason: format!("claude exited with code {}", code),
|
||||
});
|
||||
}
|
||||
|
||||
// Report successful result
|
||||
self.report_event(
|
||||
"result",
|
||||
&serde_json::json!({
|
||||
"status": "completed",
|
||||
"session_id": session_id,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
Ok(session_id)
|
||||
}
|
||||
|
||||
/// Post a job event to the orchestrator.
|
||||
async fn report_event(&self, event_type: &str, data: &serde_json::Value) {
|
||||
let payload = JobEventPayload {
|
||||
event_type: event_type.to_string(),
|
||||
data: data.clone(),
|
||||
};
|
||||
self.client.post_event(&payload).await;
|
||||
}
|
||||
|
||||
/// Poll the orchestrator for a follow-up prompt.
|
||||
async fn poll_for_prompt(&self) -> Result<Option<PromptResponse>, WorkerError> {
|
||||
self.client.poll_prompt().await
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a Claude stream event into one or more event payloads for the orchestrator.
|
||||
fn stream_event_to_payloads(event: &ClaudeStreamEvent) -> Vec<JobEventPayload> {
|
||||
let mut payloads = Vec::new();
|
||||
|
||||
match event.event_type.as_str() {
|
||||
"system" => {
|
||||
payloads.push(JobEventPayload {
|
||||
event_type: "status".to_string(),
|
||||
data: serde_json::json!({
|
||||
"message": "Claude Code session started",
|
||||
"session_id": event.session_id,
|
||||
}),
|
||||
});
|
||||
}
|
||||
"assistant" => {
|
||||
// Extract text content and tool_use blocks
|
||||
if let Some(ref blocks) = event.content {
|
||||
for block in blocks {
|
||||
match block.block_type.as_str() {
|
||||
"text" => {
|
||||
if let Some(ref text) = block.text {
|
||||
payloads.push(JobEventPayload {
|
||||
event_type: "message".to_string(),
|
||||
data: serde_json::json!({
|
||||
"role": "assistant",
|
||||
"content": text,
|
||||
}),
|
||||
});
|
||||
}
|
||||
}
|
||||
"tool_use" => {
|
||||
payloads.push(JobEventPayload {
|
||||
event_type: "tool_use".to_string(),
|
||||
data: serde_json::json!({
|
||||
"tool_name": block.name,
|
||||
"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(""),
|
||||
}),
|
||||
});
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
"result" => {
|
||||
let is_error = event
|
||||
.result
|
||||
.as_ref()
|
||||
.and_then(|r| r.is_error)
|
||||
.unwrap_or(false);
|
||||
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),
|
||||
}),
|
||||
});
|
||||
}
|
||||
_ => {
|
||||
// Forward unknown event types as status
|
||||
payloads.push(JobEventPayload {
|
||||
event_type: "status".to_string(),
|
||||
data: serde_json::json!({
|
||||
"message": format!("Claude event: {}", event.event_type),
|
||||
"raw_type": event.event_type,
|
||||
}),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
payloads
|
||||
}
|
||||
|
||||
fn truncate(s: &str, max_len: usize) -> &str {
|
||||
if s.len() <= max_len { s } else { &s[..max_len] }
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_parse_system_event() {
|
||||
let json = r#"{"type":"system","session_id":"abc-123","subtype":"init"}"#;
|
||||
let event: ClaudeStreamEvent = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(event.event_type, "system");
|
||||
assert_eq!(event.session_id.as_deref(), Some("abc-123"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_assistant_text_event() {
|
||||
let json = r#"{"type":"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();
|
||||
assert_eq!(blocks.len(), 1);
|
||||
assert_eq!(blocks[0].block_type, "text");
|
||||
assert_eq!(blocks[0].text.as_deref(), Some("Hello world"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_assistant_tool_use_event() {
|
||||
let json = r#"{"type":"assistant","content":[{"type":"tool_use","name":"Bash","input":{"command":"ls"}}]}"#;
|
||||
let event: ClaudeStreamEvent = serde_json::from_str(json).unwrap();
|
||||
let blocks = event.content.unwrap();
|
||||
assert_eq!(blocks[0].block_type, "tool_use");
|
||||
assert_eq!(blocks[0].name.as_deref(), Some("Bash"));
|
||||
assert!(blocks[0].input.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_result_event() {
|
||||
let json =
|
||||
r#"{"type":"result","result":{"is_error":false,"duration_ms":5000,"num_turns":3}}"#;
|
||||
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));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_result_error_event() {
|
||||
let json = r#"{"type":"result","result":{"is_error":true}}"#;
|
||||
let event: ClaudeStreamEvent = serde_json::from_str(json).unwrap();
|
||||
let result = event.result.unwrap();
|
||||
assert_eq!(result.is_error, Some(true));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_stream_event_to_payloads_system() {
|
||||
let event = ClaudeStreamEvent {
|
||||
event_type: "system".to_string(),
|
||||
session_id: Some("sid-123".to_string()),
|
||||
content: None,
|
||||
result: 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].event_type, "status");
|
||||
assert_eq!(payloads[0].data["session_id"], "sid-123");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_stream_event_to_payloads_assistant_text() {
|
||||
let event = ClaudeStreamEvent {
|
||||
event_type: "assistant".to_string(),
|
||||
session_id: None,
|
||||
content: Some(vec![ContentBlock {
|
||||
block_type: "text".to_string(),
|
||||
text: Some("Here's the answer".to_string()),
|
||||
name: None,
|
||||
input: None,
|
||||
content: None,
|
||||
}]),
|
||||
result: None,
|
||||
tool_name: None,
|
||||
input: None,
|
||||
output: None,
|
||||
subtype: None,
|
||||
};
|
||||
let payloads = stream_event_to_payloads(&event);
|
||||
assert_eq!(payloads.len(), 1);
|
||||
assert_eq!(payloads[0].event_type, "message");
|
||||
assert_eq!(payloads[0].data["role"], "assistant");
|
||||
assert_eq!(payloads[0].data["content"], "Here's the answer");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_stream_event_to_payloads_result_success() {
|
||||
let event = ClaudeStreamEvent {
|
||||
event_type: "result".to_string(),
|
||||
session_id: Some("s1".to_string()),
|
||||
content: None,
|
||||
result: Some(ResultInfo {
|
||||
is_error: Some(false),
|
||||
duration_ms: Some(12000),
|
||||
num_turns: Some(5),
|
||||
}),
|
||||
tool_name: None,
|
||||
input: None,
|
||||
output: None,
|
||||
subtype: None,
|
||||
};
|
||||
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");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_stream_event_to_payloads_result_error() {
|
||||
let event = ClaudeStreamEvent {
|
||||
event_type: "result".to_string(),
|
||||
session_id: None,
|
||||
content: None,
|
||||
result: Some(ResultInfo {
|
||||
is_error: Some(true),
|
||||
duration_ms: None,
|
||||
num_turns: None,
|
||||
}),
|
||||
tool_name: None,
|
||||
input: None,
|
||||
output: None,
|
||||
subtype: None,
|
||||
};
|
||||
let payloads = stream_event_to_payloads(&event);
|
||||
assert_eq!(payloads[0].data["status"], "error");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_stream_event_to_payloads_unknown_type() {
|
||||
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,
|
||||
};
|
||||
let payloads = stream_event_to_payloads(&event);
|
||||
assert_eq!(payloads.len(), 1);
|
||||
assert_eq!(payloads[0].event_type, "status");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_claude_event_payload_serde() {
|
||||
let payload = JobEventPayload {
|
||||
event_type: "message".to_string(),
|
||||
data: serde_json::json!({ "role": "assistant", "content": "hi" }),
|
||||
};
|
||||
let json = serde_json::to_string(&payload).unwrap();
|
||||
let parsed: JobEventPayload = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(parsed.event_type, "message");
|
||||
assert_eq!(parsed.data["content"], "hi");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate() {
|
||||
assert_eq!(truncate("hello", 10), "hello");
|
||||
assert_eq!(truncate("hello world", 5), "hello");
|
||||
assert_eq!(truncate("", 5), "");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
//! Worker mode for running inside Docker containers.
|
||||
//!
|
||||
//! When `ironclaw worker` is invoked, the binary starts in worker mode:
|
||||
//! - Connects to the orchestrator over HTTP
|
||||
//! - Uses a `ProxyLlmProvider` that routes LLM calls through the orchestrator
|
||||
//! - Runs container-safe tools (shell, file ops, patch)
|
||||
//! - Reports status and completion back to the orchestrator
|
||||
//!
|
||||
//! ```text
|
||||
//! ┌────────────────────────────────┐
|
||||
//! │ Docker Container │
|
||||
//! │ │
|
||||
//! │ ironclaw worker │
|
||||
//! │ ├─ ProxyLlmProvider ─────────┼──▶ Orchestrator /worker/{id}/llm/complete
|
||||
//! │ ├─ SafetyLayer │
|
||||
//! │ ├─ ToolRegistry │
|
||||
//! │ │ ├─ shell │
|
||||
//! │ │ ├─ read_file │
|
||||
//! │ │ ├─ write_file │
|
||||
//! │ │ ├─ list_dir │
|
||||
//! │ │ └─ apply_patch │
|
||||
//! │ └─ WorkerHttpClient ─────────┼──▶ Orchestrator /worker/{id}/status
|
||||
//! │ │
|
||||
//! └────────────────────────────────┘
|
||||
//! ```
|
||||
|
||||
pub mod api;
|
||||
pub mod claude_bridge;
|
||||
pub mod proxy_llm;
|
||||
pub mod runtime;
|
||||
|
||||
pub use api::WorkerHttpClient;
|
||||
pub use claude_bridge::ClaudeBridgeRuntime;
|
||||
pub use proxy_llm::ProxyLlmProvider;
|
||||
pub use runtime::WorkerRuntime;
|
||||
@@ -0,0 +1,95 @@
|
||||
//! LLM provider that proxies all calls through the orchestrator HTTP API.
|
||||
//!
|
||||
//! The worker never has direct access to API keys or session tokens.
|
||||
//! All LLM requests go through the orchestrator, which holds the real credentials.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use rust_decimal::Decimal;
|
||||
|
||||
use crate::error::LlmError;
|
||||
use crate::llm::{
|
||||
CompletionRequest, CompletionResponse, LlmProvider, ToolCompletionRequest,
|
||||
ToolCompletionResponse,
|
||||
};
|
||||
use crate::worker::api::WorkerHttpClient;
|
||||
|
||||
/// An LLM provider that routes all calls through the orchestrator's HTTP API.
|
||||
///
|
||||
/// No API keys or secrets are needed in the container. The orchestrator
|
||||
/// handles authentication and billing.
|
||||
pub struct ProxyLlmProvider {
|
||||
client: Arc<WorkerHttpClient>,
|
||||
model_name: String,
|
||||
}
|
||||
|
||||
impl ProxyLlmProvider {
|
||||
pub fn new(client: Arc<WorkerHttpClient>, model_name: String) -> Self {
|
||||
Self { client, model_name }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LlmProvider for ProxyLlmProvider {
|
||||
fn model_name(&self) -> &str {
|
||||
&self.model_name
|
||||
}
|
||||
|
||||
fn cost_per_token(&self) -> (Decimal, Decimal) {
|
||||
// Cost tracking happens on the orchestrator side
|
||||
(Decimal::ZERO, Decimal::ZERO)
|
||||
}
|
||||
|
||||
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||
self.client
|
||||
.llm_complete(&request)
|
||||
.await
|
||||
.map_err(|e| LlmError::RequestFailed {
|
||||
provider: "proxy".to_string(),
|
||||
reason: e.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn complete_with_tools(
|
||||
&self,
|
||||
request: ToolCompletionRequest,
|
||||
) -> Result<ToolCompletionResponse, LlmError> {
|
||||
self.client
|
||||
.llm_complete_with_tools(&request)
|
||||
.await
|
||||
.map_err(|e| LlmError::RequestFailed {
|
||||
provider: "proxy".to_string(),
|
||||
reason: e.to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_proxy_model_name() {
|
||||
let client = Arc::new(WorkerHttpClient::new(
|
||||
"http://localhost:50051".to_string(),
|
||||
uuid::Uuid::nil(),
|
||||
"test".to_string(),
|
||||
));
|
||||
let provider = ProxyLlmProvider::new(client, "test-model".to_string());
|
||||
assert_eq!(provider.model_name(), "test-model");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_proxy_cost_is_zero() {
|
||||
let client = Arc::new(WorkerHttpClient::new(
|
||||
"http://localhost:50051".to_string(),
|
||||
uuid::Uuid::nil(),
|
||||
"test".to_string(),
|
||||
));
|
||||
let provider = ProxyLlmProvider::new(client, "test-model".to_string());
|
||||
let (input, output) = provider.cost_per_token();
|
||||
assert_eq!(input, Decimal::ZERO);
|
||||
assert_eq!(output, Decimal::ZERO);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,491 @@
|
||||
//! Worker runtime: the main execution loop inside a container.
|
||||
//!
|
||||
//! Reuses the existing `Reasoning` and `SafetyLayer` infrastructure but
|
||||
//! connects to the orchestrator for LLM calls instead of calling APIs directly.
|
||||
//! Streams real-time events (message, tool_use, tool_result, result) through
|
||||
//! the orchestrator's job event pipeline for UI visibility.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::config::SafetyConfig;
|
||||
use crate::context::JobContext;
|
||||
use crate::error::WorkerError;
|
||||
use crate::llm::{
|
||||
ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult, ToolSelection,
|
||||
};
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::tools::ToolRegistry;
|
||||
use crate::worker::api::{CompletionReport, JobEventPayload, StatusUpdate, WorkerHttpClient};
|
||||
use crate::worker::proxy_llm::ProxyLlmProvider;
|
||||
|
||||
/// Configuration for the worker runtime.
|
||||
pub struct WorkerConfig {
|
||||
pub job_id: Uuid,
|
||||
pub orchestrator_url: String,
|
||||
pub max_iterations: u32,
|
||||
pub timeout: Duration,
|
||||
}
|
||||
|
||||
impl Default for WorkerConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
job_id: Uuid::nil(),
|
||||
orchestrator_url: String::new(),
|
||||
max_iterations: 50,
|
||||
timeout: Duration::from_secs(600),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The worker runtime runs inside a Docker container.
|
||||
///
|
||||
/// It connects to the orchestrator over HTTP, fetches its job description,
|
||||
/// then runs a tool execution loop until the job is complete. Events are
|
||||
/// streamed to the orchestrator so the UI can show real-time progress.
|
||||
pub struct WorkerRuntime {
|
||||
config: WorkerConfig,
|
||||
client: Arc<WorkerHttpClient>,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
tools: Arc<ToolRegistry>,
|
||||
}
|
||||
|
||||
impl WorkerRuntime {
|
||||
/// Create a new worker runtime.
|
||||
///
|
||||
/// Reads `IRONCLAW_WORKER_TOKEN` from the environment for auth.
|
||||
pub fn new(config: WorkerConfig) -> Result<Self, WorkerError> {
|
||||
let client = Arc::new(WorkerHttpClient::from_env(
|
||||
config.orchestrator_url.clone(),
|
||||
config.job_id,
|
||||
)?);
|
||||
|
||||
let llm: Arc<dyn LlmProvider> = Arc::new(ProxyLlmProvider::new(
|
||||
Arc::clone(&client),
|
||||
"proxied".to_string(),
|
||||
));
|
||||
|
||||
let safety = Arc::new(SafetyLayer::new(&SafetyConfig {
|
||||
max_output_length: 100_000,
|
||||
injection_check_enabled: true,
|
||||
}));
|
||||
|
||||
let tools = Arc::new(ToolRegistry::new());
|
||||
// Register only container-safe tools
|
||||
tools.register_container_tools();
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
client,
|
||||
llm,
|
||||
safety,
|
||||
tools,
|
||||
})
|
||||
}
|
||||
|
||||
/// Run the worker until the job is complete or an error occurs.
|
||||
pub async fn run(self) -> Result<(), WorkerError> {
|
||||
tracing::info!("Worker starting for job {}", self.config.job_id);
|
||||
|
||||
// Fetch job description from orchestrator
|
||||
let job = self.client.get_job().await?;
|
||||
|
||||
tracing::info!(
|
||||
"Received job: {} - {}",
|
||||
job.title,
|
||||
truncate(&job.description, 100)
|
||||
);
|
||||
|
||||
// Report that we're starting
|
||||
self.client
|
||||
.report_status(&StatusUpdate {
|
||||
state: "in_progress".to_string(),
|
||||
message: Some("Worker started, beginning execution".to_string()),
|
||||
iteration: 0,
|
||||
})
|
||||
.await?;
|
||||
|
||||
// Create reasoning engine
|
||||
let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone());
|
||||
|
||||
// Build initial context
|
||||
let mut reason_ctx = ReasoningContext::new().with_job(&job.description);
|
||||
|
||||
reason_ctx.messages.push(ChatMessage::system(format!(
|
||||
r#"You are an autonomous agent running inside a Docker container.
|
||||
|
||||
Job: {}
|
||||
Description: {}
|
||||
|
||||
You have tools for shell commands, file operations, and code editing.
|
||||
Work independently to complete this job. Report when done."#,
|
||||
job.title, job.description
|
||||
)));
|
||||
|
||||
// Run with timeout
|
||||
let result = tokio::time::timeout(self.config.timeout, async {
|
||||
self.execution_loop(&reasoning, &mut reason_ctx).await
|
||||
})
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(Ok(output)) => {
|
||||
tracing::info!("Worker completed job {} successfully", self.config.job_id);
|
||||
self.post_event(
|
||||
"result",
|
||||
serde_json::json!({
|
||||
"success": true,
|
||||
"message": truncate(&output, 2000),
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
self.client
|
||||
.report_complete(&CompletionReport {
|
||||
success: true,
|
||||
message: Some(output),
|
||||
iterations: 0,
|
||||
})
|
||||
.await?;
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
tracing::error!("Worker failed for job {}: {}", self.config.job_id, e);
|
||||
self.post_event(
|
||||
"result",
|
||||
serde_json::json!({
|
||||
"success": false,
|
||||
"message": format!("Execution failed: {}", e),
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
self.client
|
||||
.report_complete(&CompletionReport {
|
||||
success: false,
|
||||
message: Some(format!("Execution failed: {}", e)),
|
||||
iterations: 0,
|
||||
})
|
||||
.await?;
|
||||
}
|
||||
Err(_) => {
|
||||
tracing::warn!("Worker timed out for job {}", self.config.job_id);
|
||||
self.post_event(
|
||||
"result",
|
||||
serde_json::json!({
|
||||
"success": false,
|
||||
"message": "Execution timed out",
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
self.client
|
||||
.report_complete(&CompletionReport {
|
||||
success: false,
|
||||
message: Some("Execution timed out".to_string()),
|
||||
iterations: 0,
|
||||
})
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn execution_loop(
|
||||
&self,
|
||||
reasoning: &Reasoning,
|
||||
reason_ctx: &mut ReasoningContext,
|
||||
) -> Result<String, WorkerError> {
|
||||
let max_iterations = self.config.max_iterations;
|
||||
let mut last_output = String::new();
|
||||
|
||||
// Load tool definitions
|
||||
reason_ctx.available_tools = self.tools.tool_definitions().await;
|
||||
|
||||
for iteration in 1..=max_iterations {
|
||||
// Report progress
|
||||
if iteration % 5 == 1 {
|
||||
let _ = self
|
||||
.client
|
||||
.report_status(&StatusUpdate {
|
||||
state: "in_progress".to_string(),
|
||||
message: Some(format!("Iteration {}", iteration)),
|
||||
iteration,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
// Poll for follow-up prompts from the user
|
||||
self.poll_and_inject_prompt(reason_ctx).await;
|
||||
|
||||
// Refresh tools (in case WASM tools were built)
|
||||
reason_ctx.available_tools = self.tools.tool_definitions().await;
|
||||
|
||||
// Ask the LLM what to do next
|
||||
let selections = reasoning.select_tools(reason_ctx).await.map_err(|e| {
|
||||
WorkerError::ExecutionFailed {
|
||||
reason: format!("tool selection failed: {}", e),
|
||||
}
|
||||
})?;
|
||||
|
||||
if selections.is_empty() {
|
||||
// No tools selected, try direct response
|
||||
let respond_result =
|
||||
reasoning
|
||||
.respond_with_tools(reason_ctx)
|
||||
.await
|
||||
.map_err(|e| WorkerError::ExecutionFailed {
|
||||
reason: format!("respond_with_tools failed: {}", e),
|
||||
})?;
|
||||
|
||||
match respond_result {
|
||||
RespondResult::Text(response) => {
|
||||
self.post_event(
|
||||
"message",
|
||||
serde_json::json!({
|
||||
"role": "assistant",
|
||||
"content": truncate(&response, 2000),
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let response_lower = response.to_lowercase();
|
||||
if response_lower.contains("complete")
|
||||
|| response_lower.contains("finished")
|
||||
|| response_lower.contains("done")
|
||||
{
|
||||
if last_output.is_empty() {
|
||||
last_output = response.clone();
|
||||
}
|
||||
return Ok(last_output);
|
||||
}
|
||||
reason_ctx.messages.push(ChatMessage::assistant(&response));
|
||||
}
|
||||
RespondResult::ToolCalls {
|
||||
tool_calls,
|
||||
content,
|
||||
} => {
|
||||
if let Some(ref text) = content {
|
||||
self.post_event(
|
||||
"message",
|
||||
serde_json::json!({
|
||||
"role": "assistant",
|
||||
"content": truncate(text, 2000),
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
// Add assistant message with tool_calls (OpenAI protocol)
|
||||
reason_ctx
|
||||
.messages
|
||||
.push(ChatMessage::assistant_with_tool_calls(
|
||||
content,
|
||||
tool_calls.clone(),
|
||||
));
|
||||
|
||||
for tc in tool_calls {
|
||||
self.post_event(
|
||||
"tool_use",
|
||||
serde_json::json!({
|
||||
"tool_name": tc.name,
|
||||
"input": truncate(&tc.arguments.to_string(), 500),
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let result = self.execute_tool(&tc.name, &tc.arguments).await;
|
||||
|
||||
self.post_event(
|
||||
"tool_result",
|
||||
serde_json::json!({
|
||||
"tool_name": tc.name,
|
||||
"output": match &result {
|
||||
Ok(output) => truncate(output, 2000),
|
||||
Err(e) => format!("Error: {}", truncate(e, 500)),
|
||||
},
|
||||
"success": result.is_ok(),
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
if let Ok(ref output) = result {
|
||||
last_output = output.clone();
|
||||
}
|
||||
let selection = ToolSelection {
|
||||
tool_name: tc.name.clone(),
|
||||
parameters: tc.arguments.clone(),
|
||||
reasoning: String::new(),
|
||||
alternatives: vec![],
|
||||
};
|
||||
self.process_result(reason_ctx, &selection, result);
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Execute selected tools
|
||||
for selection in &selections {
|
||||
self.post_event(
|
||||
"tool_use",
|
||||
serde_json::json!({
|
||||
"tool_name": selection.tool_name,
|
||||
"input": truncate(&selection.parameters.to_string(), 500),
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let result = self
|
||||
.execute_tool(&selection.tool_name, &selection.parameters)
|
||||
.await;
|
||||
|
||||
self.post_event(
|
||||
"tool_result",
|
||||
serde_json::json!({
|
||||
"tool_name": selection.tool_name,
|
||||
"output": match &result {
|
||||
Ok(output) => truncate(output, 2000),
|
||||
Err(e) => format!("Error: {}", truncate(e, 500)),
|
||||
},
|
||||
"success": result.is_ok(),
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
if let Ok(ref output) = result {
|
||||
last_output = output.clone();
|
||||
}
|
||||
|
||||
let completed = self.process_result(reason_ctx, selection, result);
|
||||
if completed {
|
||||
return Ok(last_output);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Brief pause between iterations
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
}
|
||||
|
||||
Err(WorkerError::ExecutionFailed {
|
||||
reason: format!("max iterations ({}) exceeded", max_iterations),
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute_tool(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
params: &serde_json::Value,
|
||||
) -> Result<String, String> {
|
||||
let tool = match self.tools.get(tool_name).await {
|
||||
Some(t) => t,
|
||||
None => return Err(format!("tool '{}' not found", tool_name)),
|
||||
};
|
||||
|
||||
let ctx = JobContext::default();
|
||||
|
||||
// Validate params
|
||||
let validation = self.safety.validator().validate_tool_params(params);
|
||||
if !validation.is_valid {
|
||||
let details = validation
|
||||
.errors
|
||||
.iter()
|
||||
.map(|e| format!("{}: {}", e.field, e.message))
|
||||
.collect::<Vec<_>>()
|
||||
.join("; ");
|
||||
return Err(format!("invalid parameters: {}", details));
|
||||
}
|
||||
|
||||
// Execute with per-tool timeout
|
||||
let tool_timeout = tool.execution_timeout();
|
||||
let result = tokio::time::timeout(tool_timeout, tool.execute(params.clone(), &ctx)).await;
|
||||
|
||||
match result {
|
||||
Ok(Ok(output)) => serde_json::to_string_pretty(&output.result)
|
||||
.map_err(|e| format!("serialization error: {}", e)),
|
||||
Ok(Err(e)) => Err(e.to_string()),
|
||||
Err(_) => Err("tool execution timed out".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Process a tool result into the reasoning context. Returns true if the job is complete.
|
||||
fn process_result(
|
||||
&self,
|
||||
reason_ctx: &mut ReasoningContext,
|
||||
selection: &ToolSelection,
|
||||
result: Result<String, String>,
|
||||
) -> bool {
|
||||
match result {
|
||||
Ok(output) => {
|
||||
let sanitized = self
|
||||
.safety
|
||||
.sanitize_tool_output(&selection.tool_name, &output);
|
||||
let wrapped = self.safety.wrap_for_llm(
|
||||
&selection.tool_name,
|
||||
&sanitized.content,
|
||||
sanitized.was_modified,
|
||||
);
|
||||
|
||||
reason_ctx.messages.push(ChatMessage::tool_result(
|
||||
"tool_call_id",
|
||||
&selection.tool_name,
|
||||
wrapped,
|
||||
));
|
||||
|
||||
output.contains("TASK_COMPLETE") || output.contains("JOB_DONE")
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Tool {} failed: {}", selection.tool_name, e);
|
||||
reason_ctx.messages.push(ChatMessage::tool_result(
|
||||
"tool_call_id",
|
||||
&selection.tool_name,
|
||||
format!("Error: {}", e),
|
||||
));
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Post a job event to the orchestrator (fire-and-forget).
|
||||
async fn post_event(&self, event_type: &str, data: serde_json::Value) {
|
||||
self.client
|
||||
.post_event(&JobEventPayload {
|
||||
event_type: event_type.to_string(),
|
||||
data,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
/// Poll the orchestrator for a follow-up prompt. If one is available,
|
||||
/// inject it as a user message into the reasoning context.
|
||||
async fn poll_and_inject_prompt(&self, reason_ctx: &mut ReasoningContext) {
|
||||
match self.client.poll_prompt().await {
|
||||
Ok(Some(prompt)) => {
|
||||
tracing::info!(
|
||||
"Received follow-up prompt: {}",
|
||||
truncate(&prompt.content, 100)
|
||||
);
|
||||
self.post_event(
|
||||
"message",
|
||||
serde_json::json!({
|
||||
"role": "user",
|
||||
"content": truncate(&prompt.content, 2000),
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
reason_ctx.messages.push(ChatMessage::user(&prompt.content));
|
||||
}
|
||||
Ok(None) => {}
|
||||
Err(e) => {
|
||||
tracing::debug!("Failed to poll for prompt: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn truncate(s: &str, max: usize) -> String {
|
||||
if s.len() <= max {
|
||||
s.to_string()
|
||||
} else {
|
||||
format!("{}...", &s[..max])
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user