gitops/optimclaw
1
0
Fork 0
mirror of https://github.com/outbackdingo/optimclaw.git synced 2026-08-25 14:53:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: Sandbox jobs (#4)

Browse Source 
* 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
2026-02-11 08:31:25 +00:00
committed by GitHub
co-authored by Claude Opus 4.6
parent 202665a55c
commit ced83d5b4d
99 changed files with 17095 additions and 1162 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.dockerignore
+8
View File
@@ -0,0 +1,8 @@
target/
.git/
.env
.env.*
*.md
!CLAUDE.md
node_modules/
tools-src/
.gitignore
+3
View File
@@ -4,3 +4,6 @@
target/
# WASM build artifacts (loaded from disk, not bundled)
*.wasm
CLAUDE.md
+78 -11
View File
@@ -11,8 +11,13 @@
- **Always available** - Multi-channel access with proactive background execution
### Features
- **Multi-channel input**: TUI (Ratatui), HTTP webhooks, Telegram, WhatsApp, Slack (WASM channels)
- **Multi-channel input**: TUI (Ratatui), HTTP webhooks, WASM channels (Telegram, Slack), web gateway
- **Parallel job execution** with state machine and self-repair for stuck jobs
- **Sandbox execution**: Docker container isolation with orchestrator/worker pattern
- **Claude Code mode**: Delegate jobs to Claude CLI inside containers
- **Routines**: Scheduled (cron) and reactive (event, webhook) task execution
- **Web gateway**: Browser UI with SSE/WebSocket real-time streaming
- **Extension management**: Install, auth, activate MCP/WASM extensions
- **Extensible tools**: Built-in tools, WASM sandbox, MCP client, dynamic builder
- **Persistent memory**: Workspace with hybrid search (FTS + vector via RRF)
- **Prompt injection defense**: Sanitizer, validator, policy rules, leak detection
@@ -59,7 +64,9 @@ src/
│ ├── context_monitor.rs # Memory pressure detection
│ ├── undo.rs # Turn-based undo/redo with checkpoints
│ ├── submission.rs # Submission parsing (undo, redo, compact, clear, etc.)
│ └── task.rs # Sub-task execution framework
│ ├── task.rs # Sub-task execution framework
│ ├── routine.rs # Routine types (Trigger, Action, Guardrails)
│ └── routine_engine.rs # Routine execution (cron ticker, event matcher)
│
├── channels/ # Multi-channel input
│ ├── channel.rs # Channel trait, IncomingMessage, OutgoingResponse
@@ -72,8 +79,33 @@ src/
│ │ ├── overlay.rs # Approval overlays
│ │ └── composer.rs # Message composition
│ ├── http.rs # HTTP webhook (axum) with secret validation
│ ├── slack.rs # Stub
│ └── telegram.rs # Stub
│ ├── repl.rs # Simple REPL (for testing)
│ ├── web/ # Web gateway (browser UI)
│ │ ├── mod.rs # Gateway builder, startup
│ │ ├── server.rs # Axum router, 40+ API endpoints
│ │ ├── sse.rs # SSE broadcast manager
│ │ ├── ws.rs # WebSocket gateway + connection tracking
│ │ ├── types.rs # Request/response types, SseEvent enum
│ │ ├── auth.rs # Bearer token auth middleware
│ │ ├── log_layer.rs # Tracing layer for log streaming
│ │ └── static/ # HTML, CSS, JS (single-page app)
│ └── wasm/ # WASM channel runtime
│ ├── mod.rs
│ ├── bundled.rs # Bundled channel discovery
│ └── wrapper.rs # Channel trait wrapper for WASM modules
│
├── orchestrator/ # Internal HTTP API for sandbox containers
│ ├── mod.rs
│ ├── api.rs # Axum endpoints (LLM proxy, events, prompts)
│ ├── auth.rs # Per-job bearer token store
│ └── job_manager.rs # Container lifecycle (create, stop, cleanup)
│
├── worker/ # Runs inside Docker containers
│ ├── mod.rs
│ ├── runtime.rs # Worker execution loop (tool calls, LLM)
│ ├── claude_bridge.rs # Claude Code bridge (spawns claude CLI)
│ ├── api.rs # HTTP client to orchestrator
│ └── proxy_llm.rs # LlmProvider that proxies through orchestrator
│
├── safety/ # Prompt injection defense
│ ├── sanitizer.rs # Pattern detection, content escaping
@@ -96,6 +128,9 @@ src/
│ │ ├── file.rs # ReadFile, WriteFile, ListDir, ApplyPatch
│ │ ├── shell.rs # Shell command execution
│ │ ├── memory.rs # Memory tools (search, write, read, tree)
│ │ ├── job.rs # CreateJob, ListJobs, JobStatus, CancelJob
│ │ ├── routine.rs # routine_create/list/update/delete/history
│ │ ├── extension_tools.rs # Extension install/auth/activate/remove
│ │ └── marketplace.rs, ecommerce.rs, taskrabbit.rs, restaurant.rs (stubs)
│ ├── builder/ # Dynamic tool building
│ │ ├── core.rs # BuildRequirement, SoftwareType, Language
@@ -236,6 +271,30 @@ HEARTBEAT_ENABLED=true
HEARTBEAT_INTERVAL_SECS=1800 # 30 minutes
HEARTBEAT_NOTIFY_CHANNEL=tui
HEARTBEAT_NOTIFY_USER=default
# Web gateway
GATEWAY_ENABLED=true
GATEWAY_HOST=127.0.0.1
GATEWAY_PORT=3001
GATEWAY_AUTH_TOKEN=changeme # Required for API access
GATEWAY_USER_ID=default
# Docker sandbox
SANDBOX_ENABLED=true
SANDBOX_IMAGE=ironclaw-worker:latest
SANDBOX_MEMORY_LIMIT_MB=512
SANDBOX_TIMEOUT_SECS=1800
# Claude Code mode (runs inside sandbox containers)
CLAUDE_CODE_ENABLED=false
CLAUDE_CODE_MODEL=claude-sonnet-4-20250514
CLAUDE_CODE_MAX_TURNS=50
CLAUDE_CODE_CONFIG_DIR=/home/worker/.claude
# Routines (scheduled/reactive execution)
ROUTINES_ENABLED=true
ROUTINES_CRON_INTERVAL=60 # Tick interval in seconds
ROUTINES_MAX_CONCURRENT=3
```
### NEAR AI Provider
@@ -297,13 +356,14 @@ Key test patterns:
## Current Limitations / TODOs
1. **Slack/Telegram channels** - Stubs only, need implementation
2. **Domain-specific tools** - `marketplace.rs`, `restaurant.rs`, `taskrabbit.rs`, `ecommerce.rs` return placeholder responses; need real API integrations
3. **Integration tests** - Need testcontainers setup for PostgreSQL
4. **MCP stdio transport** - Only HTTP transport implemented
5. **WIT bindgen integration** - Auto-extract tool description/schema from WASM modules (stubbed)
6. **Capability granting after tool build** - Built tools get empty capabilities; need UX for granting HTTP/secrets access
7. **Tool versioning workflow** - No version tracking or rollback for dynamically built tools
1. **Domain-specific tools** - `marketplace.rs`, `restaurant.rs`, `taskrabbit.rs`, `ecommerce.rs` return placeholder responses; need real API integrations
2. **Integration tests** - Need testcontainers setup for PostgreSQL
3. **MCP stdio transport** - Only HTTP transport implemented
4. **WIT bindgen integration** - Auto-extract tool description/schema from WASM modules (stubbed)
5. **Capability granting after tool build** - Built tools get empty capabilities; need UX for granting HTTP/secrets access
6. **Tool versioning workflow** - No version tracking or rollback for dynamically built tools
7. **Webhook trigger endpoint** - Routines webhook trigger not yet exposed in web gateway
8. **Full channel status view** - Gateway status widget exists, but no per-channel connection dashboard
### Completed
@@ -320,6 +380,13 @@ Key test patterns:
- ✅ **Tool approval enforcement** - Tools with `requires_approval()` (shell, http, file write/patch, build_software) now gate execution, track auto-approved tools per session
- ✅ **Tool definition refresh** - Tool definitions refreshed each iteration so newly built tools become visible in same session
- ✅ **Worker tool call handling** - Uses `respond_with_tools()` to properly execute tool calls when `select_tools()` returns empty
- ✅ **Gateway control plane** - Web gateway with 40+ API endpoints, SSE/WebSocket
- ✅ **Web Control UI** - Browser-based dashboard with chat, memory, jobs, logs, extensions, routines
- ✅ **Slack/Telegram channels** - Implemented as WASM tools
- ✅ **Docker sandbox** - Orchestrator/worker containers with per-job auth
- ✅ **Claude Code mode** - Delegate jobs to Claude CLI inside containers
- ✅ **Routines system** - Cron, event, webhook, and manual triggers with guardrails
- ✅ **Extension management** - Install, auth, activate MCP/WASM extensions via CLI and web UI
## Adding a New Tool
Cargo.lock
Generated
+45
View File
@@ -996,6 +996,17 @@ dependencies = [
"syn 2.0.114",
]
[[package]]
name = "cron"
version = "0.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "eee8b2b4516038bc0f1d3c9934bcb4a13dd316e04abbc63c96757a6d75978532"
dependencies = [
"chrono",
"nom",
"once_cell",
]
[[package]]
name = "crossbeam"
version = "0.8.4"
@@ -2189,6 +2200,7 @@ dependencies = [
"bytes",
"chrono",
"clap",
"cron",
"crossterm 0.28.1",
"deadpool-postgres",
"dirs 6.0.0",
@@ -2198,6 +2210,7 @@ dependencies = [
"http-body-util",
"hyper",
"hyper-util",
"mime_guess",
"open",
"pgvector",
"postgres-types",
@@ -2495,6 +2508,16 @@ version = "0.3.17"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
[[package]]
name = "mime_guess"
version = "2.0.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f7c44f8e672c00fe5308fa235f821cb4198414e1c77935c1ab6948d3fd78550e"
dependencies = [
"mime",
"unicase",
]
[[package]]
name = "minimad"
version = "0.14.0"
@@ -2504,6 +2527,12 @@ dependencies = [
"once_cell",
]
[[package]]
name = "minimal-lexical"
version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a"
[[package]]
name = "mio"
version = "1.1.1"
@@ -2550,6 +2579,16 @@ dependencies = [
"libc",
]
[[package]]
name = "nom"
version = "7.1.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a"
dependencies = [
"memchr",
"minimal-lexical",
]
[[package]]
name = "nu-ansi-term"
version = "0.50.3"
@@ -4673,6 +4712,12 @@ dependencies = [
"winapi",
]
[[package]]
name = "unicase"
version = "2.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142"
[[package]]
name = "unicode-bidi"
version = "0.3.18"
Cargo.toml
+4
View File
@@ -58,6 +58,9 @@ axum = { version = "0.8", features = ["ws"] }
tower = "0.5"
tower-http = { version = "0.6", features = ["trace", "cors"] }
# Cron scheduling for routines
cron = "0.13"
# Safety/sanitization
regex = "1"
aho-corasick = "1"
@@ -99,6 +102,7 @@ hyper-util = { version = "0.1", features = ["server", "tokio", "http1", "http2"]
http-body-util = "0.1"
bytes = "1"
base64 = "0.22.1"
mime_guess = "2.0.5"
# macOS keychain
[target.'cfg(target_os = "macos")'.dependencies]
Dockerfile.worker
+63
View File
@@ -0,0 +1,63 @@
# Multi-stage Dockerfile for the IronClaw worker container.
#
# This image runs the ironclaw binary in worker mode inside Docker containers.
# The orchestrator creates instances of this image for sandboxed job execution.
#
# Build:
# docker build -f Dockerfile.worker -t ironclaw-worker .
#
# The image includes common development tools so workers can build software,
# run tests, and execute shell commands.
FROM rust:1.85-bookworm AS builder
WORKDIR /build
COPY . .
# Build only the ironclaw binary (release mode)
RUN cargo build --release --bin ironclaw
# ---
FROM debian:bookworm-slim
# Install common development tools
RUN apt-get update && apt-get install -y --no-install-recommends \
ca-certificates \
curl \
git \
build-essential \
pkg-config \
libssl-dev \
nodejs \
npm \
python3 \
python3-pip \
python3-venv \
&& rm -rf /var/lib/apt/lists/*
# Install Rust toolchain for the sandbox user
ENV RUSTUP_HOME=/usr/local/rustup \
CARGO_HOME=/usr/local/cargo \
PATH=/usr/local/cargo/bin:$PATH
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain 1.85.0 \
&& chmod -R a+r /usr/local/rustup /usr/local/cargo
# Install Claude Code CLI (for claude-bridge mode)
RUN npm install -g @anthropic-ai/claude-code@latest
# Copy the binary
COPY --from=builder /build/target/release/ironclaw /usr/local/bin/ironclaw
# Create non-root user (UID 1000 matches the orchestrator's container config)
RUN useradd -m -u 1000 -s /bin/bash sandbox \
&& mkdir -p /workspace \
&& chown sandbox:sandbox /workspace \
&& mkdir -p /home/sandbox/.claude \
&& chown sandbox:sandbox /home/sandbox/.claude
USER sandbox
WORKDIR /workspace
# The orchestrator passes the full command via Docker cmd.
ENTRYPOINT ["ironclaw"]
FEATURE_PARITY.md
+34 -31
View File
@@ -16,8 +16,8 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
| Feature | OpenClaw | IronClaw | Notes |
|---------|----------|----------|-------|
| Hub-and-spoke architecture | ✅ | 🚧 | IronClaw has channels but no central gateway |
| WebSocket control plane | ✅ | ❌ | Gateway with ws://127.0.0.1:18789 |
| Hub-and-spoke architecture | ✅ | ✅ | Web gateway as central hub |
| WebSocket control plane | ✅ | ✅ | Gateway with WebSocket + SSE |
| Single-user system | ✅ | ✅ | |
| Multi-agent routing | ✅ | ❌ | Workspace isolation per-agent |
| Session-based messaging | ✅ | ✅ | Per-sender sessions |
@@ -31,9 +31,9 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
| Feature | OpenClaw | IronClaw | Notes |
|---------|----------|----------|-------|
| Gateway control plane | ✅ | ❌ | Central WebSocket server |
| HTTP endpoints for Control UI | ✅ | ❌ | Web dashboard |
| Channel connection lifecycle | ✅ | 🚧 | ChannelManager handles streams |
| Gateway control plane | ✅ | ✅ | Web gateway with 40+ API endpoints |
| HTTP endpoints for Control UI | ✅ | ✅ | Web dashboard with chat, memory, jobs, logs, extensions |
| Channel connection lifecycle | ✅ | ✅ | ChannelManager + WebSocket tracker |
| Session management/routing | ✅ | ✅ | SessionManager exists |
| Configuration hot-reload | ✅ | ❌ | |
| Network modes (loopback/LAN/remote) | ✅ | 🚧 | HTTP only |
@@ -43,7 +43,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
| launchd/systemd integration | ✅ | ❌ | |
| Bonjour/mDNS discovery | ✅ | ❌ | |
| Tailscale integration | ✅ | ❌ | |
| Health check endpoints | ✅ | ❌ | |
| Health check endpoints | ✅ | ✅ | /api/health + /api/gateway/status |
| `doctor` diagnostics | ✅ | ❌ | |
### Owner: _Unassigned_
@@ -59,14 +59,14 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
| REPL (simple) | ✅ | ✅ | - | For testing |
| WASM channels | ❌ | ✅ | - | IronClaw innovation |
| WhatsApp | ✅ | ❌ | P1 | Baileys (Web) |
| Telegram | ✅ | ❌ | P1 | grammY (Bot API) |
| Telegram | ✅ | ✅ | - | WASM tool (MTProto) |
| Discord | ✅ | ❌ | P2 | discord.js |
| Signal | ✅ | ❌ | P2 | signal-cli |
| Slack | ✅ | 🚧 | P1 | Stub exists, needs implementation |
| Slack | ✅ | ✅ | - | WASM tool |
| iMessage | ✅ | ❌ | P3 | BlueBubbles recommended |
| Feishu/Lark | ✅ | ❌ | P3 | |
| LINE | ✅ | ❌ | P3 | |
| WebChat | ✅ | ❌ | P2 | Browser-based chat |
| WebChat | ✅ | ✅ | - | Web gateway chat |
| Matrix | ✅ | ❌ | P3 | E2EE support |
| Mattermost | ✅ | ❌ | P3 | |
| Google Chat | ✅ | ❌ | P3 | |
@@ -99,15 +99,15 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
| `run` (agent) | ✅ | ✅ | - | Default command |
| `tool install/list/remove` | ✅ | ✅ | - | WASM tools |
| `gateway start/stop` | ✅ | ❌ | P2 | |
| `onboard` (wizard) | ✅ | ❌ | P2 | Interactive setup |
| `onboard` (wizard) | ✅ | ✅ | - | Interactive setup |
| `tui` | ✅ | ✅ | - | Ratatui TUI |
| `config` | ✅ | ❌ | P2 | Read/write config |
| `config` | ✅ | ✅ | - | Read/write config |
| `channels` | ✅ | ❌ | P2 | Channel management |
| `models` | ✅ | 🚧 | - | Model selector in TUI |
| `status` | ✅ | ❌ | P2 | System status |
| `status` | ✅ | ✅ | - | System status |
| `agents` | ✅ | ❌ | P3 | Multi-agent management |
| `sessions` | ✅ | ❌ | P3 | Session listing |
| `memory` | ✅ | ❌ | P2 | Memory search CLI |
| `memory` | ✅ | ✅ | - | Memory search CLI |
| `skills` | ✅ | ❌ | P3 | Agent skills |
| `pairing` | ✅ | ❌ | P3 | Node pairing |
| `nodes` | ✅ | ❌ | P3 | Device management |
@@ -132,7 +132,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
| Feature | OpenClaw | IronClaw | Notes |
|---------|----------|----------|-------|
| Pi agent runtime | ✅ | ➖ | IronClaw uses custom runtime |
| RPC-based execution | ✅ | 🚧 | Worker isolation |
| RPC-based execution | ✅ | ✅ | Orchestrator/worker pattern |
| Multi-provider failover | ✅ | ❌ | Provider fallback chains |
| Per-sender sessions | ✅ | ✅ | |
| Global sessions | ✅ | ❌ | Optional shared context |
@@ -303,13 +303,13 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
| Feature | OpenClaw | IronClaw | Priority | Notes |
|---------|----------|----------|----------|-------|
| Control UI Dashboard | ✅ | ❌ | P2 | Web status/config |
| Channel status view | ✅ | ❌ | P2 | |
| Control UI Dashboard | ✅ | ✅ | - | Web gateway with chat, memory, jobs, logs, extensions |
| Channel status view | ✅ | 🚧 | P2 | Gateway status widget, full channel view pending |
| Agent management | ✅ | ❌ | P3 | |
| Model selection | ✅ | ✅ | - | TUI only |
| Config editing | ✅ | ❌ | P3 | |
| Debug/logs viewer | ✅ | ❌ | P3 | |
| WebChat interface | ✅ | ❌ | P2 | Browser chat |
| Debug/logs viewer | ✅ | ✅ | - | Real-time log streaming with level/target filters |
| WebChat interface | ✅ | ✅ | - | Web gateway chat with SSE/WebSocket |
| Canvas system (A2UI) | ✅ | ❌ | P3 | Agent-driven UI |
### Owner: _Unassigned_
@@ -320,13 +320,13 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
| Feature | OpenClaw | IronClaw | Priority | Notes |
|---------|----------|----------|----------|-------|
| Cron jobs | ✅ | ❌ | P2 | Schedule-based tasks |
| Timezone support | ✅ | ❌ | P2 | |
| One-shot/recurring jobs | ✅ | ❌ | P2 | |
| Cron jobs | ✅ | ✅ | - | Routines with cron trigger |
| Timezone support | ✅ | ✅ | - | Via cron expressions |
| One-shot/recurring jobs | ✅ | ✅ | - | Manual + cron triggers |
| `beforeInbound` hook | ✅ | ❌ | P2 | |
| `beforeOutbound` hook | ✅ | ❌ | P2 | |
| `beforeToolCall` hook | ✅ | ❌ | P2 | |
| `onMessage` hook | ✅ | ❌ | P2 | |
| `onMessage` hook | ✅ | ✅ | - | Routines with event trigger |
| `onSessionStart` hook | ✅ | ❌ | P2 | |
| `onSessionEnd` hook | ✅ | ❌ | P2 | |
| `transcribeAudio` hook | ✅ | ❌ | P3 | |
@@ -346,7 +346,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
| Feature | OpenClaw | IronClaw | Notes |
|---------|----------|----------|-------|
| Gateway token auth | ✅ | 🚧 | HTTP webhook secret |
| Gateway token auth | ✅ | ✅ | Bearer token auth on web gateway |
| Device pairing | ✅ | ❌ | |
| Tailscale identity | ✅ | ❌ | |
| OAuth flows | ✅ | 🚧 | NEAR AI OAuth |
@@ -357,7 +357,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
| TLS 1.3 minimum | ✅ | ✅ | reqwest rustls |
| SSRF protection | ✅ | ✅ | WASM allowlist |
| Loopback-first | ✅ | 🚧 | HTTP binds 0.0.0.0 |
| Docker sandbox | ✅ | ❌ | Uses WASM sandbox |
| Docker sandbox | ✅ | ✅ | Orchestrator/worker containers |
| WASM sandbox | ❌ | ✅ | IronClaw innovation |
| Tool policies | ✅ | ✅ | |
| Elevated mode | ✅ | ❌ | |
@@ -404,23 +404,26 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
- ✅ Session management
- ✅ Context compaction
- ✅ Model selection
- ✅ Gateway control plane + WebSocket
- ✅ Web Control UI (chat, memory, jobs, logs, extensions, routines)
- ✅ WebChat channel (web gateway)
- ✅ Slack channel (WASM tool)
- ✅ Telegram channel (WASM tool, MTProto)
- ✅ Docker sandbox (orchestrator/worker)
- ✅ Cron job scheduling (routines)
- ✅ CLI subcommands (onboard, config, status, memory)
- ✅ Gateway token auth
### P1 - High Priority
- ❌ Slack channel (real implementation)
- ❌ Telegram channel
- ❌ WhatsApp channel
- ❌ Multi-provider failover
- ❌ Gateway control plane + WebSocket
- ❌ Hooks system (beforeInbound, beforeToolCall, etc.)
### P2 - Medium Priority
- ❌ Cron job scheduling
- ❌ Web Control UI
- ❌ WebChat channel
- ❌ Media handling (images, PDFs)
- ❌ CLI subcommands (config, status, memory, doctor)
- ❌ Ollama/local model support
- ❌ Configuration hot-reload
- ❌ Webhook trigger endpoint in web gateway
### P3 - Lower Priority
- ❌ Discord channel
README.md
+43 -32
View File
@@ -43,7 +43,10 @@ IronClaw is the AI assistant you can actually trust with your personal and profe
### Always Available
- **Multi-channel** - REPL, HTTP webhooks, and extensible WASM channels (Telegram, Slack, and more)
- **Multi-channel** - REPL, HTTP webhooks, WASM channels (Telegram, Slack), and web gateway
- **Docker Sandbox** - Isolated container execution with per-job tokens and orchestrator/worker pattern
- **Web Gateway** - Browser UI with real-time SSE/WebSocket streaming
- **Routines** - Cron schedules, event triggers, webhook handlers for background automation
- **Heartbeat System** - Proactive background execution for monitoring and maintenance tasks
- **Parallel Jobs** - Handle multiple requests concurrently with isolated contexts
- **Self-repair** - Automatic detection and recovery of stuck operations
@@ -143,37 +146,42 @@ External content passes through multiple security layers:
## Architecture
```
┌─────────────────────────────────────────────────────────────────┐
│ Channels │
│ ┌──────┐ ┌──────┐ ┌──────────────┐ │
│ │ REPL │ │ HTTP │ │ WASM Channels│ │
│ └──┬───┘ └──┬───┘ └──────┬───────┘ │
│ └─────────┴─────────────┘ │
│ │ │
│ ┌────▼────┐ │
│ │ Router │ Intent classification │
│ └────┬────┘ │
│ │ │
│ ┌──────────▼──────────┐ │
│ │ Scheduler │ Parallel job management │
│ └──────────┬──────────┘ │
│ │ │
│ ┌───────────────┼───────────────┐ │
│ ▼ ▼ ▼ │
│ ┌─────────┐ ┌─────────┐ ┌─────────┐ │
│ │ Worker │ │ Worker │ │ Worker │ LLM reasoning │
│ └────┬────┘ └────┬────┘ └────┬────┘ │
│ └───────────────┼───────────────┘ │
│ │ │
│ ┌──────────▼──────────┐ │
│ │ Tool Registry │ │
│ │ ┌───────────────┐ │ │
│ │ │ Built-in │ │ │
│ │ │ MCP │ │ │
│ │ │ WASM Sandbox │ │ │
│ │ └───────────────┘ │ │
│ └─────────────────────┘ │
└─────────────────────────────────────────────────────────────────┘
┌────────────────────────────────────────────────────────────────────┐
│ Channels │
│ ┌──────┐ ┌──────┐ ┌─────────────┐ ┌─────────────┐ │
│ │ REPL │ │ HTTP │ │WASM Channels│ │ Web Gateway │ │
│ └──┬───┘ └──┬───┘ └──────┬──────┘ │ (SSE + WS) │ │
│ │ │ │ └──────┬──────┘ │
│ └─────────┴──────────────┴────────────────┘ │
│ │ │
│ ┌─────────▼─────────┐ │
│ │ Agent Loop │ Intent routing │
│ └────┬─────────┬────┘ │
│ │ │ │
│ ┌──────────▼───┐ ┌──▼──────────────┐ │
│ │ Scheduler │ │ Routines Engine │ │
│ │(parallel jobs)│ │(cron, event, wh) │ │
│ └──────┬───────┘ └────────┬─────────┘ │
│ │ │ │
│ ┌─────────────┼───────────────────┘ │
│ │ │ │
│ ┌───▼────┐ ┌────▼────────────────┐ │
│ │ Local │ │ Orchestrator │ │
│ │Workers │ │ ┌───────────────┐ │ │
│ │(in-proc)│ │ │ Docker Sandbox│ │ │
│ └───┬────┘ │ │ Containers │ │ │
│ │ │ │ ┌───────────┐ │ │ │
│ │ │ │ │Worker / CC│ │ │ │
│ │ │ │ └───────────┘ │ │ │
│ │ │ └───────────────┘ │ │
│ │ └─────────┬───────────┘ │
│ └──────────────────┤ │
│ │ │
│ ┌───────────▼──────────┐ │
│ │ Tool Registry │ │
│ │ Built-in, MCP, WASM │ │
│ └──────────────────────┘ │
└────────────────────────────────────────────────────────────────────┘
```
### Core Components
@@ -184,6 +192,9 @@ External content passes through multiple security layers:
| **Router** | Classifies user intent (command, query, task) |
| **Scheduler** | Manages parallel job execution with priorities |
| **Worker** | Executes jobs with LLM reasoning and tool calls |
| **Orchestrator** | Container lifecycle, LLM proxying, per-job auth |
| **Web Gateway** | Browser UI with chat, memory, jobs, logs, extensions, routines |
| **Routines Engine** | Scheduled (cron) and reactive (event, webhook) background tasks |
| **Workspace** | Persistent memory with hybrid search |
| **Safety Layer** | Prompt injection defense and content sanitization |
channels-src/slack/src/lib.rs
+24 -22
View File
@@ -108,7 +108,10 @@ struct SlackPostMessageResponse {
#[derive(Debug, Deserialize)]
struct SlackConfig {
/// Name of secret containing signing secret (for verification by host).
/// Parsed from config for forward compatibility; not yet used in WASM
/// (host handles signature verification).
#[serde(default = "default_signing_secret_name")]
#[allow(dead_code)]
signing_secret_name: String,
}
@@ -175,11 +178,7 @@ impl Guest for SlackChannel {
// Actual event callback
"event_callback" => {
if let Some(event) = event_wrapper.event {
handle_slack_event(
event,
event_wrapper.team_id,
event_wrapper.event_id,
);
handle_slack_event(event, event_wrapper.team_id, event_wrapper.event_id);
}
// Always respond 200 quickly to Slack (they have a 3s timeout)
json_response(200, serde_json::json!({"ok": true}))
@@ -230,6 +229,7 @@ impl Guest for SlackChannel {
"https://slack.com/api/chat.postMessage",
&headers.to_string(),
Some(&payload_bytes),
None,
);
match result {
@@ -243,14 +243,15 @@ impl Guest for SlackChannel {
// Parse Slack response
let slack_response: SlackPostMessageResponse =
serde_json::from_slice(&http_response.body).map_err(|e| {
format!("Failed to parse Slack response: {}", e)
})?;
serde_json::from_slice(&http_response.body)
.map_err(|e| format!("Failed to parse Slack response: {}", e))?;
if !slack_response.ok {
return Err(format!(
"Slack API error: {}",
slack_response.error.unwrap_or_else(|| "unknown".to_string())
slack_response
.error
.unwrap_or_else(|| "unknown".to_string())
));
}
@@ -277,17 +278,16 @@ impl Guest for SlackChannel {
}
/// Handle a Slack event and emit message if applicable.
fn handle_slack_event(
event: SlackEvent,
team_id: Option<String>,
_event_id: Option<String>,
) {
fn handle_slack_event(event: SlackEvent, team_id: Option<String>, _event_id: Option<String>) {
match event.event_type.as_str() {
// Direct mention of the bot
"app_mention" => {
if let (Some(user), Some(channel), Some(text), Some(ts)) =
(event.user, event.channel.clone(), event.text, event.ts.clone())
{
if let (Some(user), Some(channel), Some(text), Some(ts)) = (
event.user,
event.channel.clone(),
event.text,
event.ts.clone(),
) {
emit_message(user, text, channel, event.thread_ts.or(Some(ts)), team_id);
}
}
@@ -299,9 +299,12 @@ fn handle_slack_event(
return;
}
if let (Some(user), Some(channel), Some(text), Some(ts)) =
(event.user, event.channel.clone(), event.text, event.ts.clone())
{
if let (Some(user), Some(channel), Some(text), Some(ts)) = (
event.user,
event.channel.clone(),
event.text,
event.ts.clone(),
) {
// Only process DMs (channel IDs starting with D)
if channel.starts_with('D') {
emit_message(user, text, channel, event.thread_ts.or(Some(ts)), team_id);
@@ -335,8 +338,7 @@ fn emit_message(
team_id,
};
let metadata_json =
serde_json::to_string(&metadata).unwrap_or_else(|_| "{}".to_string());
let metadata_json = serde_json::to_string(&metadata).unwrap_or_else(|_| "{}".to_string());
// Strip @ mentions of the bot from the text for cleaner messages
let cleaned_text = strip_bot_mention(&text);
channels-src/telegram/src/lib.rs
+133 -53
View File
@@ -388,7 +388,9 @@ impl Guest for TelegramChannel {
let headers = serde_json::json!({});
let result = channel_host::http_request("GET", &url, &headers.to_string(), None);
// 35s HTTP timeout outlives Telegram's 30s server-side long-poll
let result =
channel_host::http_request("GET", &url, &headers.to_string(), None, Some(35_000));
match result {
Ok(response) => {
@@ -461,72 +463,52 @@ impl Guest for TelegramChannel {
}
fn on_respond(response: AgentResponse) -> Result<(), String> {
// Parse metadata to get chat info
let metadata: TelegramMessageMetadata = serde_json::from_str(&response.metadata_json)
.map_err(|e| format!("Failed to parse metadata: {}", e))?;
// Build sendMessage payload
let mut payload = serde_json::json!({
"chat_id": metadata.chat_id,
"text": response.content,
"parse_mode": "Markdown",
});
// Reply to the original message for context
payload["reply_to_message_id"] = serde_json::Value::Number(metadata.message_id.into());
let payload_bytes = serde_json::to_vec(&payload)
.map_err(|e| format!("Failed to serialize payload: {}", e))?;
// Make HTTP request to Telegram API
// The bot token is injected into the URL by the host
let headers = serde_json::json!({
"Content-Type": "application/json"
});
let result = channel_host::http_request(
"POST",
"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/sendMessage",
&headers.to_string(),
Some(&payload_bytes),
// Try sending with Markdown first; fall back to plain text if Telegram
// can't parse the entities (e.g. model leaked <tool_call> with underscores).
let result = send_message(
metadata.chat_id,
&response.content,
metadata.message_id,
Some("Markdown"),
);
match result {
Ok(http_response) => {
if http_response.status != 200 {
let body_str = String::from_utf8_lossy(&http_response.body);
return Err(format!(
"Telegram API returned status {}: {}",
http_response.status, body_str
));
}
// Parse Telegram response
let api_response: TelegramApiResponse<SentMessage> =
serde_json::from_slice(&http_response.body)
.map_err(|e| format!("Failed to parse Telegram response: {}", e))?;
if !api_response.ok {
return Err(format!(
"Telegram API error: {}",
api_response
.description
.unwrap_or_else(|| "unknown".to_string())
));
}
Ok(msg_id) => {
channel_host::log(
channel_host::LogLevel::Debug,
&format!(
"Sent message to chat {}: message_id={}",
metadata.chat_id,
api_response.result.map(|r| r.message_id).unwrap_or(0)
metadata.chat_id, msg_id
),
);
Ok(())
}
Err(e) => Err(format!("HTTP request failed: {}", e)),
Err(SendError::ParseEntities(detail)) => {
channel_host::log(
channel_host::LogLevel::Warn,
&format!("Markdown parse failed ({}), retrying as plain text", detail),
);
let msg_id = send_message(
metadata.chat_id,
&response.content,
metadata.message_id,
None,
)
.map_err(|e| format!("Plain-text retry also failed: {}", e))?;
channel_host::log(
channel_host::LogLevel::Debug,
&format!(
"Sent plain-text message to chat {}: message_id={}",
metadata.chat_id, msg_id
),
);
Ok(())
}
Err(e) => Err(e.to_string()),
}
}
@@ -568,6 +550,7 @@ impl Guest for TelegramChannel {
"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/sendChatAction",
&headers.to_string(),
Some(&payload_bytes),
None,
);
if let Err(e) = result {
@@ -586,6 +569,101 @@ impl Guest for TelegramChannel {
}
}
// ============================================================================
// Send Message Helper
// ============================================================================
/// Errors from send_message, split so callers can match on parse-entity failures.
enum SendError {
/// Telegram returned 400 with "can't parse entities" (Markdown issue).
ParseEntities(String),
/// Any other failure.
Other(String),
}
impl std::fmt::Display for SendError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
SendError::ParseEntities(detail) => write!(f, "parse entities error: {}", detail),
SendError::Other(msg) => write!(f, "{}", msg),
}
}
}
/// Send a message via the Telegram Bot API.
///
/// Returns the sent message_id on success. When `parse_mode` is set and
/// Telegram returns a 400 "can't parse entities" error, returns
/// `SendError::ParseEntities` so the caller can retry without formatting.
fn send_message(
chat_id: i64,
text: &str,
reply_to_message_id: i64,
parse_mode: Option<&str>,
) -> Result<i64, SendError> {
let mut payload = serde_json::json!({
"chat_id": chat_id,
"text": text,
"reply_to_message_id": reply_to_message_id,
});
if let Some(mode) = parse_mode {
payload["parse_mode"] = serde_json::Value::String(mode.to_string());
}
let payload_bytes = serde_json::to_vec(&payload)
.map_err(|e| SendError::Other(format!("Failed to serialize payload: {}", e)))?;
let headers = serde_json::json!({ "Content-Type": "application/json" });
let result = channel_host::http_request(
"POST",
"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/sendMessage",
&headers.to_string(),
Some(&payload_bytes),
None,
);
match result {
Ok(http_response) => {
if http_response.status == 400 {
let body_str = String::from_utf8_lossy(&http_response.body);
if body_str.contains("can't parse entities") {
return Err(SendError::ParseEntities(body_str.to_string()));
}
return Err(SendError::Other(format!(
"Telegram API returned 400: {}",
body_str
)));
}
if http_response.status != 200 {
let body_str = String::from_utf8_lossy(&http_response.body);
return Err(SendError::Other(format!(
"Telegram API returned status {}: {}",
http_response.status, body_str
)));
}
let api_response: TelegramApiResponse<SentMessage> =
serde_json::from_slice(&http_response.body)
.map_err(|e| SendError::Other(format!("Failed to parse response: {}", e)))?;
if !api_response.ok {
return Err(SendError::Other(format!(
"Telegram API error: {}",
api_response
.description
.unwrap_or_else(|| "unknown".to_string())
)));
}
Ok(api_response.result.map(|r| r.message_id).unwrap_or(0))
}
Err(e) => Err(SendError::Other(format!("HTTP request failed: {}", e))),
}
}
// ============================================================================
// Webhook Management
// ============================================================================
@@ -604,6 +682,7 @@ fn delete_webhook() -> Result<(), String> {
"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/deleteWebhook",
&headers.to_string(),
None,
None,
);
match result {
@@ -666,6 +745,7 @@ fn register_webhook(tunnel_url: &str, webhook_secret: Option<&str>) -> Result<()
"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/setWebhook",
&headers.to_string(),
Some(&body_bytes),
None,
);
match result {
channels-src/telegram/telegram.capabilities.json
+1
View File
@@ -26,6 +26,7 @@
"allowed_paths": ["/webhook/telegram"],
"allow_polling": true,
"min_poll_interval_ms": 30000,
"callback_timeout_secs": 45,
"workspace_prefix": "channels/telegram/",
"emit_rate_limit": {
"messages_per_minute": 100,
channels-src/telegram/telegram.wasm
BIN
View File
Binary file not shown.
channels-src/whatsapp/src/lib.rs
+1
View File
@@ -361,6 +361,7 @@ impl Guest for WhatsAppChannel {
&api_url,
&headers.to_string(),
Some(&payload_bytes),
None,
);
match result {
examples/test_heartbeat.rs
+120
View File
@@ -0,0 +1,120 @@
//! Standalone heartbeat test.
//!
//! Exercises the heartbeat system in isolation: connects to the real
//! database, reads the real HEARTBEAT.md, calls the real LLM, and prints
//! every step so you can see exactly where it breaks.
//!
//! Usage:
//! cargo run --example test_heartbeat
use std::sync::Arc;
use ironclaw::{
agent::HeartbeatRunner,
config::Config,
history::Store,
llm::{SessionConfig, create_llm_provider, create_session_manager},
workspace::Workspace,
};
#[tokio::main]
async fn main() -> anyhow::Result<()> {
// Load .env and set up logging
let _ = dotenvy::dotenv();
tracing_subscriber::fmt()
.with_env_filter("ironclaw=debug")
.init();
println!("=== Heartbeat Integration Test ===\n");
// 1. Load config
let config = Config::from_env().map_err(|e| anyhow::anyhow!("Config: {}", e))?;
println!("[1/6] Config loaded");
println!(" heartbeat.enabled = {}", config.heartbeat.enabled);
println!(
" heartbeat.interval_secs = {}",
config.heartbeat.interval_secs
);
println!(
" heartbeat.notify_channel = {:?}",
config.heartbeat.notify_channel
);
println!(
" heartbeat.notify_user = {:?}",
config.heartbeat.notify_user
);
// 2. Connect to database
let store = Store::new(&config.database).await?;
store.run_migrations().await?;
println!("[2/6] Database connected");
// 3. Create workspace
let workspace = Arc::new(Workspace::new("default", store.pool()));
println!("[3/6] Workspace created");
// 4. Read HEARTBEAT.md
let checklist = workspace.heartbeat_checklist().await;
match &checklist {
Ok(Some(content)) => {
let preview: String = content.chars().take(200).collect();
println!("[4/6] HEARTBEAT.md found ({} chars)", content.len());
println!(" Preview: {}...", preview);
}
Ok(None) => {
println!("[4/6] HEARTBEAT.md is None (no file, no seed fallback)");
println!(" Heartbeat will return Skipped.");
}
Err(e) => {
println!("[4/6] HEARTBEAT.md read error: {}", e);
}
}
// Check if the checklist would be considered "effectively empty"
if let Ok(Some(_)) = checklist {
println!(" (Will verify via runner below)");
}
// 5. Create LLM provider
let session = create_session_manager(SessionConfig {
auth_base_url: config.llm.nearai.auth_base_url.clone(),
session_path: config.llm.nearai.session_path.clone(),
..Default::default()
})
.await;
let llm = create_llm_provider(&config.llm, session)?;
println!("[5/6] LLM provider created (model: {})", llm.model_name());
// 6. Run heartbeat check
println!("[6/6] Running check_heartbeat()...\n");
let hb_config = ironclaw::agent::HeartbeatConfig::default();
let runner = HeartbeatRunner::new(hb_config, workspace, llm);
let result = runner.check_heartbeat().await;
println!("=== Result ===\n");
match &result {
ironclaw::agent::HeartbeatResult::Ok => {
println!("HeartbeatResult::Ok");
println!(" LLM responded HEARTBEAT_OK, nothing needs attention.");
}
ironclaw::agent::HeartbeatResult::NeedsAttention(msg) => {
println!("HeartbeatResult::NeedsAttention");
println!(" Message:\n{}", msg);
}
ironclaw::agent::HeartbeatResult::Skipped => {
println!("HeartbeatResult::Skipped");
println!(" No checklist found, or checklist was effectively empty.");
println!(" This means the HEARTBEAT.md either:");
println!(" - Does not exist in the workspace database");
println!(" - Contains only headers, comments, and empty checkboxes");
}
ironclaw::agent::HeartbeatResult::Failed(err) => {
println!("HeartbeatResult::Failed");
println!(" Error: {}", err);
}
}
Ok(())
}
migrations/V4__sandbox_columns.sql
+10
View File
@@ -0,0 +1,10 @@
-- Add project_dir and user_id columns for sandbox job tracking.
-- user_id was previously hardcoded to "default" in the Rust layer;
-- now it's persisted so we can filter per-user.
ALTER TABLE agent_jobs ADD COLUMN IF NOT EXISTS project_dir TEXT;
ALTER TABLE agent_jobs ADD COLUMN IF NOT EXISTS user_id TEXT NOT NULL DEFAULT 'default';
CREATE INDEX IF NOT EXISTS idx_agent_jobs_source ON agent_jobs(source);
CREATE INDEX IF NOT EXISTS idx_agent_jobs_user ON agent_jobs(user_id);
CREATE INDEX IF NOT EXISTS idx_agent_jobs_created ON agent_jobs(created_at DESC);
migrations/V5__claude_code.sql
+14
View File
@@ -0,0 +1,14 @@
-- Track which mode a sandbox job uses (worker vs claude_code).
ALTER TABLE agent_jobs ADD COLUMN IF NOT EXISTS job_mode TEXT NOT NULL DEFAULT 'worker';
-- Persist Claude Code streaming events so they survive restarts and can be
-- loaded when the frontend opens a job detail view after the fact.
CREATE TABLE IF NOT EXISTS claude_code_events (
id BIGSERIAL PRIMARY KEY,
job_id UUID NOT NULL REFERENCES agent_jobs(id),
event_type TEXT NOT NULL,
data JSONB NOT NULL,
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE INDEX IF NOT EXISTS idx_cc_events_job ON claude_code_events(job_id, id);
migrations/V6__routines.sql
+73
View File
@@ -0,0 +1,73 @@
-- Routines: scheduled and reactive job system.
--
-- A routine is a named, persistent, user-owned task with a trigger and an action.
-- Triggers fire independently (cron, event, webhook, manual) so only the
-- relevant routine's prompt hits the LLM, not the whole checklist.
CREATE TABLE routines (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
name TEXT NOT NULL,
description TEXT NOT NULL DEFAULT '',
user_id TEXT NOT NULL,
enabled BOOLEAN NOT NULL DEFAULT true,
-- Trigger definition
trigger_type TEXT NOT NULL, -- 'cron', 'event', 'webhook', 'manual'
trigger_config JSONB NOT NULL, -- type-specific config (schedule, pattern, etc.)
-- Action definition
action_type TEXT NOT NULL, -- 'lightweight', 'full_job'
action_config JSONB NOT NULL, -- prompt, context_paths, max_tokens / title, max_iterations
-- Guardrails
cooldown_secs INTEGER NOT NULL DEFAULT 300,
max_concurrent INTEGER NOT NULL DEFAULT 1,
dedup_window_secs INTEGER, -- NULL = no dedup
-- Notification preferences
notify_channel TEXT, -- NULL = use default
notify_user TEXT NOT NULL DEFAULT 'default',
notify_on_success BOOLEAN NOT NULL DEFAULT false,
notify_on_failure BOOLEAN NOT NULL DEFAULT true,
notify_on_attention BOOLEAN NOT NULL DEFAULT true,
-- Runtime state (updated by engine)
state JSONB NOT NULL DEFAULT '{}',
last_run_at TIMESTAMPTZ,
next_fire_at TIMESTAMPTZ, -- pre-computed for cron triggers
run_count BIGINT NOT NULL DEFAULT 0,
consecutive_failures INTEGER NOT NULL DEFAULT 0,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
UNIQUE (user_id, name)
);
-- Fast lookup: "which cron routines need to fire right now?"
CREATE INDEX idx_routines_next_fire
ON routines (next_fire_at)
WHERE enabled AND next_fire_at IS NOT NULL;
-- Fast lookup: event triggers for a user
CREATE INDEX idx_routines_event_triggers
ON routines (user_id)
WHERE enabled AND trigger_type = 'event';
-- Audit log of individual routine executions.
CREATE TABLE routine_runs (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
routine_id UUID NOT NULL REFERENCES routines(id) ON DELETE CASCADE,
trigger_type TEXT NOT NULL,
trigger_detail TEXT, -- e.g. matched message preview, cron expression
started_at TIMESTAMPTZ NOT NULL DEFAULT now(),
completed_at TIMESTAMPTZ,
status TEXT NOT NULL DEFAULT 'running', -- running, ok, attention, failed
result_summary TEXT,
tokens_used INTEGER,
job_id UUID REFERENCES agent_jobs(id), -- non-NULL for full_job runs
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE INDEX idx_routine_runs_routine ON routine_runs (routine_id);
CREATE INDEX idx_routine_runs_status ON routine_runs (status) WHERE status = 'running';
migrations/V7__rename_events.sql
+3
View File
@@ -0,0 +1,3 @@
-- Rename claude_code_events to job_events (generic for all sandbox job types).
ALTER TABLE claude_code_events RENAME TO job_events;
ALTER INDEX idx_cc_events_job RENAME TO idx_job_events_job;
migrations/V8__settings.sql
+16
View File
@@ -0,0 +1,16 @@
-- Settings table: key-value store for all user configuration.
--
-- Replaces ~/.ironclaw/settings.json, session.json, and mcp-servers.json.
-- Keys use dotted paths matching the existing Settings.get()/set() convention
-- (e.g., "agent.name", "sandbox.enabled", "mcp_servers").
-- One row per setting so individual values can be updated atomically.
CREATE TABLE IF NOT EXISTS settings (
user_id TEXT NOT NULL,
key TEXT NOT NULL,
value JSONB NOT NULL,
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
PRIMARY KEY (user_id, key)
);
CREATE INDEX IF NOT EXISTS idx_settings_user ON settings (user_id);
src/agent/agent_loop.rs
+640 -116
View File
@@ -9,13 +9,14 @@ use uuid::Uuid;
use crate::agent::compaction::ContextCompactor;
use crate::agent::context_monitor::ContextMonitor;
use crate::agent::heartbeat::spawn_heartbeat;
use crate::agent::routine_engine::{RoutineEngine, spawn_cron_ticker};
use crate::agent::self_repair::{DefaultSelfRepair, RepairResult, SelfRepair};
use crate::agent::session::{PendingApproval, Session, ThreadState};
use crate::agent::session_manager::SessionManager;
use crate::agent::submission::{Submission, SubmissionParser, SubmissionResult};
use crate::agent::{HeartbeatConfig as AgentHeartbeatConfig, MessageIntent, Router, Scheduler};
use crate::channels::{ChannelManager, IncomingMessage, OutgoingResponse, StatusUpdate};
use crate::config::{AgentConfig, HeartbeatConfig};
use crate::config::{AgentConfig, HeartbeatConfig, RoutineConfig};
use crate::context::ContextManager;
use crate::context::JobContext;
use crate::error::Error;
@@ -77,6 +78,7 @@ pub struct Agent {
session_manager: Arc<SessionManager>,
context_monitor: ContextMonitor,
heartbeat_config: Option<HeartbeatConfig>,
routine_config: Option<RoutineConfig>,
}
impl Agent {
@@ -89,6 +91,7 @@ impl Agent {
deps: AgentDeps,
channels: ChannelManager,
heartbeat_config: Option<HeartbeatConfig>,
routine_config: Option<RoutineConfig>,
context_manager: Option<Arc<ContextManager>>,
session_manager: Option<Arc<SessionManager>>,
) -> Self {
@@ -116,6 +119,7 @@ impl Agent {
session_manager,
context_monitor: ContextMonitor::new(),
heartbeat_config,
routine_config,
}
}
@@ -256,53 +260,28 @@ impl Agent {
let channels = self.channels.clone();
tokio::spawn(async move {
while let Some(response) = notify_rx.recv().await {
// Route notification to configured channel/user, or broadcast to all
match (&notify_channel, &notify_user) {
(Some(channel), Some(user)) => {
// Send to specific channel and user
if let Err(e) =
channels.broadcast(channel, user, response.clone()).await
{
let user = notify_user.as_deref().unwrap_or("default");
// Try the configured channel first, fall back to
// broadcasting on all channels.
let targeted_ok = if let Some(ref channel) = notify_channel {
channels
.broadcast(channel, user, response.clone())
.await
.is_ok()
} else {
false
};
if !targeted_ok {
let results = channels.broadcast_all(user, response).await;
for (ch, result) in results {
if let Err(e) = result {
tracing::warn!(
"Failed to send heartbeat to {}/{}: {}",
channel,
user,
"Failed to broadcast heartbeat to {}: {}",
ch,
e
);
} else {
tracing::debug!(
"Heartbeat notification sent to {}/{}",
channel,
user
);
}
}
(None, Some(user)) => {
// Broadcast to all channels for this user
let results = channels.broadcast_all(user, response).await;
for (ch, result) in results {
if let Err(e) = result {
tracing::warn!(
"Failed to broadcast heartbeat to {}: {}",
ch,
e
);
}
}
}
_ => {
// No explicit target, broadcast to all channels
// for the default user so notifications actually
// reach someone instead of vanishing into logs.
let results = channels.broadcast_all("default", response).await;
for (ch, result) in results {
if let Err(e) = result {
tracing::warn!(
"Failed to broadcast heartbeat to {}: {}",
ch,
e
);
}
}
}
}
@@ -330,6 +309,85 @@ impl Agent {
None
};
// Spawn routine engine if enabled
let routine_handle = if let Some(ref rt_config) = self.routine_config {
if rt_config.enabled {
if let (Some(store), Some(workspace)) = (self.store(), self.workspace()) {
// Set up notification channel (same pattern as heartbeat)
let (notify_tx, mut notify_rx) =
tokio::sync::mpsc::channel::<OutgoingResponse>(32);
let engine = Arc::new(RoutineEngine::new(
rt_config.clone(),
Arc::clone(store),
self.llm().clone(),
Arc::clone(workspace),
notify_tx,
));
// Register routine tools
self.deps
.tools
.register_routine_tools(Arc::clone(store), Arc::clone(&engine));
// Load initial event cache
engine.refresh_event_cache().await;
// Spawn notification forwarder
let channels = self.channels.clone();
tokio::spawn(async move {
while let Some(response) = notify_rx.recv().await {
let user = response
.metadata
.get("notify_user")
.and_then(|v| v.as_str())
.unwrap_or("default")
.to_string();
let results = channels.broadcast_all(&user, response).await;
for (ch, result) in results {
if let Err(e) = result {
tracing::warn!(
"Failed to broadcast routine notification to {}: {}",
ch,
e
);
}
}
}
});
// Spawn cron ticker
let cron_interval =
std::time::Duration::from_secs(rt_config.cron_check_interval_secs);
let cron_handle = spawn_cron_ticker(Arc::clone(&engine), cron_interval);
// Store engine reference for event trigger checking
// Safety: we're in run() which takes self, no other reference exists
let engine_ref = Arc::clone(&engine);
// SAFETY: self is consumed by run(), we can smuggle the engine in
// via a local to use in the message loop below.
tracing::info!(
"Routines enabled: cron ticker every {}s, max {} concurrent",
rt_config.cron_check_interval_secs,
rt_config.max_concurrent_routines
);
Some((cron_handle, engine_ref))
} else {
tracing::warn!("Routines enabled but store/workspace not available");
None
}
} else {
None
}
} else {
None
};
// Extract engine ref for use in message loop
let routine_engine_for_loop = routine_handle.as_ref().map(|(_, e)| Arc::clone(e));
// Main message loop
tracing::info!("Agent {} ready and listening", self.config.name);
@@ -374,6 +432,14 @@ impl Agent {
.await;
}
}
// Check event triggers (cheap in-memory regex, fires async if matched)
if let Some(ref engine) = routine_engine_for_loop {
let fired = engine.check_event_triggers(&message).await;
if fired > 0 {
tracing::debug!("Fired {} event-triggered routines", fired);
}
}
}
// Cleanup
@@ -383,6 +449,9 @@ impl Agent {
if let Some(handle) = heartbeat_handle {
handle.abort();
}
if let Some((cron_handle, _)) = routine_handle {
cron_handle.abort();
}
self.scheduler.stop_all().await;
self.channels.shutdown_all().await?;
@@ -393,6 +462,11 @@ impl Agent {
// Parse submission type first
let submission = SubmissionParser::parse(&message.content);
// Hydrate thread from DB if it's a historical thread not in memory
if let Some(ref external_thread_id) = message.thread_id {
self.maybe_hydrate_thread(message, external_thread_id).await;
}
// Resolve session and thread
let (session, thread_id) = self
.session_manager
@@ -444,6 +518,9 @@ impl Agent {
self.process_user_input(message, session, thread_id, &content)
.await
}
Submission::SystemCommand { command, args } => {
self.handle_system_command(&command, &args).await
}
Submission::Undo => self.process_undo(session, thread_id).await,
Submission::Redo => self.process_redo(session, thread_id).await,
Submission::Interrupt => self.process_interrupt(session, thread_id).await,
@@ -515,6 +592,107 @@ impl Agent {
}
}
/// Hydrate a historical thread from DB into memory if not already present.
///
/// Called before `resolve_thread` so that the session manager finds the
/// thread on lookup instead of creating a new one.
///
/// Creates an in-memory thread with the exact UUID the frontend sent,
/// even when the conversation has zero messages (e.g. a brand-new
/// assistant thread). Without this, `resolve_thread` would mint a
/// fresh UUID and all messages would land in the wrong conversation.
async fn maybe_hydrate_thread(&self, message: &IncomingMessage, external_thread_id: &str) {
// Only hydrate UUID-shaped thread IDs (web gateway uses UUIDs)
let thread_uuid = match Uuid::parse_str(external_thread_id) {
Ok(id) => id,
Err(_) => return,
};
// Check if already in memory
let session = self
.session_manager
.get_or_create_session(&message.user_id)
.await;
{
let sess = session.lock().await;
if sess.threads.contains_key(&thread_uuid) {
return;
}
}
// Load history from DB (may be empty for a newly created thread).
let mut chat_messages: Vec<ChatMessage> = Vec::new();
let msg_count;
if let Some(store) = self.store() {
let db_messages = store
.list_conversation_messages(thread_uuid)
.await
.unwrap_or_default();
msg_count = db_messages.len();
chat_messages = db_messages
.iter()
.filter_map(|m| match m.role.as_str() {
"user" => Some(ChatMessage::user(&m.content)),
"assistant" => Some(ChatMessage::assistant(&m.content)),
_ => None,
})
.collect();
} else {
msg_count = 0;
}
// Create thread with the historical ID and restore messages
let session_id = {
let sess = session.lock().await;
sess.id
};
let mut thread = crate::agent::session::Thread::with_id(thread_uuid, session_id);
if !chat_messages.is_empty() {
thread.restore_from_messages(chat_messages);
}
// Restore response chain from conversation metadata
if let Some(store) = self.store() {
if let Ok(Some(metadata)) = store.get_conversation_metadata(thread_uuid).await {
if let Some(rid) = metadata
.get("last_response_id")
.and_then(|v| v.as_str())
.map(String::from)
{
thread.last_response_id = Some(rid.clone());
self.llm()
.seed_response_chain(&thread_uuid.to_string(), rid);
tracing::debug!("Restored response chain for thread {}", thread_uuid);
}
}
}
// Insert into session and register with session manager
{
let mut sess = session.lock().await;
sess.threads.insert(thread_uuid, thread);
sess.active_thread = Some(thread_uuid);
sess.last_active_at = chrono::Utc::now();
}
self.session_manager
.register_thread(
&message.user_id,
&message.channel,
thread_uuid,
Arc::clone(&session),
)
.await;
tracing::debug!(
"Hydrated thread {} from DB ({} messages)",
thread_uuid,
msg_count
);
}
async fn process_user_input(
&self,
message: &IncomingMessage,
@@ -694,6 +872,7 @@ impl Agent {
match result {
Ok(AgenticLoopResult::Response(response)) => {
thread.complete_turn(&response);
self.persist_response_chain(thread);
let _ = self
.channels
.send_status(
@@ -702,6 +881,10 @@ impl Agent {
&message.metadata,
)
.await;
// Fire-and-forget: persist turn to DB
self.persist_turn(thread_id, &message.user_id, content, Some(&response));
Ok(SubmissionResult::response(response))
}
Ok(AgenticLoopResult::NeedApproval { pending }) => {
@@ -728,11 +911,95 @@ impl Agent {
}
Err(e) => {
thread.fail_turn(e.to_string());
// Persist the user message even on failure
self.persist_turn(thread_id, &message.user_id, content, None);
Ok(SubmissionResult::error(e.to_string()))
}
}
}
/// Fire-and-forget: persist a turn (user message + optional assistant response) to the DB.
fn persist_turn(
&self,
thread_id: Uuid,
user_id: &str,
user_input: &str,
response: Option<&str>,
) {
let store = match self.store() {
Some(s) => Arc::clone(s),
None => return,
};
let user_id = user_id.to_string();
let user_input = user_input.to_string();
let response = response.map(String::from);
tokio::spawn(async move {
if let Err(e) = store
.ensure_conversation(thread_id, "gateway", &user_id, None)
.await
{
tracing::warn!("Failed to ensure conversation {}: {}", thread_id, e);
return;
}
if let Err(e) = store
.add_conversation_message(thread_id, "user", &user_input)
.await
{
tracing::warn!("Failed to persist user message: {}", e);
return;
}
if let Some(ref resp) = response {
if let Err(e) = store
.add_conversation_message(thread_id, "assistant", resp)
.await
{
tracing::warn!("Failed to persist assistant message: {}", e);
}
}
});
}
/// Sync the provider's response chain ID to the thread and DB metadata.
///
/// Call after a successful agentic loop to persist the latest
/// `previous_response_id` so chaining survives restarts.
fn persist_response_chain(&self, thread: &mut crate::agent::session::Thread) {
let tid = thread.id.to_string();
let response_id = match self.llm().get_response_chain_id(&tid) {
Some(rid) => rid,
None => return,
};
// Update in-memory thread
thread.last_response_id = Some(response_id.clone());
// Fire-and-forget DB write
let store = match self.store() {
Some(s) => Arc::clone(s),
None => return,
};
let thread_id = thread.id;
tokio::spawn(async move {
let val = serde_json::json!(response_id);
if let Err(e) = store
.update_conversation_metadata_field(thread_id, "last_response_id", &val)
.await
{
tracing::warn!(
"Failed to persist response chain for thread {}: {}",
thread_id,
e
);
}
});
}
/// Run the agentic loop: call LLM, execute tools, repeat until text response.
///
/// Returns `AgenticLoopResult::Response` on completion, or
@@ -808,7 +1075,12 @@ impl Agent {
// Call LLM with current context
let context = ReasoningContext::new()
.with_messages(context_messages.clone())
.with_tools(tool_defs);
.with_tools(tool_defs)
.with_metadata({
let mut m = std::collections::HashMap::new();
m.insert("thread_id".to_string(), thread_id.to_string());
m
});
let result = reasoning.respond_with_tools(&context).await?;
@@ -832,13 +1104,16 @@ impl Agent {
// Tools have been executed or we've tried multiple times, return response
return Ok(AgenticLoopResult::Response(text));
}
RespondResult::ToolCalls(tool_calls) => {
RespondResult::ToolCalls {
tool_calls,
content,
} => {
tools_executed = true;
// Add the assistant message with tool_calls to context.
// OpenAI-compatible APIs require this before tool-result messages.
// OpenAI protocol requires this before tool-result messages.
context_messages.push(ChatMessage::assistant_with_tool_calls(
"",
content,
tool_calls.clone(),
));
@@ -960,10 +1235,26 @@ impl Agent {
if let Some((ext_name, instructions)) =
detect_auth_awaiting(&tc.name, &tool_result)
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) {
thread.enter_auth_mode(ext_name);
let auth_data = parse_auth_result(&tool_result);
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) {
thread.enter_auth_mode(ext_name.clone());
}
}
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::AuthRequired {
extension_name: ext_name,
instructions: Some(instructions.clone()),
auth_url: auth_data.auth_url,
setup_url: auth_data.setup_url,
},
&message.metadata,
)
.await;
return Ok(AgenticLoopResult::Response(instructions));
}
@@ -1024,19 +1315,59 @@ impl Agent {
.into());
}
// Execute with timeout
let result = tokio::time::timeout(std::time::Duration::from_secs(60), async {
tracing::debug!(
tool = %tool_name,
params = %params,
"Tool call started"
);
// Execute with per-tool timeout
let timeout = tool.execution_timeout();
let start = std::time::Instant::now();
let result = tokio::time::timeout(timeout, async {
tool.execute(params.clone(), job_ctx).await
})
.await
.map_err(|_| crate::error::ToolError::Timeout {
name: tool_name.to_string(),
timeout: std::time::Duration::from_secs(60),
})?
.map_err(|e| crate::error::ToolError::ExecutionFailed {
name: tool_name.to_string(),
reason: e.to_string(),
})?;
.await;
let elapsed = start.elapsed();
match &result {
Ok(Ok(output)) => {
let result_str = serde_json::to_string(&output.result)
.unwrap_or_else(|_| "<serialize error>".to_string());
tracing::debug!(
tool = %tool_name,
elapsed_ms = elapsed.as_millis() as u64,
result = %result_str,
"Tool call succeeded"
);
}
Ok(Err(e)) => {
tracing::debug!(
tool = %tool_name,
elapsed_ms = elapsed.as_millis() as u64,
error = %e,
"Tool call failed"
);
}
Err(_) => {
tracing::debug!(
tool = %tool_name,
elapsed_ms = elapsed.as_millis() as u64,
timeout_secs = timeout.as_secs(),
"Tool call timed out"
);
}
}
let result = result
.map_err(|_| crate::error::ToolError::Timeout {
name: tool_name.to_string(),
timeout,
})?
.map_err(|e| crate::error::ToolError::ExecutionFailed {
name: tool_name.to_string(),
reason: e.to_string(),
})?;
// Convert result to string
serde_json::to_string_pretty(&result.result).map_err(|e| {
@@ -1380,10 +1711,11 @@ impl Agent {
if let Some((ext_name, instructions)) =
detect_auth_awaiting(&pending.tool_name, &tool_result)
{
let auth_data = parse_auth_result(&tool_result);
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) {
thread.enter_auth_mode(ext_name);
thread.enter_auth_mode(ext_name.clone());
thread.complete_turn(&instructions);
}
}
@@ -1391,7 +1723,12 @@ impl Agent {
.channels
.send_status(
&message.channel,
StatusUpdate::Status("Awaiting token".into()),
StatusUpdate::AuthRequired {
extension_name: ext_name,
instructions: Some(instructions.clone()),
auth_url: auth_data.auth_url,
setup_url: auth_data.setup_url,
},
&message.metadata,
)
.await;
@@ -1434,6 +1771,7 @@ impl Agent {
match result {
Ok(AgenticLoopResult::Response(response)) => {
thread.complete_turn(&response);
self.persist_response_chain(thread);
let _ = self
.channels
.send_status(
@@ -1532,16 +1870,6 @@ impl Agent {
pending.extension_name
);
// Notify via channel status so the response doesn't echo the token
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::Status("Authenticated, loading tools...".into()),
&message.metadata,
)
.await;
// Auto-activate so tools are available immediately after auth
match ext_mgr.activate(&pending.extension_name).await {
Ok(activate_result) => {
@@ -1551,10 +1879,23 @@ impl Agent {
} else {
format!("\n\nTools: {}", activate_result.tools_loaded.join(", "))
};
Ok(Some(format!(
let msg = format!(
"{} authenticated and activated ({} tools loaded).{}",
pending.extension_name, tool_count, tool_list
)))
);
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::AuthCompleted {
extension_name: pending.extension_name.clone(),
success: true,
message: msg.clone(),
},
&message.metadata,
)
.await;
Ok(Some(msg))
}
Err(e) => {
tracing::warn!(
@@ -1562,16 +1903,29 @@ impl Agent {
pending.extension_name,
e
);
Ok(Some(format!(
let msg = format!(
"{} authenticated successfully, but activation failed: {}. \
Try activating manually.",
pending.extension_name, e
)))
);
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::AuthCompleted {
extension_name: pending.extension_name.clone(),
success: true,
message: msg.clone(),
},
&message.metadata,
)
.await;
Ok(Some(msg))
}
}
}
Ok(result) => {
// Unexpected state, re-enter auth mode
// Invalid token, re-enter auth mode
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) {
@@ -1580,13 +1934,43 @@ impl Agent {
}
let msg = result
.instructions
.clone()
.unwrap_or_else(|| "Invalid token. Please try again.".to_string());
// Re-emit AuthRequired so web UI re-shows the card
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::AuthRequired {
extension_name: pending.extension_name.clone(),
instructions: Some(msg.clone()),
auth_url: result.auth_url,
setup_url: result.setup_url,
},
&message.metadata,
)
.await;
Ok(Some(msg))
}
Err(e) => {
let msg = format!(
"Authentication failed for {}: {}",
pending.extension_name, e
);
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::AuthCompleted {
extension_name: pending.extension_name.clone(),
success: false,
message: msg.clone(),
},
&message.metadata,
)
.await;
Ok(Some(msg))
}
Err(e) => Ok(Some(format!(
"Authentication failed for {}: {}",
pending.extension_name, e
))),
}
}
@@ -1937,40 +2321,49 @@ impl Agent {
}
}
async fn handle_command(
/// Handle system commands that bypass thread-state checks entirely.
async fn handle_system_command(
&self,
command: &str,
_args: &[String],
) -> Result<Option<String>, Error> {
args: &[String],
) -> Result<SubmissionResult, Error> {
match command {
"help" => Ok(Some(
r#"Commands:
/job <desc> - Create a job
/status [id] - Check job status
/cancel <id> - Cancel a job
/list - List all jobs
/help <job_id> - Help a stuck job
"help" => Ok(SubmissionResult::response(concat!(
"System:\n",
" /help Show this help\n",
" /model [name] Show or switch the active model\n",
" /version Show version info\n",
" /tools List available tools\n",
" /debug Toggle debug mode\n",
" /ping Connectivity check\n",
"\n",
"Jobs:\n",
" /job <desc> Create a new job\n",
" /status [id] Check job status\n",
" /cancel <id> Cancel a job\n",
" /list List all jobs\n",
"\n",
"Session:\n",
" /undo Undo last turn\n",
" /redo Redo undone turn\n",
" /compact Compress context window\n",
" /clear Clear current thread\n",
" /interrupt Stop current operation\n",
" /new New conversation thread\n",
" /thread <id> Switch to thread\n",
" /resume <id> Resume from checkpoint\n",
"\n",
"Agent:\n",
" /heartbeat Run heartbeat check\n",
" /summarize Summarize current thread\n",
" /suggest Suggest next steps\n",
"\n",
" /quit Exit",
))),
/undo - Undo last turn
/redo - Redo undone turn
/compact - Compress context
/clear - Clear thread
/interrupt - Stop current turn
/thread new - New thread
/thread <id> - Switch thread
/resume <id> - Resume checkpoint
"ping" => Ok(SubmissionResult::response("pong!")),
/heartbeat - Run heartbeat check now
/summarize - Summarize current thread
/suggest - Suggest next steps
/quit - Exit"#
.to_string(),
)),
"ping" => Ok(Some("pong!".to_string())),
"version" => Ok(Some(format!(
"version" => Ok(SubmissionResult::response(format!(
"{} v{}",
env!("CARGO_PKG_NAME"),
env!("CARGO_PKG_VERSION")
@@ -1978,12 +2371,113 @@ impl Agent {
"tools" => {
let tools = self.tools().list().await;
Ok(Some(format!("Available tools: {}", tools.join(", "))))
Ok(SubmissionResult::response(format!(
"Available tools: {}",
tools.join(", ")
)))
}
_ => Ok(Some(format!("Unknown command: {}. Try /help", command))),
"debug" => {
// Debug toggle is handled client-side in the REPL.
// For non-REPL channels, just acknowledge.
Ok(SubmissionResult::ok_with_message(
"Debug toggle is handled by your client.",
))
}
"model" => {
if args.is_empty() {
// Show current model
let name = self.llm().active_model_name();
Ok(SubmissionResult::response(format!(
"Active model: {}",
name
)))
} else {
let requested = &args[0];
// Validate the model exists
match self.llm().list_models().await {
Ok(models) if !models.is_empty() => {
if !models.iter().any(|m| m == requested) {
return Ok(SubmissionResult::error(format!(
"Unknown model: {}. Available models:\n {}",
requested,
models.join("\n ")
)));
}
}
Ok(_) => {
// Empty model list, can't validate but try anyway
}
Err(e) => {
tracing::warn!("Could not fetch model list for validation: {}", e);
// Proceed anyway, the provider will error on the next call if invalid
}
}
match self.llm().set_model(requested) {
Ok(()) => Ok(SubmissionResult::response(format!(
"Switched model to: {}",
requested
))),
Err(e) => Ok(SubmissionResult::error(format!(
"Failed to switch model: {}",
e
))),
}
}
}
_ => Ok(SubmissionResult::error(format!(
"Unknown command: {}. Try /help",
command
))),
}
}
/// Handle legacy command routing from the Router (job commands that go through
/// process_user_input -> router -> handle_job_or_command -> here).
async fn handle_command(
&self,
command: &str,
args: &[String],
) -> Result<Option<String>, Error> {
// System commands are now handled directly via Submission::SystemCommand,
// but the router may still send us unknown /commands.
match self.handle_system_command(command, args).await? {
SubmissionResult::Response { content } => Ok(Some(content)),
SubmissionResult::Ok { message } => Ok(message),
SubmissionResult::Error { message } => Ok(Some(format!("Error: {}", message))),
_ => Ok(None),
}
}
}
/// Parsed auth result fields for emitting StatusUpdate::AuthRequired.
struct ParsedAuthData {
auth_url: Option<String>,
setup_url: Option<String>,
}
/// Extract auth_url and setup_url from a tool_auth result JSON string.
fn parse_auth_result(result: &Result<String, Error>) -> ParsedAuthData {
let parsed = result
.as_ref()
.ok()
.and_then(|s| serde_json::from_str::<serde_json::Value>(s).ok());
ParsedAuthData {
auth_url: parsed
.as_ref()
.and_then(|v| v.get("auth_url"))
.and_then(|v| v.as_str())
.map(|s| s.to_string()),
setup_url: parsed
.as_ref()
.and_then(|v| v.get("setup_url"))
.and_then(|v| v.as_str())
.map(|s| s.to_string()),
}
}
/// Check if a tool_auth result indicates the extension is awaiting a token.
@@ -1994,7 +2488,7 @@ fn detect_auth_awaiting(
tool_name: &str,
result: &Result<String, Error>,
) -> Option<(String, String)> {
if tool_name != "tool_auth" {
if tool_name != "tool_auth" && tool_name != "tool_activate" {
return None;
}
let output = result.as_ref().ok()?;
@@ -2078,4 +2572,34 @@ mod tests {
let (_, instructions) = detect_auth_awaiting("tool_auth", &result).unwrap();
assert_eq!(instructions, "Please provide your API token/key.");
}
#[test]
fn test_detect_auth_awaiting_tool_activate() {
let result: Result<String, Error> = Ok(serde_json::json!({
"name": "slack",
"kind": "McpServer",
"awaiting_token": true,
"status": "awaiting_token",
"instructions": "Provide your Slack Bot token."
})
.to_string());
let detected = detect_auth_awaiting("tool_activate", &result);
assert!(detected.is_some());
let (name, instructions) = detected.unwrap();
assert_eq!(name, "slack");
assert!(instructions.contains("Slack Bot"));
}
#[test]
fn test_detect_auth_awaiting_tool_activate_not_awaiting() {
let result: Result<String, Error> = Ok(serde_json::json!({
"name": "slack",
"tools_loaded": ["slack_post_message"],
"message": "Activated"
})
.to_string());
assert!(detect_auth_awaiting("tool_activate", &result).is_none());
}
}
src/agent/heartbeat.rs
+34 -3
View File
@@ -29,7 +29,7 @@ use std::time::Duration;
use tokio::sync::mpsc;
use crate::channels::OutgoingResponse;
use crate::llm::{ChatMessage, CompletionRequest, LlmProvider};
use crate::llm::{ChatMessage, CompletionRequest, FinishReason, LlmProvider};
use crate::workspace::Workspace;
/// Configuration for the heartbeat runner.
@@ -217,9 +217,26 @@ impl HeartbeatRunner {
]
};
// Use the model's context_length to set max_tokens. The API returns
// the total context window; we cap output at half of that (the rest is
// the prompt) with a floor of 4096.
let max_tokens = match self.llm.model_metadata().await {
Ok(meta) => {
let from_api = meta.context_length.map(|ctx| ctx / 2).unwrap_or(4096);
from_api.max(4096)
}
Err(e) => {
tracing::warn!(
"Could not fetch model metadata, using default max_tokens: {}",
e
);
4096
}
};
let request = CompletionRequest::new(messages)
.with_max_tokens(1024)
.with_temperature(0.3); // Lower temperature for more focused responses
.with_max_tokens(max_tokens)
.with_temperature(0.3);
let response = match self.llm.complete(request).await {
Ok(r) => r,
@@ -228,6 +245,20 @@ impl HeartbeatRunner {
let content = response.content.trim();
// Guard against empty content. Reasoning models (e.g. GLM-4.7) may
// burn all output tokens on chain-of-thought and return content: null.
if content.is_empty() {
return if response.finish_reason == FinishReason::Length {
HeartbeatResult::Failed(
"LLM response was truncated (finish_reason=length) with no content. \
The model may have exhausted its token budget on reasoning."
.to_string(),
)
} else {
HeartbeatResult::Failed("LLM returned empty content.".to_string())
};
}
// Check if nothing needs attention
if content == "HEARTBEAT_OK" || content.contains("HEARTBEAT_OK") {
return HeartbeatResult::Ok;
src/agent/mod.rs
+5
View File
@@ -6,6 +6,7 @@
//! - Tool invocation with safety
//! - Self-repair for stuck jobs
//! - Proactive heartbeat execution
//! - Routine-based scheduled and reactive jobs
//! - Turn-based session management with undo
//! - Context compaction for long conversations
@@ -14,6 +15,8 @@ pub mod compaction;
pub mod context_monitor;
mod heartbeat;
mod router;
pub mod routine;
pub mod routine_engine;
mod scheduler;
mod self_repair;
pub mod session;
@@ -28,6 +31,8 @@ pub use compaction::{CompactionResult, ContextCompactor};
pub use context_monitor::{CompactionStrategy, ContextBreakdown, ContextMonitor};
pub use heartbeat::{HeartbeatConfig, HeartbeatResult, HeartbeatRunner, spawn_heartbeat};
pub use router::{MessageIntent, Router};
pub use routine::{Routine, RoutineAction, RoutineRun, Trigger};
pub use routine_engine::RoutineEngine;
pub use scheduler::Scheduler;
pub use self_repair::{BrokenTool, RepairResult, RepairTask, SelfRepair, StuckJob};
pub use session::{PendingApproval, PendingAuth, Session, Thread, ThreadState, Turn, TurnState};
src/agent/routine.rs
+509
View File
@@ -0,0 +1,509 @@
//! Core types for the routines system.
//!
//! A routine is a named, persistent, user-owned task with a trigger and an action.
//! Each routine fires independently when its trigger condition is met, with only
//! that routine's prompt and context sent to the LLM.
//!
//! ```text
//! ┌──────────┐ ┌─────────┐ ┌──────────────────┐
//! │ Trigger │────▶│ Engine │────▶│ Execution Mode │
//! │ cron/event│ │guardrail│ │lightweight│full_job│
//! │ webhook │ │ check │ └──────────────────┘
//! │ manual │ └─────────┘ │
//! └──────────┘ ▼
//! ┌──────────────┐
//! │ Notify user │
//! │ if needed │
//! └──────────────┘
//! ```
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::str::FromStr;
use std::time::Duration;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use uuid::Uuid;
/// A routine is a named, persistent, user-owned task with a trigger and an action.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Routine {
pub id: Uuid,
pub name: String,
pub description: String,
pub user_id: String,
pub enabled: bool,
pub trigger: Trigger,
pub action: RoutineAction,
pub guardrails: RoutineGuardrails,
pub notify: NotifyConfig,
// Runtime state (DB-managed)
pub last_run_at: Option<DateTime<Utc>>,
pub next_fire_at: Option<DateTime<Utc>>,
pub run_count: u64,
pub consecutive_failures: u32,
pub state: serde_json::Value,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
/// When a routine should fire.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum Trigger {
/// Fire on a cron schedule (e.g. "0 9 * * MON-FRI" or "every 2h").
Cron { schedule: String },
/// Fire when a channel message matches a pattern.
Event {
/// Optional channel filter (e.g. "telegram", "slack").
channel: Option<String>,
/// Regex pattern to match against message content.
pattern: String,
},
/// Fire on incoming webhook POST to /hooks/routine/{id}.
Webhook {
/// Optional webhook path suffix (defaults to routine id).
path: Option<String>,
/// Optional shared secret for HMAC validation.
secret: Option<String>,
},
/// Only fires via tool call or CLI.
Manual,
}
impl Trigger {
/// The string tag stored in the DB trigger_type column.
pub fn type_tag(&self) -> &'static str {
match self {
Trigger::Cron { .. } => "cron",
Trigger::Event { .. } => "event",
Trigger::Webhook { .. } => "webhook",
Trigger::Manual => "manual",
}
}
/// Parse a trigger from its DB representation.
pub fn from_db(trigger_type: &str, config: serde_json::Value) -> Result<Self, String> {
match trigger_type {
"cron" => {
let schedule = config
.get("schedule")
.and_then(|v| v.as_str())
.ok_or("cron trigger missing 'schedule'")?
.to_string();
Ok(Trigger::Cron { schedule })
}
"event" => {
let pattern = config
.get("pattern")
.and_then(|v| v.as_str())
.ok_or("event trigger missing 'pattern'")?
.to_string();
let channel = config
.get("channel")
.and_then(|v| v.as_str())
.map(String::from);
Ok(Trigger::Event { channel, pattern })
}
"webhook" => {
let path = config
.get("path")
.and_then(|v| v.as_str())
.map(String::from);
let secret = config
.get("secret")
.and_then(|v| v.as_str())
.map(String::from);
Ok(Trigger::Webhook { path, secret })
}
"manual" => Ok(Trigger::Manual),
other => Err(format!("unknown trigger type: {other}")),
}
}
/// Serialize trigger-specific config to JSON for DB storage.
pub fn to_config_json(&self) -> serde_json::Value {
match self {
Trigger::Cron { schedule } => serde_json::json!({ "schedule": schedule }),
Trigger::Event { channel, pattern } => serde_json::json!({
"pattern": pattern,
"channel": channel,
}),
Trigger::Webhook { path, secret } => serde_json::json!({
"path": path,
"secret": secret,
}),
Trigger::Manual => serde_json::json!({}),
}
}
}
/// What happens when a routine fires.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum RoutineAction {
/// Single LLM call, no tools. Cheap and fast.
Lightweight {
/// The prompt sent to the LLM.
prompt: String,
/// Workspace paths to load as context (e.g. ["context/priorities.md"]).
#[serde(default)]
context_paths: Vec<String>,
/// Max output tokens (default: 4096).
#[serde(default = "default_max_tokens")]
max_tokens: u32,
},
/// Full multi-turn worker job with tool access.
FullJob {
/// Job title for the scheduler.
title: String,
/// Job description / initial prompt.
description: String,
/// Max reasoning iterations (default: 10).
#[serde(default = "default_max_iterations")]
max_iterations: u32,
},
}
fn default_max_tokens() -> u32 {
4096
}
fn default_max_iterations() -> u32 {
10
}
impl RoutineAction {
/// The string tag stored in the DB action_type column.
pub fn type_tag(&self) -> &'static str {
match self {
RoutineAction::Lightweight { .. } => "lightweight",
RoutineAction::FullJob { .. } => "full_job",
}
}
/// Parse an action from its DB representation.
pub fn from_db(action_type: &str, config: serde_json::Value) -> Result<Self, String> {
match action_type {
"lightweight" => {
let prompt = config
.get("prompt")
.and_then(|v| v.as_str())
.ok_or("lightweight action missing 'prompt'")?
.to_string();
let context_paths = config
.get("context_paths")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_str().map(String::from))
.collect()
})
.unwrap_or_default();
let max_tokens = config
.get("max_tokens")
.and_then(|v| v.as_u64())
.unwrap_or(default_max_tokens() as u64) as u32;
Ok(RoutineAction::Lightweight {
prompt,
context_paths,
max_tokens,
})
}
"full_job" => {
let title = config
.get("title")
.and_then(|v| v.as_str())
.ok_or("full_job action missing 'title'")?
.to_string();
let description = config
.get("description")
.and_then(|v| v.as_str())
.ok_or("full_job action missing 'description'")?
.to_string();
let max_iterations = config
.get("max_iterations")
.and_then(|v| v.as_u64())
.unwrap_or(default_max_iterations() as u64)
as u32;
Ok(RoutineAction::FullJob {
title,
description,
max_iterations,
})
}
other => Err(format!("unknown action type: {other}")),
}
}
/// Serialize action config to JSON for DB storage.
pub fn to_config_json(&self) -> serde_json::Value {
match self {
RoutineAction::Lightweight {
prompt,
context_paths,
max_tokens,
} => serde_json::json!({
"prompt": prompt,
"context_paths": context_paths,
"max_tokens": max_tokens,
}),
RoutineAction::FullJob {
title,
description,
max_iterations,
} => serde_json::json!({
"title": title,
"description": description,
"max_iterations": max_iterations,
}),
}
}
}
/// Guardrails to prevent runaway execution.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RoutineGuardrails {
/// Minimum time between fires.
pub cooldown: Duration,
/// Max simultaneous runs of this routine.
pub max_concurrent: u32,
/// Window for content-hash dedup (event triggers). None = no dedup.
pub dedup_window: Option<Duration>,
}
impl Default for RoutineGuardrails {
fn default() -> Self {
Self {
cooldown: Duration::from_secs(300),
max_concurrent: 1,
dedup_window: None,
}
}
}
/// Notification preferences for a routine.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NotifyConfig {
/// Channel to notify on (None = default/broadcast all).
pub channel: Option<String>,
/// User to notify.
pub user: String,
/// Notify when routine produces actionable output.
pub on_attention: bool,
/// Notify when routine errors.
pub on_failure: bool,
/// Notify when routine runs with no findings.
pub on_success: bool,
}
impl Default for NotifyConfig {
fn default() -> Self {
Self {
channel: None,
user: "default".to_string(),
on_attention: true,
on_failure: true,
on_success: false,
}
}
}
/// Status of a routine run.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RunStatus {
Running,
Ok,
Attention,
Failed,
}
impl std::fmt::Display for RunStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
RunStatus::Running => write!(f, "running"),
RunStatus::Ok => write!(f, "ok"),
RunStatus::Attention => write!(f, "attention"),
RunStatus::Failed => write!(f, "failed"),
}
}
}
impl FromStr for RunStatus {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s {
"running" => Ok(RunStatus::Running),
"ok" => Ok(RunStatus::Ok),
"attention" => Ok(RunStatus::Attention),
"failed" => Ok(RunStatus::Failed),
other => Err(format!("unknown run status: {other}")),
}
}
}
/// A single execution of a routine.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RoutineRun {
pub id: Uuid,
pub routine_id: Uuid,
pub trigger_type: String,
pub trigger_detail: Option<String>,
pub started_at: DateTime<Utc>,
pub completed_at: Option<DateTime<Utc>>,
pub status: RunStatus,
pub result_summary: Option<String>,
pub tokens_used: Option<i32>,
pub job_id: Option<Uuid>,
pub created_at: DateTime<Utc>,
}
/// Compute a content hash for event dedup.
pub fn content_hash(content: &str) -> u64 {
let mut hasher = DefaultHasher::new();
content.hash(&mut hasher);
hasher.finish()
}
/// Parse a cron expression and compute the next fire time from now.
pub fn next_cron_fire(schedule: &str) -> Result<Option<DateTime<Utc>>, String> {
let cron_schedule =
cron::Schedule::from_str(schedule).map_err(|e| format!("invalid cron: {e}"))?;
Ok(cron_schedule.upcoming(Utc).next())
}
#[cfg(test)]
mod tests {
use crate::agent::routine::{
RoutineAction, RoutineGuardrails, RunStatus, Trigger, content_hash, next_cron_fire,
};
#[test]
fn test_trigger_roundtrip() {
let trigger = Trigger::Cron {
schedule: "0 9 * * MON-FRI".to_string(),
};
let json = trigger.to_config_json();
let parsed = Trigger::from_db("cron", json).expect("parse cron");
assert!(matches!(parsed, Trigger::Cron { schedule } if schedule == "0 9 * * MON-FRI"));
}
#[test]
fn test_event_trigger_roundtrip() {
let trigger = Trigger::Event {
channel: Some("telegram".to_string()),
pattern: r"deploy\s+\w+".to_string(),
};
let json = trigger.to_config_json();
let parsed = Trigger::from_db("event", json).expect("parse event");
assert!(matches!(parsed, Trigger::Event { channel, pattern }
if channel == Some("telegram".to_string()) && pattern == r"deploy\s+\w+"));
}
#[test]
fn test_action_lightweight_roundtrip() {
let action = RoutineAction::Lightweight {
prompt: "Check PRs".to_string(),
context_paths: vec!["context/priorities.md".to_string()],
max_tokens: 2048,
};
let json = action.to_config_json();
let parsed = RoutineAction::from_db("lightweight", json).expect("parse lightweight");
assert!(
matches!(parsed, RoutineAction::Lightweight { prompt, context_paths, max_tokens }
if prompt == "Check PRs" && context_paths.len() == 1 && max_tokens == 2048)
);
}
#[test]
fn test_action_full_job_roundtrip() {
let action = RoutineAction::FullJob {
title: "Deploy review".to_string(),
description: "Review and deploy pending changes".to_string(),
max_iterations: 5,
};
let json = action.to_config_json();
let parsed = RoutineAction::from_db("full_job", json).expect("parse full_job");
assert!(
matches!(parsed, RoutineAction::FullJob { title, max_iterations, .. }
if title == "Deploy review" && max_iterations == 5)
);
}
#[test]
fn test_run_status_display_parse() {
for status in [
RunStatus::Running,
RunStatus::Ok,
RunStatus::Attention,
RunStatus::Failed,
] {
let s = status.to_string();
let parsed: RunStatus = s.parse().expect("parse status");
assert_eq!(parsed, status);
}
}
#[test]
fn test_content_hash_deterministic() {
let h1 = content_hash("deploy production");
let h2 = content_hash("deploy production");
assert_eq!(h1, h2);
let h3 = content_hash("deploy staging");
assert_ne!(h1, h3);
}
#[test]
fn test_next_cron_fire_valid() {
// Every minute should always have a next fire
let next = next_cron_fire("* * * * * *").expect("valid cron");
assert!(next.is_some());
}
#[test]
fn test_next_cron_fire_invalid() {
let result = next_cron_fire("not a cron");
assert!(result.is_err());
}
#[test]
fn test_guardrails_default() {
let g = RoutineGuardrails::default();
assert_eq!(g.cooldown.as_secs(), 300);
assert_eq!(g.max_concurrent, 1);
assert!(g.dedup_window.is_none());
}
#[test]
fn test_trigger_type_tag() {
assert_eq!(
Trigger::Cron {
schedule: String::new()
}
.type_tag(),
"cron"
);
assert_eq!(
Trigger::Event {
channel: None,
pattern: String::new()
}
.type_tag(),
"event"
);
assert_eq!(
Trigger::Webhook {
path: None,
secret: None
}
.type_tag(),
"webhook"
);
assert_eq!(Trigger::Manual.type_tag(), "manual");
}
}
src/agent/routine_engine.rs
+606
View File
@@ -0,0 +1,606 @@
//! Routine execution engine.
//!
//! Handles loading routines, checking triggers, enforcing guardrails,
//! and executing both lightweight (single LLM call) and full-job routines.
//!
//! The engine runs two independent loops:
//! - A **cron ticker** that polls the DB every N seconds for due cron routines
//! - An **event matcher** called synchronously from the agent main loop
//!
//! Lightweight routines execute inline (single LLM call, no scheduler slot).
//! Full-job routines are delegated to the existing `Scheduler`.
use std::sync::Arc;
use std::time::Duration;
use chrono::Utc;
use regex::Regex;
use tokio::sync::{RwLock, mpsc};
use uuid::Uuid;
use crate::agent::routine::{
NotifyConfig, Routine, RoutineAction, RoutineRun, RunStatus, Trigger, next_cron_fire,
};
use crate::channels::{IncomingMessage, OutgoingResponse};
use crate::config::RoutineConfig;
use crate::history::Store;
use crate::llm::{ChatMessage, CompletionRequest, FinishReason, LlmProvider};
use crate::workspace::Workspace;
/// The routine execution engine.
pub struct RoutineEngine {
config: RoutineConfig,
store: Arc<Store>,
llm: Arc<dyn LlmProvider>,
workspace: Arc<Workspace>,
/// Sender for notifications (routed to channel manager).
notify_tx: mpsc::Sender<OutgoingResponse>,
/// Currently running routine count (across all routines).
running_count: Arc<RwLock<usize>>,
/// Compiled event regex cache: routine_id -> compiled regex.
event_cache: Arc<RwLock<Vec<(Uuid, Routine, Regex)>>>,
}
impl RoutineEngine {
pub fn new(
config: RoutineConfig,
store: Arc<Store>,
llm: Arc<dyn LlmProvider>,
workspace: Arc<Workspace>,
notify_tx: mpsc::Sender<OutgoingResponse>,
) -> Self {
Self {
config,
store,
llm,
workspace,
notify_tx,
running_count: Arc::new(RwLock::new(0)),
event_cache: Arc::new(RwLock::new(Vec::new())),
}
}
/// Refresh the in-memory event trigger cache from DB.
pub async fn refresh_event_cache(&self) {
match self.store.list_event_routines().await {
Ok(routines) => {
let mut cache = Vec::new();
for routine in routines {
if let Trigger::Event { ref pattern, .. } = routine.trigger {
match Regex::new(pattern) {
Ok(re) => cache.push((routine.id, routine.clone(), re)),
Err(e) => {
tracing::warn!(
routine = %routine.name,
"Invalid event regex '{}': {}",
pattern, e
);
}
}
}
}
let count = cache.len();
*self.event_cache.write().await = cache;
tracing::debug!("Refreshed event cache: {} routines", count);
}
Err(e) => {
tracing::error!("Failed to refresh event cache: {}", e);
}
}
}
/// Check incoming message against event triggers. Returns number of routines fired.
///
/// Called synchronously from the main loop after handle_message(). The actual
/// execution is spawned async so this returns quickly.
pub async fn check_event_triggers(&self, message: &IncomingMessage) -> usize {
let cache = self.event_cache.read().await;
let mut fired = 0;
for (_, routine, re) in cache.iter() {
// Channel filter
if let Trigger::Event {
channel: Some(ch), ..
} = &routine.trigger
{
if ch != &message.channel {
continue;
}
}
// Regex match
if !re.is_match(&message.content) {
continue;
}
// Cooldown check
if !self.check_cooldown(routine) {
tracing::debug!(routine = %routine.name, "Skipped: cooldown active");
continue;
}
// Concurrent run check
if !self.check_concurrent(routine).await {
tracing::debug!(routine = %routine.name, "Skipped: max concurrent reached");
continue;
}
// Global capacity check
if *self.running_count.read().await >= self.config.max_concurrent_routines {
tracing::warn!(routine = %routine.name, "Skipped: global max concurrent reached");
continue;
}
let detail = truncate(&message.content, 200);
self.spawn_fire(routine.clone(), "event", Some(detail));
fired += 1;
}
fired
}
/// Check all due cron routines and fire them. Called by the cron ticker.
pub async fn check_cron_triggers(&self) {
let routines = match self.store.list_due_cron_routines().await {
Ok(r) => r,
Err(e) => {
tracing::error!("Failed to load due cron routines: {}", e);
return;
}
};
for routine in routines {
if *self.running_count.read().await >= self.config.max_concurrent_routines {
tracing::warn!("Global max concurrent routines reached, skipping remaining");
break;
}
if !self.check_cooldown(&routine) {
continue;
}
if !self.check_concurrent(&routine).await {
continue;
}
let detail = if let Trigger::Cron { ref schedule } = routine.trigger {
Some(schedule.clone())
} else {
None
};
self.spawn_fire(routine, "cron", detail);
}
}
/// Fire a routine manually (from tool call or CLI).
pub async fn fire_manual(&self, routine_id: Uuid) -> Result<Uuid, String> {
let routine = self
.store
.get_routine(routine_id)
.await
.map_err(|e| format!("DB error: {e}"))?
.ok_or_else(|| format!("routine {routine_id} not found"))?;
if !routine.enabled {
return Err(format!("routine '{}' is disabled", routine.name));
}
if !self.check_concurrent(&routine).await {
return Err(format!(
"routine '{}' already at max concurrent runs",
routine.name
));
}
let run_id = Uuid::new_v4();
let run = RoutineRun {
id: run_id,
routine_id: routine.id,
trigger_type: "manual".to_string(),
trigger_detail: None,
started_at: Utc::now(),
completed_at: None,
status: RunStatus::Running,
result_summary: None,
tokens_used: None,
job_id: None,
created_at: Utc::now(),
};
if let Err(e) = self.store.create_routine_run(&run).await {
return Err(format!("failed to create run record: {e}"));
}
// Execute inline for manual triggers (caller wants to wait)
let engine = EngineContext {
store: self.store.clone(),
llm: self.llm.clone(),
workspace: self.workspace.clone(),
notify_tx: self.notify_tx.clone(),
running_count: self.running_count.clone(),
max_lightweight_tokens: self.config.max_lightweight_tokens,
};
tokio::spawn(async move {
execute_routine(engine, routine, run).await;
});
Ok(run_id)
}
/// Spawn a fire in a background task.
fn spawn_fire(&self, routine: Routine, trigger_type: &str, trigger_detail: Option<String>) {
let run = RoutineRun {
id: Uuid::new_v4(),
routine_id: routine.id,
trigger_type: trigger_type.to_string(),
trigger_detail,
started_at: Utc::now(),
completed_at: None,
status: RunStatus::Running,
result_summary: None,
tokens_used: None,
job_id: None,
created_at: Utc::now(),
};
let engine = EngineContext {
store: self.store.clone(),
llm: self.llm.clone(),
workspace: self.workspace.clone(),
notify_tx: self.notify_tx.clone(),
running_count: self.running_count.clone(),
max_lightweight_tokens: self.config.max_lightweight_tokens,
};
// Record the run in DB, then spawn execution
let store = self.store.clone();
tokio::spawn(async move {
if let Err(e) = store.create_routine_run(&run).await {
tracing::error!(routine = %routine.name, "Failed to record run: {}", e);
return;
}
execute_routine(engine, routine, run).await;
});
}
fn check_cooldown(&self, routine: &Routine) -> bool {
if let Some(last_run) = routine.last_run_at {
let elapsed = Utc::now().signed_duration_since(last_run);
let cooldown = chrono::Duration::from_std(routine.guardrails.cooldown)
.unwrap_or(chrono::Duration::seconds(300));
if elapsed < cooldown {
return false;
}
}
true
}
async fn check_concurrent(&self, routine: &Routine) -> bool {
match self.store.count_running_routine_runs(routine.id).await {
Ok(count) => count < routine.guardrails.max_concurrent as i64,
Err(e) => {
tracing::error!(
routine = %routine.name,
"Failed to check concurrent runs: {}", e
);
false
}
}
}
}
/// Shared context passed to the execution function.
struct EngineContext {
store: Arc<Store>,
llm: Arc<dyn LlmProvider>,
workspace: Arc<Workspace>,
notify_tx: mpsc::Sender<OutgoingResponse>,
running_count: Arc<RwLock<usize>>,
max_lightweight_tokens: u32,
}
/// Execute a routine run. Handles both lightweight and full_job modes.
async fn execute_routine(ctx: EngineContext, routine: Routine, run: RoutineRun) {
// Increment running count
{
let mut count = ctx.running_count.write().await;
*count += 1;
}
let result = match &routine.action {
RoutineAction::Lightweight {
prompt,
context_paths,
max_tokens,
} => execute_lightweight(&ctx, &routine, prompt, context_paths, *max_tokens).await,
RoutineAction::FullJob { description, .. } => {
// Full job mode: for now, execute as lightweight with the description
// as prompt. Full scheduler integration will come as a follow-up.
tracing::info!(
routine = %routine.name,
"FullJob mode executing as lightweight (scheduler integration pending)"
);
execute_lightweight(&ctx, &routine, description, &[], ctx.max_lightweight_tokens).await
}
};
// Decrement running count
{
let mut count = ctx.running_count.write().await;
*count = count.saturating_sub(1);
}
// Process result
let (status, summary, tokens) = match result {
Ok(execution) => execution,
Err(e) => {
tracing::error!(routine = %routine.name, "Execution failed: {}", e);
(RunStatus::Failed, Some(e), None)
}
};
// Complete the run record
if let Err(e) = ctx
.store
.complete_routine_run(run.id, status, summary.as_deref(), tokens)
.await
{
tracing::error!(routine = %routine.name, "Failed to complete run record: {}", e);
}
// Update routine runtime state
let now = Utc::now();
let next_fire = if let Trigger::Cron { ref schedule } = routine.trigger {
next_cron_fire(schedule).unwrap_or(None)
} else {
None
};
let new_failures = if status == RunStatus::Failed {
routine.consecutive_failures + 1
} else {
0
};
if let Err(e) = ctx
.store
.update_routine_runtime(
routine.id,
now,
next_fire,
routine.run_count + 1,
new_failures,
&routine.state,
)
.await
{
tracing::error!(routine = %routine.name, "Failed to update runtime state: {}", e);
}
// Send notifications based on config
send_notification(
&ctx.notify_tx,
&routine.notify,
&routine.name,
status,
summary.as_deref(),
)
.await;
}
/// Execute a lightweight routine (single LLM call).
async fn execute_lightweight(
ctx: &EngineContext,
routine: &Routine,
prompt: &str,
context_paths: &[String],
max_tokens: u32,
) -> Result<(RunStatus, Option<String>, Option<i32>), String> {
// Load context from workspace
let mut context_parts = Vec::new();
for path in context_paths {
match ctx.workspace.read(path).await {
Ok(doc) => {
context_parts.push(format!("## {}\n\n{}", path, doc.content));
}
Err(e) => {
tracing::debug!(
routine = %routine.name,
"Failed to read context path {}: {}", path, e
);
}
}
}
// Load routine state from workspace
let state_path = format!("routines/{}/state.md", routine.name);
let state_content = match ctx.workspace.read(&state_path).await {
Ok(doc) => Some(doc.content),
Err(_) => None,
};
// Build the prompt
let mut full_prompt = String::new();
full_prompt.push_str(prompt);
if !context_parts.is_empty() {
full_prompt.push_str("\n\n---\n\n# Context\n\n");
full_prompt.push_str(&context_parts.join("\n\n"));
}
if let Some(state) = &state_content {
full_prompt.push_str("\n\n---\n\n# Previous State\n\n");
full_prompt.push_str(state);
}
full_prompt.push_str(
"\n\n---\n\nIf nothing needs attention, reply EXACTLY with: ROUTINE_OK\n\
If something needs attention, provide a concise summary.",
);
// Get system prompt
let system_prompt = match ctx.workspace.system_prompt().await {
Ok(p) => p,
Err(e) => {
tracing::warn!(routine = %routine.name, "Failed to get system prompt: {}", e);
String::new()
}
};
let messages = if system_prompt.is_empty() {
vec![ChatMessage::user(&full_prompt)]
} else {
vec![
ChatMessage::system(&system_prompt),
ChatMessage::user(&full_prompt),
]
};
// Determine max_tokens from model metadata with fallback
let effective_max_tokens = match ctx.llm.model_metadata().await {
Ok(meta) => {
let from_api = meta.context_length.map(|ctx| ctx / 2).unwrap_or(max_tokens);
from_api.max(max_tokens)
}
Err(_) => max_tokens,
};
let request = CompletionRequest::new(messages)
.with_max_tokens(effective_max_tokens)
.with_temperature(0.3);
let response = ctx
.llm
.complete(request)
.await
.map_err(|e| format!("LLM call failed: {e}"))?;
let content = response.content.trim();
let tokens_used = Some((response.input_tokens + response.output_tokens) as i32);
// Empty content guard (same as heartbeat)
if content.is_empty() {
return if response.finish_reason == FinishReason::Length {
Err(
"LLM response truncated (finish_reason=length) with no content. \
Model may have exhausted token budget on reasoning."
.to_string(),
)
} else {
Err("LLM returned empty content.".to_string())
};
}
// Check for the "nothing to do" sentinel
if content == "ROUTINE_OK" || content.contains("ROUTINE_OK") {
return Ok((RunStatus::Ok, None, tokens_used));
}
Ok((RunStatus::Attention, Some(content.to_string()), tokens_used))
}
/// Send a notification based on the routine's notify config and run status.
async fn send_notification(
tx: &mpsc::Sender<OutgoingResponse>,
notify: &NotifyConfig,
routine_name: &str,
status: RunStatus,
summary: Option<&str>,
) {
let should_notify = match status {
RunStatus::Ok => notify.on_success,
RunStatus::Attention => notify.on_attention,
RunStatus::Failed => notify.on_failure,
RunStatus::Running => false,
};
if !should_notify {
return;
}
let icon = match status {
RunStatus::Ok => "✅",
RunStatus::Attention => "🔔",
RunStatus::Failed => "❌",
RunStatus::Running => "⏳",
};
let message = match summary {
Some(s) => format!("{} *Routine '{}'*: {}\n\n{}", icon, routine_name, status, s),
None => format!("{} *Routine '{}'*: {}", icon, routine_name, status),
};
let response = OutgoingResponse {
content: message,
thread_id: None,
metadata: serde_json::json!({
"source": "routine",
"routine_name": routine_name,
"status": status.to_string(),
}),
};
if let Err(e) = tx.send(response).await {
tracing::error!(routine = %routine_name, "Failed to send notification: {}", e);
}
}
/// Spawn the cron ticker background task.
pub fn spawn_cron_ticker(
engine: Arc<RoutineEngine>,
interval: Duration,
) -> tokio::task::JoinHandle<()> {
tokio::spawn(async move {
let mut ticker = tokio::time::interval(interval);
// Skip immediate first tick
ticker.tick().await;
loop {
ticker.tick().await;
engine.check_cron_triggers().await;
}
})
}
fn truncate(s: &str, max: usize) -> String {
if s.len() <= max {
s.to_string()
} else {
format!("{}...", &s[..max])
}
}
#[cfg(test)]
mod tests {
use crate::agent::routine::{NotifyConfig, RunStatus};
#[test]
fn test_notification_gating() {
let config = NotifyConfig {
on_success: false,
on_failure: true,
on_attention: true,
..Default::default()
};
// on_success = false means Ok status should not notify
assert!(!config.on_success);
assert!(config.on_failure);
assert!(config.on_attention);
}
#[test]
fn test_run_status_icons() {
// Just verify the mapping doesn't panic
for status in [
RunStatus::Ok,
RunStatus::Attention,
RunStatus::Failed,
RunStatus::Running,
] {
let _ = status.to_string();
}
}
}
src/agent/scheduler.rs
+17 -17
View File
@@ -373,23 +373,23 @@ impl Scheduler {
.into());
}
// Execute with timeout
let result = tokio::time::timeout(Duration::from_secs(60), async {
tool.execute(params, &job_ctx).await
})
.await
.map_err(|_| {
Error::Tool(crate::error::ToolError::Timeout {
name: tool_name.to_string(),
timeout: Duration::from_secs(60),
})
})?
.map_err(|e| {
Error::Tool(crate::error::ToolError::ExecutionFailed {
name: tool_name.to_string(),
reason: e.to_string(),
})
})?;
// Execute with per-tool timeout
let tool_timeout = tool.execution_timeout();
let result =
tokio::time::timeout(tool_timeout, async { tool.execute(params, &job_ctx).await })
.await
.map_err(|_| {
Error::Tool(crate::error::ToolError::Timeout {
name: tool_name.to_string(),
timeout: tool_timeout,
})
})?
.map_err(|e| {
Error::Tool(crate::error::ToolError::ExecutionFailed {
name: tool_name.to_string(),
reason: e.to_string(),
})
})?;
Ok(TaskOutput::new(result.result, start.elapsed()))
}
src/agent/session.rs
+404
View File
@@ -173,6 +173,10 @@ pub struct Thread {
/// Pending auth token request (thread is in auth mode).
#[serde(default)]
pub pending_auth: Option<PendingAuth>,
/// Last NEAR AI response ID for response chaining. Persisted to DB
/// metadata so we can resume chaining across restarts.
#[serde(default)]
pub last_response_id: Option<String>,
}
impl Thread {
@@ -189,6 +193,24 @@ impl Thread {
metadata: serde_json::Value::Null,
pending_approval: None,
pending_auth: None,
last_response_id: None,
}
}
/// Create a thread with a specific ID (for DB hydration).
pub fn with_id(id: Uuid, session_id: Uuid) -> Self {
let now = Utc::now();
Self {
id,
session_id,
state: ThreadState::Idle,
turns: Vec::new(),
created_at: now,
updated_at: now,
metadata: serde_json::Value::Null,
pending_approval: None,
pending_auth: None,
last_response_id: None,
}
}
@@ -593,4 +615,386 @@ mod tests {
let restored: Thread = serde_json::from_str(&json).expect("should deserialize");
assert!(restored.pending_auth.is_none());
}
#[test]
fn test_thread_with_id() {
let specific_id = Uuid::new_v4();
let session_id = Uuid::new_v4();
let thread = Thread::with_id(specific_id, session_id);
assert_eq!(thread.id, specific_id);
assert_eq!(thread.session_id, session_id);
assert_eq!(thread.state, ThreadState::Idle);
assert!(thread.turns.is_empty());
}
#[test]
fn test_thread_with_id_restore_messages() {
let thread_id = Uuid::new_v4();
let session_id = Uuid::new_v4();
let mut thread = Thread::with_id(thread_id, session_id);
let messages = vec![
ChatMessage::user("Hello from DB"),
ChatMessage::assistant("Restored response"),
];
thread.restore_from_messages(messages);
assert_eq!(thread.id, thread_id);
assert_eq!(thread.turns.len(), 1);
assert_eq!(thread.turns[0].user_input, "Hello from DB");
assert_eq!(
thread.turns[0].response,
Some("Restored response".to_string())
);
}
#[test]
fn test_restore_from_messages_empty() {
let mut thread = Thread::new(Uuid::new_v4());
// Add a turn first, then restore with empty vec
thread.start_turn("hello");
thread.complete_turn("hi");
assert_eq!(thread.turns.len(), 1);
thread.restore_from_messages(Vec::new());
// Should clear all turns and stay idle
assert!(thread.turns.is_empty());
assert_eq!(thread.state, ThreadState::Idle);
}
#[test]
fn test_restore_from_messages_only_assistant_messages() {
let mut thread = Thread::new(Uuid::new_v4());
// Only assistant messages (no user messages to anchor turns)
let messages = vec![
ChatMessage::assistant("I'm here"),
ChatMessage::assistant("Still here"),
];
thread.restore_from_messages(messages);
// Assistant-only messages have no user turn to attach to, so
// they should be skipped entirely.
assert!(thread.turns.is_empty());
}
#[test]
fn test_restore_from_messages_multiple_user_messages_in_a_row() {
let mut thread = Thread::new(Uuid::new_v4());
// Two user messages with no assistant response between them
let messages = vec![
ChatMessage::user("first"),
ChatMessage::user("second"),
ChatMessage::assistant("reply to second"),
];
thread.restore_from_messages(messages);
// First user message becomes a turn with no response,
// second user message pairs with the assistant response.
assert_eq!(thread.turns.len(), 2);
assert_eq!(thread.turns[0].user_input, "first");
assert!(thread.turns[0].response.is_none());
assert_eq!(thread.turns[1].user_input, "second");
assert_eq!(
thread.turns[1].response,
Some("reply to second".to_string())
);
}
#[test]
fn test_thread_switch() {
let mut session = Session::new("user-1");
let t1_id = session.create_thread().id;
let t2_id = session.create_thread().id;
// After creating two threads, active should be the last one
assert_eq!(session.active_thread, Some(t2_id));
// Switch back to the first
assert!(session.switch_thread(t1_id));
assert_eq!(session.active_thread, Some(t1_id));
// Switching to a nonexistent thread should fail
let fake_id = Uuid::new_v4();
assert!(!session.switch_thread(fake_id));
// Active thread should remain unchanged
assert_eq!(session.active_thread, Some(t1_id));
}
#[test]
fn test_get_or_create_thread_idempotent() {
let mut session = Session::new("user-1");
let tid1 = session.get_or_create_thread().id;
let tid2 = session.get_or_create_thread().id;
// Should return the same thread (not create a new one each time)
assert_eq!(tid1, tid2);
assert_eq!(session.threads.len(), 1);
}
#[test]
fn test_truncate_turns() {
let mut thread = Thread::new(Uuid::new_v4());
for i in 0..5 {
thread.start_turn(format!("msg-{}", i));
thread.complete_turn(format!("resp-{}", i));
}
assert_eq!(thread.turns.len(), 5);
thread.truncate_turns(3);
assert_eq!(thread.turns.len(), 3);
// Should keep the most recent turns
assert_eq!(thread.turns[0].user_input, "msg-2");
assert_eq!(thread.turns[1].user_input, "msg-3");
assert_eq!(thread.turns[2].user_input, "msg-4");
// Turn numbers should be re-indexed
assert_eq!(thread.turns[0].turn_number, 0);
assert_eq!(thread.turns[1].turn_number, 1);
assert_eq!(thread.turns[2].turn_number, 2);
}
#[test]
fn test_truncate_turns_noop_when_fewer() {
let mut thread = Thread::new(Uuid::new_v4());
thread.start_turn("only one");
thread.complete_turn("response");
thread.truncate_turns(10);
assert_eq!(thread.turns.len(), 1);
assert_eq!(thread.turns[0].user_input, "only one");
}
#[test]
fn test_thread_interrupt_and_resume() {
let mut thread = Thread::new(Uuid::new_v4());
thread.start_turn("do something");
assert_eq!(thread.state, ThreadState::Processing);
thread.interrupt();
assert_eq!(thread.state, ThreadState::Interrupted);
let last_turn = thread.last_turn().unwrap();
assert_eq!(last_turn.state, TurnState::Interrupted);
assert!(last_turn.completed_at.is_some());
thread.resume();
assert_eq!(thread.state, ThreadState::Idle);
}
#[test]
fn test_resume_only_from_interrupted() {
let mut thread = Thread::new(Uuid::new_v4());
// Idle thread: resume should be a no-op
assert_eq!(thread.state, ThreadState::Idle);
thread.resume();
assert_eq!(thread.state, ThreadState::Idle);
// Processing thread: resume should not change state
thread.start_turn("work");
assert_eq!(thread.state, ThreadState::Processing);
thread.resume();
assert_eq!(thread.state, ThreadState::Processing);
}
#[test]
fn test_turn_fail() {
let mut thread = Thread::new(Uuid::new_v4());
thread.start_turn("risky operation");
thread.fail_turn("connection timed out");
assert_eq!(thread.state, ThreadState::Idle);
let turn = thread.last_turn().unwrap();
assert_eq!(turn.state, TurnState::Failed);
assert_eq!(turn.error, Some("connection timed out".to_string()));
assert!(turn.response.is_none());
assert!(turn.completed_at.is_some());
}
#[test]
fn test_messages_with_incomplete_last_turn() {
let mut thread = Thread::new(Uuid::new_v4());
thread.start_turn("first");
thread.complete_turn("first reply");
thread.start_turn("second (in progress)");
let messages = thread.messages();
// Should have 3 messages: user, assistant, user (no assistant for in-progress)
assert_eq!(messages.len(), 3);
assert_eq!(messages[0].content, "first");
assert_eq!(messages[1].content, "first reply");
assert_eq!(messages[2].content, "second (in progress)");
}
#[test]
fn test_thread_serialization_round_trip() {
let mut thread = Thread::new(Uuid::new_v4());
thread.start_turn("hello");
thread.complete_turn("world");
thread.last_response_id = Some("resp_abc123".to_string());
let json = serde_json::to_string(&thread).unwrap();
let restored: Thread = serde_json::from_str(&json).unwrap();
assert_eq!(restored.id, thread.id);
assert_eq!(restored.session_id, thread.session_id);
assert_eq!(restored.turns.len(), 1);
assert_eq!(restored.turns[0].user_input, "hello");
assert_eq!(restored.turns[0].response, Some("world".to_string()));
assert_eq!(restored.last_response_id, Some("resp_abc123".to_string()));
}
#[test]
fn test_session_serialization_round_trip() {
let mut session = Session::new("user-ser");
session.create_thread();
session.auto_approve_tool("echo");
let json = serde_json::to_string(&session).unwrap();
let restored: Session = serde_json::from_str(&json).unwrap();
assert_eq!(restored.user_id, "user-ser");
assert_eq!(restored.threads.len(), 1);
assert!(restored.is_tool_auto_approved("echo"));
assert!(!restored.is_tool_auto_approved("shell"));
}
#[test]
fn test_auto_approved_tools() {
let mut session = Session::new("user-1");
assert!(!session.is_tool_auto_approved("shell"));
session.auto_approve_tool("shell");
assert!(session.is_tool_auto_approved("shell"));
// Idempotent
session.auto_approve_tool("shell");
assert_eq!(session.auto_approved_tools.len(), 1);
}
#[test]
fn test_turn_tool_call_error() {
let mut turn = Turn::new(0, "test");
turn.record_tool_call("http", serde_json::json!({"url": "example.com"}));
turn.record_tool_error("timeout");
assert_eq!(turn.tool_calls.len(), 1);
assert_eq!(turn.tool_calls[0].error, Some("timeout".to_string()));
assert!(turn.tool_calls[0].result.is_none());
}
#[test]
fn test_turn_number_increments() {
let mut thread = Thread::new(Uuid::new_v4());
// Before any turns, turn_number() is 1 (1-indexed for display)
assert_eq!(thread.turn_number(), 1);
thread.start_turn("first");
thread.complete_turn("done");
assert_eq!(thread.turn_number(), 2);
thread.start_turn("second");
assert_eq!(thread.turn_number(), 3);
}
#[test]
fn test_complete_turn_on_empty_thread() {
let mut thread = Thread::new(Uuid::new_v4());
// Completing a turn when there are no turns should be a safe no-op
thread.complete_turn("phantom response");
assert_eq!(thread.state, ThreadState::Idle);
assert!(thread.turns.is_empty());
}
#[test]
fn test_fail_turn_on_empty_thread() {
let mut thread = Thread::new(Uuid::new_v4());
// Failing a turn when there are no turns should be a safe no-op
thread.fail_turn("phantom error");
assert_eq!(thread.state, ThreadState::Idle);
assert!(thread.turns.is_empty());
}
#[test]
fn test_pending_approval_flow() {
let mut thread = Thread::new(Uuid::new_v4());
let approval = PendingApproval {
request_id: Uuid::new_v4(),
tool_name: "shell".to_string(),
parameters: serde_json::json!({"command": "rm -rf /"}),
description: "dangerous command".to_string(),
tool_call_id: "call_123".to_string(),
context_messages: vec![ChatMessage::user("do it")],
};
thread.await_approval(approval);
assert_eq!(thread.state, ThreadState::AwaitingApproval);
assert!(thread.pending_approval.is_some());
let taken = thread.take_pending_approval();
assert!(taken.is_some());
assert_eq!(taken.unwrap().tool_name, "shell");
assert!(thread.pending_approval.is_none());
}
#[test]
fn test_clear_pending_approval() {
let mut thread = Thread::new(Uuid::new_v4());
let approval = PendingApproval {
request_id: Uuid::new_v4(),
tool_name: "http".to_string(),
parameters: serde_json::json!({}),
description: "test".to_string(),
tool_call_id: "call_456".to_string(),
context_messages: vec![],
};
thread.await_approval(approval);
thread.clear_pending_approval();
assert_eq!(thread.state, ThreadState::Idle);
assert!(thread.pending_approval.is_none());
}
#[test]
fn test_active_thread_accessors() {
let mut session = Session::new("user-1");
assert!(session.active_thread().is_none());
assert!(session.active_thread_mut().is_none());
let tid = session.create_thread().id;
assert!(session.active_thread().is_some());
assert_eq!(session.active_thread().unwrap().id, tid);
// Mutably modify through accessor
session.active_thread_mut().unwrap().start_turn("test");
assert_eq!(
session.active_thread().unwrap().state,
ThreadState::Processing
);
}
}
src/agent/session_manager.rs
+375
View File
@@ -110,6 +110,41 @@ impl SessionManager {
(session, thread_id)
}
/// Register a hydrated thread so subsequent `resolve_thread` calls find it.
///
/// Inserts into the thread_map and creates an undo manager for the thread.
pub async fn register_thread(
&self,
user_id: &str,
channel: &str,
thread_id: Uuid,
session: Arc<Mutex<Session>>,
) {
let key = ThreadKey {
user_id: user_id.to_string(),
channel: channel.to_string(),
external_thread_id: Some(thread_id.to_string()),
};
{
let mut thread_map = self.thread_map.write().await;
thread_map.insert(key, thread_id);
}
{
let mut undo_managers = self.undo_managers.write().await;
undo_managers
.entry(thread_id)
.or_insert_with(|| Arc::new(Mutex::new(UndoManager::new())));
}
// Ensure the session is tracked
{
let mut sessions = self.sessions.write().await;
sessions.entry(user_id.to_string()).or_insert(session);
}
}
/// Get undo manager for a thread.
pub async fn get_undo_manager(&self, thread_id: Uuid) -> Arc<Mutex<UndoManager>> {
// Fast path
@@ -296,4 +331,344 @@ mod tests {
.await;
assert_eq!(pruned, 0);
}
#[tokio::test]
async fn test_register_thread() {
use crate::agent::session::{Session, Thread};
let manager = SessionManager::new();
let thread_id = Uuid::new_v4();
// Create a session with a hydrated thread
let session = Arc::new(Mutex::new(Session::new("user-hydrate")));
{
let mut sess = session.lock().await;
let thread = Thread::with_id(thread_id, sess.id);
sess.threads.insert(thread_id, thread);
sess.active_thread = Some(thread_id);
}
// Register the thread
manager
.register_thread("user-hydrate", "gateway", thread_id, Arc::clone(&session))
.await;
// resolve_thread should find it (using the UUID as external_thread_id)
let (resolved_session, resolved_tid) = manager
.resolve_thread("user-hydrate", "gateway", Some(&thread_id.to_string()))
.await;
assert_eq!(resolved_tid, thread_id);
// Should be the same session object
let sess = resolved_session.lock().await;
assert!(sess.threads.contains_key(&thread_id));
}
#[tokio::test]
async fn test_resolve_thread_with_explicit_external_id() {
let manager = SessionManager::new();
// Two calls with the same explicit external thread ID should resolve
// to the same internal thread.
let (_, t1) = manager
.resolve_thread("user-1", "gateway", Some("ext-abc"))
.await;
let (_, t2) = manager
.resolve_thread("user-1", "gateway", Some("ext-abc"))
.await;
assert_eq!(t1, t2);
// A different external ID on the same channel/user gets a new thread.
let (_, t3) = manager
.resolve_thread("user-1", "gateway", Some("ext-xyz"))
.await;
assert_ne!(t1, t3);
}
#[tokio::test]
async fn test_resolve_thread_none_vs_some_external_id() {
let manager = SessionManager::new();
// None external_thread_id is a distinct key from Some("ext-1").
let (_, t_none) = manager.resolve_thread("user-1", "cli", None).await;
let (_, t_some) = manager.resolve_thread("user-1", "cli", Some("ext-1")).await;
assert_ne!(t_none, t_some);
}
#[tokio::test]
async fn test_resolve_thread_different_users_isolated() {
let manager = SessionManager::new();
let (_, t1) = manager
.resolve_thread("user-a", "gateway", Some("same-ext"))
.await;
let (_, t2) = manager
.resolve_thread("user-b", "gateway", Some("same-ext"))
.await;
// Same channel + same external ID but different users = different threads
assert_ne!(t1, t2);
}
#[tokio::test]
async fn test_resolve_thread_different_channels_isolated() {
let manager = SessionManager::new();
let (_, t1) = manager
.resolve_thread("user-1", "gateway", Some("thread-x"))
.await;
let (_, t2) = manager
.resolve_thread("user-1", "telegram", Some("thread-x"))
.await;
// Same user + same external ID but different channels = different threads
assert_ne!(t1, t2);
}
#[tokio::test]
async fn test_resolve_thread_stale_mapping_creates_new_thread() {
let manager = SessionManager::new();
// Create a thread normally
let (session, original_tid) = manager
.resolve_thread("user-1", "gateway", Some("ext-1"))
.await;
// Simulate the thread being removed from the session (e.g. pruned)
{
let mut sess = session.lock().await;
sess.threads.remove(&original_tid);
}
// Next resolve should detect the stale mapping and create a fresh thread
let (_, new_tid) = manager
.resolve_thread("user-1", "gateway", Some("ext-1"))
.await;
assert_ne!(original_tid, new_tid);
// The new thread should actually exist in the session
let sess = session.lock().await;
assert!(sess.threads.contains_key(&new_tid));
}
#[tokio::test]
async fn test_register_thread_preserves_uuid_on_resolve() {
use crate::agent::session::{Session, Thread};
let manager = SessionManager::new();
let known_uuid = Uuid::new_v4();
let session = Arc::new(Mutex::new(Session::new("user-web")));
let session_id = {
let sess = session.lock().await;
sess.id
};
// Simulate hydration: create thread with a known UUID
{
let mut sess = session.lock().await;
let thread = Thread::with_id(known_uuid, session_id);
sess.threads.insert(known_uuid, thread);
}
// Register it
manager
.register_thread("user-web", "gateway", known_uuid, Arc::clone(&session))
.await;
// resolve_thread with UUID as external_thread_id MUST return the same UUID,
// not mint a new one (this was the root cause of the "wrong conversation" bug)
let (_, resolved) = manager
.resolve_thread("user-web", "gateway", Some(&known_uuid.to_string()))
.await;
assert_eq!(resolved, known_uuid);
}
#[tokio::test]
async fn test_register_thread_idempotent() {
use crate::agent::session::{Session, Thread};
let manager = SessionManager::new();
let tid = Uuid::new_v4();
let session = Arc::new(Mutex::new(Session::new("user-idem")));
{
let mut sess = session.lock().await;
let thread = Thread::with_id(tid, sess.id);
sess.threads.insert(tid, thread);
}
// Register twice
manager
.register_thread("user-idem", "gateway", tid, Arc::clone(&session))
.await;
manager
.register_thread("user-idem", "gateway", tid, Arc::clone(&session))
.await;
// Should still resolve to the same thread
let (_, resolved) = manager
.resolve_thread("user-idem", "gateway", Some(&tid.to_string()))
.await;
assert_eq!(resolved, tid);
}
#[tokio::test]
async fn test_register_thread_creates_undo_manager() {
use crate::agent::session::{Session, Thread};
let manager = SessionManager::new();
let tid = Uuid::new_v4();
let session = Arc::new(Mutex::new(Session::new("user-undo")));
{
let mut sess = session.lock().await;
let thread = Thread::with_id(tid, sess.id);
sess.threads.insert(tid, thread);
}
manager
.register_thread("user-undo", "gateway", tid, Arc::clone(&session))
.await;
// Undo manager should exist for the registered thread
let undo = manager.get_undo_manager(tid).await;
let undo2 = manager.get_undo_manager(tid).await;
assert!(Arc::ptr_eq(&undo, &undo2));
}
#[tokio::test]
async fn test_register_thread_stores_session() {
use crate::agent::session::{Session, Thread};
let manager = SessionManager::new();
let tid = Uuid::new_v4();
let session = Arc::new(Mutex::new(Session::new("user-new")));
{
let mut sess = session.lock().await;
let thread = Thread::with_id(tid, sess.id);
sess.threads.insert(tid, thread);
}
// The user has no session yet in the manager
{
let sessions = manager.sessions.read().await;
assert!(!sessions.contains_key("user-new"));
}
manager
.register_thread("user-new", "gateway", tid, Arc::clone(&session))
.await;
// Now the session should be tracked
{
let sessions = manager.sessions.read().await;
assert!(sessions.contains_key("user-new"));
}
}
#[tokio::test]
async fn test_multiple_threads_per_user() {
let manager = SessionManager::new();
let (_, t1) = manager
.resolve_thread("user-1", "gateway", Some("thread-a"))
.await;
let (_, t2) = manager
.resolve_thread("user-1", "gateway", Some("thread-b"))
.await;
let (session, t3) = manager
.resolve_thread("user-1", "gateway", Some("thread-c"))
.await;
// All three should be distinct
assert_ne!(t1, t2);
assert_ne!(t2, t3);
assert_ne!(t1, t3);
// All three should exist in the same session
let sess = session.lock().await;
assert!(sess.threads.contains_key(&t1));
assert!(sess.threads.contains_key(&t2));
assert!(sess.threads.contains_key(&t3));
}
#[tokio::test]
async fn test_prune_cleans_thread_map_and_undo_managers() {
let manager = SessionManager::new();
let (stale_session, stale_tid) = manager.resolve_thread("user-stale", "cli", None).await;
// Backdate the session
{
let mut sess = stale_session.lock().await;
sess.last_active_at = chrono::Utc::now() - chrono::TimeDelta::seconds(86400 * 30);
}
// Verify thread_map and undo_managers have entries
{
let tm = manager.thread_map.read().await;
assert!(!tm.is_empty());
}
{
let um = manager.undo_managers.read().await;
assert!(um.contains_key(&stale_tid));
}
let pruned = manager
.prune_stale_sessions(std::time::Duration::from_secs(86400 * 7))
.await;
assert_eq!(pruned, 1);
// Thread map and undo managers should be cleaned up
{
let tm = manager.thread_map.read().await;
assert!(tm.is_empty());
}
{
let um = manager.undo_managers.read().await;
assert!(!um.contains_key(&stale_tid));
}
}
#[tokio::test]
async fn test_resolve_thread_active_thread_set() {
let manager = SessionManager::new();
let (session, thread_id) = manager
.resolve_thread("user-1", "gateway", Some("ext-1"))
.await;
// The resolved thread should be set as the active thread
let sess = session.lock().await;
assert_eq!(sess.active_thread, Some(thread_id));
}
#[tokio::test]
async fn test_register_then_resolve_different_channel_creates_new() {
use crate::agent::session::{Session, Thread};
let manager = SessionManager::new();
let tid = Uuid::new_v4();
let session = Arc::new(Mutex::new(Session::new("user-cross")));
{
let mut sess = session.lock().await;
let thread = Thread::with_id(tid, sess.id);
sess.threads.insert(tid, thread);
}
// Register on "gateway" channel
manager
.register_thread("user-cross", "gateway", tid, Arc::clone(&session))
.await;
// Resolve on a different channel with the same UUID string should NOT
// find the registered thread (channel is part of the key)
let (_, resolved) = manager
.resolve_thread("user-cross", "telegram", Some(&tid.to_string()))
.await;
assert_ne!(resolved, tid);
}
}
src/agent/submission.rs
+131
View File
@@ -43,6 +43,49 @@ impl SubmissionParser {
if lower == "/thread new" || lower == "/new" {
return Submission::NewThread;
}
// System commands (bypass thread-state checks)
if lower == "/help" || lower == "/?" {
return Submission::SystemCommand {
command: "help".to_string(),
args: vec![],
};
}
if lower == "/version" {
return Submission::SystemCommand {
command: "version".to_string(),
args: vec![],
};
}
if lower == "/tools" {
return Submission::SystemCommand {
command: "tools".to_string(),
args: vec![],
};
}
if lower == "/ping" {
return Submission::SystemCommand {
command: "ping".to_string(),
args: vec![],
};
}
if lower == "/debug" {
return Submission::SystemCommand {
command: "debug".to_string(),
args: vec![],
};
}
if lower.starts_with("/model") {
let args: Vec<String> = trimmed
.split_whitespace()
.skip(1)
.map(|s| s.to_string())
.collect();
return Submission::SystemCommand {
command: "model".to_string(),
args,
};
}
if lower == "/quit" || lower == "/exit" || lower == "/shutdown" {
return Submission::Quit;
}
@@ -172,6 +215,15 @@ pub enum Submission {
/// Quit the agent. Bypasses thread-state checks.
Quit,
/// System command (help, model, version, tools, ping, debug).
/// Bypasses thread-state checks and safety validation.
SystemCommand {
/// The command name (e.g. "help", "model", "version").
command: String,
/// Arguments to the command.
args: Vec<String>,
},
}
impl Submission {
@@ -238,6 +290,7 @@ impl Submission {
| Self::Heartbeat
| Self::Summarize
| Self::Suggest
| Self::SystemCommand { .. }
)
}
}
@@ -504,6 +557,84 @@ mod tests {
);
}
#[test]
fn test_parser_system_command_help() {
let submission = SubmissionParser::parse("/help");
assert!(
matches!(submission, Submission::SystemCommand { command, args } if command == "help" && args.is_empty())
);
let submission = SubmissionParser::parse("/?");
assert!(
matches!(submission, Submission::SystemCommand { command, .. } if command == "help")
);
let submission = SubmissionParser::parse("/HELP");
assert!(
matches!(submission, Submission::SystemCommand { command, .. } if command == "help")
);
}
#[test]
fn test_parser_system_command_model() {
// No args: show current model
let submission = SubmissionParser::parse("/model");
assert!(
matches!(submission, Submission::SystemCommand { command, args } if command == "model" && args.is_empty())
);
// With args: switch model
let submission = SubmissionParser::parse("/model gpt-4o");
assert!(
matches!(submission, Submission::SystemCommand { command, args } if command == "model" && args == vec!["gpt-4o"])
);
// Case insensitive command, preserves arg case
let submission = SubmissionParser::parse("/MODEL Claude-3.5");
assert!(
matches!(submission, Submission::SystemCommand { command, args } if command == "model" && args == vec!["Claude-3.5"])
);
}
#[test]
fn test_parser_system_command_version() {
let submission = SubmissionParser::parse("/version");
assert!(
matches!(submission, Submission::SystemCommand { command, args } if command == "version" && args.is_empty())
);
}
#[test]
fn test_parser_system_command_tools() {
let submission = SubmissionParser::parse("/tools");
assert!(
matches!(submission, Submission::SystemCommand { command, args } if command == "tools" && args.is_empty())
);
}
#[test]
fn test_parser_system_command_ping() {
let submission = SubmissionParser::parse("/ping");
assert!(
matches!(submission, Submission::SystemCommand { command, args } if command == "ping" && args.is_empty())
);
}
#[test]
fn test_parser_system_command_debug() {
let submission = SubmissionParser::parse("/debug");
assert!(
matches!(submission, Submission::SystemCommand { command, args } if command == "debug" && args.is_empty())
);
}
#[test]
fn test_parser_system_command_is_control() {
let submission = SubmissionParser::parse("/help");
assert!(submission.is_control());
assert!(!submission.starts_turn());
}
#[test]
fn test_parser_quit() {
assert!(matches!(SubmissionParser::parse("/quit"), Submission::Quit));
src/agent/worker.rs
+52 -4
View File
@@ -272,7 +272,10 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
));
}
}
RespondResult::ToolCalls(tool_calls) => {
RespondResult::ToolCalls {
tool_calls,
content,
} => {
// Model returned tool calls - execute them
tracing::debug!(
"Job {} respond_with_tools returned {} tool calls",
@@ -280,6 +283,14 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
tool_calls.len()
);
// 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 {
let result = self.execute_tool(&tc.name, &tc.arguments).await;
@@ -417,14 +428,51 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
.into());
}
// Execute with timeout and timing
tracing::debug!(
tool = %tool_name,
params = %params,
job = %job_id,
"Tool call started"
);
// Execute with per-tool timeout and timing
let tool_timeout = tool.execution_timeout();
let start = std::time::Instant::now();
let result = tokio::time::timeout(Duration::from_secs(60), async {
let result = tokio::time::timeout(tool_timeout, async {
tool.execute(params.clone(), &job_ctx).await
})
.await;
let elapsed = start.elapsed();
match &result {
Ok(Ok(output)) => {
let result_str = serde_json::to_string(&output.result)
.unwrap_or_else(|_| "<serialize error>".to_string());
tracing::debug!(
tool = %tool_name,
elapsed_ms = elapsed.as_millis() as u64,
result = %result_str,
"Tool call succeeded"
);
}
Ok(Err(e)) => {
tracing::debug!(
tool = %tool_name,
elapsed_ms = elapsed.as_millis() as u64,
error = %e,
"Tool call failed"
);
}
Err(_) => {
tracing::debug!(
tool = %tool_name,
elapsed_ms = elapsed.as_millis() as u64,
timeout_secs = tool_timeout.as_secs(),
"Tool call timed out"
);
}
}
// Record action in memory and get the ActionRecord for persistence
let action = match &result {
Ok(Ok(output)) => {
@@ -479,7 +527,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
let output = result
.map_err(|_| crate::error::ToolError::Timeout {
name: tool_name.to_string(),
timeout: Duration::from_secs(60),
timeout: tool_timeout,
})?
.map_err(|e| crate::error::ToolError::ExecutionFailed {
name: tool_name.to_string(),
src/bootstrap.rs
+325
View File
@@ -0,0 +1,325 @@
//! Bootstrap configuration for IronClaw.
//!
//! These are the only settings that MUST live on disk because they're needed
//! before the database connection is established. Everything else lives in the
//! `settings` table in PostgreSQL.
//!
//! File: `~/.ironclaw/bootstrap.json`
use std::path::PathBuf;
use serde::{Deserialize, Serialize};
use crate::settings::KeySource;
/// Minimal config needed to connect to the database and decrypt secrets.
///
/// This is the only JSON file IronClaw reads from disk at startup.
/// All other configuration lives in the `settings` table in PostgreSQL.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BootstrapConfig {
/// Database connection URL (postgres://...).
#[serde(default)]
pub database_url: Option<String>,
/// Database connection pool size.
#[serde(default)]
pub database_pool_size: Option<usize>,
/// Source for the secrets master key.
#[serde(default)]
pub secrets_master_key_source: KeySource,
/// Whether onboarding wizard has been completed.
#[serde(default)]
pub onboard_completed: bool,
}
impl Default for BootstrapConfig {
fn default() -> Self {
Self {
database_url: None,
database_pool_size: None,
secrets_master_key_source: KeySource::None,
onboard_completed: false,
}
}
}
impl BootstrapConfig {
/// Default bootstrap file path: `~/.ironclaw/bootstrap.json`.
pub fn default_path() -> PathBuf {
dirs::home_dir()
.unwrap_or_else(|| PathBuf::from("."))
.join(".ironclaw")
.join("bootstrap.json")
}
/// Legacy settings.json path (for migration detection).
pub fn legacy_settings_path() -> PathBuf {
dirs::home_dir()
.unwrap_or_else(|| PathBuf::from("."))
.join(".ironclaw")
.join("settings.json")
}
/// Load from the default path, falling back to legacy settings.json,
/// then to defaults if neither exists.
pub fn load() -> Self {
let bootstrap_path = Self::default_path();
if bootstrap_path.exists() {
return Self::load_from(&bootstrap_path);
}
// Fall back to legacy settings.json (extract just the 4 bootstrap fields)
let legacy_path = Self::legacy_settings_path();
if legacy_path.exists() {
return Self::load_from_legacy(&legacy_path);
}
Self::default()
}
/// Load from a specific path.
pub fn load_from(path: &PathBuf) -> Self {
match std::fs::read_to_string(path) {
Ok(data) => serde_json::from_str(&data).unwrap_or_default(),
Err(_) => Self::default(),
}
}
/// Extract bootstrap fields from a legacy settings.json.
fn load_from_legacy(path: &PathBuf) -> Self {
match std::fs::read_to_string(path) {
Ok(data) => {
// The legacy Settings struct is a superset; serde will ignore extra fields.
serde_json::from_str(&data).unwrap_or_default()
}
Err(_) => Self::default(),
}
}
/// Save to the default path.
pub fn save(&self) -> std::io::Result<()> {
self.save_to(&Self::default_path())
}
/// Save to a specific path.
pub fn save_to(&self, path: &PathBuf) -> std::io::Result<()> {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let json = serde_json::to_string_pretty(self)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string()))?;
std::fs::write(path, json)
}
}
/// One-time migration from disk config files to the database settings table.
///
/// On first boot after upgrade, checks if:
/// 1. `~/.ironclaw/settings.json` exists
/// 2. The DB settings table is empty for this user
///
/// If both conditions hold, migrates settings, MCP servers, and session data
/// to the database, writes `bootstrap.json`, and renames old files to `.migrated`.
pub async fn migrate_disk_to_db(
store: &crate::history::Store,
user_id: &str,
) -> Result<(), MigrationError> {
let legacy_settings_path = BootstrapConfig::legacy_settings_path();
if !legacy_settings_path.exists() {
tracing::debug!("No legacy settings.json found, skipping disk-to-DB migration");
return Ok(());
}
// Only migrate if DB is empty for this user
let has_settings = store.has_settings(user_id).await.map_err(|e| {
MigrationError::Database(format!("Failed to check existing settings: {}", e))
})?;
if has_settings {
tracing::debug!(
"DB already has settings for user '{}', skipping migration",
user_id
);
return Ok(());
}
tracing::info!("Migrating disk settings to database...");
// 1. Load and migrate settings.json
let settings = crate::settings::Settings::load_from(&legacy_settings_path);
let db_map = settings.to_db_map();
if !db_map.is_empty() {
store
.set_all_settings(user_id, &db_map)
.await
.map_err(|e| {
MigrationError::Database(format!("Failed to write settings to DB: {}", e))
})?;
tracing::info!("Migrated {} settings to database", db_map.len());
}
// 2. Write bootstrap.json with the 4 essential fields
let bootstrap = BootstrapConfig {
database_url: settings.database_url.clone(),
database_pool_size: settings.database_pool_size,
secrets_master_key_source: settings.secrets_master_key_source,
onboard_completed: settings.onboard_completed,
};
bootstrap
.save()
.map_err(|e| MigrationError::Io(format!("Failed to write bootstrap.json: {}", e)))?;
tracing::info!("Wrote bootstrap.json");
// 3. Migrate mcp-servers.json if it exists
let ironclaw_dir = dirs::home_dir()
.unwrap_or_else(|| PathBuf::from("."))
.join(".ironclaw");
let mcp_path = ironclaw_dir.join("mcp-servers.json");
if mcp_path.exists() {
match std::fs::read_to_string(&mcp_path) {
Ok(content) => match serde_json::from_str::<serde_json::Value>(&content) {
Ok(value) => {
store
.set_setting(user_id, "mcp_servers", &value)
.await
.map_err(|e| {
MigrationError::Database(format!(
"Failed to write MCP servers to DB: {}",
e
))
})?;
tracing::info!("Migrated mcp-servers.json to database");
rename_to_migrated(&mcp_path);
}
Err(e) => {
tracing::warn!("Failed to parse mcp-servers.json: {}", e);
}
},
Err(e) => {
tracing::warn!("Failed to read mcp-servers.json: {}", e);
}
}
}
// 4. Migrate session.json if it exists
let session_path = ironclaw_dir.join("session.json");
if session_path.exists() {
match std::fs::read_to_string(&session_path) {
Ok(content) => match serde_json::from_str::<serde_json::Value>(&content) {
Ok(value) => {
store
.set_setting(user_id, "nearai.session", &value)
.await
.map_err(|e| {
MigrationError::Database(format!(
"Failed to write session to DB: {}",
e
))
})?;
tracing::info!("Migrated session.json to database");
rename_to_migrated(&session_path);
}
Err(e) => {
tracing::warn!("Failed to parse session.json: {}", e);
}
},
Err(e) => {
tracing::warn!("Failed to read session.json: {}", e);
}
}
}
// 5. Rename settings.json to .migrated (don't delete, safety net)
rename_to_migrated(&legacy_settings_path);
tracing::info!("Disk-to-DB migration complete");
Ok(())
}
/// Rename a file to `<name>.migrated` as a safety net.
fn rename_to_migrated(path: &PathBuf) {
let mut migrated = path.as_os_str().to_owned();
migrated.push(".migrated");
if let Err(e) = std::fs::rename(path, &migrated) {
tracing::warn!("Failed to rename {} to .migrated: {}", path.display(), e);
}
}
/// Errors that can occur during disk-to-DB migration.
#[derive(Debug, thiserror::Error)]
pub enum MigrationError {
#[error("Database error: {0}")]
Database(String),
#[error("IO error: {0}")]
Io(String),
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
#[test]
fn test_bootstrap_save_load() {
let dir = tempdir().unwrap();
let path = dir.path().join("bootstrap.json");
let config = BootstrapConfig {
database_url: Some("postgres://localhost/test".to_string()),
database_pool_size: Some(5),
secrets_master_key_source: KeySource::Keychain,
onboard_completed: true,
};
config.save_to(&path).unwrap();
let loaded = BootstrapConfig::load_from(&path);
assert_eq!(
loaded.database_url,
Some("postgres://localhost/test".to_string())
);
assert_eq!(loaded.database_pool_size, Some(5));
assert_eq!(loaded.secrets_master_key_source, KeySource::Keychain);
assert!(loaded.onboard_completed);
}
#[test]
fn test_bootstrap_from_legacy_settings() {
let dir = tempdir().unwrap();
let path = dir.path().join("settings.json");
// Write a legacy settings.json with many extra fields
let legacy = serde_json::json!({
"database_url": "postgres://localhost/ironclaw",
"database_pool_size": 10,
"secrets_master_key_source": "keychain",
"onboard_completed": true,
"selected_model": "claude-3-5-sonnet",
"agent": { "name": "testbot", "max_parallel_jobs": 3 },
"heartbeat": { "enabled": true }
});
std::fs::write(&path, serde_json::to_string_pretty(&legacy).unwrap()).unwrap();
let config = BootstrapConfig::load_from_legacy(&path);
assert_eq!(
config.database_url,
Some("postgres://localhost/ironclaw".to_string())
);
assert_eq!(config.database_pool_size, Some(10));
assert_eq!(config.secrets_master_key_source, KeySource::Keychain);
assert!(config.onboard_completed);
}
#[test]
fn test_bootstrap_defaults() {
let config = BootstrapConfig::default();
assert!(config.database_url.is_none());
assert!(config.database_pool_size.is_none());
assert_eq!(config.secrets_master_key_source, KeySource::None);
assert!(!config.onboard_completed);
}
}
src/channels/channel.rs
+19
View File
@@ -114,6 +114,12 @@ pub enum StatusUpdate {
StreamChunk(String),
/// General status message.
Status(String),
/// A sandbox job has started (shown as a clickable card in the UI).
JobStarted {
job_id: String,
title: String,
browse_url: String,
},
/// Tool requires user approval before execution.
ApprovalNeeded {
request_id: String,
@@ -121,6 +127,19 @@ pub enum StatusUpdate {
description: String,
parameters: serde_json::Value,
},
/// Extension needs user authentication (token or OAuth).
AuthRequired {
extension_name: String,
instructions: Option<String>,
auth_url: Option<String>,
setup_url: Option<String>,
},
/// Extension authentication completed.
AuthCompleted {
extension_name: String,
success: bool,
message: String,
},
}
/// Trait for message channels.
src/channels/repl.rs
+52 -2
View File
@@ -42,12 +42,25 @@ const SLASH_COMMANDS: &[&str] = &[
"/quit",
"/exit",
"/debug",
"/model",
"/undo",
"/redo",
"/clear",
"/compact",
"/new",
"/interrupt",
"/version",
"/tools",
"/ping",
"/job",
"/status",
"/cancel",
"/list",
"/heartbeat",
"/summarize",
"/suggest",
"/thread",
"/resume",
];
/// Rustyline helper for slash-command tab completion.
@@ -295,10 +308,11 @@ impl Channel for ReplChannel {
continue;
}
// Handle local REPL commands
// Handle local REPL commands (only commands that need
// immediate local handling stay here)
match line.to_lowercase().as_str() {
"/quit" | "/exit" => break,
"/help" | "/?" => {
"/help" => {
print_help();
continue;
}
@@ -413,6 +427,15 @@ impl Channel for ReplChannel {
print!("{chunk}");
let _ = io::stdout().flush();
}
StatusUpdate::JobStarted {
job_id,
title,
browse_url,
} => {
eprintln!(
" \x1b[36m[job]\x1b[0m {title} \x1b[90m({job_id})\x1b[0m \x1b[4m{browse_url}\x1b[0m"
);
}
StatusUpdate::Status(msg) => {
if debug || msg.contains("approval") || msg.contains("Approval") {
eprintln!(" \x1b[90m{msg}\x1b[0m");
@@ -472,6 +495,33 @@ impl Channel for ReplChannel {
eprintln!(" {bot_border}");
eprintln!();
}
StatusUpdate::AuthRequired {
extension_name,
instructions,
setup_url,
..
} => {
eprintln!();
eprintln!("\x1b[33m Authentication required for {extension_name}\x1b[0m");
if let Some(ref instr) = instructions {
eprintln!(" {instr}");
}
if let Some(ref url) = setup_url {
eprintln!(" \x1b[4m{url}\x1b[0m");
}
eprintln!();
}
StatusUpdate::AuthCompleted {
extension_name,
success,
message,
} => {
if success {
eprintln!("\x1b[32m {extension_name}: {message}\x1b[0m");
} else {
eprintln!("\x1b[31m {extension_name}: {message}\x1b[0m");
}
}
}
Ok(())
}
src/channels/wasm/bundled.rs
+109 -48
View File
@@ -1,68 +1,125 @@
//! Bundled WASM channels that can be installed locally.
//! Known WASM channels that can be installed from build artifacts.
//!
//! Instead of embedding WASM binaries in the host binary via include_bytes!,
//! channels are compiled separately and installed from their build output
//! directories during onboarding.
//!
//! Channel source layout:
//! channels-src/<name>/
//! target/wasm32-wasip2/release/<name>_channel.wasm
//! <name>.capabilities.json
use std::path::Path;
use std::path::{Path, PathBuf};
use tokio::fs;
#[derive(Clone, Copy)]
struct BundledChannel {
name: &'static str,
wasm: &'static [u8],
capabilities: &'static [u8],
/// Compile-time project root, used to locate channels-src/ in dev builds.
const CARGO_MANIFEST_DIR: &str = env!("CARGO_MANIFEST_DIR");
/// Known channel names and their crate names (for locating build artifacts).
const KNOWN_CHANNELS: &[(&str, &str)] = &[
("telegram", "telegram_channel"),
("slack", "slack_channel"),
("whatsapp", "whatsapp_channel"),
];
/// Names of known channels that can be installed.
pub fn bundled_channel_names() -> Vec<&'static str> {
KNOWN_CHANNELS.iter().map(|(name, _)| *name).collect()
}
/// Names of bundled channels shipped with IronClaw.
pub fn bundled_channel_names() -> &'static [&'static str] {
&["telegram"]
/// Resolve the channels source directory.
///
/// Checks (in order):
/// 1. `IRONCLAW_CHANNELS_SRC` env var
/// 2. `<CARGO_MANIFEST_DIR>/channels-src/` (dev builds)
fn channels_src_dir() -> PathBuf {
if let Ok(dir) = std::env::var("IRONCLAW_CHANNELS_SRC") {
return PathBuf::from(dir);
}
PathBuf::from(CARGO_MANIFEST_DIR).join("channels-src")
}
/// Install a bundled channel into a channels directory.
/// Locate the build artifacts for a channel.
///
/// Returns (wasm_path, capabilities_path) or an error if files are missing.
fn locate_channel_artifacts(name: &str) -> Result<(PathBuf, PathBuf), String> {
let (_, crate_name) = KNOWN_CHANNELS
.iter()
.find(|(n, _)| *n == name)
.ok_or_else(|| format!("Unknown channel '{}'", name))?;
let src_dir = channels_src_dir();
let channel_dir = src_dir.join(name);
let wasm_path = channel_dir
.join("target/wasm32-wasip2/release")
.join(format!("{}.wasm", crate_name));
let caps_path = channel_dir.join(format!("{}.capabilities.json", name));
if !wasm_path.exists() {
return Err(format!(
"Channel '{}' WASM not found at {}. Build it first:\n \
cd {} && cargo build --target wasm32-wasip2 --release",
name,
wasm_path.display(),
channel_dir.display()
));
}
if !caps_path.exists() {
return Err(format!(
"Channel '{}' capabilities not found at {}",
name,
caps_path.display()
));
}
Ok((wasm_path, caps_path))
}
/// Install a channel from build artifacts into the channels directory.
pub async fn install_bundled_channel(
name: &str,
target_dir: &Path,
force: bool,
) -> Result<(), String> {
let channel = bundled_channel(name)
.ok_or_else(|| format!("Unknown bundled channel '{}'", name.to_lowercase()))?;
let (wasm_src, caps_src) = locate_channel_artifacts(name)?;
fs::create_dir_all(target_dir)
.await
.map_err(|e| format!("Failed to create channels directory: {}", e))?;
let wasm_path = target_dir.join(format!("{}.wasm", channel.name));
let caps_path = target_dir.join(format!("{}.capabilities.json", channel.name));
let wasm_dst = target_dir.join(format!("{}.wasm", name));
let caps_dst = target_dir.join(format!("{}.capabilities.json", name));
let has_existing = wasm_path.exists() || caps_path.exists();
let has_existing = wasm_dst.exists() || caps_dst.exists();
if has_existing && !force {
return Err(format!(
"Channel '{}' already exists at {}",
channel.name,
name,
target_dir.display()
));
}
fs::write(&wasm_path, channel.wasm)
fs::copy(&wasm_src, &wasm_dst)
.await
.map_err(|e| format!("Failed to write {}: {}", wasm_path.display(), e))?;
fs::write(&caps_path, channel.capabilities)
.map_err(|e| format!("Failed to copy {}: {}", wasm_src.display(), e))?;
fs::copy(&caps_src, &caps_dst)
.await
.map_err(|e| format!("Failed to write {}: {}", caps_path.display(), e))?;
.map_err(|e| format!("Failed to copy {}: {}", caps_src.display(), e))?;
Ok(())
}
fn bundled_channel(name: &str) -> Option<BundledChannel> {
if name.eq_ignore_ascii_case("telegram") {
Some(BundledChannel {
name: "telegram",
wasm: include_bytes!("../../../channels-src/telegram/telegram.wasm"),
capabilities: include_bytes!(
"../../../channels-src/telegram/telegram.capabilities.json"
),
})
} else {
None
}
/// Check which known channels have build artifacts available.
pub fn available_channel_names() -> Vec<&'static str> {
KNOWN_CHANNELS
.iter()
.filter(|(name, _)| locate_channel_artifacts(name).is_ok())
.map(|(name, _)| *name)
.collect()
}
#[cfg(test)]
@@ -73,31 +130,35 @@ mod tests {
use super::*;
#[test]
fn test_bundled_channel_names_contains_telegram() {
assert!(bundled_channel_names().contains(&"telegram"));
fn test_known_channels_includes_all_three() {
let names = bundled_channel_names();
assert!(names.contains(&"telegram"));
assert!(names.contains(&"slack"));
assert!(names.contains(&"whatsapp"));
}
#[test]
fn test_channels_src_dir_default() {
let dir = channels_src_dir();
assert!(dir.ends_with("channels-src"));
}
#[test]
fn test_locate_unknown_channel_errors() {
assert!(locate_channel_artifacts("nonexistent").is_err());
}
#[tokio::test]
async fn test_install_bundled_channel_writes_files() {
let dir = tempdir().unwrap();
install_bundled_channel("telegram", dir.path(), false)
.await
.unwrap();
assert!(dir.path().join("telegram.wasm").exists());
assert!(dir.path().join("telegram.capabilities.json").exists());
}
#[tokio::test]
async fn test_install_bundled_channel_refuses_overwrite_without_force() {
async fn test_install_refuses_overwrite_without_force() {
let dir = tempdir().unwrap();
let wasm_path = dir.path().join("telegram.wasm");
fs::write(&wasm_path, b"custom").await.unwrap();
let result = install_bundled_channel("telegram", dir.path(), false).await;
// Either fails because artifacts missing OR because file exists
assert!(result.is_err());
// Original file should be untouched
let existing = fs::read(&wasm_path).await.unwrap();
assert_eq!(existing, b"custom");
}
src/channels/wasm/mod.rs
+1 -1
View File
@@ -89,7 +89,7 @@ mod schema;
mod wrapper;
// Core types
pub use bundled::{bundled_channel_names, install_bundled_channel};
pub use bundled::{available_channel_names, bundled_channel_names, install_bundled_channel};
pub use capabilities::{ChannelCapabilities, EmitRateLimitConfig, HttpEndpointConfig, PollConfig};
pub use error::WasmChannelError;
pub use host::{ChannelEmitRateLimiter, ChannelHostState, EmittedMessage};
src/channels/wasm/wrapper.rs
+134 -6
View File
@@ -141,6 +141,22 @@ impl ChannelStoreData {
result
}
/// Replace injected credential values with `[REDACTED]` in text.
///
/// Prevents credentials from leaking through error messages, logs, or
/// return values to WASM. reqwest::Error includes the full URL in its
/// Display output, so any error from an injected-URL request will
/// contain the raw credential unless we scrub it.
fn redact_credentials(&self, text: &str) -> String {
let mut result = text.to_string();
for (name, value) in &self.credentials {
if !value.is_empty() {
result = result.replace(value, &format!("[REDACTED:{}]", name));
}
}
result
}
}
// Implement WasiView to provide WASI context and resource table
@@ -187,6 +203,7 @@ impl near::agent::channel_host::Host for ChannelStoreData {
url: String,
headers_json: String,
body: Option<Vec<u8>>,
timeout_ms: Option<u32>,
) -> Result<near::agent::channel_host::HttpResponse, String> {
tracing::info!(
method = %method,
@@ -276,12 +293,21 @@ impl near::agent::channel_host::Host for ChannelStoreData {
request = request.body(body_bytes);
}
// Send request with timeout
let response = request
.timeout(std::time::Duration::from_secs(30))
.send()
.await
.map_err(|e| format!("HTTP request failed: {}", e))?;
// Send request with caller-specified timeout (default 30s).
// Cap at callback_timeout to prevent outliving the host wrapper.
let timeout = std::time::Duration::from_millis(timeout_ms.unwrap_or(30_000) as u64);
let response = request.timeout(timeout).send().await.map_err(|e| {
// Walk the full error chain so we get the actual root cause
// (DNS, TLS, connection refused, etc.) instead of just
// "error sending request for url (...)".
let mut chain = format!("HTTP request failed: {}", e);
let mut source = std::error::Error::source(&e);
while let Some(cause) = source {
chain.push_str(&format!(" -> {}", cause));
source = cause.source();
}
chain
})?;
let status = response.status().as_u16();
let response_headers: std::collections::HashMap<String, String> = response
@@ -330,6 +356,11 @@ impl near::agent::channel_host::Host for ChannelStoreData {
})
});
// Scrub credential values from error messages before logging or returning
// to WASM. reqwest::Error includes the full URL (with injected credentials)
// in its Display output.
let result = result.map_err(|e| self.redact_credentials(&e));
match &result {
Ok(resp) => {
tracing::info!(status = resp.status, "http_request completed successfully");
@@ -1858,6 +1889,29 @@ fn status_to_wit(status: &StatusUpdate, metadata: &serde_json::Value) -> wit_cha
message: format!("Approval needed: {} - {}", tool_name, description),
metadata_json,
},
StatusUpdate::JobStarted { job_id, title, .. } => wit_channel::StatusUpdate {
status: wit_channel::StatusType::Thinking,
message: format!("Job started: {} ({})", title, job_id),
metadata_json,
},
StatusUpdate::AuthRequired { extension_name, .. } => wit_channel::StatusUpdate {
status: wit_channel::StatusType::Thinking,
message: format!("Auth required: {}", extension_name),
metadata_json,
},
StatusUpdate::AuthCompleted {
extension_name,
success,
..
} => wit_channel::StatusUpdate {
status: wit_channel::StatusType::Thinking,
message: format!(
"Auth {}: {}",
if *success { "completed" } else { "failed" },
extension_name
),
metadata_json,
},
}
}
@@ -2350,4 +2404,78 @@ mod tests {
assert_eq!(cloned.message, "hello");
assert_eq!(cloned.metadata_json, "{\"a\":1}");
}
#[test]
fn test_redact_credentials_replaces_values() {
use super::ChannelStoreData;
let mut creds = std::collections::HashMap::new();
creds.insert(
"TELEGRAM_BOT_TOKEN".to_string(),
"8218490433:AAEZeUxwqZ5OO3mOCXv7fKvpdhDgsmBBNis".to_string(),
);
creds.insert("OTHER_SECRET".to_string(), "s3cret".to_string());
let store =
ChannelStoreData::new(1024 * 1024, "test", ChannelCapabilities::default(), creds);
let error = "HTTP request failed: error sending request for url \
(https://api.telegram.org/bot8218490433:AAEZeUxwqZ5OO3mOCXv7fKvpdhDgsmBBNis/getUpdates)";
let redacted = store.redact_credentials(error);
assert!(
!redacted.contains("8218490433:AAEZeUxwqZ5OO3mOCXv7fKvpdhDgsmBBNis"),
"credential value should be redacted"
);
assert!(
redacted.contains("[REDACTED:TELEGRAM_BOT_TOKEN]"),
"redacted text should contain placeholder name"
);
assert!(
!redacted.contains("s3cret"),
"other credentials should also be redacted"
);
}
#[test]
fn test_redact_credentials_no_op_without_credentials() {
use super::ChannelStoreData;
let store = ChannelStoreData::new(
1024 * 1024,
"test",
ChannelCapabilities::default(),
std::collections::HashMap::new(),
);
let input = "some error message";
assert_eq!(store.redact_credentials(input), input);
}
#[test]
fn test_redact_credentials_skips_empty_values() {
use super::ChannelStoreData;
let mut creds = std::collections::HashMap::new();
creds.insert("EMPTY_TOKEN".to_string(), String::new());
let store =
ChannelStoreData::new(1024 * 1024, "test", ChannelCapabilities::default(), creds);
let input = "should not match anything";
assert_eq!(store.redact_credentials(input), input);
}
/// Verify that the block_on-inside-spawn_blocking pattern used by the WASM
/// channel HTTP host function doesn't deadlock or panic.
#[tokio::test]
async fn test_block_on_inside_spawn_blocking_does_not_deadlock() {
let result = tokio::task::spawn_blocking(|| {
tokio::runtime::Handle::current().block_on(async { 42 })
})
.await
.expect("spawn_blocking panicked");
assert_eq!(result, 42);
}
}
src/channels/web/mod.rs
+98 -24
View File
@@ -10,7 +10,7 @@
//! ◄── GET /api/chat/events ── SSE stream
//! ─── GET /api/chat/ws ─────► WebSocket (bidirectional)
//! ─── GET /api/memory/* ────► Workspace
//! ─── GET /api/jobs/* ──────► ContextManager
//! ─── GET /api/jobs/* ──────► Database
//! ◄── GET / ───────────────── Static HTML/CSS/JS
//! ```
@@ -31,9 +31,10 @@ use tokio_stream::wrappers::ReceiverStream;
use crate::agent::SessionManager;
use crate::channels::{Channel, IncomingMessage, MessageStream, OutgoingResponse, StatusUpdate};
use crate::config::GatewayConfig;
use crate::context::ContextManager;
use crate::error::ChannelError;
use crate::extensions::ExtensionManager;
use crate::history::Store;
use crate::orchestrator::job_manager::ContainerJobManager;
use crate::tools::ToolRegistry;
use crate::workspace::Workspace;
@@ -70,11 +71,13 @@ impl GatewayChannel {
msg_tx: tokio::sync::RwLock::new(None),
sse: SseManager::new(),
workspace: None,
context_manager: None,
session_manager: None,
log_broadcaster: None,
extension_manager: None,
tool_registry: None,
store: None,
job_manager: None,
prompt_queue: None,
user_id: config.user_id.clone(),
shutdown_tx: tokio::sync::RwLock::new(None),
ws_tracker: Some(Arc::new(ws::WsConnectionTracker::new())),
@@ -93,11 +96,13 @@ impl GatewayChannel {
msg_tx: tokio::sync::RwLock::new(None),
sse: SseManager::new(),
workspace: self.state.workspace.clone(),
context_manager: self.state.context_manager.clone(),
session_manager: self.state.session_manager.clone(),
log_broadcaster: self.state.log_broadcaster.clone(),
extension_manager: self.state.extension_manager.clone(),
tool_registry: self.state.tool_registry.clone(),
store: self.state.store.clone(),
job_manager: self.state.job_manager.clone(),
prompt_queue: self.state.prompt_queue.clone(),
user_id: self.state.user_id.clone(),
shutdown_tx: tokio::sync::RwLock::new(None),
ws_tracker: self.state.ws_tracker.clone(),
@@ -112,12 +117,6 @@ impl GatewayChannel {
self
}
/// Inject the context manager for the jobs API.
pub fn with_context_manager(mut self, cm: Arc<ContextManager>) -> Self {
self.rebuild_state(|s| s.context_manager = Some(cm));
self
}
/// Inject the session manager for thread/session info.
pub fn with_session_manager(mut self, sm: Arc<SessionManager>) -> Self {
self.rebuild_state(|s| s.session_manager = Some(sm));
@@ -142,6 +141,34 @@ impl GatewayChannel {
self
}
/// Inject the database store for sandbox job persistence.
pub fn with_store(mut self, store: Arc<Store>) -> Self {
self.rebuild_state(|s| s.store = Some(store));
self
}
/// Inject the container job manager for sandbox operations.
pub fn with_job_manager(mut self, jm: Arc<ContainerJobManager>) -> Self {
self.rebuild_state(|s| s.job_manager = Some(jm));
self
}
/// Inject the prompt queue for Claude Code follow-up prompts.
pub fn with_prompt_queue(
mut self,
pq: Arc<
tokio::sync::Mutex<
std::collections::HashMap<
uuid::Uuid,
std::collections::VecDeque<crate::orchestrator::api::PendingPrompt>,
>,
>,
>,
) -> Self {
self.rebuild_state(|s| s.prompt_queue = Some(pq));
self
}
/// Get the auth token (for printing to console on startup).
pub fn auth_token(&self) -> &str {
&self.auth_token
@@ -173,11 +200,7 @@ impl Channel for GatewayChannel {
),
})?;
let bound_addr =
server::start_server(addr, self.state.clone(), self.auth_token.clone()).await?;
tracing::info!("Web gateway listening on http://{}", bound_addr);
tracing::info!("Auth token: {}", self.auth_token);
server::start_server(addr, self.state.clone(), self.auth_token.clone()).await?;
Ok(Box::pin(ReceiverStream::new(rx)))
}
@@ -200,17 +223,48 @@ impl Channel for GatewayChannel {
async fn send_status(
&self,
status: StatusUpdate,
_metadata: &serde_json::Value,
metadata: &serde_json::Value,
) -> Result<(), ChannelError> {
let thread_id = metadata
.get("thread_id")
.and_then(|v| v.as_str())
.map(String::from);
let event = match status {
StatusUpdate::Thinking(msg) => SseEvent::Thinking { message: msg },
StatusUpdate::ToolStarted { name } => SseEvent::ToolStarted { name },
StatusUpdate::ToolCompleted { name, success } => {
SseEvent::ToolCompleted { name, success }
}
StatusUpdate::ToolResult { name, preview } => SseEvent::ToolResult { name, preview },
StatusUpdate::StreamChunk(content) => SseEvent::StreamChunk { content },
StatusUpdate::Status(msg) => SseEvent::Status { message: msg },
StatusUpdate::Thinking(msg) => SseEvent::Thinking {
message: msg,
thread_id: thread_id.clone(),
},
StatusUpdate::ToolStarted { name } => SseEvent::ToolStarted {
name,
thread_id: thread_id.clone(),
},
StatusUpdate::ToolCompleted { name, success } => SseEvent::ToolCompleted {
name,
success,
thread_id: thread_id.clone(),
},
StatusUpdate::ToolResult { name, preview } => SseEvent::ToolResult {
name,
preview,
thread_id: thread_id.clone(),
},
StatusUpdate::StreamChunk(content) => SseEvent::StreamChunk {
content,
thread_id: thread_id.clone(),
},
StatusUpdate::Status(msg) => SseEvent::Status {
message: msg,
thread_id: thread_id.clone(),
},
StatusUpdate::JobStarted {
job_id,
title,
browse_url,
} => SseEvent::JobStarted {
job_id,
title,
browse_url,
},
StatusUpdate::ApprovalNeeded {
request_id,
tool_name,
@@ -223,6 +277,26 @@ impl Channel for GatewayChannel {
parameters: serde_json::to_string_pretty(&parameters)
.unwrap_or_else(|_| parameters.to_string()),
},
StatusUpdate::AuthRequired {
extension_name,
instructions,
auth_url,
setup_url,
} => SseEvent::AuthRequired {
extension_name,
instructions,
auth_url,
setup_url,
},
StatusUpdate::AuthCompleted {
extension_name,
success,
message,
} => SseEvent::AuthCompleted {
extension_name,
success,
message,
},
};
self.state.sse.broadcast(event);
src/channels/web/server.rs
+1333 -96
View File
File diff suppressed because it is too large Load Diff
src/channels/web/sse.rs
+12 -2
View File
@@ -79,7 +79,15 @@ impl SseManager {
SseEvent::StreamChunk { .. } => "stream_chunk",
SseEvent::Status { .. } => "status",
SseEvent::ApprovalNeeded { .. } => "approval_needed",
SseEvent::AuthRequired { .. } => "auth_required",
SseEvent::AuthCompleted { .. } => "auth_completed",
SseEvent::Error { .. } => "error",
SseEvent::JobStarted { .. } => "job_started",
SseEvent::JobMessage { .. } => "job_message",
SseEvent::JobToolUse { .. } => "job_tool_use",
SseEvent::JobToolResult { .. } => "job_tool_result",
SseEvent::JobStatus { .. } => "job_status",
SseEvent::JobResult { .. } => "job_result",
SseEvent::Heartbeat => "heartbeat",
};
Ok(Event::default().event(event_type).data(data))
@@ -152,13 +160,14 @@ mod tests {
manager.broadcast(SseEvent::Status {
message: "test".to_string(),
thread_id: None,
});
let event = rx.next().await;
assert!(event.is_some());
let event = event.unwrap().unwrap();
match event {
SseEvent::Status { message } => assert_eq!(message, "test"),
SseEvent::Status { message, .. } => assert_eq!(message, "test"),
_ => panic!("unexpected event type"),
}
}
@@ -172,11 +181,12 @@ mod tests {
manager.broadcast(SseEvent::Thinking {
message: "working".to_string(),
thread_id: None,
});
let event = stream.next().await.unwrap();
match event {
SseEvent::Thinking { message } => assert_eq!(message, "working"),
SseEvent::Thinking { message, .. } => assert_eq!(message, "working"),
_ => panic!("Expected Thinking event"),
}
}
src/channels/web/static/app.js
+1358 -39
View File
File diff suppressed because it is too large Load Diff
src/channels/web/static/index.html
+82 -7
View File
@@ -10,12 +10,19 @@
<body>
<!-- Auth Screen -->
<div id="auth-screen">
<h1>IronClaw</h1>
<div class="auth-form">
<input type="password" id="token-input" placeholder="Auth token" autofocus>
<button onclick="authenticate()">Connect</button>
<div class="auth-card-login">
<div class="auth-brand">
<h1>IronClaw</h1>
<p class="auth-tagline">Secure AI Assistant</p>
</div>
<div class="auth-form">
<label for="token-input">Gateway Token</label>
<input type="password" id="token-input" placeholder="Paste your auth token" autofocus>
<button onclick="authenticate()">Connect</button>
</div>
<div id="auth-error"></div>
<p class="auth-hint">Enter the GATEWAY_AUTH_TOKEN from your .env configuration.</p>
</div>
<div id="auth-error"></div>
</div>
<!-- Main App (hidden until authenticated) -->
@@ -26,16 +33,33 @@
<button data-tab="memory">Memory</button>
<button data-tab="jobs">Jobs</button>
<button data-tab="logs">Logs</button>
<button data-tab="routines">Routines</button>
<button data-tab="extensions">Extensions</button>
<div class="spacer"></div>
<div class="status">
<div class="status" id="gateway-status-trigger">
<div class="dot" id="sse-dot"></div>
<span id="sse-status">Connected</span>
<div class="gateway-popover" id="gateway-popover"></div>
</div>
</div>
<!-- Chat Tab -->
<div class="tab-panel active" id="tab-chat">
<div class="thread-sidebar" id="thread-sidebar">
<div class="thread-sidebar-header">
<span>Threads</span>
<button class="thread-new-btn" onclick="createNewThread()" title="New thread (Ctrl/Cmd+N)">+</button>
<button class="thread-toggle-btn" id="thread-toggle-btn" onclick="toggleThreadSidebar()" title="Toggle sidebar">&laquo;</button>
</div>
<div class="assistant-item" id="assistant-thread" onclick="switchToAssistant()">
<span class="assistant-label">Assistant</span>
<span class="assistant-meta" id="assistant-meta"></span>
</div>
<div class="threads-section-header">
<span>Conversations</span>
</div>
<div class="thread-list" id="thread-list"></div>
</div>
<div class="chat-container">
<div class="chat-messages" id="chat-messages"></div>
<div class="chat-status" id="chat-status"></div>
@@ -56,10 +80,20 @@
<div class="memory-tree" id="memory-tree"></div>
</div>
<div class="memory-content">
<div class="memory-breadcrumb" id="memory-breadcrumb">workspace /</div>
<div class="memory-breadcrumb" id="memory-breadcrumb">
<span id="memory-breadcrumb-path">workspace /</span>
<button class="memory-edit-btn" id="memory-edit-btn" style="display:none" onclick="startMemoryEdit()">Edit</button>
</div>
<div class="memory-viewer" id="memory-viewer">
<div class="empty">Select a file to view its contents</div>
</div>
<div class="memory-editor" id="memory-editor" style="display:none">
<textarea id="memory-edit-textarea"></textarea>
<div class="memory-editor-actions">
<button class="btn-save" onclick="saveMemoryEdit()">Save</button>
<button class="btn-cancel-edit" onclick="cancelMemoryEdit()">Cancel</button>
</div>
</div>
</div>
</div>
</div>
@@ -73,6 +107,7 @@
<tr>
<th>ID</th>
<th>Title</th>
<th>Source</th>
<th>Status</th>
<th>Created</th>
<th>Actions</th>
@@ -104,9 +139,48 @@
</div>
</div>
<!-- Routines Tab -->
<div class="tab-panel" id="tab-routines">
<div class="routines-container">
<div class="routines-summary" id="routines-summary"></div>
<table class="routines-table" id="routines-table">
<thead>
<tr>
<th>Name</th>
<th>Trigger</th>
<th>Action</th>
<th>Last Run</th>
<th>Next Run</th>
<th>Runs</th>
<th>Status</th>
<th>Actions</th>
</tr>
</thead>
<tbody id="routines-tbody"></tbody>
</table>
<div class="empty-state" id="routines-empty" style="display:none">
No routines configured. Ask the assistant to create one.
</div>
<div class="routine-detail" id="routine-detail" style="display:none"></div>
</div>
</div>
<!-- Extensions Tab -->
<div class="tab-panel" id="tab-extensions">
<div class="extensions-container">
<div class="extensions-section">
<h3>Install Extension</h3>
<div class="ext-install-form" id="ext-install-form">
<input type="text" id="ext-install-name" placeholder="Extension name (required)">
<input type="text" id="ext-install-url" placeholder="URL (optional)">
<select id="ext-install-kind">
<option value="mcp_server">MCP Server</option>
<option value="wasm_tool">WASM Tool</option>
<option value="wasm_channel">WASM Channel</option>
</select>
<button onclick="installExtension()">Install</button>
</div>
</div>
<div class="extensions-section">
<h3>Installed Extensions</h3>
<div class="extensions-list" id="extensions-list">
@@ -130,6 +204,7 @@
</div>
</div>
<div id="toasts"></div>
<script src="/app.js"></script>
</body>
</html>
src/channels/web/static/style.css
+1460 -9
View File
File diff suppressed because it is too large Load Diff
src/channels/web/types.rs
+413 -9
View File
@@ -24,10 +24,17 @@ pub struct ThreadInfo {
pub turn_count: usize,
pub created_at: String,
pub updated_at: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub title: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub thread_type: Option<String>,
}
#[derive(Debug, Serialize)]
pub struct ThreadListResponse {
/// The pinned assistant thread (always present after first load).
pub assistant_thread: Option<ThreadInfo>,
/// Regular conversation threads.
pub threads: Vec<ThreadInfo>,
pub active_thread: Option<Uuid>,
}
@@ -54,6 +61,12 @@ pub struct ToolCallInfo {
pub struct HistoryResponse {
pub thread_id: Uuid,
pub turns: Vec<TurnInfo>,
/// Whether there are older messages available.
#[serde(default)]
pub has_more: bool,
/// Cursor for the next page (ISO8601 timestamp of the oldest message returned).
#[serde(skip_serializing_if = "Option::is_none")]
pub oldest_timestamp: Option<String>,
}
// --- Approval ---
@@ -63,6 +76,8 @@ pub struct ApprovalRequest {
pub request_id: String,
/// "approve", "always", or "deny"
pub action: String,
/// Thread that owns the pending approval (so the agent loop finds the right session).
pub thread_id: Option<String>,
}
// --- SSE Event Types ---
@@ -73,17 +88,49 @@ pub enum SseEvent {
#[serde(rename = "response")]
Response { content: String, thread_id: String },
#[serde(rename = "thinking")]
Thinking { message: String },
Thinking {
message: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "tool_started")]
ToolStarted { name: String },
ToolStarted {
name: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "tool_completed")]
ToolCompleted { name: String, success: bool },
ToolCompleted {
name: String,
success: bool,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "tool_result")]
ToolResult { name: String, preview: String },
ToolResult {
name: String,
preview: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "stream_chunk")]
StreamChunk { content: String },
StreamChunk {
content: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "status")]
Status { message: String },
Status {
message: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "job_started")]
JobStarted {
job_id: String,
title: String,
browse_url: String,
},
#[serde(rename = "approval_needed")]
ApprovalNeeded {
request_id: String,
@@ -91,10 +138,59 @@ pub enum SseEvent {
description: String,
parameters: String,
},
#[serde(rename = "auth_required")]
AuthRequired {
extension_name: String,
#[serde(skip_serializing_if = "Option::is_none")]
instructions: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
auth_url: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
setup_url: Option<String>,
},
#[serde(rename = "auth_completed")]
AuthCompleted {
extension_name: String,
success: bool,
message: String,
},
#[serde(rename = "error")]
Error { message: String },
Error {
message: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "heartbeat")]
Heartbeat,
// Sandbox job streaming events (worker + Claude Code bridge)
#[serde(rename = "job_message")]
JobMessage {
job_id: String,
role: String,
content: String,
},
#[serde(rename = "job_tool_use")]
JobToolUse {
job_id: String,
tool_name: String,
input: serde_json::Value,
},
#[serde(rename = "job_tool_result")]
JobToolResult {
job_id: String,
tool_name: String,
output: String,
},
#[serde(rename = "job_status")]
JobStatus { job_id: String, message: String },
#[serde(rename = "job_result")]
JobResult {
job_id: String,
status: String,
#[serde(skip_serializing_if = "Option::is_none")]
session_id: Option<String>,
},
}
// --- Memory ---
@@ -188,6 +284,54 @@ pub struct JobSummaryResponse {
pub stuck: usize,
}
#[derive(Debug, Serialize)]
pub struct JobDetailResponse {
pub id: Uuid,
pub title: String,
pub description: String,
pub state: String,
pub user_id: String,
pub created_at: String,
pub started_at: Option<String>,
pub completed_at: Option<String>,
pub elapsed_secs: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub project_dir: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub browse_url: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub job_mode: Option<String>,
pub transitions: Vec<TransitionInfo>,
}
// --- Project Files ---
#[derive(Debug, Serialize)]
pub struct ProjectFileEntry {
pub name: String,
pub path: String,
pub is_dir: bool,
}
#[derive(Debug, Serialize)]
pub struct ProjectFilesResponse {
pub entries: Vec<ProjectFileEntry>,
}
#[derive(Debug, Serialize)]
pub struct ProjectFileReadResponse {
pub path: String,
pub content: String,
}
#[derive(Debug, Serialize)]
pub struct TransitionInfo {
pub from: String,
pub to: String,
pub timestamp: String,
pub reason: Option<String>,
}
// --- Extensions ---
#[derive(Debug, Serialize)]
@@ -262,6 +406,21 @@ impl ActionResponse {
}
}
// --- Auth Token ---
/// Request to submit an auth token for an extension (dedicated endpoint).
#[derive(Debug, Deserialize)]
pub struct AuthTokenRequest {
pub extension_name: String,
pub token: String,
}
/// Request to cancel an in-progress auth flow.
#[derive(Debug, Deserialize)]
pub struct AuthCancelRequest {
pub extension_name: String,
}
// --- WebSocket ---
/// Message sent by a WebSocket client to the server.
@@ -280,7 +439,18 @@ pub enum WsClientMessage {
request_id: String,
/// "approve", "always", or "deny"
action: String,
/// Thread that owns the pending approval.
thread_id: Option<String>,
},
/// Submit an auth token for an extension (bypasses message pipeline).
#[serde(rename = "auth_token")]
AuthToken {
extension_name: String,
token: String,
},
/// Cancel an in-progress auth flow.
#[serde(rename = "auth_cancel")]
AuthCancel { extension_name: String },
/// Client heartbeat ping.
#[serde(rename = "ping")]
Ping,
@@ -317,9 +487,17 @@ impl WsServerMessage {
SseEvent::ToolResult { .. } => "tool_result",
SseEvent::StreamChunk { .. } => "stream_chunk",
SseEvent::Status { .. } => "status",
SseEvent::JobStarted { .. } => "job_started",
SseEvent::ApprovalNeeded { .. } => "approval_needed",
SseEvent::AuthRequired { .. } => "auth_required",
SseEvent::AuthCompleted { .. } => "auth_completed",
SseEvent::Error { .. } => "error",
SseEvent::Heartbeat => "heartbeat",
SseEvent::JobMessage { .. } => "job_message",
SseEvent::JobToolUse { .. } => "job_tool_use",
SseEvent::JobToolResult { .. } => "job_tool_result",
SseEvent::JobStatus { .. } => "job_status",
SseEvent::JobResult { .. } => "job_result",
};
let data = serde_json::to_value(event).unwrap_or(serde_json::Value::Null);
WsServerMessage::Event {
@@ -329,6 +507,96 @@ impl WsServerMessage {
}
}
// --- Routines ---
#[derive(Debug, Serialize)]
pub struct RoutineInfo {
pub id: Uuid,
pub name: String,
pub description: String,
pub enabled: bool,
pub trigger_type: String,
pub trigger_summary: String,
pub action_type: String,
pub last_run_at: Option<String>,
pub next_fire_at: Option<String>,
pub run_count: u64,
pub consecutive_failures: u32,
pub status: String,
}
#[derive(Debug, Serialize)]
pub struct RoutineListResponse {
pub routines: Vec<RoutineInfo>,
}
#[derive(Debug, Serialize)]
pub struct RoutineSummaryResponse {
pub total: u64,
pub enabled: u64,
pub disabled: u64,
pub failing: u64,
pub runs_today: u64,
}
#[derive(Debug, Serialize)]
pub struct RoutineDetailResponse {
pub id: Uuid,
pub name: String,
pub description: String,
pub enabled: bool,
pub trigger: serde_json::Value,
pub action: serde_json::Value,
pub guardrails: serde_json::Value,
pub notify: serde_json::Value,
pub last_run_at: Option<String>,
pub next_fire_at: Option<String>,
pub run_count: u64,
pub consecutive_failures: u32,
pub created_at: String,
pub recent_runs: Vec<RoutineRunInfo>,
}
#[derive(Debug, Serialize)]
pub struct RoutineRunInfo {
pub id: Uuid,
pub trigger_type: String,
pub started_at: String,
pub completed_at: Option<String>,
pub status: String,
pub result_summary: Option<String>,
pub tokens_used: Option<i32>,
}
// --- Settings ---
#[derive(Debug, Serialize)]
pub struct SettingResponse {
pub key: String,
pub value: serde_json::Value,
pub updated_at: String,
}
#[derive(Debug, Serialize)]
pub struct SettingsListResponse {
pub settings: Vec<SettingResponse>,
}
#[derive(Debug, Deserialize)]
pub struct SettingWriteRequest {
pub value: serde_json::Value,
}
#[derive(Debug, Deserialize)]
pub struct SettingsImportRequest {
pub settings: std::collections::HashMap<String, serde_json::Value>,
}
#[derive(Debug, Serialize)]
pub struct SettingsExportResponse {
pub settings: std::collections::HashMap<String, serde_json::Value>,
}
// --- Health ---
#[derive(Debug, Serialize)]
@@ -371,12 +639,36 @@ mod tests {
#[test]
fn test_ws_client_approval_parse() {
let json = r#"{"type":"approval","request_id":"abc-123","action":"approve"}"#;
let json =
r#"{"type":"approval","request_id":"abc-123","action":"approve","thread_id":"t1"}"#;
let msg: WsClientMessage = serde_json::from_str(json).unwrap();
match msg {
WsClientMessage::Approval { request_id, action } => {
WsClientMessage::Approval {
request_id,
action,
thread_id,
} => {
assert_eq!(request_id, "abc-123");
assert_eq!(action, "approve");
assert_eq!(thread_id.as_deref(), Some("t1"));
}
_ => panic!("Expected Approval variant"),
}
}
#[test]
fn test_ws_client_approval_parse_no_thread() {
let json = r#"{"type":"approval","request_id":"abc-123","action":"deny"}"#;
let msg: WsClientMessage = serde_json::from_str(json).unwrap();
match msg {
WsClientMessage::Approval {
request_id,
action,
thread_id,
} => {
assert_eq!(request_id, "abc-123");
assert_eq!(action, "deny");
assert!(thread_id.is_none());
}
_ => panic!("Expected Approval variant"),
}
@@ -437,6 +729,7 @@ mod tests {
fn test_ws_server_from_sse_thinking() {
let sse = SseEvent::Thinking {
message: "reasoning...".to_string(),
thread_id: None,
};
let ws = WsServerMessage::from_sse_event(&sse);
match ws {
@@ -477,4 +770,115 @@ mod tests {
_ => panic!("Expected Event variant"),
}
}
// ---- Auth type tests ----
#[test]
fn test_ws_client_auth_token_parse() {
let json = r#"{"type":"auth_token","extension_name":"notion","token":"sk-123"}"#;
let msg: WsClientMessage = serde_json::from_str(json).unwrap();
match msg {
WsClientMessage::AuthToken {
extension_name,
token,
} => {
assert_eq!(extension_name, "notion");
assert_eq!(token, "sk-123");
}
_ => panic!("Expected AuthToken variant"),
}
}
#[test]
fn test_ws_client_auth_cancel_parse() {
let json = r#"{"type":"auth_cancel","extension_name":"notion"}"#;
let msg: WsClientMessage = serde_json::from_str(json).unwrap();
match msg {
WsClientMessage::AuthCancel { extension_name } => {
assert_eq!(extension_name, "notion");
}
_ => panic!("Expected AuthCancel variant"),
}
}
#[test]
fn test_sse_auth_required_serialize() {
let event = SseEvent::AuthRequired {
extension_name: "notion".to_string(),
instructions: Some("Get your token from...".to_string()),
auth_url: None,
setup_url: Some("https://notion.so/integrations".to_string()),
};
let json = serde_json::to_string(&event).unwrap();
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
assert_eq!(parsed["type"], "auth_required");
assert_eq!(parsed["extension_name"], "notion");
assert_eq!(parsed["instructions"], "Get your token from...");
assert!(parsed.get("auth_url").is_none());
assert_eq!(parsed["setup_url"], "https://notion.so/integrations");
}
#[test]
fn test_sse_auth_completed_serialize() {
let event = SseEvent::AuthCompleted {
extension_name: "notion".to_string(),
success: true,
message: "notion authenticated (3 tools loaded)".to_string(),
};
let json = serde_json::to_string(&event).unwrap();
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
assert_eq!(parsed["type"], "auth_completed");
assert_eq!(parsed["extension_name"], "notion");
assert_eq!(parsed["success"], true);
}
#[test]
fn test_ws_server_from_sse_auth_required() {
let sse = SseEvent::AuthRequired {
extension_name: "openai".to_string(),
instructions: Some("Enter API key".to_string()),
auth_url: None,
setup_url: None,
};
let ws = WsServerMessage::from_sse_event(&sse);
match ws {
WsServerMessage::Event { event_type, data } => {
assert_eq!(event_type, "auth_required");
assert_eq!(data["extension_name"], "openai");
}
_ => panic!("Expected Event variant"),
}
}
#[test]
fn test_ws_server_from_sse_auth_completed() {
let sse = SseEvent::AuthCompleted {
extension_name: "slack".to_string(),
success: false,
message: "Invalid token".to_string(),
};
let ws = WsServerMessage::from_sse_event(&sse);
match ws {
WsServerMessage::Event { event_type, data } => {
assert_eq!(event_type, "auth_completed");
assert_eq!(data["success"], false);
}
_ => panic!("Expected Event variant"),
}
}
#[test]
fn test_auth_token_request_deserialize() {
let json = r#"{"extension_name":"telegram","token":"bot12345"}"#;
let req: AuthTokenRequest = serde_json::from_str(json).unwrap();
assert_eq!(req.extension_name, "telegram");
assert_eq!(req.token, "bot12345");
}
#[test]
fn test_auth_cancel_request_deserialize() {
let json = r#"{"extension_name":"telegram"}"#;
let req: AuthCancelRequest = serde_json::from_str(json).unwrap();
assert_eq!(req.extension_name, "telegram");
}
}
src/channels/web/ws.rs
+73 -3
View File
@@ -170,7 +170,11 @@ async fn handle_client_message(
.await;
}
}
WsClientMessage::Approval { request_id, action } => {
WsClientMessage::Approval {
request_id,
action,
thread_id,
} => {
let (approved, always) = match action.as_str() {
"approve" => (true, false),
"always" => (true, true),
@@ -214,12 +218,71 @@ async fn handle_client_message(
}
};
let msg = IncomingMessage::new("gateway", user_id, content);
let mut msg = IncomingMessage::new("gateway", user_id, content);
if let Some(ref tid) = thread_id {
msg = msg.with_thread(tid);
}
let tx_guard = state.msg_tx.read().await;
if let Some(ref tx) = *tx_guard {
let _ = tx.send(msg).await;
}
}
WsClientMessage::AuthToken {
extension_name,
token,
} => {
if let Some(ref ext_mgr) = state.extension_manager {
match ext_mgr.auth(&extension_name, Some(&token)).await {
Ok(result) if result.status == "authenticated" => {
let msg = match ext_mgr.activate(&extension_name).await {
Ok(r) => format!(
"{} authenticated ({} tools loaded)",
extension_name,
r.tools_loaded.len()
),
Err(e) => format!(
"{} authenticated but activation failed: {}",
extension_name, e
),
};
crate::channels::web::server::clear_auth_mode(state).await;
state
.sse
.broadcast(crate::channels::web::types::SseEvent::AuthCompleted {
extension_name,
success: true,
message: msg,
});
}
Ok(result) => {
state
.sse
.broadcast(crate::channels::web::types::SseEvent::AuthRequired {
extension_name,
instructions: result.instructions,
auth_url: result.auth_url,
setup_url: result.setup_url,
});
}
Err(e) => {
let _ = direct_tx
.send(WsServerMessage::Error {
message: format!("Auth failed: {}", e),
})
.await;
}
}
} else {
let _ = direct_tx
.send(WsServerMessage::Error {
message: "Extension manager not available".to_string(),
})
.await;
}
}
WsClientMessage::AuthCancel { .. } => {
crate::channels::web::server::clear_auth_mode(state).await;
}
WsClientMessage::Ping => {
let _ = direct_tx.send(WsServerMessage::Pong).await;
}
@@ -328,6 +391,7 @@ mod tests {
WsClientMessage::Approval {
request_id: request_id.to_string(),
action: "approve".to_string(),
thread_id: Some("thread-42".to_string()),
},
&state,
"user1",
@@ -338,6 +402,8 @@ mod tests {
let incoming = agent_rx.recv().await.unwrap();
// The content should be a serialized ExecApproval
assert!(incoming.content.contains("ExecApproval"));
// Thread should be forwarded onto the IncomingMessage.
assert_eq!(incoming.thread_id.as_deref(), Some("thread-42"));
}
#[tokio::test]
@@ -349,6 +415,7 @@ mod tests {
WsClientMessage::Approval {
request_id: Uuid::new_v4().to_string(),
action: "maybe".to_string(),
thread_id: None,
},
&state,
"user1",
@@ -374,6 +441,7 @@ mod tests {
WsClientMessage::Approval {
request_id: "not-a-uuid".to_string(),
action: "approve".to_string(),
thread_id: None,
},
&state,
"user1",
@@ -398,11 +466,13 @@ mod tests {
msg_tx: tokio::sync::RwLock::new(msg_tx),
sse: SseManager::new(),
workspace: None,
context_manager: None,
session_manager: None,
log_broadcaster: None,
extension_manager: None,
tool_registry: None,
store: None,
job_manager: None,
prompt_queue: None,
user_id: "test".to_string(),
shutdown_tx: tokio::sync::RwLock::new(None),
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
src/cli/config.rs
+119 -54
View File
@@ -1,6 +1,7 @@
//! Configuration management CLI commands.
//!
//! Commands for viewing and modifying settings.
//! Settings are stored in PostgreSQL (env > DB > default).
use clap::Subcommand;
@@ -36,41 +37,80 @@ pub enum ConfigCommand {
path: String,
},
/// Show the settings file path
/// Show the settings storage info
Path,
}
/// Run a config command.
pub fn run_config_command(cmd: ConfigCommand) -> anyhow::Result<()> {
///
/// Connects to the database to read/write settings. Falls back to disk
/// if the database is not available.
pub async fn run_config_command(cmd: ConfigCommand) -> anyhow::Result<()> {
let _ = dotenvy::dotenv();
// Try to connect to the DB for settings access
let store = match connect_store().await {
Ok(s) => Some(s),
Err(e) => {
eprintln!(
"Warning: Could not connect to database ({}), using disk fallback",
e
);
None
}
};
match cmd {
ConfigCommand::List { filter } => list_settings(filter),
ConfigCommand::Get { path } => get_setting(&path),
ConfigCommand::Set { path, value } => set_setting(&path, &value),
ConfigCommand::Reset { path } => reset_setting(&path),
ConfigCommand::Path => show_path(),
ConfigCommand::List { filter } => list_settings(store.as_ref(), filter).await,
ConfigCommand::Get { path } => get_setting(store.as_ref(), &path).await,
ConfigCommand::Set { path, value } => set_setting(store.as_ref(), &path, &value).await,
ConfigCommand::Reset { path } => reset_setting(store.as_ref(), &path).await,
ConfigCommand::Path => show_path(store.is_some()),
}
}
/// Bootstrap a DB connection for config commands.
async fn connect_store() -> anyhow::Result<crate::history::Store> {
let config = crate::config::Config::from_env().map_err(|e| anyhow::anyhow!("{}", e))?;
let store = crate::history::Store::new(&config.database).await?;
store.run_migrations().await?;
Ok(store)
}
const DEFAULT_USER_ID: &str = "default";
/// Load settings: DB if available, else disk.
async fn load_settings(store: Option<&crate::history::Store>) -> Settings {
if let Some(store) = store {
match store.get_all_settings(DEFAULT_USER_ID).await {
Ok(map) if !map.is_empty() => return Settings::from_db_map(&map),
_ => {}
}
}
Settings::load()
}
/// List all settings.
fn list_settings(filter: Option<String>) -> anyhow::Result<()> {
let settings = Settings::load();
async fn list_settings(
store: Option<&crate::history::Store>,
filter: Option<String>,
) -> anyhow::Result<()> {
let settings = load_settings(store).await;
let all = settings.list();
// Find the longest key for alignment
let max_key_len = all.iter().map(|(k, _)| k.len()).max().unwrap_or(0);
println!("Settings:");
let source = if store.is_some() { "database" } else { "disk" };
println!("Settings (source: {}):", source);
println!();
for (key, value) in all {
// Skip if filter is set and doesn't match
if let Some(ref f) = filter {
if !key.starts_with(f) {
continue;
}
}
// Truncate long values for display
let display_value = if value.len() > 60 {
format!("{}...", &value[..57])
} else {
@@ -84,8 +124,8 @@ fn list_settings(filter: Option<String>) -> anyhow::Result<()> {
}
/// Get a specific setting.
fn get_setting(path: &str) -> anyhow::Result<()> {
let settings = Settings::load();
async fn get_setting(store: Option<&crate::history::Store>, path: &str) -> anyhow::Result<()> {
let settings = load_settings(store).await;
match settings.get(path) {
Some(value) => {
@@ -99,67 +139,92 @@ fn get_setting(path: &str) -> anyhow::Result<()> {
}
/// Set a setting value.
fn set_setting(path: &str, value: &str) -> anyhow::Result<()> {
let mut settings = Settings::load();
async fn set_setting(
store: Option<&crate::history::Store>,
path: &str,
value: &str,
) -> anyhow::Result<()> {
let mut settings = load_settings(store).await;
// Try to set the value
settings
.set(path, value)
.map_err(|e| anyhow::anyhow!("{}", e))?;
// Save to disk
settings.save()?;
// Save to DB if available, otherwise disk
if let Some(store) = store {
let json_value = match serde_json::from_str::<serde_json::Value>(value) {
Ok(v) => v,
Err(_) => serde_json::Value::String(value.to_string()),
};
store
.set_setting(DEFAULT_USER_ID, path, &json_value)
.await
.map_err(|e| anyhow::anyhow!("Failed to save to database: {}", e))?;
} else {
settings.save()?;
}
println!("Set {} = {}", path, value);
Ok(())
}
/// Reset a setting to default.
fn reset_setting(path: &str) -> anyhow::Result<()> {
let mut settings = Settings::load();
// Get the default value for display
async fn reset_setting(store: Option<&crate::history::Store>, path: &str) -> anyhow::Result<()> {
let default = Settings::default();
let default_value = default
.get(path)
.ok_or_else(|| anyhow::anyhow!("Unknown setting: {}", path))?;
// Reset it
settings.reset(path).map_err(|e| anyhow::anyhow!("{}", e))?;
// Save to disk
settings.save()?;
// Delete from DB (falling back to default) or reset on disk
if let Some(store) = store {
store
.delete_setting(DEFAULT_USER_ID, path)
.await
.map_err(|e| anyhow::anyhow!("Failed to delete setting from database: {}", e))?;
} else {
let mut settings = Settings::load();
settings.reset(path).map_err(|e| anyhow::anyhow!("{}", e))?;
settings.save()?;
}
println!("Reset {} to default: {}", path, default_value);
Ok(())
}
/// Show the settings file path.
fn show_path() -> anyhow::Result<()> {
let path = Settings::default_path();
println!("{}", path.display());
if path.exists() {
let metadata = std::fs::metadata(&path)?;
println!(" Size: {} bytes", metadata.len());
if let Ok(modified) = metadata.modified() {
use std::time::SystemTime;
let duration = SystemTime::now()
.duration_since(modified)
.unwrap_or_default();
let secs = duration.as_secs();
if secs < 60 {
println!(" Modified: {} seconds ago", secs);
} else if secs < 3600 {
println!(" Modified: {} minutes ago", secs / 60);
} else if secs < 86400 {
println!(" Modified: {} hours ago", secs / 3600);
} else {
println!(" Modified: {} days ago", secs / 86400);
}
}
/// Show the settings storage info.
fn show_path(has_db: bool) -> anyhow::Result<()> {
if has_db {
println!("Settings stored in: PostgreSQL (settings table)");
println!(
"Bootstrap config: {}",
crate::bootstrap::BootstrapConfig::default_path().display()
);
} else {
println!(" (does not exist, using defaults)");
let path = Settings::default_path();
println!("Settings stored in: {} (disk fallback)", path.display());
if path.exists() {
let metadata = std::fs::metadata(&path)?;
println!(" Size: {} bytes", metadata.len());
if let Ok(modified) = metadata.modified() {
use std::time::SystemTime;
let duration = SystemTime::now()
.duration_since(modified)
.unwrap_or_default();
let secs = duration.as_secs();
if secs < 60 {
println!(" Modified: {} seconds ago", secs);
} else if secs < 3600 {
println!(" Modified: {} minutes ago", secs / 60);
} else if secs < 86400 {
println!(" Modified: {} hours ago", secs / 3600);
} else {
println!(" Modified: {} days ago", secs / 86400);
}
}
} else {
println!(" (does not exist, using defaults)");
}
}
Ok(())
src/cli/mcp.rs
+60 -11
View File
@@ -13,9 +13,7 @@ use crate::secrets::{PostgresSecretsStore, SecretsCrypto, SecretsStore};
use crate::tools::mcp::{
McpClient, McpServerConfig, McpSessionManager, OAuthConfig,
auth::{authorize_mcp_server, is_authenticated},
config::{
add_mcp_server, get_mcp_server, load_mcp_servers, remove_mcp_server, save_mcp_servers,
},
config::{self, McpServersFile},
};
#[derive(Subcommand, Debug, Clone)]
@@ -173,8 +171,11 @@ async fn add_server(
// Validate
config.validate()?;
// Save
add_mcp_server(config).await?;
// Save (DB if available, else disk)
let store = connect_store().await;
let mut servers = load_servers(store.as_ref()).await?;
servers.upsert(config);
save_servers(store.as_ref(), &servers).await?;
println!();
println!(" ✓ Added MCP server '{}'", name);
@@ -192,7 +193,12 @@ async fn add_server(
/// Remove an MCP server.
async fn remove_server(name: String) -> anyhow::Result<()> {
remove_mcp_server(&name).await?;
let store = connect_store().await;
let mut servers = load_servers(store.as_ref()).await?;
if !servers.remove(&name) {
anyhow::bail!("Server '{}' not found", name);
}
save_servers(store.as_ref(), &servers).await?;
println!();
println!(" ✓ Removed MCP server '{}'", name);
@@ -203,7 +209,8 @@ async fn remove_server(name: String) -> anyhow::Result<()> {
/// List configured MCP servers.
async fn list_servers(verbose: bool) -> anyhow::Result<()> {
let servers = load_mcp_servers().await?;
let store = connect_store().await;
let servers = load_servers(store.as_ref()).await?;
if servers.servers.is_empty() {
println!();
@@ -261,7 +268,12 @@ async fn list_servers(verbose: bool) -> anyhow::Result<()> {
/// Authenticate with an MCP server.
async fn auth_server(name: String, user_id: String) -> anyhow::Result<()> {
// Get server config
let server = get_mcp_server(&name).await?;
let store = connect_store().await;
let servers = load_servers(store.as_ref()).await?;
let server = servers
.get(&name)
.cloned()
.ok_or_else(|| anyhow::anyhow!("Server '{}' not found", name))?;
// Initialize secrets store
let secrets = get_secrets_store().await?;
@@ -329,7 +341,12 @@ async fn auth_server(name: String, user_id: String) -> anyhow::Result<()> {
/// Test connection to an MCP server.
async fn test_server(name: String, user_id: String) -> anyhow::Result<()> {
// Get server config
let server = get_mcp_server(&name).await?;
let store = connect_store().await;
let servers = load_servers(store.as_ref()).await?;
let server = servers
.get(&name)
.cloned()
.ok_or_else(|| anyhow::anyhow!("Server '{}' not found", name))?;
println!();
println!(" Testing connection to '{}'...", name);
@@ -420,7 +437,8 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> {
/// Toggle server enabled/disabled state.
async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Result<()> {
let mut servers = load_mcp_servers().await?;
let store = connect_store().await;
let mut servers = load_servers(store.as_ref()).await?;
let server = servers
.get_mut(&name)
@@ -435,7 +453,7 @@ async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Res
};
server.enabled = new_state;
save_mcp_servers(&servers).await?;
save_servers(store.as_ref(), &servers).await?;
let status = if new_state { "enabled" } else { "disabled" };
println!();
@@ -445,6 +463,37 @@ async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Res
Ok(())
}
const DEFAULT_USER_ID: &str = "default";
/// Try to connect to the database store for DB-backed config.
async fn connect_store() -> Option<Store> {
let config = Config::from_env().ok()?;
let store = Store::new(&config.database).await.ok()?;
store.run_migrations().await.ok()?;
Some(store)
}
/// Load MCP servers (DB if available, else disk).
async fn load_servers(store: Option<&Store>) -> Result<McpServersFile, config::ConfigError> {
if let Some(store) = store {
config::load_mcp_servers_from_db(store, DEFAULT_USER_ID).await
} else {
config::load_mcp_servers().await
}
}
/// Save MCP servers (DB if available, else disk).
async fn save_servers(
store: Option<&Store>,
servers: &McpServersFile,
) -> Result<(), config::ConfigError> {
if let Some(store) = store {
config::save_mcp_servers_to_db(store, DEFAULT_USER_ID, servers).await
} else {
config::save_mcp_servers(servers).await
}
}
/// Initialize and return the secrets store.
async fn get_secrets_store() -> anyhow::Result<Arc<dyn SecretsStore + Send + Sync>> {
let config = Config::from_env()?;
src/cli/mod.rs
+36
View File
@@ -88,6 +88,42 @@ pub enum Command {
/// Show system health and diagnostics
Status,
/// Run as a sandboxed worker inside a Docker container (internal use).
/// This is invoked automatically by the orchestrator, not by users directly.
Worker {
/// Job ID to execute.
#[arg(long)]
job_id: uuid::Uuid,
/// URL of the orchestrator's internal API.
#[arg(long, default_value = "http://host.docker.internal:50051")]
orchestrator_url: String,
/// Maximum iterations before stopping.
#[arg(long, default_value = "50")]
max_iterations: u32,
},
/// Run as a Claude Code bridge inside a Docker container (internal use).
/// Spawns the `claude` CLI and streams output back to the orchestrator.
ClaudeBridge {
/// Job ID to execute.
#[arg(long)]
job_id: uuid::Uuid,
/// URL of the orchestrator's internal API.
#[arg(long, default_value = "http://host.docker.internal:50051")]
orchestrator_url: String,
/// Maximum agentic turns for Claude Code.
#[arg(long, default_value = "50")]
max_turns: u32,
/// Claude model to use (e.g. "sonnet", "opus").
#[arg(long, default_value = "sonnet")]
model: String,
},
}
impl Cli {
src/config.rs
+219 -121
View File
@@ -1,4 +1,9 @@
//! Configuration for IronClaw.
//!
//! Settings are loaded with priority: env var > database > default.
//! The database replaces the old `settings.json` file for all settings
//! except the 4 bootstrap fields (database_url, pool_size, secrets key
//! source, onboard_completed) which live in `~/.ironclaw/bootstrap.json`.
use std::path::PathBuf;
use std::time::Duration;
@@ -6,6 +11,7 @@ use std::time::Duration;
use secrecy::{ExposeSecret, SecretString};
use crate::error::ConfigError;
use crate::settings::Settings;
/// Main configuration for the agent.
#[derive(Debug, Clone)]
@@ -21,28 +27,67 @@ pub struct Config {
pub secrets: SecretsConfig,
pub builder: BuilderModeConfig,
pub heartbeat: HeartbeatConfig,
pub routines: RoutineConfig,
pub sandbox: SandboxModeConfig,
pub claude_code: ClaudeCodeConfig,
}
impl Config {
/// Load configuration from environment variables.
pub fn from_env() -> Result<Self, ConfigError> {
// Load .env file if present (ignore errors if not found)
/// Load configuration from environment variables and the database.
///
/// Priority: env var > DB settings > default.
/// This is the primary way to load config after DB is connected.
pub async fn from_db(
store: &crate::history::Store,
user_id: &str,
bootstrap: &crate::bootstrap::BootstrapConfig,
) -> Result<Self, ConfigError> {
let _ = dotenvy::dotenv();
// Load all settings from DB into a Settings struct
let db_settings = match store.get_all_settings(user_id).await {
Ok(map) => Settings::from_db_map(&map),
Err(e) => {
tracing::warn!("Failed to load settings from DB, using defaults: {}", e);
Settings::default()
}
};
Self::build(bootstrap, &db_settings)
}
/// Load configuration from environment variables only (no database).
///
/// Used during early startup before the database is connected,
/// and by CLI commands that don't have DB access.
/// Falls back to legacy `settings.json` on disk if present.
pub fn from_env() -> Result<Self, ConfigError> {
let _ = dotenvy::dotenv();
let bootstrap = crate::bootstrap::BootstrapConfig::load();
let settings = Settings::load();
Self::build(&bootstrap, &settings)
}
/// Build config from bootstrap + settings (shared by from_env and from_db).
fn build(
bootstrap: &crate::bootstrap::BootstrapConfig,
settings: &Settings,
) -> Result<Self, ConfigError> {
Ok(Self {
database: DatabaseConfig::from_env()?,
llm: LlmConfig::from_env()?,
embeddings: EmbeddingsConfig::from_env()?,
tunnel: TunnelConfig::from_env()?,
channels: ChannelsConfig::from_env()?,
agent: AgentConfig::from_env()?,
safety: SafetyConfig::from_env()?,
wasm: WasmConfig::from_env()?,
secrets: SecretsConfig::from_env()?,
builder: BuilderModeConfig::from_env()?,
heartbeat: HeartbeatConfig::from_env()?,
sandbox: SandboxModeConfig::from_env()?,
database: DatabaseConfig::resolve(bootstrap)?,
llm: LlmConfig::resolve(settings)?,
embeddings: EmbeddingsConfig::resolve(settings)?,
tunnel: TunnelConfig::resolve(settings)?,
channels: ChannelsConfig::resolve(settings)?,
agent: AgentConfig::resolve(settings)?,
safety: SafetyConfig::resolve()?,
wasm: WasmConfig::resolve()?,
secrets: SecretsConfig::resolve(bootstrap)?,
builder: BuilderModeConfig::resolve()?,
heartbeat: HeartbeatConfig::resolve(settings)?,
routines: RoutineConfig::resolve()?,
sandbox: SandboxModeConfig::resolve()?,
claude_code: ClaudeCodeConfig::resolve()?,
})
}
}
@@ -51,48 +96,17 @@ impl Config {
///
/// Used by channels and tools that need public webhook endpoints.
/// The tunnel URL is shared across all channels (Telegram, Slack, etc.).
///
/// # Security Notes
///
/// **Webhook endpoints** (e.g., `/webhook/telegram`) should NOT use tunnel-level
/// authentication because webhook providers (Telegram, Slack, GitHub) need
/// unauthenticated access to POST updates. Security for webhooks comes from:
/// - Webhook signature verification (provider-specific secrets)
/// - IP allowlisting (if supported by provider)
///
/// **Non-webhook endpoints** (admin APIs, health checks) CAN be protected using
/// tunnel provider features:
/// - ngrok: Basic Auth, OAuth, IP restrictions
/// - Cloudflare: Access policies, mTLS
///
/// These protections are configured in the tunnel provider, not here.
///
/// # Supported Providers
///
/// - **ngrok**: `ngrok http 8080` -> `https://abc123.ngrok.io`
/// - **Cloudflare Tunnel**: `cloudflared tunnel --url http://localhost:8080`
/// - **localtunnel**: `lt --port 8080`
/// - Any service that provides a public HTTPS URL to localhost
#[derive(Debug, Clone, Default)]
pub struct TunnelConfig {
/// Public URL from tunnel provider (e.g., "https://abc123.ngrok.io").
///
/// When set, channels that support webhooks will register their endpoints
/// with this base URL instead of using polling.
pub public_url: Option<String>,
}
impl TunnelConfig {
fn from_env() -> Result<Self, ConfigError> {
// Priority: env var > settings file
let public_url = optional_env("TUNNEL_URL")?.or_else(|| {
crate::settings::Settings::load()
.tunnel
.public_url
.filter(|s| !s.is_empty())
});
fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
let public_url = optional_env("TUNNEL_URL")?
.or_else(|| settings.tunnel.public_url.clone().filter(|s| !s.is_empty()));
// Validate URL format if provided
if let Some(ref url) = public_url {
if !url.starts_with("https://") {
return Err(ConfigError::InvalidValue {
@@ -111,8 +125,6 @@ impl TunnelConfig {
}
/// Get the webhook URL for a given path.
///
/// Returns `None` if no tunnel is configured.
pub fn webhook_url(&self, path: &str) -> Option<String> {
self.public_url.as_ref().map(|base| {
let base = base.trim_end_matches('/');
@@ -130,18 +142,14 @@ pub struct DatabaseConfig {
}
impl DatabaseConfig {
fn from_env() -> Result<Self, ConfigError> {
let settings = crate::settings::Settings::load();
// Priority: env var > settings > error (required)
fn resolve(bootstrap: &crate::bootstrap::BootstrapConfig) -> Result<Self, ConfigError> {
let url = optional_env("DATABASE_URL")?
.or(settings.database_url.clone())
.or_else(|| bootstrap.database_url.clone())
.ok_or_else(|| ConfigError::MissingRequired {
key: "database_url".to_string(),
hint: "Run 'ironclaw onboard' or set DATABASE_URL environment variable".to_string(),
})?;
// Priority: env var > settings > default
let pool_size = optional_env("DATABASE_POOL_SIZE")?
.map(|s| s.parse())
.transpose()
@@ -149,7 +157,7 @@ impl DatabaseConfig {
key: "DATABASE_POOL_SIZE".to_string(),
message: format!("must be a positive integer: {e}"),
})?
.or(settings.database_pool_size)
.or(bootstrap.database_pool_size)
.unwrap_or(10);
Ok(Self {
@@ -215,17 +223,15 @@ pub struct NearAiConfig {
}
impl LlmConfig {
fn from_env() -> Result<Self, ConfigError> {
fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
let api_key = optional_env("NEARAI_API_KEY")?.map(SecretString::from);
// Determine API mode: explicit setting, or infer from API key presence
let api_mode = if let Some(mode_str) = optional_env("NEARAI_API_MODE")? {
mode_str.parse().map_err(|e| ConfigError::InvalidValue {
key: "NEARAI_API_MODE".to_string(),
message: e,
})?
} else if api_key.is_some() {
// If API key is provided, default to chat_completions mode
NearAiApiMode::ChatCompletions
} else {
NearAiApiMode::Responses
@@ -233,10 +239,8 @@ impl LlmConfig {
Ok(Self {
nearai: NearAiConfig {
// Load model from saved settings first, then env, then default
model: crate::settings::Settings::load()
.selected_model
.or_else(|| optional_env("NEARAI_MODEL").ok().flatten())
model: optional_env("NEARAI_MODEL")?
.or_else(|| settings.selected_model.clone())
.unwrap_or_else(|| {
"fireworks::accounts/fireworks/models/llama4-maverick-instruct-basic"
.to_string()
@@ -265,8 +269,6 @@ pub struct EmbeddingsConfig {
/// OpenAI API key (for OpenAI provider).
pub openai_api_key: Option<SecretString>,
/// Model to use for embeddings.
/// For OpenAI: "text-embedding-3-small", "text-embedding-3-large", "text-embedding-ada-002"
/// For NEAR AI: Uses the configured session for auth.
pub model: String,
}
@@ -282,18 +284,15 @@ impl Default for EmbeddingsConfig {
}
impl EmbeddingsConfig {
fn from_env() -> Result<Self, ConfigError> {
let settings = crate::settings::Settings::load();
fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
let openai_api_key = optional_env("OPENAI_API_KEY")?.map(SecretString::from);
// Priority: env var > settings > default
let provider = optional_env("EMBEDDING_PROVIDER")?
.unwrap_or_else(|| settings.embeddings.provider.clone());
let model =
optional_env("EMBEDDING_MODEL")?.unwrap_or_else(|| settings.embeddings.model.clone());
// Priority: env var > settings > auto-detect from API key
let enabled = optional_env("EMBEDDING_ENABLED")?
.map(|s| s.parse())
.transpose()
@@ -301,10 +300,7 @@ impl EmbeddingsConfig {
key: "EMBEDDING_ENABLED".to_string(),
message: format!("must be 'true' or 'false': {e}"),
})?
.unwrap_or_else(|| {
// Check settings, or auto-enable if API key present
settings.embeddings.enabled || openai_api_key.is_some()
});
.unwrap_or_else(|| settings.embeddings.enabled || openai_api_key.is_some());
Ok(Self {
enabled,
@@ -338,6 +334,8 @@ pub struct ChannelsConfig {
pub wasm_channels_dir: std::path::PathBuf,
/// Whether WASM channels are enabled.
pub wasm_channels_enabled: bool,
/// Telegram owner user ID. When set, the bot only responds to this user.
pub telegram_owner_id: Option<i64>,
}
#[derive(Debug, Clone)]
@@ -364,7 +362,7 @@ pub struct GatewayConfig {
}
impl ChannelsConfig {
fn from_env() -> Result<Self, ConfigError> {
fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
let http = if optional_env("HTTP_PORT")?.is_some() || optional_env("HTTP_HOST")?.is_some() {
Some(HttpConfig {
host: optional_env("HTTP_HOST")?.unwrap_or_else(|| "0.0.0.0".to_string()),
@@ -385,7 +383,7 @@ impl ChannelsConfig {
let gateway = if optional_env("GATEWAY_ENABLED")?
.map(|s| s.to_lowercase() == "true" || s == "1")
.unwrap_or(false)
.unwrap_or(true)
{
Some(GatewayConfig {
host: optional_env("GATEWAY_HOST")?.unwrap_or_else(|| "127.0.0.1".to_string()),
@@ -425,6 +423,14 @@ impl ChannelsConfig {
message: format!("must be 'true' or 'false': {e}"),
})?
.unwrap_or(true),
telegram_owner_id: optional_env("TELEGRAM_OWNER_ID")?
.map(|s| s.parse())
.transpose()
.map_err(|e| ConfigError::InvalidValue {
key: "TELEGRAM_OWNER_ID".to_string(),
message: format!("must be an integer: {e}"),
})?
.or(settings.channels.telegram_owner_id),
})
}
}
@@ -450,14 +456,13 @@ pub struct AgentConfig {
pub use_planning: bool,
/// Session idle timeout. Sessions inactive longer than this are pruned.
pub session_idle_timeout: Duration,
/// Allow chat to use filesystem/shell tools directly (bypass sandbox).
pub allow_local_tools: bool,
}
impl AgentConfig {
fn from_env() -> Result<Self, ConfigError> {
let settings = crate::settings::Settings::load();
fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
Ok(Self {
// Priority: env var > settings > default
name: optional_env("AGENT_NAME")?.unwrap_or_else(|| settings.agent.name.clone()),
max_parallel_jobs: optional_env("AGENT_MAX_PARALLEL_JOBS")?
.map(|s| s.parse())
@@ -523,6 +528,14 @@ impl AgentConfig {
})?
.unwrap_or(settings.agent.session_idle_timeout_secs),
),
allow_local_tools: optional_env("ALLOW_LOCAL_TOOLS")?
.map(|s| s.parse())
.transpose()
.map_err(|e| ConfigError::InvalidValue {
key: "ALLOW_LOCAL_TOOLS".to_string(),
message: format!("must be 'true' or 'false': {e}"),
})?
.unwrap_or(false),
})
}
}
@@ -535,7 +548,7 @@ pub struct SafetyConfig {
}
impl SafetyConfig {
fn from_env() -> Result<Self, ConfigError> {
fn resolve() -> Result<Self, ConfigError> {
Ok(Self {
max_output_length: parse_optional_env("SAFETY_MAX_OUTPUT_LENGTH", 100_000)?,
injection_check_enabled: optional_env("SAFETY_INJECTION_CHECK_ENABLED")?
@@ -573,7 +586,6 @@ pub struct WasmConfig {
#[derive(Clone, Default)]
pub struct SecretsConfig {
/// Master key for encrypting secrets.
/// Source determined by KeySource in settings.
pub master_key: Option<SecretString>,
/// Whether secrets management is enabled.
pub enabled: bool,
@@ -592,38 +604,28 @@ impl std::fmt::Debug for SecretsConfig {
}
impl SecretsConfig {
fn from_env() -> Result<Self, ConfigError> {
fn resolve(bootstrap: &crate::bootstrap::BootstrapConfig) -> Result<Self, ConfigError> {
use crate::settings::KeySource;
let settings = crate::settings::Settings::load();
// Priority: env var > keychain (based on settings) > disabled
let (master_key, source) = if let Some(env_key) = optional_env("SECRETS_MASTER_KEY")? {
// Env var takes priority (for CI/Docker)
(Some(SecretString::from(env_key)), KeySource::Env)
} else {
match settings.secrets_master_key_source {
KeySource::Keychain => {
// Try to load from OS keychain
match crate::secrets::keychain::get_master_key() {
Ok(key_bytes) => {
let key_hex: String =
key_bytes.iter().map(|b| format!("{:02x}", b)).collect();
(Some(SecretString::from(key_hex)), KeySource::Keychain)
}
Err(_) => {
// Keychain configured but key not found
// This might happen if keychain was cleared
tracing::warn!(
"Secrets configured for keychain but key not found. \
Run 'ironclaw onboard' to reconfigure."
);
(None, KeySource::None)
}
match bootstrap.secrets_master_key_source {
KeySource::Keychain => match crate::secrets::keychain::get_master_key() {
Ok(key_bytes) => {
let key_hex: String =
key_bytes.iter().map(|b| format!("{:02x}", b)).collect();
(Some(SecretString::from(key_hex)), KeySource::Keychain)
}
}
Err(_) => {
tracing::warn!(
"Secrets configured for keychain but key not found. \
Run 'ironclaw onboard' to reconfigure."
);
(None, KeySource::None)
}
},
KeySource::Env => {
// Settings say env, but no env var found
tracing::warn!(
"Secrets configured for env var but SECRETS_MASTER_KEY not set."
);
@@ -635,7 +637,6 @@ impl SecretsConfig {
let enabled = master_key.is_some();
// Validate master key length if provided
if let Some(ref key) = master_key {
if key.expose_secret().len() < 32 {
return Err(ConfigError::InvalidValue {
@@ -681,7 +682,7 @@ fn default_tools_dir() -> PathBuf {
}
impl WasmConfig {
fn from_env() -> Result<Self, ConfigError> {
fn resolve() -> Result<Self, ConfigError> {
Ok(Self {
enabled: optional_env("WASM_ENABLED")?
.map(|s| s.parse())
@@ -752,7 +753,7 @@ pub struct BuilderModeConfig {
impl Default for BuilderModeConfig {
fn default() -> Self {
Self {
enabled: true, // Builder enabled by default
enabled: true,
build_dir: None,
max_iterations: 20,
timeout_secs: 600,
@@ -762,7 +763,7 @@ impl Default for BuilderModeConfig {
}
impl BuilderModeConfig {
fn from_env() -> Result<Self, ConfigError> {
fn resolve() -> Result<Self, ConfigError> {
Ok(Self {
enabled: optional_env("BUILDER_ENABLED")?
.map(|s| s.parse())
@@ -771,7 +772,7 @@ impl BuilderModeConfig {
key: "BUILDER_ENABLED".to_string(),
message: format!("must be 'true' or 'false': {e}"),
})?
.unwrap_or(true), // Builder enabled by default
.unwrap_or(true),
build_dir: optional_env("BUILDER_DIR")?.map(PathBuf::from),
max_iterations: parse_optional_env("BUILDER_MAX_ITERATIONS", 20)?,
timeout_secs: parse_optional_env("BUILDER_TIMEOUT_SECS", 600)?,
@@ -826,11 +827,8 @@ impl Default for HeartbeatConfig {
}
impl HeartbeatConfig {
fn from_env() -> Result<Self, ConfigError> {
let settings = crate::settings::Settings::load();
fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
Ok(Self {
// Priority: env var > settings > default
enabled: optional_env("HEARTBEAT_ENABLED")?
.map(|s| s.parse())
.transpose()
@@ -848,9 +846,55 @@ impl HeartbeatConfig {
})?
.unwrap_or(settings.heartbeat.interval_secs),
notify_channel: optional_env("HEARTBEAT_NOTIFY_CHANNEL")?
.or(settings.heartbeat.notify_channel.clone()),
.or_else(|| settings.heartbeat.notify_channel.clone()),
notify_user: optional_env("HEARTBEAT_NOTIFY_USER")?
.or(settings.heartbeat.notify_user.clone()),
.or_else(|| settings.heartbeat.notify_user.clone()),
})
}
}
/// Routines configuration.
#[derive(Debug, Clone)]
pub struct RoutineConfig {
/// Whether the routines system is enabled.
pub enabled: bool,
/// How often (seconds) to poll for cron routines that need firing.
pub cron_check_interval_secs: u64,
/// Max routines executing concurrently across all users.
pub max_concurrent_routines: usize,
/// Default cooldown between fires (seconds).
pub default_cooldown_secs: u64,
/// Max output tokens for lightweight routine LLM calls.
pub max_lightweight_tokens: u32,
}
impl Default for RoutineConfig {
fn default() -> Self {
Self {
enabled: true,
cron_check_interval_secs: 15,
max_concurrent_routines: 10,
default_cooldown_secs: 300,
max_lightweight_tokens: 4096,
}
}
}
impl RoutineConfig {
fn resolve() -> Result<Self, ConfigError> {
Ok(Self {
enabled: optional_env("ROUTINES_ENABLED")?
.map(|s| s.parse())
.transpose()
.map_err(|e| ConfigError::InvalidValue {
key: "ROUTINES_ENABLED".to_string(),
message: format!("must be 'true' or 'false': {e}"),
})?
.unwrap_or(true),
cron_check_interval_secs: parse_optional_env("ROUTINES_CRON_INTERVAL", 15)?,
max_concurrent_routines: parse_optional_env("ROUTINES_MAX_CONCURRENT", 10)?,
default_cooldown_secs: parse_optional_env("ROUTINES_DEFAULT_COOLDOWN", 300)?,
max_lightweight_tokens: parse_optional_env("ROUTINES_MAX_TOKENS", 4096)?,
})
}
}
@@ -879,7 +923,7 @@ pub struct SandboxModeConfig {
impl Default for SandboxModeConfig {
fn default() -> Self {
Self {
enabled: true, // Enabled by default
enabled: true,
policy: "readonly".to_string(),
timeout_secs: 120,
memory_limit_mb: 2048,
@@ -892,7 +936,7 @@ impl Default for SandboxModeConfig {
}
impl SandboxModeConfig {
fn from_env() -> Result<Self, ConfigError> {
fn resolve() -> Result<Self, ConfigError> {
let extra_domains = optional_env("SANDBOX_EXTRA_DOMAINS")?
.map(|s| s.split(',').map(|d| d.trim().to_string()).collect())
.unwrap_or_default();
@@ -948,6 +992,60 @@ impl SandboxModeConfig {
}
}
/// Claude Code sandbox configuration.
#[derive(Debug, Clone)]
pub struct ClaudeCodeConfig {
/// Whether Claude Code sandbox mode is available.
pub enabled: bool,
/// Host directory containing Claude auth session (mounted read-only).
pub config_dir: std::path::PathBuf,
/// Claude model to use (e.g. "sonnet", "opus").
pub model: String,
/// Maximum agentic turns before stopping.
pub max_turns: u32,
/// Memory limit in MB for Claude Code containers (heavier than workers).
pub memory_limit_mb: u64,
}
impl Default for ClaudeCodeConfig {
fn default() -> Self {
Self {
enabled: false,
config_dir: dirs::home_dir()
.unwrap_or_else(|| std::path::PathBuf::from("."))
.join(".claude"),
model: "sonnet".to_string(),
max_turns: 50,
memory_limit_mb: 4096,
}
}
}
impl ClaudeCodeConfig {
fn resolve() -> Result<Self, ConfigError> {
let defaults = Self::default();
Ok(Self {
enabled: optional_env("CLAUDE_CODE_ENABLED")?
.map(|s| s.parse())
.transpose()
.map_err(|e| ConfigError::InvalidValue {
key: "CLAUDE_CODE_ENABLED".to_string(),
message: format!("must be 'true' or 'false': {e}"),
})?
.unwrap_or(defaults.enabled),
config_dir: optional_env("CLAUDE_CONFIG_DIR")?
.map(std::path::PathBuf::from)
.unwrap_or(defaults.config_dir),
model: optional_env("CLAUDE_CODE_MODEL")?.unwrap_or(defaults.model),
max_turns: parse_optional_env("CLAUDE_CODE_MAX_TURNS", defaults.max_turns)?,
memory_limit_mb: parse_optional_env(
"CLAUDE_CODE_MEMORY_LIMIT_MB",
defaults.memory_limit_mb,
)?,
})
}
}
// Helper functions
fn optional_env(key: &str) -> Result<Option<String>, ConfigError> {
src/error.rs
+53
View File
@@ -39,6 +39,12 @@ pub enum Error {
#[error("Workspace error: {0}")]
Workspace(#[from] WorkspaceError),
#[error("Orchestrator error: {0}")]
Orchestrator(#[from] OrchestratorError),
#[error("Worker error: {0}")]
Worker(#[from] WorkerError),
}
/// Configuration-related errors.
@@ -308,5 +314,52 @@ pub enum WorkspaceError {
HeartbeatError { reason: String },
}
/// Orchestrator errors (internal API, container management).
#[derive(Debug, thiserror::Error)]
pub enum OrchestratorError {
#[error("Container creation failed for job {job_id}: {reason}")]
ContainerCreationFailed { job_id: Uuid, reason: String },
#[error("Container not found for job {job_id}")]
ContainerNotFound { job_id: Uuid },
#[error("Container for job {job_id} is in unexpected state: {state}")]
InvalidContainerState { job_id: Uuid, state: String },
#[error("Worker authentication failed: {reason}")]
AuthFailed { reason: String },
#[error("Internal API error: {reason}")]
ApiError { reason: String },
#[error("Docker error: {reason}")]
Docker { reason: String },
#[error("Job {job_id} timed out in container")]
ContainerTimeout { job_id: Uuid },
}
/// Worker errors (container-side execution).
#[derive(Debug, thiserror::Error)]
pub enum WorkerError {
#[error("Failed to connect to orchestrator at {url}: {reason}")]
ConnectionFailed { url: String, reason: String },
#[error("LLM proxy request failed: {reason}")]
LlmProxyFailed { reason: String },
#[error("Secret resolution failed for {secret_name}: {reason}")]
SecretResolveFailed { secret_name: String, reason: String },
#[error("Orchestrator returned error for job {job_id}: {reason}")]
OrchestratorRejected { job_id: Uuid, reason: String },
#[error("Worker execution failed: {reason}")]
ExecutionFailed { reason: String },
#[error("Missing worker token (IRONCLAW_WORKER_TOKEN not set)")]
MissingToken,
}
/// Result type alias for the agent.
pub type Result<T> = std::result::Result<T, Error>;
src/extensions/manager.rs
+62 -10
View File
@@ -23,9 +23,7 @@ use crate::tools::mcp::auth::{
PkceChallenge, authorize_mcp_server, build_authorization_url, discover_full_oauth_metadata,
find_available_port, is_authenticated, register_client,
};
use crate::tools::mcp::config::{
McpServerConfig, add_mcp_server, get_mcp_server, load_mcp_servers, remove_mcp_server,
};
use crate::tools::mcp::config::McpServerConfig;
use crate::tools::mcp::session::McpSessionManager;
use crate::tools::wasm::{WasmToolLoader, WasmToolRuntime, discover_tools};
@@ -58,6 +56,8 @@ pub struct ExtensionManager {
/// Tunnel URL for remote OAuth callbacks (used in future iterations).
_tunnel_url: Option<String>,
user_id: String,
/// Optional database store for DB-backed MCP config.
store: Option<Arc<crate::history::Store>>,
}
impl ExtensionManager {
@@ -71,6 +71,7 @@ impl ExtensionManager {
wasm_channels_dir: PathBuf,
tunnel_url: Option<String>,
user_id: String,
store: Option<Arc<crate::history::Store>>,
) -> Self {
Self {
registry: ExtensionRegistry::new(),
@@ -85,6 +86,7 @@ impl ExtensionManager {
pending_auth: RwLock::new(HashMap::new()),
_tunnel_url: tunnel_url,
user_id,
store,
}
}
@@ -191,7 +193,7 @@ impl ExtensionManager {
// List MCP servers
if kind_filter.is_none() || kind_filter == Some(ExtensionKind::McpServer) {
match load_mcp_servers().await {
match self.load_mcp_servers().await {
Ok(servers) => {
for server in &servers.servers {
let authenticated =
@@ -304,7 +306,7 @@ impl ExtensionManager {
self.mcp_clients.write().await.remove(name);
// Remove from config
remove_mcp_server(name)
self.remove_mcp_server(name)
.await
.map_err(|e| ExtensionError::Config(e.to_string()))?;
@@ -342,6 +344,54 @@ impl ExtensionManager {
}
}
// ── MCP config helpers (DB with disk fallback) ─────────────────────
async fn load_mcp_servers(
&self,
) -> Result<crate::tools::mcp::config::McpServersFile, crate::tools::mcp::config::ConfigError>
{
if let Some(ref store) = self.store {
crate::tools::mcp::config::load_mcp_servers_from_db(store, &self.user_id).await
} else {
crate::tools::mcp::config::load_mcp_servers().await
}
}
async fn get_mcp_server(
&self,
name: &str,
) -> Result<McpServerConfig, crate::tools::mcp::config::ConfigError> {
let servers = self.load_mcp_servers().await?;
servers.get(name).cloned().ok_or_else(|| {
crate::tools::mcp::config::ConfigError::ServerNotFound {
name: name.to_string(),
}
})
}
async fn add_mcp_server(
&self,
config: McpServerConfig,
) -> Result<(), crate::tools::mcp::config::ConfigError> {
config.validate()?;
if let Some(ref store) = self.store {
crate::tools::mcp::config::add_mcp_server_db(store, &self.user_id, config).await
} else {
crate::tools::mcp::config::add_mcp_server(config).await
}
}
async fn remove_mcp_server(
&self,
name: &str,
) -> Result<(), crate::tools::mcp::config::ConfigError> {
if let Some(ref store) = self.store {
crate::tools::mcp::config::remove_mcp_server_db(store, &self.user_id, name).await
} else {
crate::tools::mcp::config::remove_mcp_server(name).await
}
}
// ── Private helpers ──────────────────────────────────────────────────
async fn install_from_entry(
@@ -381,7 +431,7 @@ impl ExtensionManager {
url: &str,
) -> Result<InstallResult, ExtensionError> {
// Check if already installed
if get_mcp_server(name).await.is_ok() {
if self.get_mcp_server(name).await.is_ok() {
return Err(ExtensionError::AlreadyInstalled(name.to_string()));
}
@@ -390,7 +440,7 @@ impl ExtensionManager {
.validate()
.map_err(|e| ExtensionError::InvalidUrl(e.to_string()))?;
add_mcp_server(config)
self.add_mcp_server(config)
.await
.map_err(|e| ExtensionError::Config(e.to_string()))?;
@@ -465,7 +515,8 @@ impl ExtensionManager {
name: &str,
token: Option<&str>,
) -> Result<AuthResult, ExtensionError> {
let server = get_mcp_server(name)
let server = self
.get_mcp_server(name)
.await
.map_err(|e| ExtensionError::NotInstalled(e.to_string()))?;
@@ -784,7 +835,8 @@ impl ExtensionManager {
}
}
let server = get_mcp_server(name)
let server = self
.get_mcp_server(name)
.await
.map_err(|e| ExtensionError::NotInstalled(e.to_string()))?;
@@ -893,7 +945,7 @@ impl ExtensionManager {
/// Determine what kind of installed extension this is.
async fn determine_installed_kind(&self, name: &str) -> Result<ExtensionKind, ExtensionError> {
// Check MCP servers first
if get_mcp_server(name).await.is_ok() {
if self.get_mcp_server(name).await.is_ok() {
return Ok(ExtensionKind::McpServer);
}
src/history/mod.rs
+4 -1
View File
@@ -9,4 +9,7 @@ mod analytics;
mod store;
pub use analytics::{JobStats, ToolStats};
pub use store::{LlmCallRecord, Store};
pub use store::{
ConversationMessage, ConversationSummary, JobEventRecord, LlmCallRecord, SandboxJobRecord,
SandboxJobSummary, Store,
};
src/history/store.rs
+1117 -6
View File
File diff suppressed because it is too large Load Diff
src/lib.rs
+3
View File
@@ -39,6 +39,7 @@
//! - **Continuous learning** - Improve estimates from historical data
pub mod agent;
pub mod bootstrap;
pub mod channels;
pub mod cli;
pub mod config;
@@ -49,12 +50,14 @@ pub mod evaluation;
pub mod extensions;
pub mod history;
pub mod llm;
pub mod orchestrator;
pub mod safety;
pub mod sandbox;
pub mod secrets;
pub mod settings;
pub mod setup;
pub mod tools;
pub mod worker;
pub mod workspace;
pub use config::Config;
src/llm/mod.rs
+2 -2
View File
@@ -13,8 +13,8 @@ pub mod session;
pub use nearai::{ModelInfo, NearAiProvider};
pub use nearai_chat::NearAiChatProvider;
pub use provider::{
ChatMessage, CompletionRequest, CompletionResponse, LlmProvider, Role, ToolCall,
ToolCompletionRequest, ToolCompletionResponse, ToolDefinition, ToolResult,
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, ModelMetadata,
Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse, ToolDefinition, ToolResult,
};
pub use reasoning::{ActionPlan, Reasoning, ReasoningContext, RespondResult, ToolSelection};
pub use session::{SessionConfig, SessionManager, create_session_manager};
src/llm/nearai.rs
+308 -20
View File
@@ -3,6 +3,7 @@
//! This provider uses the NEAR AI chat-api which provides a unified interface
//! to multiple LLM models (OpenAI, Anthropic, etc.) with user authentication.
use std::collections::HashMap;
use std::sync::Arc;
use async_trait::async_trait;
@@ -31,11 +32,23 @@ pub struct ModelInfo {
pub provider: Option<String>,
}
/// Per-thread chaining state: the last response ID and how many input
/// messages were included in that request. This lets subsequent calls send
/// only the delta (new messages since last call).
struct ChainState {
response_id: String,
input_count: usize,
}
/// NEAR AI Chat API provider.
pub struct NearAiProvider {
client: Client,
config: NearAiConfig,
session: Arc<SessionManager>,
active_model: std::sync::RwLock<String>,
/// Per-thread response ID chaining state.
/// Key is thread_id from request metadata.
response_chains: std::sync::RwLock<HashMap<String, ChainState>>,
}
impl NearAiProvider {
@@ -46,13 +59,64 @@ impl NearAiProvider {
.build()
.unwrap_or_else(|_| Client::new());
let active_model = std::sync::RwLock::new(config.model.clone());
Self {
client,
config,
session,
active_model,
response_chains: std::sync::RwLock::new(HashMap::new()),
}
}
/// Seed a response chain for a thread (e.g. when restoring from DB).
pub fn seed_response_id(&self, thread_id: &str, response_id: String) {
let mut chains = self
.response_chains
.write()
.expect("response_chains lock poisoned");
chains.insert(
thread_id.to_string(),
ChainState {
response_id,
input_count: 0,
},
);
}
/// Get the last response ID for a thread (for persistence).
pub fn get_response_id(&self, thread_id: &str) -> Option<String> {
let chains = self
.response_chains
.read()
.expect("response_chains lock poisoned");
chains.get(thread_id).map(|c| c.response_id.clone())
}
/// Store a response chain state after a successful call.
fn store_chain(&self, thread_id: &str, response_id: String, input_count: usize) {
let mut chains = self
.response_chains
.write()
.expect("response_chains lock poisoned");
chains.insert(
thread_id.to_string(),
ChainState {
response_id,
input_count,
},
);
}
/// Clear the chain for a thread (on error / fallback).
fn clear_chain(&self, thread_id: &str) {
let mut chains = self
.response_chains
.write()
.expect("response_chains lock poisoned");
chains.remove(thread_id);
}
fn api_url(&self, path: &str) -> String {
format!(
"{}/v1/{}",
@@ -291,18 +355,34 @@ impl NearAiProvider {
}
}
/// Split messages into system instructions and non-system input messages.
/// Split messages into system instructions and non-system input items.
/// The OpenAI Responses API expects system prompts in an `instructions` field,
/// not as a message with role "system" in the input array.
fn split_messages(messages: Vec<ChatMessage>) -> (Option<String>, Vec<NearAiMessage>) {
///
/// When `chaining` is true, tool result messages (role=tool) are converted to
/// `NearAiInputItem::FunctionCallOutput` for the Responses API protocol.
fn split_messages(
messages: Vec<ChatMessage>,
chaining: bool,
) -> (Option<String>, Vec<NearAiInputItem>) {
let mut instructions: Vec<String> = Vec::new();
let mut input: Vec<NearAiMessage> = Vec::new();
let mut input: Vec<NearAiInputItem> = Vec::new();
for msg in messages {
if msg.role == Role::System {
instructions.push(msg.content);
} else if chaining && msg.role == Role::Tool {
if let Some(ref call_id) = msg.tool_call_id {
input.push(NearAiInputItem::FunctionCallOutput {
item_type: "function_call_output".to_string(),
call_id: call_id.clone(),
output: msg.content,
});
} else {
input.push(NearAiInputItem::Message(msg.into()));
}
} else {
input.push(msg.into());
input.push(NearAiInputItem::Message(msg.into()));
}
}
@@ -318,12 +398,14 @@ fn split_messages(messages: Vec<ChatMessage>) -> (Option<String>, Vec<NearAiMess
#[async_trait]
impl LlmProvider for NearAiProvider {
async fn complete(&self, req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
let (instructions, input) = split_messages(req.messages);
let thread_id = req.metadata.get("thread_id").cloned();
let (instructions, input) = split_messages(req.messages, false);
let request = NearAiRequest {
model: self.config.model.clone(),
model: self.active_model_name(),
instructions,
input,
previous_response_id: None,
temperature: req.temperature,
max_output_tokens: req.max_tokens,
stream: Some(false),
@@ -350,6 +432,7 @@ impl LlmProvider for NearAiProvider {
finish_reason: FinishReason::Stop,
input_tokens: usage.input_tokens,
output_tokens: usage.output_tokens,
response_id: None,
});
}
@@ -367,6 +450,7 @@ impl LlmProvider for NearAiProvider {
finish_reason: FinishReason::Stop,
input_tokens: 0,
output_tokens: 0,
response_id: None,
});
}
Err(e) => return Err(e),
@@ -423,11 +507,17 @@ impl LlmProvider for NearAiProvider {
);
}
// Store response ID for chaining
if let Some(ref tid) = thread_id {
self.store_chain(tid, response.id.clone(), 0);
}
Ok(CompletionResponse {
content: text,
finish_reason: FinishReason::Stop,
input_tokens: response.usage.input_tokens,
output_tokens: response.usage.output_tokens,
response_id: Some(response.id),
})
}
@@ -435,7 +525,33 @@ impl LlmProvider for NearAiProvider {
&self,
req: ToolCompletionRequest,
) -> Result<ToolCompletionResponse, LlmError> {
let (instructions, input) = split_messages(req.messages);
let thread_id = req.metadata.get("thread_id").cloned();
// Look up chaining state for this thread
let chain_state = thread_id.as_ref().and_then(|tid| {
let chains = self
.response_chains
.read()
.expect("response_chains lock poisoned");
chains
.get(tid)
.map(|c| (c.response_id.clone(), c.input_count))
});
let chaining = chain_state.is_some();
let (previous_response_id, prev_input_count) = chain_state
.map(|(rid, count)| (Some(rid), count))
.unwrap_or((None, 0));
// When chaining, only send new messages (the delta since last call).
// Tool results are converted to function_call_output items.
let (instructions, all_input) = split_messages(req.messages, chaining);
let input = if chaining && all_input.len() > prev_input_count {
all_input[prev_input_count..].to_vec()
} else {
all_input.clone()
};
let total_input_count = all_input.len();
let tools: Vec<NearAiTool> = req
.tools
@@ -449,18 +565,58 @@ impl LlmProvider for NearAiProvider {
.collect();
let request = NearAiRequest {
model: self.config.model.clone(),
instructions,
model: self.active_model_name(),
instructions: if chaining { None } else { instructions.clone() },
input,
previous_response_id: previous_response_id.clone(),
temperature: req.temperature,
max_output_tokens: req.max_tokens,
stream: Some(false),
tools: if tools.is_empty() { None } else { Some(tools) },
tools: if tools.is_empty() {
None
} else {
Some(tools.clone())
},
};
// Try to get structured response, fall back to alternative formats
// Try to get structured response, fall back to alternative formats.
// If chaining fails (bad previous_response_id), retry with full history.
let response: NearAiResponse = match self.send_request("responses", &request).await {
Ok(r) => r,
Err(ref e) if chaining && is_chain_error(e) => {
tracing::warn!(
"Response chaining failed, retrying with full history: {}",
e
);
if let Some(ref tid) = thread_id {
self.clear_chain(tid);
}
let (instructions_full, input_full) = split_messages(
// Rebuild from the original input (non-chaining mode)
{
let mut msgs = Vec::new();
if let Some(ref instr) = instructions {
msgs.push(ChatMessage::system(instr.clone()));
}
for item in &all_input {
msgs.push(item.to_chat_message());
}
msgs
},
false,
);
let retry_request = NearAiRequest {
model: self.active_model_name(),
instructions: instructions_full,
input: input_full,
previous_response_id: None,
temperature: request.temperature,
max_output_tokens: request.max_output_tokens,
stream: Some(false),
tools: request.tools.clone(),
};
self.send_request("responses", &retry_request).await?
}
Err(LlmError::InvalidResponse { reason, .. }) if reason.contains("Raw: ") => {
let raw_text = reason.split("Raw: ").nth(1).unwrap_or("");
@@ -490,6 +646,7 @@ impl LlmProvider for NearAiProvider {
finish_reason,
input_tokens: usage.input_tokens,
output_tokens: usage.output_tokens,
response_id: None,
});
}
@@ -507,6 +664,7 @@ impl LlmProvider for NearAiProvider {
finish_reason: FinishReason::Stop,
input_tokens: 0,
output_tokens: 0,
response_id: None,
});
}
Err(e) => return Err(e),
@@ -560,12 +718,18 @@ impl LlmProvider for NearAiProvider {
FinishReason::ToolUse
};
// Store response ID for chaining on subsequent calls
if let Some(ref tid) = thread_id {
self.store_chain(tid, response.id.clone(), total_input_count);
}
Ok(ToolCompletionResponse {
content: if text.is_empty() { None } else { Some(text) },
tool_calls,
finish_reason,
input_tokens: response.usage.input_tokens,
output_tokens: response.usage.output_tokens,
response_id: Some(response.id),
})
}
@@ -584,6 +748,30 @@ impl LlmProvider for NearAiProvider {
let models = NearAiProvider::list_models(self).await?;
Ok(models.into_iter().map(|m| m.name).collect())
}
fn active_model_name(&self) -> String {
self.active_model
.read()
.expect("active_model lock poisoned")
.clone()
}
fn set_model(&self, model: &str) -> Result<(), LlmError> {
let mut guard = self
.active_model
.write()
.expect("active_model lock poisoned");
*guard = model.to_string();
Ok(())
}
fn seed_response_chain(&self, thread_id: &str, response_id: String) {
self.seed_response_id(thread_id, response_id);
}
fn get_response_chain_id(&self, thread_id: &str) -> Option<String> {
self.get_response_id(thread_id)
}
}
// NEAR AI API types
@@ -597,8 +785,11 @@ struct NearAiRequest {
/// System instructions (replaces sending system role in input)
#[serde(skip_serializing_if = "Option::is_none")]
instructions: Option<String>,
/// Input messages (user/assistant/tool only, NOT system)
input: Vec<NearAiMessage>,
/// Input items: messages and/or function_call_output entries.
input: Vec<NearAiInputItem>,
/// Chain this request to a previous response (avoids resending full context).
#[serde(skip_serializing_if = "Option::is_none")]
previous_response_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
temperature: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
@@ -609,7 +800,7 @@ struct NearAiRequest {
tools: Option<Vec<NearAiTool>>,
}
#[derive(Debug, Serialize, Deserialize)]
#[derive(Debug, Serialize, Deserialize, Clone)]
struct NearAiMessage {
role: String,
content: String,
@@ -630,7 +821,68 @@ impl From<ChatMessage> for NearAiMessage {
}
}
#[derive(Debug, Serialize)]
/// Input item for the Responses API. Either a regular message or a
/// function_call_output (for returning tool results when chaining).
#[derive(Debug, Serialize, Deserialize, Clone)]
#[serde(untagged)]
enum NearAiInputItem {
Message(NearAiMessage),
FunctionCallOutput {
#[serde(rename = "type")]
item_type: String,
call_id: String,
output: String,
},
}
impl NearAiInputItem {
/// Convert back to a ChatMessage (used for fallback retry).
fn to_chat_message(&self) -> ChatMessage {
match self {
NearAiInputItem::Message(msg) => {
let role = match msg.role.as_str() {
"system" => Role::System,
"user" => Role::User,
"assistant" => Role::Assistant,
"tool" => Role::Tool,
_ => Role::User,
};
ChatMessage {
role,
content: msg.content.clone(),
tool_call_id: None,
name: None,
tool_calls: None,
}
}
NearAiInputItem::FunctionCallOutput {
call_id, output, ..
} => ChatMessage {
role: Role::Tool,
content: output.clone(),
tool_call_id: Some(call_id.clone()),
name: None,
tool_calls: None,
},
}
}
}
/// Check if an LLM error is likely caused by an invalid previous_response_id.
fn is_chain_error(err: &LlmError) -> bool {
match err {
LlmError::RequestFailed { reason, .. } => {
let lower = reason.to_lowercase();
lower.contains("previous_response_id")
|| lower.contains("previous response")
|| lower.contains("not found")
|| lower.contains("invalid response id")
}
_ => false,
}
}
#[derive(Debug, Clone, Serialize)]
struct NearAiTool {
#[serde(rename = "type")]
tool_type: String,
@@ -833,14 +1085,17 @@ mod tests {
ChatMessage::user("Hello"),
ChatMessage::assistant("Hi there!"),
];
let (instructions, input) = split_messages(messages);
let (instructions, input) = split_messages(messages, false);
assert_eq!(
instructions,
Some("You are a helpful assistant".to_string())
);
assert_eq!(input.len(), 2);
assert_eq!(input[0].role, "user");
assert_eq!(input[1].role, "assistant");
// Verify the input items are messages
match &input[0] {
NearAiInputItem::Message(m) => assert_eq!(m.role, "user"),
_ => panic!("expected Message"),
}
}
#[test]
@@ -849,7 +1104,7 @@ mod tests {
ChatMessage::user("Hello"),
ChatMessage::assistant("Hi there!"),
];
let (instructions, input) = split_messages(messages);
let (instructions, input) = split_messages(messages, false);
assert!(instructions.is_none());
assert_eq!(input.len(), 2);
}
@@ -861,11 +1116,44 @@ mod tests {
ChatMessage::system("Second instruction"),
ChatMessage::user("Hello"),
];
let (instructions, input) = split_messages(messages);
let (instructions, input) = split_messages(messages, false);
assert_eq!(
instructions,
Some("First instruction\n\nSecond instruction".to_string())
);
assert_eq!(input.len(), 1);
}
#[test]
fn test_split_messages_chaining_converts_tool_results() {
let messages = vec![
ChatMessage::user("Hello"),
ChatMessage::tool_result("call_123", "my_tool", "result data"),
];
let (_, input) = split_messages(messages, true);
assert_eq!(input.len(), 2);
match &input[1] {
NearAiInputItem::FunctionCallOutput {
call_id, output, ..
} => {
assert_eq!(call_id, "call_123");
assert_eq!(output, "result data");
}
_ => panic!("expected FunctionCallOutput"),
}
}
#[test]
fn test_split_messages_no_chaining_keeps_tool_as_message() {
let messages = vec![
ChatMessage::user("Hello"),
ChatMessage::tool_result("call_123", "my_tool", "result data"),
];
let (_, input) = split_messages(messages, false);
assert_eq!(input.len(), 2);
match &input[1] {
NearAiInputItem::Message(m) => assert_eq!(m.role, "tool"),
_ => panic!("expected Message"),
}
}
}
src/llm/nearai_chat.rs
+69 -18
View File
@@ -13,14 +13,15 @@ use serde::{Deserialize, Serialize};
use crate::config::NearAiConfig;
use crate::error::LlmError;
use crate::llm::provider::{
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, Role, ToolCall,
ToolCompletionRequest, ToolCompletionResponse,
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, ModelMetadata,
Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse,
};
/// NEAR AI Chat Completions API provider.
pub struct NearAiChatProvider {
client: Client,
config: NearAiConfig,
active_model: std::sync::RwLock<String>,
}
impl NearAiChatProvider {
@@ -37,7 +38,12 @@ impl NearAiChatProvider {
.build()
.unwrap_or_else(|_| Client::new());
Ok(Self { client, config })
let active_model = std::sync::RwLock::new(config.model.clone());
Ok(Self {
client,
config,
active_model,
})
}
fn api_url(&self, path: &str) -> String {
@@ -65,6 +71,11 @@ impl NearAiChatProvider {
tracing::debug!("Sending request to NEAR AI Chat: {}", url);
// Log the request body for debugging tool call issues
if let Ok(json) = serde_json::to_string(body) {
tracing::debug!("NEAR AI Chat request body: {}", json);
}
let response = self
.client
.post(&url)
@@ -111,8 +122,8 @@ impl NearAiChatProvider {
})
}
/// Fetch available models.
pub async fn list_models(&self) -> Result<Vec<String>, LlmError> {
/// Fetch available models with full metadata from the `/v1/models` endpoint.
async fn fetch_models(&self) -> Result<Vec<ApiModelEntry>, LlmError> {
let url = self.api_url("models");
let response = self
@@ -138,12 +149,7 @@ impl NearAiChatProvider {
#[derive(Deserialize)]
struct ModelsResponse {
data: Vec<ModelEntry>,
}
#[derive(Deserialize)]
struct ModelEntry {
id: String,
data: Vec<ApiModelEntry>,
}
let resp: ModelsResponse =
@@ -152,10 +158,18 @@ impl NearAiChatProvider {
reason: format!("JSON parse error: {}", e),
})?;
Ok(resp.data.into_iter().map(|m| m.id).collect())
Ok(resp.data)
}
}
/// Model entry as returned by the `/v1/models` API.
#[derive(Debug, Deserialize)]
struct ApiModelEntry {
id: String,
#[serde(default)]
context_length: Option<u32>,
}
#[async_trait]
impl LlmProvider for NearAiChatProvider {
async fn complete(&self, req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
@@ -163,7 +177,7 @@ impl LlmProvider for NearAiChatProvider {
req.messages.into_iter().map(|m| m.into()).collect();
let request = ChatCompletionRequest {
model: self.config.model.clone(),
model: self.active_model_name(),
messages,
temperature: req.temperature,
max_tokens: req.max_tokens,
@@ -197,6 +211,7 @@ impl LlmProvider for NearAiChatProvider {
finish_reason,
input_tokens: response.usage.prompt_tokens,
output_tokens: response.usage.completion_tokens,
response_id: None,
})
}
@@ -221,7 +236,7 @@ impl LlmProvider for NearAiChatProvider {
.collect();
let request = ChatCompletionRequest {
model: self.config.model.clone(),
model: self.active_model_name(),
messages,
temperature: req.temperature,
max_tokens: req.max_tokens,
@@ -278,6 +293,7 @@ impl LlmProvider for NearAiChatProvider {
finish_reason,
input_tokens: response.usage.prompt_tokens,
output_tokens: response.usage.completion_tokens,
response_id: None,
})
}
@@ -291,7 +307,34 @@ impl LlmProvider for NearAiChatProvider {
}
async fn list_models(&self) -> Result<Vec<String>, LlmError> {
NearAiChatProvider::list_models(self).await
let models = self.fetch_models().await?;
Ok(models.into_iter().map(|m| m.id).collect())
}
async fn model_metadata(&self) -> Result<ModelMetadata, LlmError> {
let active = self.active_model_name();
let models = self.fetch_models().await?;
let current = models.iter().find(|m| m.id == active);
Ok(ModelMetadata {
id: active,
context_length: current.and_then(|m| m.context_length),
})
}
fn active_model_name(&self) -> String {
self.active_model
.read()
.expect("active_model lock poisoned")
.clone()
}
fn set_model(&self, model: &str) -> Result<(), crate::error::LlmError> {
let mut guard = self
.active_model
.write()
.expect("active_model lock poisoned");
*guard = model.to_string();
Ok(())
}
}
@@ -332,6 +375,7 @@ impl From<ChatMessage> for ChatCompletionMessage {
Role::Assistant => "assistant",
Role::Tool => "tool",
};
let tool_calls = msg.tool_calls.map(|calls| {
calls
.into_iter()
@@ -345,9 +389,16 @@ impl From<ChatMessage> for ChatCompletionMessage {
})
.collect()
});
let content = if role == "assistant" && tool_calls.is_some() && msg.content.is_empty() {
None
} else {
Some(msg.content)
};
Self {
role: role.to_string(),
content: Some(msg.content),
content,
tool_call_id: msg.tool_call_id,
name: msg.name,
tool_calls,
@@ -454,7 +505,7 @@ mod tests {
},
];
let msg = ChatMessage::assistant_with_tool_calls("", tool_calls);
let msg = ChatMessage::assistant_with_tool_calls(None, tool_calls);
let chat_msg: ChatCompletionMessage = msg.into();
assert_eq!(chat_msg.role, "assistant");
@@ -484,7 +535,7 @@ mod tests {
name: "test".to_string(),
arguments: serde_json::json!({"key": "value"}),
};
let msg = ChatMessage::assistant_with_tool_calls("", vec![tc]);
let msg = ChatMessage::assistant_with_tool_calls(None, vec![tc]);
let chat_msg: ChatCompletionMessage = msg.into();
let calls = chat_msg.tool_calls.unwrap();
src/llm/provider.rs
+64 -11
View File
@@ -27,9 +27,8 @@ pub struct ChatMessage {
/// Name of the tool for tool results.
#[serde(skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
/// Tool calls requested by the assistant (for conversation replay).
/// OpenAI-compatible APIs require the assistant message to include
/// tool_calls when followed by tool result messages.
/// Tool calls made by the assistant (OpenAI protocol requires these
/// to appear on the assistant message preceding tool result messages).
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_calls: Option<Vec<ToolCall>>,
}
@@ -68,17 +67,14 @@ impl ChatMessage {
}
}
/// Create an assistant message that requested tool calls.
/// Create an assistant message that includes tool calls.
///
/// OpenAI-compatible APIs require the assistant message to carry the
/// `tool_calls` array when followed by tool-result messages.
pub fn assistant_with_tool_calls(
content: impl Into<String>,
tool_calls: Vec<ToolCall>,
) -> Self {
/// Per the OpenAI protocol, an assistant message with tool_calls must
/// precede the corresponding tool result messages in the conversation.
pub fn assistant_with_tool_calls(content: Option<String>, tool_calls: Vec<ToolCall>) -> Self {
Self {
role: Role::Assistant,
content: content.into(),
content: content.unwrap_or_default(),
tool_call_id: None,
name: None,
tool_calls: if tool_calls.is_empty() {
@@ -112,6 +108,8 @@ pub struct CompletionRequest {
pub max_tokens: Option<u32>,
pub temperature: Option<f32>,
pub stop_sequences: Option<Vec<String>>,
/// Opaque metadata passed through to the provider (e.g. thread_id for chaining).
pub metadata: std::collections::HashMap<String, String>,
}
impl CompletionRequest {
@@ -122,6 +120,7 @@ impl CompletionRequest {
max_tokens: None,
temperature: None,
stop_sequences: None,
metadata: std::collections::HashMap::new(),
}
}
@@ -145,6 +144,8 @@ pub struct CompletionResponse {
pub input_tokens: u32,
pub output_tokens: u32,
pub finish_reason: FinishReason,
/// Provider-specific response ID (e.g. for NEAR AI response chaining).
pub response_id: Option<String>,
}
/// Why the completion finished.
@@ -191,6 +192,8 @@ pub struct ToolCompletionRequest {
pub temperature: Option<f32>,
/// How to handle tool use: "auto", "required", or "none".
pub tool_choice: Option<String>,
/// Opaque metadata passed through to the provider (e.g. thread_id for chaining).
pub metadata: std::collections::HashMap<String, String>,
}
impl ToolCompletionRequest {
@@ -202,6 +205,7 @@ impl ToolCompletionRequest {
max_tokens: None,
temperature: None,
tool_choice: None,
metadata: std::collections::HashMap::new(),
}
}
@@ -234,6 +238,16 @@ pub struct ToolCompletionResponse {
pub input_tokens: u32,
pub output_tokens: u32,
pub finish_reason: FinishReason,
/// Provider-specific response ID (e.g. for NEAR AI response chaining).
pub response_id: Option<String>,
}
/// Metadata about a model returned by the provider's API.
#[derive(Debug, Clone)]
pub struct ModelMetadata {
pub id: String,
/// Total context window size in tokens.
pub context_length: Option<u32>,
}
/// Trait for LLM providers.
@@ -260,6 +274,45 @@ pub trait LlmProvider: Send + Sync {
Ok(Vec::new())
}
/// Fetch metadata for the current model (context length, etc.).
/// Default returns the model name with no size info.
async fn model_metadata(&self) -> Result<ModelMetadata, LlmError> {
Ok(ModelMetadata {
id: self.model_name().to_string(),
context_length: None,
})
}
/// Get the currently active model name.
///
/// May differ from `model_name()` if the model was switched at runtime
/// via `set_model()`. Default returns `model_name()`.
fn active_model_name(&self) -> String {
self.model_name().to_string()
}
/// Switch the active model at runtime. Not all providers support this.
fn set_model(&self, _model: &str) -> Result<(), LlmError> {
Err(LlmError::RequestFailed {
provider: "unknown".to_string(),
reason: "Runtime model switching not supported by this provider".to_string(),
})
}
/// Seed a response chain for a thread (e.g. restoring from DB).
///
/// Providers that support response chaining (e.g. NEAR AI `previous_response_id`)
/// store this so subsequent calls send only delta messages.
fn seed_response_chain(&self, _thread_id: &str, _response_id: String) {}
/// Get the last response chain ID for a thread.
///
/// Returns `None` if the provider doesn't support chaining or has no
/// stored state for this thread.
fn get_response_chain_id(&self, _thread_id: &str) -> Option<String> {
None
}
/// Calculate cost for a completion.
fn calculate_cost(&self, input_tokens: u32, output_tokens: u32) -> Decimal {
let (input_cost, output_cost) = self.cost_per_token();
src/llm/reasoning.rs
+353 -29
View File
@@ -21,6 +21,8 @@ pub struct ReasoningContext {
pub job_description: Option<String>,
/// Current state description.
pub current_state: Option<String>,
/// Opaque metadata forwarded to the LLM provider (e.g. thread_id for chaining).
pub metadata: std::collections::HashMap<String, String>,
}
impl ReasoningContext {
@@ -31,6 +33,7 @@ impl ReasoningContext {
available_tools: Vec::new(),
job_description: None,
current_state: None,
metadata: std::collections::HashMap::new(),
}
}
@@ -57,6 +60,12 @@ impl ReasoningContext {
self.job_description = Some(description.into());
self
}
/// Set metadata (forwarded to the LLM provider).
pub fn with_metadata(mut self, metadata: std::collections::HashMap<String, String>) -> Self {
self.metadata = metadata;
self
}
}
impl Default for ReasoningContext {
@@ -114,7 +123,12 @@ pub enum RespondResult {
/// A text response (no tools needed).
Text(String),
/// The model wants to call tools. Caller should execute them and call back.
ToolCalls(Vec<ToolCall>),
/// Includes the optional content from the assistant message (some models
/// include explanatory text alongside tool calls).
ToolCalls {
tool_calls: Vec<ToolCall>,
content: Option<String>,
},
}
/// Reasoning engine for the agent.
@@ -192,10 +206,11 @@ impl Reasoning {
return Ok(vec![]);
}
let request =
let mut request =
ToolCompletionRequest::new(context.messages.clone(), context.available_tools.clone())
.with_max_tokens(1024)
.with_tool_choice("auto");
request.metadata = context.metadata.clone();
let response = self.llm.complete_with_tools(request).await?;
@@ -271,7 +286,9 @@ Respond in JSON format:
pub async fn respond(&self, context: &ReasoningContext) -> Result<String, LlmError> {
match self.respond_with_tools(context).await? {
RespondResult::Text(text) => Ok(text),
RespondResult::ToolCalls(calls) => {
RespondResult::ToolCalls {
tool_calls: calls, ..
} => {
// Format tool calls as text (legacy behavior for non-agentic callers)
let tool_info: Vec<String> = calls
.iter()
@@ -298,28 +315,49 @@ Respond in JSON format:
// If we have tools, use tool completion mode
if !context.available_tools.is_empty() {
let request = ToolCompletionRequest::new(messages, context.available_tools.clone())
let mut request = ToolCompletionRequest::new(messages, context.available_tools.clone())
.with_max_tokens(4096)
.with_temperature(0.7)
.with_tool_choice("auto");
request.metadata = context.metadata.clone();
let response = self.llm.complete_with_tools(request).await?;
// If there were tool calls, return them for execution
if !response.tool_calls.is_empty() {
return Ok(RespondResult::ToolCalls(response.tool_calls));
return Ok(RespondResult::ToolCalls {
tool_calls: response.tool_calls,
content: response.content,
});
}
let content = response
.content
.unwrap_or_else(|| "I'm not sure how to respond to that.".to_string());
// Some models (e.g. GLM-4.7) emit tool calls as XML tags in content
// instead of using the structured tool_calls field. Try to recover
// them before giving up and returning plain text.
let recovered = recover_tool_calls_from_content(&content, &context.available_tools);
if !recovered.is_empty() {
let cleaned = clean_response(&content);
return Ok(RespondResult::ToolCalls {
tool_calls: recovered,
content: if cleaned.is_empty() {
None
} else {
Some(cleaned)
},
});
}
Ok(RespondResult::Text(clean_response(&content)))
} else {
// No tools, use simple completion
let request = CompletionRequest::new(messages)
let mut request = CompletionRequest::new(messages)
.with_max_tokens(4096)
.with_temperature(0.7);
request.metadata = context.metadata.clone();
let response = self.llm.complete(request).await?;
Ok(RespondResult::Text(clean_response(&response.content)))
@@ -462,47 +500,178 @@ fn extract_json(text: &str) -> Option<&str> {
}
}
/// Clean up LLM response by stripping thinking tags and reasoning patterns.
/// Clean up LLM response by stripping model-internal tags and reasoning patterns.
///
/// Some models (GLM-4.7, etc.) emit XML-tagged internal state like
/// Try to extract tool calls from content text where the model emitted them
/// as XML tags instead of using the structured tool_calls field.
///
/// Handles these formats:
/// - `<tool_call>tool_name</tool_call>` (bare name)
/// - `<tool_call>{"name":"x","arguments":{}}</tool_call>` (JSON)
/// - `<|tool_call|>...<|/tool_call|>` (pipe-delimited variant)
/// - `<function_call>...</function_call>` (function_call variant)
///
/// Only returns calls whose name matches an available tool.
fn recover_tool_calls_from_content(
content: &str,
available_tools: &[ToolDefinition],
) -> Vec<ToolCall> {
let tool_names: std::collections::HashSet<&str> =
available_tools.iter().map(|t| t.name.as_str()).collect();
let mut calls = Vec::new();
for (open, close) in &[
("<tool_call>", "</tool_call>"),
("<|tool_call|>", "<|/tool_call|>"),
("<function_call>", "</function_call>"),
("<|function_call|>", "<|/function_call|>"),
] {
let mut remaining = content;
while let Some(start) = remaining.find(open) {
let inner_start = start + open.len();
let after = &remaining[inner_start..];
let Some(end) = after.find(close) else {
break;
};
let inner = after[..end].trim();
remaining = &after[end + close.len()..];
if inner.is_empty() {
continue;
}
// Try JSON first: {"name":"x","arguments":{}}
if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(inner) {
if let Some(name) = parsed.get("name").and_then(|v| v.as_str()) {
if tool_names.contains(name) {
let arguments = parsed
.get("arguments")
.cloned()
.unwrap_or(serde_json::Value::Object(Default::default()));
calls.push(ToolCall {
id: format!("recovered_{}", calls.len()),
name: name.to_string(),
arguments,
});
continue;
}
}
}
// Bare tool name (e.g. "<tool_call>tool_list</tool_call>")
let name = inner.trim();
if tool_names.contains(name) {
calls.push(ToolCall {
id: format!("recovered_{}", calls.len()),
name: name.to_string(),
arguments: serde_json::Value::Object(Default::default()),
});
}
}
}
calls
}
/// `<tool_call>tool_list</tool_call>` or `<|tool_call|>` in the content field
/// instead of using the standard OpenAI tool_calls array. We strip all of
/// these before the response reaches channels/users.
fn clean_response(text: &str) -> String {
let text = strip_thinking_tags(text);
let text = strip_internal_tags(text);
strip_reasoning_patterns(&text)
}
/// Strip `<thinking>...</thinking>` blocks from LLM output.
/// Tags that are model-internal and should never reach users.
const INTERNAL_TAGS: &[&str] = &["thinking", "tool_call", "function_call", "tool_calls"];
/// Strip all model-internal XML tags from LLM output.
///
/// Some models (especially Claude with extended thinking) include internal
/// reasoning in thinking tags. We strip these before showing to users.
fn strip_thinking_tags(text: &str) -> String {
/// Handles standard XML tags (`<tag>...</tag>`) and pipe-delimited variants
/// (`<|tag|>...<|/tag|>`) used by some models (e.g. GLM-4.7).
fn strip_internal_tags(text: &str) -> String {
let mut result = text.to_string();
for tag in INTERNAL_TAGS {
result = strip_xml_tag(&result, tag);
result = strip_pipe_tag(&result, tag);
}
// Collapse triple+ newlines left behind by removed blocks
while result.contains("\n\n\n") {
result = result.replace("\n\n\n", "\n\n");
}
result.trim().to_string()
}
/// Strip `<tag>...</tag>` and `<tag ...>...</tag>` blocks from text.
fn strip_xml_tag(text: &str, tag: &str) -> String {
let open_exact = format!("<{}>", tag);
let open_prefix = format!("<{} ", tag); // for <tag attr="...">
let close = format!("</{}>", tag);
let mut result = String::with_capacity(text.len());
let mut remaining = text;
while let Some(start) = remaining.find("<thinking>") {
loop {
// Find the next opening tag (exact or with attributes)
let exact_pos = remaining.find(&open_exact);
let prefix_pos = remaining.find(&open_prefix);
let start = match (exact_pos, prefix_pos) {
(Some(a), Some(b)) => a.min(b),
(Some(a), None) => a,
(None, Some(b)) => b,
(None, None) => break,
};
// Add everything before the tag
result.push_str(&remaining[..start]);
// Find the end of the opening tag (the closing >)
let after_open = &remaining[start..];
let open_end = match after_open.find('>') {
Some(pos) => start + pos + 1,
None => break, // malformed, stop
};
// Find the closing tag
if let Some(end_offset) = remaining[start..].find("</thinking>") {
// Skip past the closing tag (start + offset + tag length)
let end = start + end_offset + "</thinking>".len();
if let Some(close_offset) = remaining[open_end..].find(&close) {
let end = open_end + close_offset + close.len();
remaining = &remaining[end..];
} else {
// No closing tag found, discard everything from here
// (malformed, but handle gracefully by not including the unclosed tag)
// No closing tag, discard from here (malformed)
remaining = "";
break;
}
}
// Add any remaining content after the last thinking block
result.push_str(remaining);
result
}
// Clean up any double newlines left behind
let mut cleaned = result.trim().to_string();
while cleaned.contains("\n\n\n") {
cleaned = cleaned.replace("\n\n\n", "\n\n");
/// Strip `<|tag|>...<|/tag|>` pipe-delimited blocks from text.
///
/// Some models (e.g. certain Chinese LLMs) use this format instead of
/// standard XML tags.
fn strip_pipe_tag(text: &str, tag: &str) -> String {
let open = format!("<|{}|>", tag);
let close = format!("<|/{}|>", tag);
let mut result = String::with_capacity(text.len());
let mut remaining = text;
while let Some(start) = remaining.find(&open) {
result.push_str(&remaining[..start]);
if let Some(close_offset) = remaining[start..].find(&close) {
let end = start + close_offset + close.len();
remaining = &remaining[end..];
} else {
remaining = "";
break;
}
}
cleaned
result.push_str(remaining);
result
}
/// Strip any remaining reasoning that wasn't in proper <thinking> tags.
@@ -574,7 +743,7 @@ That's my plan."#;
#[test]
fn test_strip_thinking_tags_basic() {
let input = "<thinking>Let me think about this...</thinking>Hello, user!";
let output = strip_thinking_tags(input);
let output = strip_internal_tags(input);
assert_eq!(output, "Hello, user!");
}
@@ -582,7 +751,7 @@ That's my plan."#;
fn test_strip_thinking_tags_multiple() {
let input =
"<thinking>First thought</thinking>Hello<thinking>Second thought</thinking> world!";
let output = strip_thinking_tags(input);
let output = strip_internal_tags(input);
assert_eq!(output, "Hello world!");
}
@@ -594,14 +763,14 @@ I need to consider:
2. How to respond
</thinking>
Here is my response to your question."#;
let output = strip_thinking_tags(input);
let output = strip_internal_tags(input);
assert_eq!(output, "Here is my response to your question.");
}
#[test]
fn test_strip_thinking_tags_no_tags() {
let input = "Just a normal response without thinking tags.";
let output = strip_thinking_tags(input);
let output = strip_internal_tags(input);
assert_eq!(output, "Just a normal response without thinking tags.");
}
@@ -609,10 +778,77 @@ Here is my response to your question."#;
fn test_strip_thinking_tags_unclosed() {
// Malformed: unclosed tag should strip from there to end
let input = "Hello <thinking>this never closes";
let output = strip_thinking_tags(input);
let output = strip_internal_tags(input);
assert_eq!(output, "Hello");
}
#[test]
fn test_strip_tool_call_tags() {
// GLM-4.7 emits this garbage instead of using the tool_calls array
let input = "<tool_call>tool_list</tool_call>";
let output = strip_internal_tags(input);
assert_eq!(output, "");
}
#[test]
fn test_strip_tool_call_with_surrounding_text() {
let input = "Here is my answer.\n\n<tool_call>\n{\"name\": \"search\", \"arguments\": {}}\n</tool_call>";
let output = strip_internal_tags(input);
assert_eq!(output, "Here is my answer.");
}
#[test]
fn test_strip_multiple_internal_tags() {
let input = "<thinking>Let me think</thinking>Hello!\n<tool_call>some_tool</tool_call>";
let output = strip_internal_tags(input);
assert_eq!(output, "Hello!");
}
#[test]
fn test_strip_function_call_tags() {
let input = "Response text<function_call>{\"name\": \"foo\"}</function_call>";
let output = strip_internal_tags(input);
assert_eq!(output, "Response text");
}
#[test]
fn test_strip_tool_calls_plural() {
let input = "<tool_calls>[{\"id\": \"1\"}]</tool_calls>Actual response.";
let output = strip_internal_tags(input);
assert_eq!(output, "Actual response.");
}
#[test]
fn test_strip_pipe_delimited_tags() {
let input = "<|tool_call|>{\"name\": \"search\"}<|/tool_call|>Hello!";
let output = strip_internal_tags(input);
assert_eq!(output, "Hello!");
}
#[test]
fn test_strip_pipe_delimited_thinking() {
let input = "<|thinking|>reasoning here<|/thinking|>The answer is 42.";
let output = strip_internal_tags(input);
assert_eq!(output, "The answer is 42.");
}
#[test]
fn test_strip_xml_tag_with_attributes() {
let input = "<tool_call type=\"function\">search()</tool_call>Done.";
let output = strip_internal_tags(input);
assert_eq!(output, "Done.");
}
#[test]
fn test_clean_response_preserves_normal_content() {
let input = "The function tool_call_handler works great. No tags here!";
let output = clean_response(input);
assert_eq!(
output,
"The function tool_call_handler works great. No tags here!"
);
}
#[test]
fn test_strip_reasoning_paragraph_break() {
// Content after paragraph break with "here" marker
@@ -660,4 +896,92 @@ Here is my response to your question."#;
let output = clean_response(input);
assert_eq!(output, "Here's the answer.");
}
// -- recover_tool_calls_from_content tests --
fn make_tools(names: &[&str]) -> Vec<ToolDefinition> {
names
.iter()
.map(|n| ToolDefinition {
name: n.to_string(),
description: String::new(),
parameters: serde_json::json!({}),
})
.collect()
}
#[test]
fn test_recover_bare_tool_name() {
let tools = make_tools(&["tool_list", "tool_auth"]);
let content = "<tool_call>tool_list</tool_call>";
let calls = recover_tool_calls_from_content(content, &tools);
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].name, "tool_list");
assert_eq!(calls[0].arguments, serde_json::json!({}));
}
#[test]
fn test_recover_json_tool_call() {
let tools = make_tools(&["memory_search"]);
let content =
r#"<tool_call>{"name": "memory_search", "arguments": {"query": "test"}}</tool_call>"#;
let calls = recover_tool_calls_from_content(content, &tools);
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].name, "memory_search");
assert_eq!(calls[0].arguments, serde_json::json!({"query": "test"}));
}
#[test]
fn test_recover_pipe_delimited() {
let tools = make_tools(&["tool_list"]);
let content = "<|tool_call|>tool_list<|/tool_call|>";
let calls = recover_tool_calls_from_content(content, &tools);
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].name, "tool_list");
}
#[test]
fn test_recover_unknown_tool_ignored() {
let tools = make_tools(&["tool_list"]);
let content = "<tool_call>nonexistent_tool</tool_call>";
let calls = recover_tool_calls_from_content(content, &tools);
assert!(calls.is_empty());
}
#[test]
fn test_recover_no_tags() {
let tools = make_tools(&["tool_list"]);
let content = "Just a normal response.";
let calls = recover_tool_calls_from_content(content, &tools);
assert!(calls.is_empty());
}
#[test]
fn test_recover_multiple_tool_calls() {
let tools = make_tools(&["tool_list", "tool_auth"]);
let content = "<tool_call>tool_list</tool_call>\n<tool_call>tool_auth</tool_call>";
let calls = recover_tool_calls_from_content(content, &tools);
assert_eq!(calls.len(), 2);
assert_eq!(calls[0].name, "tool_list");
assert_eq!(calls[1].name, "tool_auth");
}
#[test]
fn test_recover_function_call_variant() {
let tools = make_tools(&["shell"]);
let content =
r#"<function_call>{"name": "shell", "arguments": {"cmd": "ls"}}</function_call>"#;
let calls = recover_tool_calls_from_content(content, &tools);
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].name, "shell");
}
#[test]
fn test_recover_with_surrounding_text() {
let tools = make_tools(&["tool_list"]);
let content = "Let me check.\n\n<tool_call>tool_list</tool_call>\n\nDone.";
let calls = recover_tool_calls_from_content(content, &tools);
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].name, "tool_list");
}
}
src/llm/session.rs
+77 -2
View File
@@ -61,6 +61,10 @@ pub struct SessionManager {
token: RwLock<Option<SecretString>>,
/// Prevents thundering herd during concurrent 401s.
renewal_lock: Mutex<()>,
/// Optional database store for persisting session to the settings table.
store: RwLock<Option<Arc<crate::history::Store>>>,
/// User ID for DB settings (default: "default").
user_id: RwLock<String>,
}
impl SessionManager {
@@ -74,6 +78,8 @@ impl SessionManager {
.unwrap_or_else(|_| Client::new()),
token: RwLock::new(None),
renewal_lock: Mutex::new(()),
store: RwLock::new(None),
user_id: RwLock::new("default".to_string()),
};
// Try to load existing session synchronously during construction
@@ -103,6 +109,8 @@ impl SessionManager {
.unwrap_or_else(|_| Client::new()),
token: RwLock::new(None),
renewal_lock: Mutex::new(()),
store: RwLock::new(None),
user_id: RwLock::new("default".to_string()),
};
if let Err(e) = manager.load_session().await {
@@ -112,6 +120,21 @@ impl SessionManager {
manager
}
/// Attach a database store for persisting session tokens.
///
/// When a store is attached, session tokens are saved to the `settings`
/// table (key: `nearai.session_token`) in addition to the disk file.
/// On load, DB is preferred over disk.
pub async fn attach_store(&self, store: Arc<crate::history::Store>, user_id: &str) {
*self.store.write().await = Some(store);
*self.user_id.write().await = user_id.to_string();
// Try to load from DB (may have been saved by a previous run)
if let Err(e) = self.load_session_from_db().await {
tracing::debug!("No session in DB: {}", e);
}
}
/// Get the current session token, returning an error if not authenticated.
pub async fn get_token(&self) -> Result<SecretString, LlmError> {
let guard = self.token.read().await;
@@ -460,7 +483,7 @@ impl SessionManager {
Ok(())
}
/// Save session data to disk.
/// Save session data to disk and (if available) to the database.
async fn save_session(&self, token: &str, auth_provider: Option<&str>) -> Result<(), LlmError> {
let session = SessionData {
session_token: token.to_string(),
@@ -468,7 +491,7 @@ impl SessionManager {
auth_provider: auth_provider.map(String::from),
};
// Ensure parent directory exists
// Save to disk (always, as bootstrap fallback)
if let Some(parent) = self.config.session_path.parent() {
tokio::fs::create_dir_all(parent).await.map_err(|e| {
LlmError::Io(std::io::Error::new(
@@ -498,6 +521,58 @@ impl SessionManager {
})?;
tracing::debug!("Session saved to {}", self.config.session_path.display());
// Also save to DB if a store is attached
if let Some(ref store) = *self.store.read().await {
let user_id = self.user_id.read().await.clone();
let session_json = serde_json::to_value(&session)
.unwrap_or(serde_json::Value::String(token.to_string()));
if let Err(e) = store
.set_setting(&user_id, "nearai.session_token", &session_json)
.await
{
tracing::warn!("Failed to save session to DB: {}", e);
} else {
tracing::debug!("Session also saved to DB settings");
}
}
Ok(())
}
/// Try to load session from the database.
async fn load_session_from_db(&self) -> Result<(), LlmError> {
let store_guard = self.store.read().await;
let store = store_guard
.as_ref()
.ok_or_else(|| LlmError::SessionRenewalFailed {
provider: "nearai".to_string(),
reason: "No DB store attached".to_string(),
})?;
let user_id = self.user_id.read().await.clone();
let value = store
.get_setting(&user_id, "nearai.session_token")
.await
.map_err(|e| LlmError::SessionRenewalFailed {
provider: "nearai".to_string(),
reason: format!("DB query failed: {}", e),
})?
.ok_or_else(|| LlmError::SessionRenewalFailed {
provider: "nearai".to_string(),
reason: "No session in DB".to_string(),
})?;
let session: SessionData =
serde_json::from_value(value).map_err(|e| LlmError::SessionRenewalFailed {
provider: "nearai".to_string(),
reason: format!("Failed to parse DB session: {}", e),
})?;
let mut guard = self.token.write().await;
*guard = Some(SecretString::from(session.session_token));
tracing::info!("Loaded session from DB settings");
Ok(())
}
src/main.rs
+290 -20
View File
@@ -24,14 +24,17 @@ use ironclaw::{
extensions::ExtensionManager,
history::Store,
llm::{SessionConfig, create_llm_provider, create_session_manager},
orchestrator::{
ContainerJobConfig, ContainerJobManager, OrchestratorApi, TokenStore,
api::OrchestratorState,
},
safety::SafetyLayer,
secrets::{PostgresSecretsStore, SecretsCrypto, SecretsStore},
settings::Settings,
setup::{SetupConfig, SetupWizard},
tools::{
ToolRegistry,
mcp::{McpClient, McpSessionManager, config::load_mcp_servers, is_authenticated},
wasm::{WasmToolLoader, WasmToolRuntime},
mcp::{McpClient, McpSessionManager, config::load_mcp_servers_from_db, is_authenticated},
wasm::{WasmToolLoader, WasmToolRuntime, load_dev_tools},
},
workspace::{EmbeddingProvider, NearAiEmbeddings, OpenAiEmbeddings, Workspace},
};
@@ -53,9 +56,14 @@ async fn main() -> anyhow::Result<()> {
return run_tool_command(tool_cmd.clone()).await;
}
Some(Command::Config(config_cmd)) => {
// Config commands don't need logging setup
return ironclaw::cli::run_config_command(config_cmd.clone())
.map_err(|e| anyhow::anyhow!("{}", e));
// Config commands need DB access for settings
tracing_subscriber::fmt()
.with_env_filter(
EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")),
)
.init();
return ironclaw::cli::run_config_command(config_cmd.clone()).await;
}
Some(Command::Mcp(mcp_cmd)) => {
// Simple logging for MCP commands
@@ -130,6 +138,83 @@ async fn main() -> anyhow::Result<()> {
return run_status_command().await;
}
Some(Command::Worker {
job_id,
orchestrator_url,
max_iterations,
}) => {
// Worker mode: runs inside a Docker container.
// Simple logging (no TUI, no DB, no channels).
tracing_subscriber::fmt()
.with_env_filter(
EnvFilter::try_from_default_env()
.unwrap_or_else(|_| EnvFilter::new("ironclaw=info")),
)
.init();
tracing::info!(
"Starting worker for job {} (orchestrator: {})",
job_id,
orchestrator_url
);
let config = ironclaw::worker::runtime::WorkerConfig {
job_id: *job_id,
orchestrator_url: orchestrator_url.clone(),
max_iterations: *max_iterations,
timeout: std::time::Duration::from_secs(600),
};
let runtime = ironclaw::worker::WorkerRuntime::new(config)
.map_err(|e| anyhow::anyhow!("Worker init failed: {}", e))?;
runtime
.run()
.await
.map_err(|e| anyhow::anyhow!("Worker failed: {}", e))?;
return Ok(());
}
Some(Command::ClaudeBridge {
job_id,
orchestrator_url,
max_turns,
model,
}) => {
// Claude Code bridge mode: runs inside a Docker container.
// Spawns the `claude` CLI and streams output to the orchestrator.
tracing_subscriber::fmt()
.with_env_filter(
EnvFilter::try_from_default_env()
.unwrap_or_else(|_| EnvFilter::new("ironclaw=info")),
)
.init();
tracing::info!(
"Starting Claude Code bridge for job {} (orchestrator: {}, model: {})",
job_id,
orchestrator_url,
model
);
let config = ironclaw::worker::claude_bridge::ClaudeBridgeConfig {
job_id: *job_id,
orchestrator_url: orchestrator_url.clone(),
max_turns: *max_turns,
model: model.clone(),
timeout: std::time::Duration::from_secs(1800),
};
let runtime = ironclaw::worker::ClaudeBridgeRuntime::new(config)
.map_err(|e| anyhow::anyhow!("Claude bridge init failed: {}", e))?;
runtime
.run()
.await
.map_err(|e| anyhow::anyhow!("Claude bridge failed: {}", e))?;
return Ok(());
}
Some(Command::Onboard {
skip_auth,
channels_only,
@@ -163,8 +248,11 @@ async fn main() -> anyhow::Result<()> {
}
}
// Load configuration (after potential setup)
let config = match Config::from_env() {
// Load bootstrap config (4 fields that must live on disk)
let bootstrap = ironclaw::bootstrap::BootstrapConfig::load();
// Load initial config from env + disk (before DB is available)
let mut config = match Config::from_env() {
Ok(c) => c,
Err(ironclaw::error::ConfigError::MissingRequired { key, hint }) => {
eprintln!("Configuration error: Missing required setting '{}'", key);
@@ -224,7 +312,38 @@ async fn main() -> anyhow::Result<()> {
let store = Store::new(&config.database).await?;
store.run_migrations().await?;
tracing::info!("Database connected and migrations applied");
Some(Arc::new(store))
// One-time migration: move disk config files into the DB settings table.
if let Err(e) = ironclaw::bootstrap::migrate_disk_to_db(&store, "default").await {
tracing::warn!("Disk-to-DB settings migration failed: {}", e);
}
// Reload config from DB now that we have a connection.
// Priority: env var > DB setting > default.
match Config::from_db(&store, "default", &bootstrap).await {
Ok(db_config) => {
config = db_config;
tracing::info!("Configuration reloaded from database");
}
Err(e) => {
tracing::warn!(
"Failed to reload config from DB, keeping env-based config: {}",
e
);
}
}
let store = Arc::new(store);
// Attach store to session manager so tokens save to DB too
session.attach_store(Arc::clone(&store), "default").await;
// Mark any jobs left in "running" or "creating" state as "interrupted".
if let Err(e) = store.cleanup_stale_sandbox_jobs().await {
tracing::warn!("Failed to cleanup stale sandbox jobs: {}", e);
}
Some(store)
};
// Initialize LLM provider (clone session so we can reuse it for embeddings)
@@ -289,8 +408,11 @@ async fn main() -> anyhow::Result<()> {
tools.register_memory_tools(workspace);
}
// Register builder tool if enabled
if config.builder.enabled {
// Register builder tool if enabled.
// When sandbox is enabled and allow_local_tools is false, skip builder registration
// because register_builder_tool also registers dev tools (shell, file ops) that would
// bypass the sandbox. The builder runs inside containers instead.
if config.builder.enabled && (config.agent.allow_local_tools || !config.sandbox.enabled) {
tools
.register_builder_tool(
llm.clone(),
@@ -339,6 +461,8 @@ async fn main() -> anyhow::Result<()> {
let wasm_tools_future = async {
if let Some(ref runtime) = wasm_tool_runtime {
let loader = WasmToolLoader::new(Arc::clone(runtime), Arc::clone(&tools));
// Load installed tools from ~/.ironclaw/tools/
match loader.load_from_dir(&config.wasm.tools_dir).await {
Ok(results) => {
if !results.loaded.is_empty() {
@@ -356,12 +480,32 @@ async fn main() -> anyhow::Result<()> {
tracing::warn!("Failed to scan WASM tools directory: {}", e);
}
}
// Load dev tools from build artifacts (overrides installed if newer)
match load_dev_tools(&loader, &config.wasm.tools_dir).await {
Ok(results) => {
if !results.loaded.is_empty() {
tracing::info!(
"Loaded {} dev WASM tools from build artifacts",
results.loaded.len()
);
}
}
Err(e) => {
tracing::debug!("No dev WASM tools found: {}", e);
}
}
}
};
let mcp_servers_future = async {
if let Some(ref secrets) = secrets_store {
match load_mcp_servers().await {
let servers_result = if let Some(ref s) = store {
load_mcp_servers_from_db(s, "default").await
} else {
ironclaw::tools::mcp::config::load_mcp_servers().await
};
match servers_result {
Ok(servers) => {
let enabled: Vec<_> = servers.enabled_servers().cloned().collect();
if !enabled.is_empty() {
@@ -470,6 +614,7 @@ async fn main() -> anyhow::Result<()> {
config.channels.wasm_channels_dir.clone(),
config.tunnel.public_url.clone(),
"default".to_string(),
store.clone(),
));
tools.register_extension_tools(Arc::clone(&manager));
tracing::info!("Extension manager initialized with in-chat discovery tools");
@@ -482,6 +627,77 @@ async fn main() -> anyhow::Result<()> {
None
};
// Set up orchestrator for sandboxed job execution
// When allow_local_tools is false (default), the LLM uses create_job for FS/shell work.
// When allow_local_tools is true, dev tools are also registered directly (current behavior).
if config.agent.allow_local_tools {
tools.register_dev_tools();
tracing::info!(
"Local tools enabled (allow_local_tools=true), dev tools registered directly"
);
}
// Shared state for job events (used by both orchestrator and web gateway)
let job_event_tx: Option<
tokio::sync::broadcast::Sender<(uuid::Uuid, ironclaw::channels::web::types::SseEvent)>,
> = if config.sandbox.enabled {
let (tx, _) = tokio::sync::broadcast::channel(256);
Some(tx)
} else {
None
};
let prompt_queue = Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::<
uuid::Uuid,
std::collections::VecDeque<ironclaw::orchestrator::api::PendingPrompt>,
>::new()));
let container_job_manager: Option<Arc<ContainerJobManager>> = if config.sandbox.enabled {
let token_store = TokenStore::new();
let job_config = ContainerJobConfig {
image: config.sandbox.image.clone(),
memory_limit_mb: config.sandbox.memory_limit_mb,
cpu_shares: config.sandbox.cpu_shares,
orchestrator_port: 50051,
claude_config_dir: if config.claude_code.enabled {
Some(config.claude_code.config_dir.clone())
} else {
None
},
claude_code_model: config.claude_code.model.clone(),
claude_code_max_turns: config.claude_code.max_turns,
claude_code_memory_limit_mb: config.claude_code.memory_limit_mb,
};
let jm = Arc::new(ContainerJobManager::new(job_config, token_store.clone()));
// Start the orchestrator internal API in the background
let orchestrator_state = OrchestratorState {
llm: llm.clone(),
job_manager: Arc::clone(&jm),
token_store,
job_event_tx: job_event_tx.clone(),
prompt_queue: Arc::clone(&prompt_queue),
store: store.clone(),
};
tokio::spawn(async move {
if let Err(e) = OrchestratorApi::start(orchestrator_state, 50051).await {
tracing::error!("Orchestrator API failed: {}", e);
}
});
tracing::info!("Orchestrator API started on :50051, sandbox delegation enabled");
if config.claude_code.enabled {
tracing::info!(
"Claude Code sandbox mode available (model: {}, max_turns: {})",
config.claude_code.model,
config.claude_code.max_turns
);
}
Some(jm)
} else {
None
};
tracing::info!(
"Tool registry initialized with {} total tools",
tools.count()
@@ -563,6 +779,17 @@ async fn main() -> anyhow::Result<()> {
);
}
// Inject owner_id for Telegram so the bot only responds
// to the bound user account.
if channel_name == "telegram" {
if let Some(owner_id) = config.channels.telegram_owner_id {
config_updates.insert(
"owner_id".to_string(),
serde_json::json!(owner_id),
);
}
}
if !config_updates.is_empty() {
channel_arc.update_config(config_updates).await;
tracing::info!(
@@ -621,7 +848,7 @@ async fn main() -> anyhow::Result<()> {
channels.add(Box::new(SharedWasmChannel::new(channel_arc)));
}
if has_webhook_channels && config.tunnel.public_url.is_some() {
if has_webhook_channels {
webhook_routes.push(create_wasm_channel_router(
wasm_router,
extension_manager.as_ref().map(Arc::clone),
@@ -693,6 +920,19 @@ async fn main() -> anyhow::Result<()> {
Arc::new(ws)
});
// Seed workspace with core identity files on first boot
if let Some(ref ws) = workspace {
match ws.seed_if_empty().await {
Ok(count) if count > 0 => {
tracing::info!("Workspace seeded with {} core files", count);
}
Ok(_) => {}
Err(e) => {
tracing::warn!("Failed to seed workspace: {}", e);
}
}
}
// Backfill embeddings if we just enabled the provider
if let (Some(ws), Some(_)) = (&workspace, &embeddings) {
match ws.backfill_embeddings().await {
@@ -712,8 +952,12 @@ async fn main() -> anyhow::Result<()> {
// Create session manager (shared between agent and web gateway)
let session_manager = Arc::new(SessionManager::new());
// Register job tools
tools.register_job_tools(Arc::clone(&context_manager));
// Register job tools (sandbox deps auto-injected when container_job_manager is available)
tools.register_job_tools(
Arc::clone(&context_manager),
container_job_manager.clone(),
store.clone(),
);
// Add web gateway channel if configured
if let Some(ref gw_config) = config.channels.gateway {
@@ -721,19 +965,44 @@ async fn main() -> anyhow::Result<()> {
if let Some(ref ws) = workspace {
gw = gw.with_workspace(Arc::clone(ws));
}
gw = gw.with_context_manager(Arc::clone(&context_manager));
gw = gw.with_session_manager(Arc::clone(&session_manager));
gw = gw.with_log_broadcaster(Arc::clone(&log_broadcaster));
gw = gw.with_tool_registry(Arc::clone(&tools));
if let Some(ref ext_mgr) = extension_manager {
gw = gw.with_extension_manager(Arc::clone(ext_mgr));
}
if let Some(ref s) = store {
gw = gw.with_store(Arc::clone(s));
}
if let Some(ref jm) = container_job_manager {
gw = gw.with_job_manager(Arc::clone(jm));
}
if config.sandbox.enabled {
gw = gw.with_prompt_queue(Arc::clone(&prompt_queue));
// Spawn a task to forward job events from the broadcast channel to SSE
if let Some(ref tx) = job_event_tx {
let mut rx = tx.subscribe();
let gw_state = Arc::clone(gw.state());
tokio::spawn(async move {
while let Ok((_job_id, event)) = rx.recv().await {
gw_state.sse.broadcast(event);
}
});
}
}
tracing::info!(
"Web gateway enabled on {}:{}",
gw_config.host,
gw_config.port
);
tracing::info!(
"Web UI: http://{}:{}/?token={}",
gw_config.host,
gw_config.port,
gw.auth_token()
);
channels.add(Box::new(gw));
}
@@ -752,6 +1021,7 @@ async fn main() -> anyhow::Result<()> {
deps,
channels,
Some(config.heartbeat.clone()),
Some(config.routines.clone()),
Some(context_manager),
Some(session_manager),
);
@@ -774,15 +1044,15 @@ async fn main() -> anyhow::Result<()> {
///
/// Returns `Some(reason)` if onboarding should be triggered, `None` otherwise.
fn check_onboard_needed() -> Option<&'static str> {
let settings = Settings::load();
let bootstrap = ironclaw::bootstrap::BootstrapConfig::load();
// Database not configured (and not in env)
if settings.database_url.is_none() && std::env::var("DATABASE_URL").is_err() {
if bootstrap.database_url.is_none() && std::env::var("DATABASE_URL").is_err() {
return Some("Database not configured");
}
// Secrets not configured (and not in env)
if settings.secrets_master_key_source == ironclaw::settings::KeySource::None
if bootstrap.secrets_master_key_source == ironclaw::settings::KeySource::None
&& std::env::var("SECRETS_MASTER_KEY").is_err()
&& !ironclaw::secrets::keychain::has_master_key()
{
@@ -792,7 +1062,7 @@ fn check_onboard_needed() -> Option<&'static str> {
// First run (onboarding never completed and no session)
let session_path = ironclaw::llm::session::default_session_path();
if !settings.onboard_completed && !session_path.exists() {
if !bootstrap.onboard_completed && !session_path.exists() {
return Some("First run");
}
src/orchestrator/api.rs
+511
View File
@@ -0,0 +1,511 @@
//! Internal HTTP API for worker-to-orchestrator communication.
//!
//! This runs on a separate port (default 50051) from the web gateway.
//! All endpoints are authenticated via per-job bearer tokens.
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use axum::extract::{Path, State};
use axum::http::StatusCode;
use axum::routing::{get, post};
use axum::{Json, Router};
use serde::{Deserialize, Serialize};
use tokio::sync::{Mutex, broadcast};
use uuid::Uuid;
use crate::channels::web::types::SseEvent;
use crate::history::Store;
use crate::llm::{CompletionRequest, LlmProvider, ToolCompletionRequest};
use crate::orchestrator::auth::{TokenStore, worker_auth_middleware};
use crate::orchestrator::job_manager::ContainerJobManager;
use crate::worker::api::JobEventPayload;
use crate::worker::api::{
CompletionReport, JobDescription, ProxyCompletionRequest, ProxyCompletionResponse,
ProxyToolCompletionRequest, ProxyToolCompletionResponse, StatusUpdate,
};
/// A follow-up prompt queued for a Claude Code bridge.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PendingPrompt {
pub content: String,
pub done: bool,
}
/// Shared state for the orchestrator API.
#[derive(Clone)]
pub struct OrchestratorState {
pub llm: Arc<dyn LlmProvider>,
pub job_manager: Arc<ContainerJobManager>,
pub token_store: TokenStore,
/// Broadcast channel for job events (consumed by the web gateway SSE).
pub job_event_tx: Option<broadcast::Sender<(Uuid, SseEvent)>>,
/// Buffered follow-up prompts for sandbox jobs, keyed by job_id.
pub prompt_queue: Arc<Mutex<HashMap<Uuid, VecDeque<PendingPrompt>>>>,
/// Database handle for persisting job events.
pub store: Option<Arc<Store>>,
}
/// The orchestrator's internal API server.
pub struct OrchestratorApi;
impl OrchestratorApi {
/// Build the axum router for the internal API.
pub fn router(state: OrchestratorState) -> Router {
Router::new()
// Worker routes: authenticated via route_layer middleware.
.route("/worker/{job_id}/job", get(get_job))
.route("/worker/{job_id}/llm/complete", post(llm_complete))
.route(
"/worker/{job_id}/llm/complete_with_tools",
post(llm_complete_with_tools),
)
.route("/worker/{job_id}/status", post(report_status))
.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_layer(axum::middleware::from_fn_with_state(
state.token_store.clone(),
worker_auth_middleware,
))
// Unauthenticated routes (added after the layer).
.route("/health", get(health_check))
.with_state(state)
}
/// Start the internal API server on the given port.
///
/// On macOS/Windows (Docker Desktop), binds to loopback only because
/// Docker Desktop routes `host.docker.internal` through its VM to the
/// host's `127.0.0.1`.
///
/// On Linux, containers reach the host via the docker bridge gateway
/// (`172.17.0.1`), which is NOT loopback. Binding to `127.0.0.1`
/// would reject container traffic. We bind to all interfaces instead
/// and rely on `worker_auth_middleware` (applied as a route_layer on
/// every `/worker/` endpoint) to reject unauthenticated requests.
pub async fn start(
state: OrchestratorState,
port: u16,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
let router = Self::router(state);
let addr = if cfg!(target_os = "linux") {
std::net::SocketAddr::from(([0, 0, 0, 0], port))
} else {
std::net::SocketAddr::from(([127, 0, 0, 1], port))
};
tracing::info!("Orchestrator internal API listening on {}", addr);
let listener = tokio::net::TcpListener::bind(addr).await?;
axum::serve(listener, router).await?;
Ok(())
}
}
// -- Handlers --
//
// All /worker/ handlers below are behind the worker_auth_middleware route_layer,
// so they don't need to validate tokens themselves.
async fn health_check() -> &'static str {
"ok"
}
async fn get_job(
State(state): State<OrchestratorState>,
Path(job_id): Path<Uuid>,
) -> Result<Json<JobDescription>, StatusCode> {
let handle = state
.job_manager
.get_handle(job_id)
.await
.ok_or(StatusCode::NOT_FOUND)?;
Ok(Json(JobDescription {
title: format!("Job {}", job_id),
description: handle.task_description,
project_dir: handle.project_dir.map(|p| p.display().to_string()),
}))
}
async fn llm_complete(
State(state): State<OrchestratorState>,
Path(job_id): Path<Uuid>,
Json(req): Json<ProxyCompletionRequest>,
) -> Result<Json<ProxyCompletionResponse>, StatusCode> {
let completion_req = CompletionRequest {
messages: req.messages,
max_tokens: req.max_tokens,
temperature: req.temperature,
stop_sequences: req.stop_sequences,
metadata: std::collections::HashMap::new(),
};
let resp = state.llm.complete(completion_req).await.map_err(|e| {
tracing::error!("LLM completion failed for job {}: {}", job_id, e);
StatusCode::BAD_GATEWAY
})?;
Ok(Json(ProxyCompletionResponse {
content: resp.content,
input_tokens: resp.input_tokens,
output_tokens: resp.output_tokens,
finish_reason: format_finish_reason(resp.finish_reason),
}))
}
async fn llm_complete_with_tools(
State(state): State<OrchestratorState>,
Path(job_id): Path<Uuid>,
Json(req): Json<ProxyToolCompletionRequest>,
) -> Result<Json<ProxyToolCompletionResponse>, StatusCode> {
let tool_req = ToolCompletionRequest {
messages: req.messages,
tools: req.tools,
max_tokens: req.max_tokens,
temperature: req.temperature,
tool_choice: req.tool_choice,
metadata: std::collections::HashMap::new(),
};
let resp = state.llm.complete_with_tools(tool_req).await.map_err(|e| {
tracing::error!("LLM tool completion failed for job {}: {}", job_id, e);
StatusCode::BAD_GATEWAY
})?;
Ok(Json(ProxyToolCompletionResponse {
content: resp.content,
tool_calls: resp.tool_calls,
input_tokens: resp.input_tokens,
output_tokens: resp.output_tokens,
finish_reason: format_finish_reason(resp.finish_reason),
}))
}
async fn report_status(
Path(job_id): Path<Uuid>,
Json(update): Json<StatusUpdate>,
) -> Result<StatusCode, StatusCode> {
tracing::debug!(
job_id = %job_id,
state = %update.state,
iteration = update.iteration,
"Worker status update"
);
Ok(StatusCode::OK)
}
async fn report_complete(
State(state): State<OrchestratorState>,
Path(job_id): Path<Uuid>,
Json(report): Json<CompletionReport>,
) -> Result<StatusCode, StatusCode> {
if report.success {
tracing::info!(
job_id = %job_id,
"Worker reported job complete"
);
} else {
tracing::warn!(
job_id = %job_id,
message = ?report.message,
"Worker reported job failure"
);
}
// Store the result and clean up the container
let result = crate::orchestrator::job_manager::CompletionResult {
success: report.success,
message: report.message.clone(),
};
let _ = state.job_manager.complete_job(job_id, result).await;
Ok(StatusCode::OK)
}
// -- Sandbox job event handlers --
/// Receive a job event from a worker or Claude Code bridge and broadcast + persist it.
async fn job_event_handler(
State(state): State<OrchestratorState>,
Path(job_id): Path<Uuid>,
Json(payload): Json<JobEventPayload>,
) -> Result<StatusCode, StatusCode> {
tracing::debug!(
job_id = %job_id,
event_type = %payload.event_type,
"Job event received"
);
// Persist to DB (fire-and-forget)
if let Some(ref store) = state.store {
let store = Arc::clone(store);
let event_type = payload.event_type.clone();
let data = payload.data.clone();
tokio::spawn(async move {
if let Err(e) = store.save_job_event(job_id, &event_type, &data).await {
tracing::warn!(job_id = %job_id, "Failed to persist job event: {}", e);
}
});
}
// Convert to SSE event and broadcast
let job_id_str = job_id.to_string();
let sse_event = match payload.event_type.as_str() {
"message" => SseEvent::JobMessage {
job_id: job_id_str,
role: payload
.data
.get("role")
.and_then(|v| v.as_str())
.unwrap_or("assistant")
.to_string(),
content: payload
.data
.get("content")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
},
"tool_use" => SseEvent::JobToolUse {
job_id: job_id_str,
tool_name: payload
.data
.get("tool_name")
.and_then(|v| v.as_str())
.unwrap_or("unknown")
.to_string(),
input: payload
.data
.get("input")
.cloned()
.unwrap_or(serde_json::Value::Null),
},
"tool_result" => SseEvent::JobToolResult {
job_id: job_id_str,
tool_name: payload
.data
.get("tool_name")
.and_then(|v| v.as_str())
.unwrap_or("unknown")
.to_string(),
output: payload
.data
.get("output")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
},
"result" => SseEvent::JobResult {
job_id: job_id_str,
status: payload
.data
.get("status")
.and_then(|v| v.as_str())
.unwrap_or("unknown")
.to_string(),
session_id: payload
.data
.get("session_id")
.and_then(|v| v.as_str())
.map(|s| s.to_string()),
},
_ => SseEvent::JobStatus {
job_id: job_id_str,
message: payload
.data
.get("message")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
},
};
// Broadcast via the channel (if configured)
if let Some(ref tx) = state.job_event_tx {
let _ = tx.send((job_id, sse_event));
}
Ok(StatusCode::OK)
}
/// Return the next queued follow-up prompt for a Claude Code bridge.
/// Returns 204 No Content if no prompt is available.
async fn get_prompt_handler(
State(state): State<OrchestratorState>,
Path(job_id): Path<Uuid>,
) -> Result<(StatusCode, Json<serde_json::Value>), StatusCode> {
let mut queue = state.prompt_queue.lock().await;
if let Some(prompts) = queue.get_mut(&job_id) {
if let Some(prompt) = prompts.pop_front() {
return Ok((
StatusCode::OK,
Json(serde_json::json!({
"content": prompt.content,
"done": prompt.done,
})),
));
}
}
// Return 204 with an empty body. The Json wrapper requires some value
// but the status code signals "nothing here".
Ok((StatusCode::NO_CONTENT, Json(serde_json::Value::Null)))
}
fn format_finish_reason(reason: crate::llm::FinishReason) -> String {
match reason {
crate::llm::FinishReason::Stop => "stop".to_string(),
crate::llm::FinishReason::Length => "length".to_string(),
crate::llm::FinishReason::ToolUse => "tool_use".to_string(),
crate::llm::FinishReason::ContentFilter => "content_filter".to_string(),
crate::llm::FinishReason::Unknown => "unknown".to_string(),
}
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use axum::body::Body;
use axum::http::Request;
use tower::ServiceExt;
use uuid::Uuid;
use crate::error::LlmError;
use crate::llm::{
CompletionRequest, CompletionResponse, ToolCompletionRequest, ToolCompletionResponse,
};
use crate::orchestrator::auth::TokenStore;
use crate::orchestrator::job_manager::{ContainerJobConfig, ContainerJobManager};
use super::*;
/// Stub LLM provider that panics if called (tests only exercise routing/auth).
struct StubLlm;
#[async_trait::async_trait]
impl crate::llm::LlmProvider for StubLlm {
fn model_name(&self) -> &str {
"stub"
}
fn cost_per_token(&self) -> (rust_decimal::Decimal, rust_decimal::Decimal) {
(rust_decimal::Decimal::ZERO, rust_decimal::Decimal::ZERO)
}
async fn complete(&self, _req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
Err(LlmError::RequestFailed {
provider: "stub".into(),
reason: "not implemented".into(),
})
}
async fn complete_with_tools(
&self,
_req: ToolCompletionRequest,
) -> Result<ToolCompletionResponse, LlmError> {
Err(LlmError::RequestFailed {
provider: "stub".into(),
reason: "not implemented".into(),
})
}
}
fn test_state() -> OrchestratorState {
let token_store = TokenStore::new();
let jm = ContainerJobManager::new(ContainerJobConfig::default(), token_store.clone());
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,
}
}
#[tokio::test]
async fn health_requires_no_auth() {
let state = test_state();
let router = OrchestratorApi::router(state);
let req = Request::builder()
.uri("/health")
.body(Body::empty())
.unwrap();
let resp = router.oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
}
#[tokio::test]
async fn worker_route_rejects_missing_token() {
let state = test_state();
let router = OrchestratorApi::router(state);
let job_id = Uuid::new_v4();
let req = Request::builder()
.uri(format!("/worker/{}/job", job_id))
.body(Body::empty())
.unwrap();
let resp = router.oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
}
#[tokio::test]
async fn worker_route_rejects_wrong_token() {
let state = test_state();
let router = OrchestratorApi::router(state);
let job_id = Uuid::new_v4();
let req = Request::builder()
.uri(format!("/worker/{}/job", job_id))
.header("Authorization", "Bearer totally-bogus")
.body(Body::empty())
.unwrap();
let resp = router.oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
}
#[tokio::test]
async fn worker_route_accepts_valid_token() {
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/{}/job", job_id))
.header("Authorization", format!("Bearer {}", token))
.body(Body::empty())
.unwrap();
let resp = router.oneshot(req).await.unwrap();
// 404 because no container exists for this job_id, but NOT 401.
assert_eq!(resp.status(), StatusCode::NOT_FOUND);
}
#[tokio::test]
async fn token_for_job_a_rejected_on_job_b() {
let state = test_state();
let job_a = Uuid::new_v4();
let job_b = Uuid::new_v4();
let token_a = state.token_store.create_token(job_a).await;
let router = OrchestratorApi::router(state);
// Use job_a's token to hit job_b's endpoint
let req = Request::builder()
.uri(format!("/worker/{}/job", job_b))
.header("Authorization", format!("Bearer {}", token_a))
.body(Body::empty())
.unwrap();
let resp = router.oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
}
}
src/orchestrator/auth.rs
+162
View File
@@ -0,0 +1,162 @@
//! Per-job bearer token authentication for worker-to-orchestrator communication.
//!
//! Security properties:
//! - Tokens are cryptographically random (32 bytes, hex-encoded)
//! - 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
use std::collections::HashMap;
use std::sync::Arc;
use axum::extract::{Request, State};
use axum::http::StatusCode;
use axum::middleware::Next;
use axum::response::Response;
use rand::Rng;
use tokio::sync::RwLock;
use uuid::Uuid;
/// In-memory store for per-job authentication tokens.
#[derive(Clone)]
pub struct TokenStore {
/// Maps job_id -> bearer token. Never logged or persisted.
tokens: Arc<RwLock<HashMap<Uuid, String>>>,
}
impl TokenStore {
pub fn new() -> Self {
Self {
tokens: Arc::new(RwLock::new(HashMap::new())),
}
}
/// Generate and store a new token for a job.
pub async fn create_token(&self, job_id: Uuid) -> String {
let token = generate_token();
self.tokens.write().await.insert(job_id, token.clone());
token
}
/// Validate a token for a specific job.
pub async fn validate(&self, job_id: Uuid, token: &str) -> bool {
self.tokens
.read()
.await
.get(&job_id)
.map(|stored| stored == token)
.unwrap_or(false)
}
/// Remove a token (on container cleanup).
pub async fn revoke(&self, job_id: Uuid) {
self.tokens.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()
}
}
impl Default for TokenStore {
fn default() -> Self {
Self::new()
}
}
/// Generate a cryptographically random token (32 bytes, hex-encoded = 64 chars).
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()
}
/// Axum middleware that validates worker bearer tokens.
///
/// Extracts the job_id from the path (`/worker/{job_id}/...`) and validates
/// the `Authorization: Bearer <token>` header against the token store.
///
/// Wire up with `axum::middleware::from_fn_with_state(token_store, worker_auth_middleware)`.
pub async fn worker_auth_middleware(
State(token_store): State<TokenStore>,
request: Request,
next: Next,
) -> Result<Response, StatusCode> {
let path = request.uri().path().to_string();
let job_id = extract_job_id_from_path(&path).ok_or(StatusCode::BAD_REQUEST)?;
let token = request
.headers()
.get("authorization")
.and_then(|v| v.to_str().ok())
.and_then(|v| v.strip_prefix("Bearer "))
.ok_or(StatusCode::UNAUTHORIZED)?;
if !token_store.validate(job_id, token).await {
return Err(StatusCode::UNAUTHORIZED);
}
Ok(next.run(request).await)
}
/// Extract job UUID from a path like `/worker/{uuid}/...`
fn extract_job_id_from_path(path: &str) -> Option<Uuid> {
let parts: Vec<&str> = path.trim_start_matches('/').split('/').collect();
if parts.len() >= 2 && parts[0] == "worker" {
Uuid::parse_str(parts[1]).ok()
} else {
None
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_token_create_and_validate() {
let store = TokenStore::new();
let job_id = Uuid::new_v4();
let token = store.create_token(job_id).await;
assert_eq!(token.len(), 64); // 32 bytes hex = 64 chars
assert!(store.validate(job_id, &token).await);
assert!(!store.validate(job_id, "wrong-token").await);
assert!(!store.validate(Uuid::new_v4(), &token).await);
}
#[tokio::test]
async fn test_token_revoke() {
let store = TokenStore::new();
let job_id = Uuid::new_v4();
let token = store.create_token(job_id).await;
assert!(store.validate(job_id, &token).await);
store.revoke(job_id).await;
assert!(!store.validate(job_id, &token).await);
}
#[test]
fn test_extract_job_id() {
let id = Uuid::new_v4();
let path = format!("/worker/{}/llm/complete", id);
assert_eq!(extract_job_id_from_path(&path), Some(id));
assert_eq!(extract_job_id_from_path("/other/path"), None);
assert_eq!(extract_job_id_from_path("/worker/not-a-uuid/foo"), None);
}
#[test]
fn test_token_is_random() {
let t1 = generate_token();
let t2 = generate_token();
assert_ne!(t1, t2);
}
}
src/orchestrator/job_manager.rs
+474
View File
@@ -0,0 +1,474 @@
//! Container lifecycle management for sandboxed jobs.
//!
//! Extends the existing `SandboxManager` infrastructure to support persistent
//! containers with their own agent loops (as opposed to ephemeral per-command containers).
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use chrono::{DateTime, Utc};
use tokio::sync::RwLock;
use uuid::Uuid;
use crate::error::OrchestratorError;
use crate::orchestrator::auth::TokenStore;
use crate::sandbox::connect_docker;
/// Which mode a sandbox container runs in.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JobMode {
/// Standard IronClaw worker with proxied LLM calls.
Worker,
/// Claude Code bridge that spawns the `claude` CLI directly.
ClaudeCode,
}
impl JobMode {
pub fn as_str(&self) -> &'static str {
match self {
Self::Worker => "worker",
Self::ClaudeCode => "claude_code",
}
}
}
impl std::fmt::Display for JobMode {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.as_str())
}
}
/// Configuration for the container job manager.
#[derive(Debug, Clone)]
pub struct ContainerJobConfig {
/// Docker image for worker containers.
pub image: String,
/// Default memory limit in MB.
pub memory_limit_mb: u64,
/// Default CPU shares.
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>,
/// Claude model to use in ClaudeCode mode.
pub claude_code_model: String,
/// Maximum turns for Claude Code.
pub claude_code_max_turns: u32,
/// Memory limit in MB for Claude Code containers (heavier than workers).
pub claude_code_memory_limit_mb: u64,
}
impl Default for ContainerJobConfig {
fn default() -> Self {
Self {
image: "ironclaw-worker:latest".to_string(),
memory_limit_mb: 2048,
cpu_shares: 1024,
orchestrator_port: 50051,
claude_config_dir: None,
claude_code_model: "sonnet".to_string(),
claude_code_max_turns: 50,
claude_code_memory_limit_mb: 4096,
}
}
}
/// State of a container.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ContainerState {
Creating,
Running,
Stopped,
Failed,
}
impl std::fmt::Display for ContainerState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Creating => write!(f, "creating"),
Self::Running => write!(f, "running"),
Self::Stopped => write!(f, "stopped"),
Self::Failed => write!(f, "failed"),
}
}
}
/// Handle to a running container job.
#[derive(Debug, Clone)]
pub struct ContainerHandle {
pub job_id: Uuid,
pub container_id: String,
pub state: ContainerState,
pub mode: JobMode,
pub created_at: DateTime<Utc>,
pub project_dir: Option<PathBuf>,
pub task_description: String,
/// 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.
// It lives only in the TokenStore (never logged, serialized, or persisted).
}
/// Result reported by a worker on completion.
#[derive(Debug, Clone)]
pub struct CompletionResult {
pub success: bool,
pub message: Option<String>,
}
/// Manages the lifecycle of Docker containers for sandboxed job execution.
pub struct ContainerJobManager {
config: ContainerJobConfig,
token_store: TokenStore,
containers: Arc<RwLock<HashMap<Uuid, ContainerHandle>>>,
}
impl ContainerJobManager {
pub fn new(config: ContainerJobConfig, token_store: TokenStore) -> Self {
Self {
config,
token_store,
containers: Arc::new(RwLock::new(HashMap::new())),
}
}
/// 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.
pub async fn create_job(
&self,
job_id: Uuid,
task: &str,
project_dir: Option<PathBuf>,
mode: JobMode,
) -> Result<String, OrchestratorError> {
// Generate auth token (stored in TokenStore, never logged)
let token = self.token_store.create_token(job_id).await;
// Record the handle
let handle = ContainerHandle {
job_id,
container_id: String::new(), // set after container creation
state: ContainerState::Creating,
mode,
created_at: Utc::now(),
project_dir: project_dir.clone(),
task_description: task.to_string(),
completion_result: None,
};
self.containers.write().await.insert(job_id, handle);
// Connect to Docker
let docker = connect_docker()
.await
.map_err(|e| OrchestratorError::Docker {
reason: e.to_string(),
})?;
// Build container configuration
let orchestrator_host = if cfg!(target_os = "linux") {
"172.17.0.1"
} else {
"host.docker.internal"
};
let orchestrator_url = format!(
"http://{}:{}",
orchestrator_host, self.config.orchestrator_port
);
let mut env_vec = vec![
format!("IRONCLAW_WORKER_TOKEN={}", token),
format!("IRONCLAW_JOB_ID={}", job_id),
format!("IRONCLAW_ORCHESTRATOR_URL={}", orchestrator_url),
];
// 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() {
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()
),
});
}
}
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
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()));
}
}
// Memory limit: Claude Code gets more memory
let memory_mb = match mode {
JobMode::ClaudeCode => self.config.claude_code_memory_limit_mb,
JobMode::Worker => self.config.memory_limit_mb,
};
// Create the container
use bollard::container::{Config, CreateContainerOptions};
use bollard::models::HostConfig;
let host_config = HostConfig {
binds: if binds.is_empty() { None } else { Some(binds) },
memory: Some((memory_mb * 1024 * 1024) as i64),
cpu_shares: Some(self.config.cpu_shares as i64),
network_mode: Some("bridge".to_string()),
extra_hosts: Some(vec!["host.docker.internal:host-gateway".to_string()]),
cap_drop: Some(vec!["ALL".to_string()]),
cap_add: Some(vec![
"CHOWN".to_string(),
"SETUID".to_string(),
"SETGID".to_string(),
]),
security_opt: Some(vec!["no-new-privileges:true".to_string()]),
tmpfs: Some(
[("/tmp".to_string(), "size=512M".to_string())]
.into_iter()
.collect(),
),
..Default::default()
};
// Build CMD based on mode
let cmd = match mode {
JobMode::Worker => vec![
"worker".to_string(),
"--job-id".to_string(),
job_id.to_string(),
"--orchestrator-url".to_string(),
orchestrator_url,
],
JobMode::ClaudeCode => vec![
"claude-bridge".to_string(),
"--job-id".to_string(),
job_id.to_string(),
"--orchestrator-url".to_string(),
orchestrator_url,
"--max-turns".to_string(),
self.config.claude_code_max_turns.to_string(),
"--model".to_string(),
self.config.claude_code_model.clone(),
],
};
let container_config = Config {
image: Some(self.config.image.clone()),
cmd: Some(cmd),
env: Some(env_vec),
host_config: Some(host_config),
user: Some("1000:1000".to_string()),
working_dir: Some("/workspace".to_string()),
..Default::default()
};
let container_name = match mode {
JobMode::Worker => format!("ironclaw-worker-{}", job_id),
JobMode::ClaudeCode => format!("ironclaw-claude-{}", job_id),
};
let options = CreateContainerOptions {
name: container_name,
..Default::default()
};
let response = docker
.create_container(Some(options), container_config)
.await
.map_err(|e| OrchestratorError::ContainerCreationFailed {
job_id,
reason: e.to_string(),
})?;
let container_id = response.id;
// Start the container
docker
.start_container::<String>(&container_id, None)
.await
.map_err(|e| OrchestratorError::ContainerCreationFailed {
job_id,
reason: format!("failed to start container: {}", e),
})?;
// Update handle with container ID
if let Some(handle) = self.containers.write().await.get_mut(&job_id) {
handle.container_id = container_id;
handle.state = ContainerState::Running;
}
tracing::info!(
job_id = %job_id,
"Created and started worker container"
);
Ok(token)
}
/// Stop a running container job.
pub async fn stop_job(&self, job_id: Uuid) -> Result<(), OrchestratorError> {
let container_id = {
let containers = self.containers.read().await;
containers
.get(&job_id)
.map(|h| h.container_id.clone())
.ok_or(OrchestratorError::ContainerNotFound { job_id })?
};
if container_id.is_empty() {
return Err(OrchestratorError::InvalidContainerState {
job_id,
state: "creating (no container ID yet)".to_string(),
});
}
let docker = connect_docker()
.await
.map_err(|e| OrchestratorError::Docker {
reason: e.to_string(),
})?;
// Stop the container (10 second grace period)
let _ = docker
.stop_container(
&container_id,
Some(bollard::container::StopContainerOptions { t: 10 }),
)
.await;
// Remove the container
let _ = docker
.remove_container(
&container_id,
Some(bollard::container::RemoveContainerOptions {
force: true,
..Default::default()
}),
)
.await;
// Update state
if let Some(handle) = self.containers.write().await.get_mut(&job_id) {
handle.state = ContainerState::Stopped;
}
// Revoke the auth token
self.token_store.revoke(job_id).await;
tracing::info!(job_id = %job_id, "Stopped worker container");
Ok(())
}
/// Mark a job as complete with a result. The container is stopped but the
/// handle is kept so `CreateJobTool` can read the completion message.
pub async fn complete_job(
&self,
job_id: Uuid,
result: CompletionResult,
) -> Result<(), OrchestratorError> {
// Store the result before stopping
{
let mut containers = self.containers.write().await;
if let Some(handle) = containers.get_mut(&job_id) {
handle.completion_result = Some(result);
handle.state = ContainerState::Stopped;
}
}
// Stop container and revoke token (but keep handle in map)
let container_id = {
let containers = self.containers.read().await;
containers.get(&job_id).map(|h| h.container_id.clone())
};
if let Some(cid) = container_id {
if !cid.is_empty() {
if let Ok(docker) = connect_docker().await {
let _ = docker
.stop_container(
&cid,
Some(bollard::container::StopContainerOptions { t: 5 }),
)
.await;
let _ = docker
.remove_container(
&cid,
Some(bollard::container::RemoveContainerOptions {
force: true,
..Default::default()
}),
)
.await;
}
}
}
self.token_store.revoke(job_id).await;
tracing::info!(job_id = %job_id, "Completed worker container");
Ok(())
}
/// Remove a completed job handle from memory (called after result is read).
pub async fn cleanup_job(&self, job_id: Uuid) {
self.containers.write().await.remove(&job_id);
}
/// 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()
}
/// List all active container jobs.
pub async fn list_jobs(&self) -> Vec<ContainerHandle> {
self.containers.read().await.values().cloned().collect()
}
/// Get a reference to the token store.
pub fn token_store(&self) -> &TokenStore {
&self.token_store
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_container_job_config_default() {
let config = ContainerJobConfig::default();
assert_eq!(config.orchestrator_port, 50051);
assert_eq!(config.memory_limit_mb, 2048);
}
#[test]
fn test_container_state_display() {
assert_eq!(ContainerState::Running.to_string(), "running");
assert_eq!(ContainerState::Stopped.to_string(), "stopped");
}
}
src/orchestrator/mod.rs
+37
View File
@@ -0,0 +1,37 @@
//! Orchestrator for managing sandboxed worker containers.
//!
//! The orchestrator runs in the main agent process and provides:
//! - An internal HTTP API for worker communication (LLM proxy, status, secrets)
//! - Per-job bearer token authentication
//! - Container lifecycle management (create, monitor, stop)
//!
//! ```text
//! ┌───────────────────────────────────────────────┐
//! │ Orchestrator │
//! │ │
//! │ Internal API (:50051) │
//! │ POST /worker/{id}/llm/complete │
//! │ POST /worker/{id}/llm/complete_with_tools │
//! │ GET /worker/{id}/job │
//! │ POST /worker/{id}/status │
//! │ POST /worker/{id}/complete │
//! │ │
//! │ ContainerJobManager │
//! │ create_job() -> container + token │
//! │ stop_job() │
//! │ list_jobs() │
//! │ │
//! │ TokenStore │
//! │ per-job bearer tokens (in-memory only) │
//! └───────────────────────────────────────────────┘
//! ```
pub mod api;
pub mod auth;
pub mod job_manager;
pub use api::OrchestratorApi;
pub use auth::TokenStore;
pub use job_manager::{
CompletionResult, ContainerHandle, ContainerJobConfig, ContainerJobManager, JobMode,
};
src/sandbox/container.rs
+29 -2
View File
@@ -491,9 +491,36 @@ impl ContainerRunner {
}
/// Connect to the Docker daemon.
///
/// Tries these locations in order:
/// 1. `DOCKER_HOST` env var (bollard default)
/// 2. `/var/run/docker.sock` (Linux default)
/// 3. `~/.docker/run/docker.sock` (Docker Desktop on macOS)
pub async fn connect_docker() -> Result<Docker> {
Docker::connect_with_local_defaults().map_err(|e| SandboxError::DockerNotAvailable {
reason: e.to_string(),
// First try bollard defaults (checks DOCKER_HOST, then /var/run/docker.sock)
if let Ok(docker) = Docker::connect_with_local_defaults() {
if docker.ping().await.is_ok() {
return Ok(docker);
}
}
// Try Docker Desktop socket (macOS)
if let Some(home) = std::env::var_os("HOME") {
let desktop_sock = std::path::Path::new(&home).join(".docker/run/docker.sock");
if desktop_sock.exists() {
let sock_str = desktop_sock.to_string_lossy();
if let Ok(docker) =
Docker::connect_with_socket(&sock_str, 120, bollard::API_DEFAULT_VERSION)
{
if docker.ping().await.is_ok() {
return Ok(docker);
}
}
}
}
Err(SandboxError::DockerNotAvailable {
reason: "Socket not found: /var/run/docker.sock".to_string(),
})
}
src/settings.rs
+101
View File
@@ -149,6 +149,11 @@ pub struct ChannelSettings {
#[serde(default)]
pub http_host: Option<String>,
/// Telegram owner user ID. When set, the bot only responds to this user.
/// Captured during setup by having the user message the bot.
#[serde(default)]
pub telegram_owner_id: Option<i64>,
/// Enabled WASM channels by name.
/// Channels not in this list but present in the channels directory will still load.
/// This is primarily used by the setup wizard to track which channels were configured.
@@ -490,6 +495,51 @@ impl Settings {
.join("settings.json")
}
/// Reconstruct Settings from a flat key-value map (as stored in the DB).
///
/// Each key is a dotted path (e.g., "agent.name"), value is a JSONB value.
/// Missing keys get their default value.
pub fn from_db_map(map: &std::collections::HashMap<String, serde_json::Value>) -> Self {
// Start with defaults, then overlay each DB setting
let mut settings = Self::default();
for (key, value) in map {
// Convert the JSONB value to a string for the existing set() method
let value_str = match value {
serde_json::Value::String(s) => s.clone(),
serde_json::Value::Bool(b) => b.to_string(),
serde_json::Value::Number(n) => n.to_string(),
serde_json::Value::Null => "null".to_string(),
other => other.to_string(),
};
if let Err(e) = settings.set(key, &value_str) {
tracing::warn!(
"Failed to apply DB setting '{}' = '{}': {}",
key,
value_str,
e
);
}
}
settings
}
/// Flatten Settings into a key-value map suitable for DB storage.
///
/// Each entry is a (dotted_path, JSONB value) pair.
pub fn to_db_map(&self) -> std::collections::HashMap<String, serde_json::Value> {
let json = match serde_json::to_value(self) {
Ok(v) => v,
Err(_) => return std::collections::HashMap::new(),
};
let mut map = std::collections::HashMap::new();
collect_settings_json(&json, String::new(), &mut map);
map
}
/// Load settings from disk, returning default if not found.
pub fn load() -> Self {
Self::load_from(&Self::default_path())
@@ -656,6 +706,29 @@ impl Settings {
}
}
/// Recursively collect settings paths with their JSON values (for DB storage).
fn collect_settings_json(
value: &serde_json::Value,
prefix: String,
results: &mut std::collections::HashMap<String, serde_json::Value>,
) {
match value {
serde_json::Value::Object(obj) => {
for (key, val) in obj {
let path = if prefix.is_empty() {
key.clone()
} else {
format!("{}.{}", prefix, key)
};
collect_settings_json(val, path, results);
}
}
other => {
results.insert(prefix, other.clone());
}
}
}
/// Recursively collect settings paths and values.
fn collect_settings(
value: &serde_json::Value,
@@ -799,4 +872,32 @@ mod tests {
assert_eq!(settings.embeddings.provider, "nearai");
assert_eq!(settings.embeddings.model, "text-embedding-3-small");
}
#[test]
fn test_telegram_owner_id_round_trip() {
let dir = tempdir().unwrap();
let path = dir.path().join("settings.json");
let mut settings = Settings::default();
settings.channels.telegram_owner_id = Some(123456789);
settings.save_to(&path).unwrap();
let loaded = Settings::load_from(&path);
assert_eq!(loaded.channels.telegram_owner_id, Some(123456789));
}
#[test]
fn test_telegram_owner_id_default_none() {
let settings = Settings::default();
assert_eq!(settings.channels.telegram_owner_id, None);
}
#[test]
fn test_telegram_owner_id_via_set() {
let mut settings = Settings::default();
settings
.set("channels.telegram_owner_id", "987654321")
.unwrap();
assert_eq!(settings.channels.telegram_owner_id, Some(987654321));
}
}
src/setup/channels.rs
+160 -1
View File
@@ -52,6 +52,16 @@ impl SecretsContext {
.await
.unwrap_or(false)
}
/// Read a secret from the database (decrypted).
pub async fn get_secret(&self, name: &str) -> Result<SecretString, String> {
let decrypted = self
.store
.get_decrypted(&self.user_id, name)
.await
.map_err(|e| format!("Failed to read secret: {}", e))?;
Ok(SecretString::from(decrypted.expose().to_string()))
}
}
/// Result of Telegram setup.
@@ -60,6 +70,7 @@ pub struct TelegramSetupResult {
pub enabled: bool,
pub bot_username: Option<String>,
pub webhook_secret: Option<String>,
pub owner_id: Option<i64>,
}
/// Telegram Bot API response for getMe.
@@ -76,6 +87,32 @@ struct TelegramUser {
first_name: String,
}
/// Telegram Bot API response for getUpdates.
#[derive(Debug, Deserialize)]
struct TelegramGetUpdatesResponse {
ok: bool,
result: Vec<TelegramUpdate>,
}
#[derive(Debug, Deserialize)]
struct TelegramUpdate {
#[allow(dead_code)]
update_id: i64,
message: Option<TelegramUpdateMessage>,
}
#[derive(Debug, Deserialize)]
struct TelegramUpdateMessage {
from: Option<TelegramUpdateUser>,
}
#[derive(Debug, Deserialize)]
struct TelegramUpdateUser {
id: i64,
first_name: String,
username: Option<String>,
}
/// Set up Telegram bot channel.
///
/// Guides the user through:
@@ -96,12 +133,14 @@ pub async fn setup_telegram(secrets: &SecretsContext) -> Result<TelegramSetupRes
if secrets.secret_exists("telegram_bot_token").await {
print_info("Existing Telegram token found in database.");
if !confirm("Replace existing token?", false).map_err(|e| e.to_string())? {
// Still offer to configure webhook secret if not already done
// Still offer to configure webhook secret and owner binding
let webhook_secret = setup_telegram_webhook_secret(secrets).await?;
let owner_id = bind_telegram_owner_flow(secrets).await?;
return Ok(TelegramSetupResult {
enabled: true,
bot_username: None,
webhook_secret,
owner_id,
});
}
}
@@ -122,6 +161,9 @@ pub async fn setup_telegram(secrets: &SecretsContext) -> Result<TelegramSetupRes
secrets.save_secret("telegram_bot_token", &token).await?;
print_success("Token saved to database");
// Bind bot to owner's Telegram account
let owner_id = bind_telegram_owner(&token).await?;
// Offer webhook secret configuration
let webhook_secret = setup_telegram_webhook_secret(secrets).await?;
@@ -129,6 +171,7 @@ pub async fn setup_telegram(secrets: &SecretsContext) -> Result<TelegramSetupRes
enabled: true,
bot_username: username,
webhook_secret,
owner_id,
})
}
Err(e) => {
@@ -141,12 +184,128 @@ pub async fn setup_telegram(secrets: &SecretsContext) -> Result<TelegramSetupRes
enabled: false,
bot_username: None,
webhook_secret: None,
owner_id: None,
})
}
}
}
}
/// Bind the bot to the owner's Telegram account by having them send a message.
///
/// Polls `getUpdates` until a message arrives, then captures the sender's user ID.
/// Returns `None` if the user declines or the flow times out.
async fn bind_telegram_owner(token: &SecretString) -> Result<Option<i64>, String> {
println!();
print_info("Account Binding (recommended):");
print_info("Binding restricts the bot so only YOU can use it.");
print_info("Without this, anyone who finds your bot can send it messages.");
println!();
if !confirm("Bind bot to your Telegram account?", true).map_err(|e| e.to_string())? {
print_info("Skipping account binding. Bot will accept messages from all users.");
return Ok(None);
}
print_info("Send any message (e.g. /start) to your bot in Telegram.");
print_info("Waiting for your message (up to 120 seconds)...");
let client = Client::builder()
.timeout(std::time::Duration::from_secs(35))
.build()
.map_err(|e| format!("Failed to create HTTP client: {}", e))?;
// Clear any existing webhook so getUpdates works
let delete_url = format!(
"https://api.telegram.org/bot{}/deleteWebhook",
token.expose_secret()
);
let _ = client.post(&delete_url).send().await;
let updates_url = format!(
"https://api.telegram.org/bot{}/getUpdates",
token.expose_secret()
);
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(120);
while std::time::Instant::now() < deadline {
let response = client
.get(&updates_url)
.query(&[("timeout", "30"), ("allowed_updates", "[\"message\"]")])
.send()
.await
.map_err(|e| format!("getUpdates request failed: {}", e))?;
if !response.status().is_success() {
return Err(format!("getUpdates returned status {}", response.status()));
}
let body: TelegramGetUpdatesResponse = response
.json()
.await
.map_err(|e| format!("Failed to parse getUpdates response: {}", e))?;
if !body.ok {
return Err("Telegram API returned error for getUpdates".to_string());
}
// Find the first message with a sender
for update in &body.result {
if let Some(ref msg) = update.message {
if let Some(ref from) = msg.from {
let display_name = from
.username
.as_ref()
.map(|u| format!("@{}", u))
.unwrap_or_else(|| from.first_name.clone());
print_success(&format!(
"Received message from {} (ID: {})",
display_name, from.id
));
// Acknowledge the update so it doesn't pile up
let ack_url = format!(
"https://api.telegram.org/bot{}/getUpdates",
token.expose_secret()
);
let _ = client
.get(&ack_url)
.query(&[("offset", &(update.update_id + 1).to_string())])
.send()
.await;
return Ok(Some(from.id));
}
}
}
}
print_error("Timed out waiting for a message. You can re-run setup to try again.");
print_info("Bot will accept messages from all users until owner is bound.");
Ok(None)
}
/// Bind flow when the token already exists (reads from secrets store).
///
/// Retrieves the saved bot token and delegates to `bind_telegram_owner`.
async fn bind_telegram_owner_flow(secrets: &SecretsContext) -> Result<Option<i64>, String> {
// Check current settings first
let settings = Settings::load();
if settings.channels.telegram_owner_id.is_some() {
print_info("Bot is already bound to a Telegram account.");
if !confirm("Re-bind to a different account?", false).map_err(|e| e.to_string())? {
return Ok(settings.channels.telegram_owner_id);
}
}
// We need the token to poll getUpdates
let token = secrets.get_secret("telegram_bot_token").await?;
bind_telegram_owner(&token).await
}
/// Set up a tunnel for exposing the agent to the internet.
///
/// This is shared across all channels that need webhook endpoints.
src/setup/wizard.rs
+24 -8
View File
@@ -17,7 +17,7 @@ use secrecy::SecretString;
use tokio_postgres::NoTls;
use crate::channels::wasm::{
ChannelCapabilitiesFile, bundled_channel_names, install_bundled_channel,
ChannelCapabilitiesFile, available_channel_names, install_bundled_channel,
};
use crate::llm::{SessionConfig, SessionManager};
use crate::secrets::SecretsCrypto;
@@ -689,6 +689,9 @@ impl SetupWizard {
} else if channel_name == "telegram" {
let telegram_result =
setup_telegram(ctx).await.map_err(SetupError::Channel)?;
if let Some(owner_id) = telegram_result.owner_id {
self.settings.channels.telegram_owner_id = Some(owner_id);
}
crate::setup::channels::WasmChannelSetupResult {
enabled: telegram_result.enabled,
channel_name: "telegram".to_string(),
@@ -978,7 +981,7 @@ async fn install_missing_bundled_channels(
) -> Result<Vec<String>, SetupError> {
let mut installed = Vec::new();
for name in bundled_channel_names().iter().copied() {
for name in available_channel_names().iter().copied() {
if already_installed.contains(name) {
continue;
}
@@ -995,7 +998,7 @@ async fn install_missing_bundled_channels(
fn wasm_channel_option_names(discovered: &[(String, ChannelCapabilitiesFile)]) -> Vec<String> {
let mut names: Vec<String> = discovered.iter().map(|(name, _)| name.clone()).collect();
for bundled in bundled_channel_names().iter().copied() {
for bundled in available_channel_names().iter().copied() {
if !names.iter().any(|name| name == bundled) {
names.push(bundled.to_string());
}
@@ -1009,7 +1012,7 @@ async fn install_selected_bundled_channels(
selected_channels: &[String],
already_installed: &HashSet<String>,
) -> Result<Option<Vec<String>>, SetupError> {
let bundled: HashSet<&str> = bundled_channel_names().iter().copied().collect();
let bundled: HashSet<&str> = available_channel_names().iter().copied().collect();
let selected_missing: HashSet<String> = selected_channels
.iter()
.filter(|name| bundled.contains(name.as_str()) && !already_installed.contains(*name))
@@ -1092,16 +1095,29 @@ mod tests {
}
#[test]
fn test_wasm_channel_option_names_includes_bundled_when_missing() {
fn test_wasm_channel_option_names_includes_available_when_missing() {
let discovered = Vec::new();
let options = wasm_channel_option_names(&discovered);
assert_eq!(options, vec!["telegram".to_string()]);
let available = available_channel_names();
// All available (built) channels should appear
for name in &available {
assert!(
options.contains(&name.to_string()),
"expected '{}' in options",
name
);
}
}
#[test]
fn test_wasm_channel_option_names_dedupes_bundled() {
fn test_wasm_channel_option_names_dedupes_available() {
let discovered = vec![(String::from("telegram"), ChannelCapabilitiesFile::default())];
let options = wasm_channel_option_names(&discovered);
assert_eq!(options, vec!["telegram".to_string()]);
// telegram should appear exactly once despite being both discovered and available
assert_eq!(
options.iter().filter(|n| *n == "telegram").count(),
1,
"telegram should not be duplicated"
);
}
}
src/tools/builder/core.rs
+12 -1
View File
@@ -689,9 +689,20 @@ Create alongside the .wasm file to grant capabilities:
.messages
.push(ChatMessage::user("Continue with the next step."));
}
RespondResult::ToolCalls(tool_calls) => {
RespondResult::ToolCalls {
tool_calls,
content,
} => {
tools_executed = true;
// Add assistant message with tool_calls (OpenAI protocol)
reason_ctx
.messages
.push(ChatMessage::assistant_with_tool_calls(
content,
tool_calls.clone(),
));
// Execute each tool call
for tc in tool_calls {
logs.push(BuildLog {
src/tools/builtin/extension_tools.rs
+55 -18
View File
@@ -185,8 +185,9 @@ impl Tool for ToolAuthTool {
}
fn description(&self) -> &str {
"Authenticate an installed extension. For MCP servers, starts OAuth flow. \
For WASM tools with manual auth, returns instructions; call again with token param to complete."
"Initiate authentication for an extension. For OAuth, returns a URL. \
For manual auth, returns instructions. The user provides their token \
through a secure channel, never through this tool."
}
fn parameters_schema(&self) -> serde_json::Value {
@@ -196,10 +197,6 @@ impl Tool for ToolAuthTool {
"name": {
"type": "string",
"description": "Extension name to authenticate"
},
"token": {
"type": "string",
"description": "API token/key for manual auth (WASM tools). Provide after user gives you the token."
}
},
"required": ["name"]
@@ -218,11 +215,9 @@ impl Tool for ToolAuthTool {
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("name is required".to_string()))?;
let token = params.get("token").and_then(|v| v.as_str());
let result = self
.manager
.auth(name, token)
.auth(name, None)
.await
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
@@ -316,16 +311,53 @@ impl Tool for ToolActivateTool {
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("name is required".to_string()))?;
let result = self
.manager
.activate(name)
.await
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
match self.manager.activate(name).await {
Ok(result) => {
let output = serde_json::to_value(&result)
.unwrap_or_else(|_| serde_json::json!({"error": "serialization failed"}));
Ok(ToolOutput::success(output, start.elapsed()))
}
Err(activate_err) => {
let err_str = activate_err.to_string();
let needs_auth = err_str.contains("authentication")
|| err_str.contains("401")
|| err_str.contains("Unauthorized")
|| err_str.contains("not authenticated");
let output = serde_json::to_value(&result)
.unwrap_or_else(|_| serde_json::json!({"error": "serialization failed"}));
if !needs_auth {
return Err(ToolError::ExecutionFailed(err_str));
}
Ok(ToolOutput::success(output, start.elapsed()))
// Activation failed due to missing auth; initiate auth flow
// so the agent loop can show the auth card.
match self.manager.auth(name, None).await {
Ok(auth_result) if auth_result.status == "authenticated" => {
// Auth succeeded (e.g. env var was set); retry activation.
let result = self
.manager
.activate(name)
.await
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
let output = serde_json::to_value(&result).unwrap_or_else(
|_| serde_json::json!({"error": "serialization failed"}),
);
Ok(ToolOutput::success(output, start.elapsed()))
}
Ok(auth_result) => {
// Auth needs user input (awaiting_token). Return the auth
// result so detect_auth_awaiting picks it up.
let output = serde_json::to_value(&auth_result).unwrap_or_else(
|_| serde_json::json!({"error": "serialization failed"}),
);
Ok(ToolOutput::success(output, start.elapsed()))
}
Err(auth_err) => Err(ToolError::ExecutionFailed(format!(
"Activation failed ({}), and authentication also failed: {}",
err_str, auth_err
))),
}
}
}
}
}
@@ -499,7 +531,11 @@ mod tests {
assert!(tool.requires_approval());
let schema = tool.parameters_schema();
assert!(schema["properties"].get("name").is_some());
assert!(schema["properties"].get("token").is_some());
// token param must NOT be in schema (security: tokens never go through LLM)
assert!(
schema["properties"].get("token").is_none(),
"tool_auth must not have a token parameter"
);
}
#[test]
@@ -550,6 +586,7 @@ mod tests {
std::path::PathBuf::from("/tmp/ironclaw-test-channels"),
None,
"test".to_string(),
None,
))
}
}
src/tools/builtin/file.rs
+17 -1
View File
@@ -11,7 +11,7 @@ use async_trait::async_trait;
use tokio::fs;
use crate::context::JobContext;
use crate::tools::tool::{Tool, ToolError, ToolOutput};
use crate::tools::tool::{Tool, ToolDomain, ToolError, ToolOutput};
use crate::workspace::paths as ws_paths;
/// Well-known workspace filenames that must go through memory_write, not write_file.
@@ -246,6 +246,10 @@ impl Tool for ReadFileTool {
fn requires_approval(&self) -> bool {
true // Reading local files should require approval
}
fn domain(&self) -> ToolDomain {
ToolDomain::Container
}
}
/// Write file contents tool.
@@ -359,6 +363,10 @@ impl Tool for WriteFileTool {
fn requires_sanitization(&self) -> bool {
false // We're writing, not reading external data
}
fn domain(&self) -> ToolDomain {
ToolDomain::Container
}
}
/// List directory contents tool.
@@ -467,6 +475,10 @@ impl Tool for ListDirTool {
fn requires_approval(&self) -> bool {
true // Directory listings can leak filesystem structure
}
fn domain(&self) -> ToolDomain {
ToolDomain::Container
}
}
/// Recursively list directory contents.
@@ -685,6 +697,10 @@ impl Tool for ApplyPatchTool {
fn requires_sanitization(&self) -> bool {
false // We're writing, not reading external data
}
fn domain(&self) -> ToolDomain {
ToolDomain::Container
}
}
#[cfg(test)]
src/tools/builtin/job.rs
+538 -42
View File
@@ -1,77 +1,108 @@
//! Job management tools.
//!
//! These tools allow the LLM to manage jobs:
//! - Create new jobs/tasks
//! - Create new jobs/tasks (with optional sandbox delegation)
//! - List existing jobs
//! - Check job status
//! - Cancel running jobs
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use chrono::Utc;
use uuid::Uuid;
use crate::context::{ContextManager, JobContext, JobState};
use crate::history::{SandboxJobRecord, Store};
use crate::orchestrator::job_manager::{ContainerJobManager, JobMode};
use crate::tools::tool::{Tool, ToolError, ToolOutput};
/// Tool for creating a new job.
///
/// When sandbox deps are injected (via `with_sandbox`), the tool automatically
/// delegates execution to a Docker container. Otherwise it creates an in-memory
/// job via the ContextManager. The LLM never needs to know the difference.
pub struct CreateJobTool {
context_manager: Arc<ContextManager>,
job_manager: Option<Arc<ContainerJobManager>>,
store: Option<Arc<Store>>,
}
impl CreateJobTool {
pub fn new(context_manager: Arc<ContextManager>) -> Self {
Self { context_manager }
}
}
#[async_trait]
impl Tool for CreateJobTool {
fn name(&self) -> &str {
"create_job"
Self {
context_manager,
job_manager: None,
store: None,
}
}
fn description(&self) -> &str {
"Create a new job or task for the agent to work on. Use this when the user wants \
you to do something substantial that should be tracked as a separate job."
/// Inject sandbox dependencies so `create_job` delegates to Docker containers.
pub fn with_sandbox(
mut self,
job_manager: Arc<ContainerJobManager>,
store: Option<Arc<Store>>,
) -> Self {
self.job_manager = Some(job_manager);
self.store = store;
self
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"title": {
"type": "string",
"description": "A short title for the job (max 100 chars)"
},
"description": {
"type": "string",
"description": "Full description of what needs to be done"
fn sandbox_enabled(&self) -> bool {
self.job_manager.is_some()
}
/// Persist a sandbox job record (fire-and-forget).
fn persist_job(&self, record: SandboxJobRecord) {
if let Some(store) = self.store.clone() {
tokio::spawn(async move {
if let Err(e) = store.save_sandbox_job(&record).await {
tracing::warn!(job_id = %record.id, "Failed to persist sandbox job: {}", e);
}
},
"required": ["title", "description"]
})
});
}
}
async fn execute(
/// Update sandbox job status in DB (fire-and-forget).
fn update_status(
&self,
params: serde_json::Value,
job_id: Uuid,
status: &str,
success: Option<bool>,
message: Option<String>,
started_at: Option<chrono::DateTime<Utc>>,
completed_at: Option<chrono::DateTime<Utc>>,
) {
if let Some(store) = self.store.clone() {
let status = status.to_string();
tokio::spawn(async move {
if let Err(e) = store
.update_sandbox_job_status(
job_id,
&status,
success,
message.as_deref(),
started_at,
completed_at,
)
.await
{
tracing::warn!(job_id = %job_id, "Failed to update sandbox job status: {}", e);
}
});
}
}
/// Execute via in-memory ContextManager (no sandbox).
async fn execute_local(
&self,
title: &str,
description: &str,
ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let start = std::time::Instant::now();
let title = params
.get("title")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("missing 'title' parameter".into()))?;
let description = params
.get("description")
.and_then(|v| v.as_str())
.ok_or_else(|| {
ToolError::InvalidParameters("missing 'description' parameter".into())
})?;
match self
.context_manager
.create_job_for_user(&ctx.user_id, title, description)
@@ -95,6 +126,374 @@ impl Tool for CreateJobTool {
}
}
/// Execute via sandboxed Docker container.
async fn execute_sandbox(
&self,
task: &str,
explicit_dir: Option<PathBuf>,
wait: bool,
mode: JobMode,
ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let start = std::time::Instant::now();
let jm = self.job_manager.as_ref().expect("sandbox deps required");
let job_id = Uuid::new_v4();
let (project_dir, browse_id) = resolve_project_dir(explicit_dir, job_id)?;
let project_dir_str = project_dir.display().to_string();
// Persist the job to DB before creating the container.
self.persist_job(SandboxJobRecord {
id: job_id,
task: task.to_string(),
status: "creating".to_string(),
user_id: ctx.user_id.clone(),
project_dir: project_dir_str.clone(),
success: None,
failure_reason: None,
created_at: Utc::now(),
started_at: None,
completed_at: None,
});
// Persist the job mode to DB
if mode == JobMode::ClaudeCode {
if let Some(store) = self.store.clone() {
let job_id_copy = job_id;
tokio::spawn(async move {
if let Err(e) = store
.update_sandbox_job_mode(job_id_copy, "claude_code")
.await
{
tracing::warn!(job_id = %job_id_copy, "Failed to set job mode: {}", e);
}
});
}
}
// Create the container job with the pre-determined job_id.
let _token = jm
.create_job(job_id, task, Some(project_dir), mode)
.await
.map_err(|e| {
self.update_status(
job_id,
"failed",
Some(false),
Some(e.to_string()),
None,
Some(Utc::now()),
);
ToolError::ExecutionFailed(format!("failed to create container: {}", e))
})?;
// Container started successfully.
let now = Utc::now();
self.update_status(job_id, "running", None, None, Some(now), None);
if !wait {
let result = serde_json::json!({
"job_id": job_id.to_string(),
"status": "started",
"message": "Container started. Use job tools to check status.",
"project_dir": project_dir_str,
"browse_url": format!("/projects/{}", browse_id),
});
return Ok(ToolOutput::success(result, start.elapsed()));
}
// Wait for completion by polling the container state.
let timeout = Duration::from_secs(600);
let poll_interval = Duration::from_secs(2);
let deadline = tokio::time::Instant::now() + timeout;
loop {
if tokio::time::Instant::now() > deadline {
let _ = jm.stop_job(job_id).await;
jm.cleanup_job(job_id).await;
self.update_status(
job_id,
"failed",
Some(false),
Some("Timed out (10 minutes)".to_string()),
None,
Some(Utc::now()),
);
return Err(ToolError::ExecutionFailed(
"container execution timed out (10 minutes)".to_string(),
));
}
match jm.get_handle(job_id).await {
Some(handle) => match handle.state {
crate::orchestrator::job_manager::ContainerState::Running
| crate::orchestrator::job_manager::ContainerState::Creating => {
tokio::time::sleep(poll_interval).await;
}
crate::orchestrator::job_manager::ContainerState::Stopped => {
let message = handle
.completion_result
.as_ref()
.and_then(|r| r.message.clone())
.unwrap_or_else(|| "Container job completed".to_string());
let success = handle
.completion_result
.as_ref()
.map(|r| r.success)
.unwrap_or(true);
jm.cleanup_job(job_id).await;
let finished_at = Utc::now();
if success {
self.update_status(
job_id,
"completed",
Some(true),
None,
None,
Some(finished_at),
);
let result = serde_json::json!({
"job_id": job_id.to_string(),
"status": "completed",
"output": message,
"project_dir": project_dir_str,
"browse_url": format!("/projects/{}", browse_id),
});
return Ok(ToolOutput::success(result, start.elapsed()));
} else {
self.update_status(
job_id,
"failed",
Some(false),
Some(message.clone()),
None,
Some(finished_at),
);
return Err(ToolError::ExecutionFailed(format!(
"container job failed: {}",
message
)));
}
}
crate::orchestrator::job_manager::ContainerState::Failed => {
let message = handle
.completion_result
.as_ref()
.and_then(|r| r.message.clone())
.unwrap_or_else(|| "unknown failure".to_string());
jm.cleanup_job(job_id).await;
self.update_status(
job_id,
"failed",
Some(false),
Some(message.clone()),
None,
Some(Utc::now()),
);
return Err(ToolError::ExecutionFailed(format!(
"container job failed: {}",
message
)));
}
},
None => {
self.update_status(
job_id,
"completed",
Some(true),
None,
None,
Some(Utc::now()),
);
let result = serde_json::json!({
"job_id": job_id.to_string(),
"status": "completed",
"output": "Container job completed",
"project_dir": project_dir_str,
"browse_url": format!("/projects/{}", browse_id),
});
return Ok(ToolOutput::success(result, start.elapsed()));
}
}
}
}
}
/// The base directory where all project directories must live.
fn projects_base() -> PathBuf {
dirs::home_dir()
.unwrap_or_else(|| PathBuf::from("."))
.join(".ironclaw")
.join("projects")
}
/// Resolve the project directory, creating it if it doesn't exist.
///
/// Auto-creates `~/.ironclaw/projects/{project_id}/` so every sandbox job has a
/// persistent bind mount that survives container teardown.
///
/// When an explicit path is provided (e.g. job restarts reusing the old dir),
/// it is validated to fall within `~/.ironclaw/projects/` after canonicalization.
fn resolve_project_dir(
explicit: Option<PathBuf>,
project_id: Uuid,
) -> Result<(PathBuf, String), ToolError> {
let base = projects_base();
std::fs::create_dir_all(&base).map_err(|e| {
ToolError::ExecutionFailed(format!(
"failed to create projects base {}: {}",
base.display(),
e
))
})?;
let canonical_base = base.canonicalize().map_err(|e| {
ToolError::ExecutionFailed(format!("failed to canonicalize projects base: {}", e))
})?;
let dir = match explicit {
Some(d) => d,
None => canonical_base.join(project_id.to_string()),
};
std::fs::create_dir_all(&dir).map_err(|e| {
ToolError::ExecutionFailed(format!(
"failed to create project dir {}: {}",
dir.display(),
e
))
})?;
// Canonicalize resolves symlinks, `..`, etc. so we can do a reliable prefix check.
let canonical_dir = dir.canonicalize().map_err(|e| {
ToolError::ExecutionFailed(format!(
"failed to canonicalize project dir {}: {}",
dir.display(),
e
))
})?;
if !canonical_dir.starts_with(&canonical_base) {
return Err(ToolError::InvalidParameters(format!(
"project directory must be under {}",
canonical_base.display()
)));
}
let browse_id = canonical_dir
.file_name()
.map(|n| n.to_string_lossy().to_string())
.unwrap_or_else(|| project_id.to_string());
Ok((canonical_dir, browse_id))
}
#[async_trait]
impl Tool for CreateJobTool {
fn name(&self) -> &str {
"create_job"
}
fn description(&self) -> &str {
if self.sandbox_enabled() {
"Create and execute a job. The job runs in a sandboxed Docker container with its own \
sub-agent that has shell, file read/write, list_dir, and apply_patch tools. Use this \
whenever the user asks you to build, create, or work on something. The task \
description should be detailed enough for the sub-agent to work independently. \
Set wait=false to start immediately while continuing the conversation. Set mode \
to 'claude_code' for complex software engineering tasks."
} else {
"Create a new job or task for the agent to work on. Use this when the user wants \
you to do something substantial that should be tracked as a separate job."
}
}
fn parameters_schema(&self) -> serde_json::Value {
if self.sandbox_enabled() {
serde_json::json!({
"type": "object",
"properties": {
"title": {
"type": "string",
"description": "Clear description of what to accomplish"
},
"description": {
"type": "string",
"description": "Full description of what needs to be done"
},
"wait": {
"type": "boolean",
"description": "If true (default), wait for the container to complete and return results. \
If false, start the container and return the job_id immediately."
},
"mode": {
"type": "string",
"enum": ["worker", "claude_code"],
"description": "Execution mode. 'worker' (default) uses the IronClaw sub-agent. \
'claude_code' uses Claude Code CLI for full agentic software engineering."
}
},
"required": ["title", "description"]
})
} else {
serde_json::json!({
"type": "object",
"properties": {
"title": {
"type": "string",
"description": "A short title for the job (max 100 chars)"
},
"description": {
"type": "string",
"description": "Full description of what needs to be done"
}
},
"required": ["title", "description"]
})
}
}
fn execution_timeout(&self) -> Duration {
if self.sandbox_enabled() {
// Sandbox polls for up to 10 min internally; give an extra 60s buffer.
Duration::from_secs(660)
} else {
Duration::from_secs(30)
}
}
async fn execute(
&self,
params: serde_json::Value,
ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let title = params
.get("title")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("missing 'title' parameter".into()))?;
let description = params
.get("description")
.and_then(|v| v.as_str())
.ok_or_else(|| {
ToolError::InvalidParameters("missing 'description' parameter".into())
})?;
if self.sandbox_enabled() {
let wait = params.get("wait").and_then(|v| v.as_bool()).unwrap_or(true);
let mode = match params.get("mode").and_then(|v| v.as_str()) {
Some("claude_code") => JobMode::ClaudeCode,
_ => JobMode::Worker,
};
// Combine title and description into the task prompt for the sub-agent.
let task = format!("{}\n\n{}", title, description);
self.execute_sandbox(&task, None, wait, mode, ctx).await
} else {
self.execute_local(title, description, ctx).await
}
}
fn requires_sanitization(&self) -> bool {
false
}
@@ -377,10 +776,13 @@ mod tests {
use super::*;
#[tokio::test]
async fn test_create_job_tool() {
async fn test_create_job_tool_local() {
let manager = Arc::new(ContextManager::new(5));
let tool = CreateJobTool::new(manager.clone());
// Without sandbox deps, it should use the local path
assert!(!tool.sandbox_enabled());
let params = serde_json::json!({
"title": "Test Job",
"description": "A test job description"
@@ -391,6 +793,37 @@ mod tests {
let job_id = result.result.get("job_id").unwrap().as_str().unwrap();
assert!(!job_id.is_empty());
assert_eq!(
result.result.get("status").unwrap().as_str().unwrap(),
"pending"
);
}
#[test]
fn test_schema_changes_with_sandbox() {
let manager = Arc::new(ContextManager::new(5));
// Without sandbox
let tool = CreateJobTool::new(Arc::clone(&manager));
let schema = tool.parameters_schema();
let props = schema.get("properties").unwrap().as_object().unwrap();
assert!(props.contains_key("title"));
assert!(props.contains_key("description"));
assert!(
!props.contains_key("project_dir"),
"project_dir must not be exposed to the LLM"
);
assert!(!props.contains_key("wait"));
assert!(!props.contains_key("mode"));
}
#[test]
fn test_execution_timeout_sandbox() {
let manager = Arc::new(ContextManager::new(5));
// Without sandbox: default timeout
let tool = CreateJobTool::new(Arc::clone(&manager));
assert_eq!(tool.execution_timeout(), Duration::from_secs(30));
}
#[tokio::test]
@@ -429,4 +862,67 @@ mod tests {
"Test Job"
);
}
#[test]
fn test_resolve_project_dir_auto() {
let project_id = Uuid::new_v4();
let (dir, browse_id) = resolve_project_dir(None, project_id).unwrap();
assert!(dir.exists());
assert!(dir.ends_with(project_id.to_string()));
assert_eq!(browse_id, project_id.to_string());
// Must be under the projects base
let base = projects_base().canonicalize().unwrap();
assert!(dir.starts_with(&base));
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn test_resolve_project_dir_explicit_under_base() {
let base = projects_base();
std::fs::create_dir_all(&base).unwrap();
let explicit = base.join("test_explicit_project");
let project_id = Uuid::new_v4();
let (dir, browse_id) = resolve_project_dir(Some(explicit.clone()), project_id).unwrap();
assert!(dir.exists());
assert_eq!(browse_id, "test_explicit_project");
let canonical_base = base.canonicalize().unwrap();
assert!(dir.starts_with(&canonical_base));
let _ = std::fs::remove_dir_all(&explicit);
}
#[test]
fn test_resolve_project_dir_rejects_outside_base() {
let tmp = tempfile::tempdir().unwrap();
let escape_attempt = tmp.path().join("evil_project");
let result = resolve_project_dir(Some(escape_attempt), Uuid::new_v4());
assert!(result.is_err());
let err = result.unwrap_err().to_string();
assert!(
err.contains("must be under"),
"expected 'must be under' error, got: {}",
err
);
}
#[test]
fn test_resolve_project_dir_rejects_traversal() {
// Attempt to escape via `..` components
let base = projects_base();
let traversal = base.join("legit").join("..").join("..").join(".ssh");
let result = resolve_project_dir(Some(traversal), Uuid::new_v4());
assert!(result.is_err());
let err = result.unwrap_err().to_string();
assert!(
err.contains("must be under"),
"expected 'must be under' error, got: {}",
err
);
}
}
src/tools/builtin/mod.rs
+4
View File
@@ -10,6 +10,7 @@ mod json;
mod marketplace;
mod memory;
mod restaurant;
pub mod routine;
mod shell;
mod taskrabbit;
mod time;
@@ -26,6 +27,9 @@ pub use json::JsonTool;
pub use marketplace::MarketplaceTool;
pub use memory::{MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool};
pub use restaurant::RestaurantTool;
pub use routine::{
RoutineCreateTool, RoutineDeleteTool, RoutineHistoryTool, RoutineListTool, RoutineUpdateTool,
};
pub use shell::ShellTool;
pub use taskrabbit::TaskRabbitTool;
pub use time::TimeTool;
src/tools/builtin/routine.rs
+654
View File
@@ -0,0 +1,654 @@
//! LLM-facing tools for managing routines.
//!
//! Five tools let the agent manage routines conversationally:
//! - `routine_create` - Create a new routine
//! - `routine_list` - List all routines with status
//! - `routine_update` - Modify or toggle a routine
//! - `routine_delete` - Remove a routine
//! - `routine_history` - View past runs
use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use chrono::Utc;
use uuid::Uuid;
use crate::agent::routine::{
NotifyConfig, Routine, RoutineAction, RoutineGuardrails, Trigger, next_cron_fire,
};
use crate::agent::routine_engine::RoutineEngine;
use crate::context::JobContext;
use crate::history::Store;
use crate::tools::tool::{Tool, ToolError, ToolOutput};
// ==================== routine_create ====================
pub struct RoutineCreateTool {
store: Arc<Store>,
engine: Arc<RoutineEngine>,
}
impl RoutineCreateTool {
pub fn new(store: Arc<Store>, engine: Arc<RoutineEngine>) -> Self {
Self { store, engine }
}
}
#[async_trait]
impl Tool for RoutineCreateTool {
fn name(&self) -> &str {
"routine_create"
}
fn description(&self) -> &str {
"Create a new routine (scheduled or event-driven task). \
Supports cron schedules, event pattern matching, webhooks, and manual triggers. \
Use this when the user wants something to happen periodically or reactively."
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"name": {
"type": "string",
"description": "Unique name for the routine (e.g. 'daily-pr-review')"
},
"description": {
"type": "string",
"description": "What this routine does"
},
"trigger_type": {
"type": "string",
"enum": ["cron", "event", "webhook", "manual"],
"description": "When the routine fires"
},
"schedule": {
"type": "string",
"description": "Cron expression (for cron trigger). E.g. '0 9 * * MON-FRI' for weekdays at 9am. Uses 6-field cron (sec min hour day month weekday)."
},
"event_pattern": {
"type": "string",
"description": "Regex pattern to match messages (for event trigger)"
},
"event_channel": {
"type": "string",
"description": "Optional channel filter for event trigger (e.g. 'telegram')"
},
"prompt": {
"type": "string",
"description": "The prompt/instructions for the routine"
},
"context_paths": {
"type": "array",
"items": { "type": "string" },
"description": "Workspace paths to load as context (e.g. ['context/priorities.md'])"
},
"action_type": {
"type": "string",
"enum": ["lightweight", "full_job"],
"description": "Execution mode: 'lightweight' (single LLM call, default) or 'full_job' (multi-turn with tools)"
},
"cooldown_secs": {
"type": "integer",
"description": "Minimum seconds between fires (default: 300)"
}
},
"required": ["name", "trigger_type", "prompt"]
})
}
async fn execute(
&self,
params: serde_json::Value,
ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let start = std::time::Instant::now();
let name = params
.get("name")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("missing 'name'".to_string()))?;
let description = params
.get("description")
.and_then(|v| v.as_str())
.unwrap_or("");
let trigger_type = params
.get("trigger_type")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("missing 'trigger_type'".to_string()))?;
let prompt = params
.get("prompt")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("missing 'prompt'".to_string()))?;
// Build trigger
let trigger = match trigger_type {
"cron" => {
let schedule =
params
.get("schedule")
.and_then(|v| v.as_str())
.ok_or_else(|| {
ToolError::InvalidParameters(
"cron trigger requires 'schedule'".to_string(),
)
})?;
// Validate cron expression
next_cron_fire(schedule).map_err(|e| {
ToolError::InvalidParameters(format!("invalid cron schedule: {e}"))
})?;
Trigger::Cron {
schedule: schedule.to_string(),
}
}
"event" => {
let pattern = params
.get("event_pattern")
.and_then(|v| v.as_str())
.ok_or_else(|| {
ToolError::InvalidParameters(
"event trigger requires 'event_pattern'".to_string(),
)
})?;
// Validate regex
regex::Regex::new(pattern)
.map_err(|e| ToolError::InvalidParameters(format!("invalid regex: {e}")))?;
let channel = params
.get("event_channel")
.and_then(|v| v.as_str())
.map(String::from);
Trigger::Event {
channel,
pattern: pattern.to_string(),
}
}
"webhook" => Trigger::Webhook {
path: None,
secret: None,
},
"manual" => Trigger::Manual,
other => {
return Err(ToolError::InvalidParameters(format!(
"unknown trigger_type: {other}"
)));
}
};
// Build action
let action_type = params
.get("action_type")
.and_then(|v| v.as_str())
.unwrap_or("lightweight");
let context_paths: Vec<String> = params
.get("context_paths")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_str().map(String::from))
.collect()
})
.unwrap_or_default();
let action = match action_type {
"lightweight" => RoutineAction::Lightweight {
prompt: prompt.to_string(),
context_paths,
max_tokens: 4096,
},
"full_job" => RoutineAction::FullJob {
title: name.to_string(),
description: prompt.to_string(),
max_iterations: 10,
},
other => {
return Err(ToolError::InvalidParameters(format!(
"unknown action_type: {other}"
)));
}
};
let cooldown_secs = params
.get("cooldown_secs")
.and_then(|v| v.as_u64())
.unwrap_or(300);
// Compute next fire time for cron
let next_fire = if let Trigger::Cron { ref schedule } = trigger {
next_cron_fire(schedule).unwrap_or(None)
} else {
None
};
let routine = Routine {
id: Uuid::new_v4(),
name: name.to_string(),
description: description.to_string(),
user_id: ctx.user_id.clone(),
enabled: true,
trigger,
action,
guardrails: RoutineGuardrails {
cooldown: Duration::from_secs(cooldown_secs),
max_concurrent: 1,
dedup_window: None,
},
notify: NotifyConfig::default(),
last_run_at: None,
next_fire_at: next_fire,
run_count: 0,
consecutive_failures: 0,
state: serde_json::json!({}),
created_at: Utc::now(),
updated_at: Utc::now(),
};
self.store
.create_routine(&routine)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("failed to create routine: {e}")))?;
// Refresh event cache if this is an event trigger
if routine.trigger.type_tag() == "event" {
self.engine.refresh_event_cache().await;
}
let result = serde_json::json!({
"id": routine.id.to_string(),
"name": routine.name,
"trigger_type": routine.trigger.type_tag(),
"next_fire_at": routine.next_fire_at.map(|t| t.to_rfc3339()),
"status": "created",
});
Ok(ToolOutput::success(result, start.elapsed()))
}
fn requires_sanitization(&self) -> bool {
false
}
}
// ==================== routine_list ====================
pub struct RoutineListTool {
store: Arc<Store>,
}
impl RoutineListTool {
pub fn new(store: Arc<Store>) -> Self {
Self { store }
}
}
#[async_trait]
impl Tool for RoutineListTool {
fn name(&self) -> &str {
"routine_list"
}
fn description(&self) -> &str {
"List all routines with their status, trigger info, and next fire time."
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {},
"required": []
})
}
async fn execute(
&self,
_params: serde_json::Value,
ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let start = std::time::Instant::now();
let routines = self
.store
.list_routines(&ctx.user_id)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("failed to list routines: {e}")))?;
let list: Vec<serde_json::Value> = routines
.iter()
.map(|r| {
serde_json::json!({
"id": r.id.to_string(),
"name": r.name,
"description": r.description,
"enabled": r.enabled,
"trigger_type": r.trigger.type_tag(),
"action_type": r.action.type_tag(),
"last_run_at": r.last_run_at.map(|t| t.to_rfc3339()),
"next_fire_at": r.next_fire_at.map(|t| t.to_rfc3339()),
"run_count": r.run_count,
"consecutive_failures": r.consecutive_failures,
})
})
.collect();
let result = serde_json::json!({
"count": list.len(),
"routines": list,
});
Ok(ToolOutput::success(result, start.elapsed()))
}
fn requires_sanitization(&self) -> bool {
false
}
}
// ==================== routine_update ====================
pub struct RoutineUpdateTool {
store: Arc<Store>,
engine: Arc<RoutineEngine>,
}
impl RoutineUpdateTool {
pub fn new(store: Arc<Store>, engine: Arc<RoutineEngine>) -> Self {
Self { store, engine }
}
}
#[async_trait]
impl Tool for RoutineUpdateTool {
fn name(&self) -> &str {
"routine_update"
}
fn description(&self) -> &str {
"Update an existing routine. Can modify trigger, prompt, schedule, or toggle enabled state. \
Pass the routine name and only the fields you want to change."
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"name": {
"type": "string",
"description": "Name of the routine to update"
},
"enabled": {
"type": "boolean",
"description": "Enable or disable the routine"
},
"prompt": {
"type": "string",
"description": "New prompt/instructions"
},
"schedule": {
"type": "string",
"description": "New cron schedule (for cron triggers)"
},
"description": {
"type": "string",
"description": "New description"
}
},
"required": ["name"]
})
}
async fn execute(
&self,
params: serde_json::Value,
ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let start = std::time::Instant::now();
let name = params
.get("name")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("missing 'name'".to_string()))?;
let mut routine = self
.store
.get_routine_by_name(&ctx.user_id, name)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("DB error: {e}")))?
.ok_or_else(|| ToolError::ExecutionFailed(format!("routine '{}' not found", name)))?;
// Apply updates
if let Some(enabled) = params.get("enabled").and_then(|v| v.as_bool()) {
routine.enabled = enabled;
}
if let Some(desc) = params.get("description").and_then(|v| v.as_str()) {
routine.description = desc.to_string();
}
if let Some(prompt) = params.get("prompt").and_then(|v| v.as_str()) {
match &mut routine.action {
RoutineAction::Lightweight { prompt: p, .. } => *p = prompt.to_string(),
RoutineAction::FullJob { description: d, .. } => *d = prompt.to_string(),
}
}
if let Some(schedule) = params.get("schedule").and_then(|v| v.as_str()) {
// Validate
next_cron_fire(schedule)
.map_err(|e| ToolError::InvalidParameters(format!("invalid cron schedule: {e}")))?;
routine.trigger = Trigger::Cron {
schedule: schedule.to_string(),
};
routine.next_fire_at = next_cron_fire(schedule).unwrap_or(None);
}
self.store
.update_routine(&routine)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("failed to update: {e}")))?;
// Refresh event cache in case trigger changed
self.engine.refresh_event_cache().await;
let result = serde_json::json!({
"name": routine.name,
"enabled": routine.enabled,
"trigger_type": routine.trigger.type_tag(),
"next_fire_at": routine.next_fire_at.map(|t| t.to_rfc3339()),
"status": "updated",
});
Ok(ToolOutput::success(result, start.elapsed()))
}
fn requires_sanitization(&self) -> bool {
false
}
}
// ==================== routine_delete ====================
pub struct RoutineDeleteTool {
store: Arc<Store>,
engine: Arc<RoutineEngine>,
}
impl RoutineDeleteTool {
pub fn new(store: Arc<Store>, engine: Arc<RoutineEngine>) -> Self {
Self { store, engine }
}
}
#[async_trait]
impl Tool for RoutineDeleteTool {
fn name(&self) -> &str {
"routine_delete"
}
fn description(&self) -> &str {
"Delete a routine permanently. This also removes all run history."
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"name": {
"type": "string",
"description": "Name of the routine to delete"
}
},
"required": ["name"]
})
}
async fn execute(
&self,
params: serde_json::Value,
ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let start = std::time::Instant::now();
let name = params
.get("name")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("missing 'name'".to_string()))?;
let routine = self
.store
.get_routine_by_name(&ctx.user_id, name)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("DB error: {e}")))?
.ok_or_else(|| ToolError::ExecutionFailed(format!("routine '{}' not found", name)))?;
let deleted = self
.store
.delete_routine(routine.id)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("failed to delete: {e}")))?;
// Refresh event cache
self.engine.refresh_event_cache().await;
let result = serde_json::json!({
"name": name,
"deleted": deleted,
});
Ok(ToolOutput::success(result, start.elapsed()))
}
fn requires_sanitization(&self) -> bool {
false
}
}
// ==================== routine_history ====================
pub struct RoutineHistoryTool {
store: Arc<Store>,
}
impl RoutineHistoryTool {
pub fn new(store: Arc<Store>) -> Self {
Self { store }
}
}
#[async_trait]
impl Tool for RoutineHistoryTool {
fn name(&self) -> &str {
"routine_history"
}
fn description(&self) -> &str {
"View the execution history of a routine. Shows recent runs with status, duration, and results."
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"name": {
"type": "string",
"description": "Name of the routine"
},
"limit": {
"type": "integer",
"description": "Max runs to return (default: 10)",
"default": 10
}
},
"required": ["name"]
})
}
async fn execute(
&self,
params: serde_json::Value,
ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let start = std::time::Instant::now();
let name = params
.get("name")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("missing 'name'".to_string()))?;
let limit = params
.get("limit")
.and_then(|v| v.as_i64())
.unwrap_or(10)
.min(50);
let routine = self
.store
.get_routine_by_name(&ctx.user_id, name)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("DB error: {e}")))?
.ok_or_else(|| ToolError::ExecutionFailed(format!("routine '{}' not found", name)))?;
let runs = self
.store
.list_routine_runs(routine.id, limit)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("failed to list runs: {e}")))?;
let run_list: Vec<serde_json::Value> = runs
.iter()
.map(|r| {
let duration_secs = r
.completed_at
.map(|c| c.signed_duration_since(r.started_at).num_seconds());
serde_json::json!({
"id": r.id.to_string(),
"trigger_type": r.trigger_type,
"trigger_detail": r.trigger_detail,
"started_at": r.started_at.to_rfc3339(),
"completed_at": r.completed_at.map(|t| t.to_rfc3339()),
"duration_secs": duration_secs,
"status": r.status.to_string(),
"result_summary": r.result_summary,
"tokens_used": r.tokens_used,
})
})
.collect();
let result = serde_json::json!({
"routine": name,
"total_runs": routine.run_count,
"runs": run_list,
});
Ok(ToolOutput::success(result, start.elapsed()))
}
fn requires_sanitization(&self) -> bool {
false
}
}
src/tools/builtin/shell.rs
+5 -1
View File
@@ -30,7 +30,7 @@ use tokio::process::Command;
use crate::context::JobContext;
use crate::sandbox::{SandboxManager, SandboxPolicy};
use crate::tools::tool::{Tool, ToolError, ToolOutput};
use crate::tools::tool::{Tool, ToolDomain, ToolError, ToolOutput};
/// Maximum output size before truncation (64KB).
const MAX_OUTPUT_SIZE: usize = 64 * 1024;
@@ -386,6 +386,10 @@ impl Tool for ShellTool {
fn requires_sanitization(&self) -> bool {
true // Shell output could contain anything
}
fn domain(&self) -> ToolDomain {
ToolDomain::Container
}
}
/// Truncate output to fit within limits.
src/tools/mcp/config.rs
+80
View File
@@ -327,6 +327,86 @@ pub async fn get_mcp_server(name: &str) -> Result<McpServerConfig, ConfigError>
})
}
// ==================== Database-backed MCP server config ====================
/// Load MCP server configurations from the database settings table.
///
/// Falls back to the disk file if DB has no entry.
pub async fn load_mcp_servers_from_db(
store: &crate::history::Store,
user_id: &str,
) -> Result<McpServersFile, ConfigError> {
match store.get_setting(user_id, "mcp_servers").await {
Ok(Some(value)) => {
let config: McpServersFile = serde_json::from_value(value)?;
Ok(config)
}
Ok(None) => {
// No entry in DB, fall back to disk
load_mcp_servers().await
}
Err(e) => {
tracing::warn!(
"Failed to load MCP servers from DB: {}, falling back to disk",
e
);
load_mcp_servers().await
}
}
}
/// Save MCP server configurations to the database settings table.
pub async fn save_mcp_servers_to_db(
store: &crate::history::Store,
user_id: &str,
config: &McpServersFile,
) -> Result<(), ConfigError> {
let value = serde_json::to_value(config)?;
store
.set_setting(user_id, "mcp_servers", &value)
.await
.map_err(|e| {
ConfigError::Io(std::io::Error::new(
std::io::ErrorKind::Other,
e.to_string(),
))
})?;
Ok(())
}
/// Add a new MCP server configuration (DB-backed).
pub async fn add_mcp_server_db(
store: &crate::history::Store,
user_id: &str,
config: McpServerConfig,
) -> Result<(), ConfigError> {
config.validate()?;
let mut servers = load_mcp_servers_from_db(store, user_id).await?;
servers.upsert(config);
save_mcp_servers_to_db(store, user_id, &servers).await?;
Ok(())
}
/// Remove an MCP server by name (DB-backed).
pub async fn remove_mcp_server_db(
store: &crate::history::Store,
user_id: &str,
name: &str,
) -> Result<(), ConfigError> {
let mut servers = load_mcp_servers_from_db(store, user_id).await?;
if !servers.remove(name) {
return Err(ConfigError::ServerNotFound {
name: name.to_string(),
});
}
save_mcp_servers_to_db(store, user_id, &servers).await?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
src/tools/mod.rs
+1 -1
View File
@@ -23,4 +23,4 @@ pub use builder::{
};
pub use registry::ToolRegistry;
pub use sandbox::ToolSandbox;
pub use tool::{Tool, ToolError, ToolOutput};
pub use tool::{Tool, ToolDomain, ToolError, ToolOutput};
src/tools/registry.rs
+79 -4
View File
@@ -7,7 +7,9 @@ use tokio::sync::RwLock;
use crate::context::ContextManager;
use crate::extensions::ExtensionManager;
use crate::history::Store;
use crate::llm::{LlmProvider, ToolDefinition};
use crate::orchestrator::job_manager::ContainerJobManager;
use crate::safety::SafetyLayer;
use crate::tools::builder::{BuildSoftwareTool, BuilderConfig, LlmSoftwareBuilder};
use crate::tools::builtin::{
@@ -16,7 +18,7 @@ use crate::tools::builtin::{
ReadFileTool, ShellTool, TimeTool, ToolActivateTool, ToolAuthTool, ToolInstallTool,
ToolListTool, ToolRemoveTool, ToolSearchTool, WriteFileTool,
};
use crate::tools::tool::Tool;
use crate::tools::tool::{Tool, ToolDomain};
use crate::tools::wasm::{
Capabilities, ResourceLimits, WasmError, WasmStorageError, WasmToolRuntime, WasmToolStore,
WasmToolWrapper,
@@ -120,6 +122,39 @@ impl ToolRegistry {
tracing::info!("Registered {} built-in tools", self.count());
}
/// Register only orchestrator-domain tools (safe for the main process).
///
/// This registers tools that don't touch the filesystem or run shell commands:
/// echo, time, json, http. Use this when `allow_local_tools = false` and
/// container-domain tools should only be available inside sandboxed containers.
pub fn register_orchestrator_tools(&self) {
self.register_builtin_tools();
// register_builtin_tools already only registers orchestrator-domain tools
}
/// Register container-domain tools (filesystem, shell, code).
///
/// These tools are intended to run inside sandboxed Docker containers.
/// Call this in the worker process, not the orchestrator (unless `allow_local_tools = true`).
pub fn register_container_tools(&self) {
self.register_dev_tools();
}
/// Get tool definitions filtered by domain.
pub async fn tool_definitions_for_domain(&self, domain: ToolDomain) -> Vec<ToolDefinition> {
self.tools
.read()
.await
.values()
.filter(|tool| tool.domain() == domain)
.map(|tool| ToolDefinition {
name: tool.name().to_string(),
description: tool.description().to_string(),
parameters: tool.parameters_schema(),
})
.collect()
}
/// Register development tools for building software.
///
/// These tools provide shell access, file operations, and code editing
@@ -151,9 +186,19 @@ impl ToolRegistry {
/// Register job management tools.
///
/// Job tools allow the LLM to create, list, check status, and cancel jobs.
/// These enable natural language job management without hardcoded intent parsing.
pub fn register_job_tools(&self, context_manager: Arc<ContextManager>) {
self.register_sync(Arc::new(CreateJobTool::new(Arc::clone(&context_manager))));
/// When sandbox deps are provided, `create_job` automatically delegates to
/// Docker containers. Otherwise it creates in-memory jobs via ContextManager.
pub fn register_job_tools(
&self,
context_manager: Arc<ContextManager>,
job_manager: Option<Arc<ContainerJobManager>>,
store: Option<Arc<Store>>,
) {
let mut create_tool = CreateJobTool::new(Arc::clone(&context_manager));
if let Some(jm) = job_manager {
create_tool = create_tool.with_sandbox(jm, store);
}
self.register_sync(Arc::new(create_tool));
self.register_sync(Arc::new(ListJobsTool::new(Arc::clone(&context_manager))));
self.register_sync(Arc::new(JobStatusTool::new(Arc::clone(&context_manager))));
self.register_sync(Arc::new(CancelJobTool::new(context_manager)));
@@ -174,6 +219,36 @@ impl ToolRegistry {
tracing::info!("Registered 6 extension management tools");
}
/// Register routine management tools.
///
/// These allow the LLM to create, list, update, delete, and view history
/// of routines (scheduled and event-driven tasks).
pub fn register_routine_tools(
&self,
store: Arc<Store>,
engine: Arc<crate::agent::routine_engine::RoutineEngine>,
) {
use crate::tools::builtin::{
RoutineCreateTool, RoutineDeleteTool, RoutineHistoryTool, RoutineListTool,
RoutineUpdateTool,
};
self.register_sync(Arc::new(RoutineCreateTool::new(
Arc::clone(&store),
Arc::clone(&engine),
)));
self.register_sync(Arc::new(RoutineListTool::new(Arc::clone(&store))));
self.register_sync(Arc::new(RoutineUpdateTool::new(
Arc::clone(&store),
Arc::clone(&engine),
)));
self.register_sync(Arc::new(RoutineDeleteTool::new(
Arc::clone(&store),
Arc::clone(&engine),
)));
self.register_sync(Arc::new(RoutineHistoryTool::new(store)));
tracing::info!("Registered 5 routine management tools");
}
/// Register the software builder tool.
///
/// The builder tool allows the agent to create new software including WASM tools,
src/tools/tool.rs
+35
View File
@@ -9,6 +9,18 @@ use thiserror::Error;
use crate::context::JobContext;
/// Where a tool should execute: orchestrator process or inside a container.
///
/// Orchestrator tools run in the main agent process (memory access, job mgmt, etc).
/// Container tools run inside Docker containers (shell, file ops, code mods).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ToolDomain {
/// Safe to run in the orchestrator (pure functions, memory, job management).
Orchestrator,
/// Must run inside a sandboxed container (filesystem, shell, code).
Container,
}
/// Error type for tool execution.
#[derive(Debug, Error)]
pub enum ToolError {
@@ -160,6 +172,23 @@ pub trait Tool: Send + Sync {
false
}
/// Maximum time this tool is allowed to run before the caller kills it.
/// Override for long-running tools like sandbox execution.
/// Default: 60 seconds.
fn execution_timeout(&self) -> Duration {
Duration::from_secs(60)
}
/// Where this tool should execute.
///
/// `Orchestrator` tools run in the main agent process (safe, no FS access).
/// `Container` tools run inside Docker containers (shell, file ops).
///
/// Default: `Orchestrator` (safe for the main process).
fn domain(&self) -> ToolDomain {
ToolDomain::Orchestrator
}
/// Get the tool schema for LLM function calling.
fn schema(&self) -> ToolSchema {
ToolSchema {
@@ -242,4 +271,10 @@ mod tests {
assert_eq!(schema.name, "echo");
assert!(!schema.description.is_empty());
}
#[test]
fn test_execution_timeout_default() {
let tool = EchoTool;
assert_eq!(tool.execution_timeout(), Duration::from_secs(60));
}
}
src/tools/wasm/loader.rs
+188
View File
@@ -2,6 +2,7 @@
//!
//! This module provides a way to load WASM tools dynamically at runtime from:
//! - A directory containing `<name>.wasm` and `<name>.capabilities.json`
//! - Build artifacts in `tools-src/` (dev mode, auto-detected)
//! - Database storage (via [`WasmToolStore`])
//!
//! # Example: Loading from Directory
@@ -19,6 +20,13 @@
//! loader.load_from_dir(Path::new("~/.ironclaw/tools/")).await?;
//! ```
//!
//! # Dev Mode
//!
//! When `load_dev_tools()` is called, the loader scans `tools-src/*/` for build
//! artifacts. Tools found there are loaded directly from the build output,
//! skipping the install directory. This means during development you just
//! rebuild the WASM and restart the host, no manual copy step needed.
//!
//! # Security
//!
//! Tools loaded from files are assigned `TrustLevel::User` by default, meaning
@@ -312,6 +320,159 @@ impl LoadResults {
}
}
/// Compile-time project root, used to locate tools-src/ in dev builds.
const CARGO_MANIFEST_DIR: &str = env!("CARGO_MANIFEST_DIR");
/// Resolve the tools source directory.
///
/// Checks (in order):
/// 1. `IRONCLAW_TOOLS_SRC` env var
/// 2. `<CARGO_MANIFEST_DIR>/tools-src/` (dev builds)
fn tools_src_dir() -> PathBuf {
if let Ok(dir) = std::env::var("IRONCLAW_TOOLS_SRC") {
return PathBuf::from(dir);
}
PathBuf::from(CARGO_MANIFEST_DIR).join("tools-src")
}
/// Discover WASM tools available as build artifacts in `tools-src/`.
///
/// Scans each subdirectory for:
/// - `tools-src/<name>/target/wasm32-wasip2/release/<crate_name>_tool.wasm`
/// - `tools-src/<name>/<name>-tool.capabilities.json`
///
/// Returns a map of install-name (e.g. "gmail-tool") to paths.
pub async fn discover_dev_tools() -> Result<HashMap<String, DiscoveredTool>, std::io::Error> {
let src_dir = tools_src_dir();
let mut tools = HashMap::new();
if !src_dir.is_dir() {
return Ok(tools);
}
let mut entries = fs::read_dir(&src_dir).await?;
while let Some(entry) = entries.next_entry().await? {
let path = entry.path();
if !path.is_dir() {
continue;
}
let dir_name = match path.file_name().and_then(|n| n.to_str()) {
Some(n) => n.to_string(),
None => continue,
};
// Convention: crate name uses underscores, directory uses hyphens
let crate_name = dir_name.replace('-', "_");
let install_name = format!("{}-tool", dir_name);
let wasm_path = path
.join("target/wasm32-wasip2/release")
.join(format!("{}_tool.wasm", crate_name));
if !wasm_path.exists() {
continue;
}
let caps_path = path.join(format!("{}-tool.capabilities.json", dir_name));
tools.insert(
install_name,
DiscoveredTool {
wasm_path,
capabilities_path: if caps_path.exists() {
Some(caps_path)
} else {
None
},
},
);
}
Ok(tools)
}
/// Load WASM tools from build artifacts in `tools-src/`.
///
/// In dev mode, tools can be loaded directly from their build output without
/// needing to install them to `~/.ironclaw/tools/` first. Build artifacts
/// that are newer than installed copies take priority.
///
/// Set `IRONCLAW_TOOLS_SRC` env var to override the source directory.
pub async fn load_dev_tools(
loader: &WasmToolLoader,
install_dir: &Path,
) -> Result<LoadResults, WasmLoadError> {
let dev_tools = discover_dev_tools().await?;
let mut results = LoadResults::default();
if dev_tools.is_empty() {
return Ok(results);
}
for (name, discovered) in &dev_tools {
// Check if the build artifact is newer than the installed copy
let installed_path = install_dir.join(format!("{}.wasm", name));
let should_load = if installed_path.exists() {
// Compare modification times: prefer fresher build artifact
match (
fs::metadata(&discovered.wasm_path).await,
fs::metadata(&installed_path).await,
) {
(Ok(dev_meta), Ok(inst_meta)) => {
let dev_modified = dev_meta.modified().unwrap_or(std::time::UNIX_EPOCH);
let inst_modified = inst_meta.modified().unwrap_or(std::time::UNIX_EPOCH);
dev_modified > inst_modified
}
_ => true,
}
} else {
true
};
if !should_load {
continue;
}
tracing::info!(
name = name,
wasm_path = %discovered.wasm_path.display(),
"Loading dev tool from build artifacts (newer than installed)"
);
match loader
.load_from_files(
name,
&discovered.wasm_path,
discovered.capabilities_path.as_deref(),
)
.await
{
Ok(()) => {
results.loaded.push(name.clone());
}
Err(e) => {
tracing::error!(
name = name,
error = %e,
"Failed to load dev tool"
);
results.errors.push((discovered.wasm_path.clone(), e));
}
}
}
if !results.loaded.is_empty() {
tracing::info!(
count = results.loaded.len(),
tools = ?results.loaded,
"Loaded dev tools from build artifacts"
);
}
Ok(results)
}
/// Discover WASM tool files in a directory without loading them.
///
/// Returns a map of tool name -> (wasm_path, capabilities_path).
@@ -430,4 +591,31 @@ mod tests {
let err = WasmLoadError::WasmNotFound(std::path::PathBuf::from("/foo/bar.wasm"));
assert!(err.to_string().contains("/foo/bar.wasm"));
}
#[test]
fn test_tools_src_dir_default() {
let dir = super::tools_src_dir();
assert!(dir.ends_with("tools-src"));
}
#[tokio::test]
async fn test_discover_dev_tools_finds_build_artifacts() {
// This test relies on the actual tools-src/ directory in the repo.
// If build artifacts exist, they should be discovered.
let tools = super::discover_dev_tools().await.unwrap();
// If any tools have been built, they should appear with "-tool" suffix
for (name, discovered) in &tools {
assert!(
name.ends_with("-tool"),
"Dev tool name should end with -tool: {}",
name
);
assert!(
discovered.wasm_path.exists(),
"WASM should exist: {:?}",
discovered.wasm_path
);
}
}
}
src/tools/wasm/mod.rs
+4 -1
View File
@@ -115,7 +115,10 @@ pub use storage::{
};
// Loader
pub use loader::{DiscoveredTool, LoadResults, WasmLoadError, WasmToolLoader, discover_tools};
pub use loader::{
DiscoveredTool, LoadResults, WasmLoadError, WasmToolLoader, discover_dev_tools, discover_tools,
load_dev_tools,
};
// Capabilities schema (for parsing *.capabilities.json files)
pub use capabilities_schema::{
src/tools/wasm/wrapper.rs
+306 -154
View File
@@ -1,16 +1,23 @@
//! WASM tool wrapper implementing the Tool trait.
//!
//! Uses wasmtime::component::bindgen! to generate typed bindings from the WIT
//! interface, ensuring all host functions are properly registered under the
//! correct `near:agent/host` namespace.
//!
//! Each execution creates a fresh instance (NEAR pattern) to ensure
//! isolation and deterministic behavior.
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use async_trait::async_trait;
use wasmtime::Store;
use wasmtime::component::{Component, Linker, Val};
use wasmtime::component::{Component, Linker};
use wasmtime_wasi::{ResourceTable, WasiCtx, WasiCtxBuilder, WasiView};
use crate::context::JobContext;
use crate::safety::LeakDetector;
use crate::tools::tool::{Tool, ToolError, ToolOutput};
use crate::tools::wasm::capabilities::Capabilities;
use crate::tools::wasm::error::WasmError;
@@ -18,21 +25,259 @@ use crate::tools::wasm::host::{HostState, LogLevel};
use crate::tools::wasm::limits::{ResourceLimits, WasmResourceLimiter};
use crate::tools::wasm::runtime::{PreparedModule, WasmToolRuntime};
/// Store data for WASM execution.
// Generate component model bindings from the WIT file.
//
// This creates:
// - `near::agent::host::Host` trait + `add_to_linker()` for the import interface
// - `SandboxedTool` struct with `instantiate()` for the world
// - `exports::near::agent::tool::*` types for the export interface
wasmtime::component::bindgen!({
path: "wit/tool.wit",
world: "sandboxed-tool",
async: false,
with: {},
});
// Alias the export interface types for convenience.
use exports::near::agent::tool as wit_tool;
/// Store data for WASM tool execution.
///
/// Contains both the resource limiter and host state.
/// Contains the resource limiter, host state, WASI context, and injected
/// credentials. Fresh instance created per execution (NEAR pattern).
struct StoreData {
limiter: WasmResourceLimiter,
host_state: HostState,
wasi: WasiCtx,
table: ResourceTable,
/// Injected credentials for URL/header substitution.
/// Keys are placeholder names like "GOOGLE_ACCESS_TOKEN".
credentials: HashMap<String, String>,
}
impl StoreData {
fn new(memory_limit: u64, capabilities: Capabilities) -> Self {
fn new(
memory_limit: u64,
capabilities: Capabilities,
credentials: HashMap<String, String>,
) -> Self {
// Minimal WASI context: no filesystem, no env vars (security)
let wasi = WasiCtxBuilder::new().build();
Self {
limiter: WasmResourceLimiter::new(memory_limit),
host_state: HostState::new(capabilities),
wasi,
table: ResourceTable::new(),
credentials,
}
}
/// Inject credentials into a string by replacing placeholders.
///
/// Replaces patterns like `{GOOGLE_ACCESS_TOKEN}` with actual values.
/// WASM tools reference credentials by placeholder, never seeing real values.
fn inject_credentials(&self, input: &str, context: &str) -> String {
let mut result = input.to_string();
for (name, value) in &self.credentials {
let placeholder = format!("{{{}}}", name);
if result.contains(&placeholder) {
tracing::debug!(
placeholder = %placeholder,
context = %context,
"Replacing credential placeholder in tool request"
);
result = result.replace(&placeholder, value);
}
}
result
}
/// Replace injected credential values with `[REDACTED]` in text.
///
/// Prevents credentials from leaking through error messages or logs.
/// reqwest::Error includes the full URL in its Display output, so any
/// error from an injected-URL request will contain the raw credential
/// unless we scrub it.
fn redact_credentials(&self, text: &str) -> String {
let mut result = text.to_string();
for (name, value) in &self.credentials {
if !value.is_empty() {
result = result.replace(value, &format!("[REDACTED:{}]", name));
}
}
result
}
}
// Provide WASI context for the WASM component.
// Required because tools are compiled with wasm32-wasip2 target.
impl WasiView for StoreData {
fn ctx(&mut self) -> &mut WasiCtx {
&mut self.wasi
}
fn table(&mut self) -> &mut ResourceTable {
&mut self.table
}
}
// Implement the generated Host trait from bindgen.
//
// This registers all 6 host functions under the `near:agent/host` namespace:
// log, now-millis, workspace-read, http-request, secret-exists, tool-invoke
impl near::agent::host::Host for StoreData {
fn log(&mut self, level: near::agent::host::LogLevel, message: String) {
let log_level = match level {
near::agent::host::LogLevel::Trace => LogLevel::Trace,
near::agent::host::LogLevel::Debug => LogLevel::Debug,
near::agent::host::LogLevel::Info => LogLevel::Info,
near::agent::host::LogLevel::Warn => LogLevel::Warn,
near::agent::host::LogLevel::Error => LogLevel::Error,
};
let _ = self.host_state.log(log_level, message);
}
fn now_millis(&mut self) -> u64 {
self.host_state.now_millis()
}
fn workspace_read(&mut self, path: String) -> Option<String> {
self.host_state.workspace_read(&path).ok().flatten()
}
fn http_request(
&mut self,
method: String,
url: String,
headers_json: String,
body: Option<Vec<u8>>,
timeout_ms: Option<u32>,
) -> Result<near::agent::host::HttpResponse, String> {
// Inject credentials into URL (e.g., replace {TELEGRAM_BOT_TOKEN})
let injected_url = self.inject_credentials(&url, "url");
// Check HTTP allowlist
self.host_state
.check_http_allowed(&injected_url, &method)
.map_err(|e| format!("HTTP not allowed: {}", e))?;
// Record for rate limiting
self.host_state
.record_http_request()
.map_err(|e| format!("Rate limit exceeded: {}", e))?;
// Parse headers and inject credentials into header values
let raw_headers: HashMap<String, String> =
serde_json::from_str(&headers_json).unwrap_or_default();
let headers: HashMap<String, String> = raw_headers
.into_iter()
.map(|(k, v)| {
(
k.clone(),
self.inject_credentials(&v, &format!("header:{}", k)),
)
})
.collect();
let url = injected_url;
let leak_detector = LeakDetector::new();
let header_vec: Vec<(String, String)> = headers
.iter()
.map(|(k, v)| (k.clone(), v.clone()))
.collect();
leak_detector
.scan_http_request(&url, &header_vec, body.as_deref())
.map_err(|e| format!("Potential secret leak blocked: {}", e))?;
// Make HTTP request using blocking I/O.
// We're inside a spawn_blocking context, so use block_on.
let result = tokio::runtime::Handle::current().block_on(async {
let client = reqwest::Client::new();
let mut request = match method.to_uppercase().as_str() {
"GET" => client.get(&url),
"POST" => client.post(&url),
"PUT" => client.put(&url),
"DELETE" => client.delete(&url),
"PATCH" => client.patch(&url),
"HEAD" => client.head(&url),
_ => return Err(format!("Unsupported HTTP method: {}", method)),
};
for (key, value) in headers {
request = request.header(&key, &value);
}
if let Some(body_bytes) = body {
request = request.body(body_bytes);
}
// Caller-specified timeout (default 30s)
let timeout = Duration::from_millis(timeout_ms.unwrap_or(30_000) as u64);
let response = request.timeout(timeout).send().await.map_err(|e| {
// Walk the full error chain for the actual root cause
let mut chain = format!("HTTP request failed: {}", e);
let mut source = std::error::Error::source(&e);
while let Some(cause) = source {
chain.push_str(&format!(" -> {}", cause));
source = cause.source();
}
chain
})?;
let status = response.status().as_u16();
let response_headers: HashMap<String, String> = response
.headers()
.iter()
.filter_map(|(k, v)| {
v.to_str()
.ok()
.map(|v| (k.as_str().to_string(), v.to_string()))
})
.collect();
let headers_json = serde_json::to_string(&response_headers).unwrap_or_default();
let body = response
.bytes()
.await
.map_err(|e| format!("Failed to read response body: {}", e))?
.to_vec();
// Leak detection on response body
if let Ok(body_str) = std::str::from_utf8(&body) {
leak_detector
.scan_and_clean(body_str)
.map_err(|e| format!("Potential secret leak in response: {}", e))?;
}
Ok(near::agent::host::HttpResponse {
status,
headers_json,
body,
})
});
// Redact credentials from error messages before returning to WASM
result.map_err(|e| self.redact_credentials(&e))
}
fn tool_invoke(&mut self, alias: String, _params_json: String) -> Result<String, String> {
// Validate capability and resolve alias
let _real_name = self.host_state.check_tool_invoke_allowed(&alias)?;
self.host_state.record_tool_invoke()?;
// Tool invocation requires async context and access to the tool registry,
// which aren't available inside a synchronous WASM callback.
Err("Tool invocation from WASM tools is not yet supported".to_string())
}
fn secret_exists(&mut self, name: String) -> bool {
self.host_state.secret_exists(&name)
}
}
/// A Tool implementation backed by a WASM component.
@@ -49,6 +294,9 @@ pub struct WasmToolWrapper {
description: String,
/// Cached schema (from PreparedModule or override).
schema: serde_json::Value,
/// Injected credentials for HTTP requests (e.g., OAuth tokens).
/// Keys are placeholder names like "GOOGLE_ACCESS_TOKEN".
credentials: HashMap<String, String>,
}
impl WasmToolWrapper {
@@ -64,6 +312,7 @@ impl WasmToolWrapper {
runtime,
prepared,
capabilities,
credentials: HashMap::new(),
}
}
@@ -79,11 +328,34 @@ impl WasmToolWrapper {
self
}
/// Set credentials for HTTP request injection.
pub fn with_credentials(mut self, credentials: HashMap<String, String>) -> Self {
self.credentials = credentials;
self
}
/// Get the resource limits for this tool.
pub fn limits(&self) -> &ResourceLimits {
&self.prepared.limits
}
/// Add all host functions to the linker using generated bindings.
///
/// Uses the bindgen-generated `add_to_linker` function to properly register
/// all host functions with correct component model signatures under the
/// `near:agent/host` namespace.
fn add_host_functions(linker: &mut Linker<StoreData>) -> Result<(), WasmError> {
// Add WASI support (required by components built with wasm32-wasip2)
wasmtime_wasi::add_to_linker_sync(linker)
.map_err(|e| WasmError::ConfigError(format!("Failed to add WASI functions: {}", e)))?;
// Add our custom host interface using the generated add_to_linker
near::agent::host::add_to_linker(linker, |state| state)
.map_err(|e| WasmError::ConfigError(format!("Failed to add host functions: {}", e)))?;
Ok(())
}
/// Execute the WASM tool synchronously (called from spawn_blocking).
fn execute_sync(
&self,
@@ -94,7 +366,11 @@ impl WasmToolWrapper {
let limits = &self.prepared.limits;
// Create store with fresh state (NEAR pattern: fresh instance per call)
let store_data = StoreData::new(limits.memory_bytes, self.capabilities.clone());
let store_data = StoreData::new(
limits.memory_bytes,
self.capabilities.clone(),
self.credentials.clone(),
);
let mut store = Store::new(engine, store_data);
// Configure fuel if enabled
@@ -115,172 +391,46 @@ impl WasmToolWrapper {
let component = Component::new(engine, self.prepared.component_bytes())
.map_err(|e| WasmError::CompilationFailed(e.to_string()))?;
// Create linker and add host functions
// Create linker with all host functions properly namespaced
let mut linker = Linker::new(engine);
self.add_host_functions(&mut linker)?;
Self::add_host_functions(&mut linker)?;
// Instantiate the component
let instance = linker
.instantiate(&mut store, &component)
// Instantiate using the generated bindings
let instance = SandboxedTool::instantiate(&mut store, &component, &linker)
.map_err(|e| WasmError::InstantiationFailed(e.to_string()))?;
// Get the execute function
let execute_func = instance
.get_func(&mut store, "execute")
.ok_or_else(|| WasmError::MissingExport("execute".to_string()))?;
// Prepare request
// Prepare the request
let params_json = serde_json::to_string(&params)
.map_err(|e| WasmError::InvalidResponseJson(e.to_string()))?;
// Build request record
// Note: The exact calling convention depends on how WIT records are lowered.
// With component model, we'd use typed bindings from wit-bindgen.
// For now, we use the lower-level Val API.
let request_params = Val::String(params_json);
let request_context = match context_json {
Some(ctx) => Val::Option(Some(Box::new(Val::String(ctx)))),
None => Val::Option(None),
let request = wit_tool::Request {
params: params_json,
context: context_json,
};
// Create request record (params, context)
let request = Val::Record(vec![
("params".to_string(), request_params),
("context".to_string(), request_context),
]);
// Call the function
let mut results = vec![Val::Bool(false)]; // Placeholder for response
execute_func
.call(&mut store, &[request], &mut results)
.map_err(|e| {
// Check for specific trap types
let error_str = e.to_string();
if error_str.contains("out of fuel") {
WasmError::FuelExhausted { limit: limits.fuel }
} else if error_str.contains("unreachable") {
WasmError::Trapped("unreachable code executed".to_string())
} else {
WasmError::Trapped(error_str)
}
})?;
// Post-call completion (cleanup)
execute_func
.post_return(&mut store)
.map_err(|e| WasmError::Trapped(format!("post_return failed: {}", e)))?;
// Extract response
let response = &results[0];
let (result_str, error_str) = extract_response(response)?;
// Call execute using the generated typed interface
let tool_iface = instance.near_agent_tool();
let response = tool_iface.call_execute(&mut store, &request).map_err(|e| {
let error_str = e.to_string();
if error_str.contains("out of fuel") {
WasmError::FuelExhausted { limit: limits.fuel }
} else if error_str.contains("unreachable") {
WasmError::Trapped("unreachable code executed".to_string())
} else {
WasmError::Trapped(error_str)
}
})?;
// Get logs from host state
let logs = store.data_mut().host_state.take_logs();
// Check for tool-level error
if let Some(err) = error_str {
if let Some(err) = response.error {
return Err(WasmError::ToolReturnedError(err));
}
// Return result (or empty string if none)
Ok((result_str.unwrap_or_default(), logs))
}
/// Add host functions to the linker.
fn add_host_functions(&self, linker: &mut Linker<StoreData>) -> Result<(), WasmError> {
// Note: With WIT bindgen, these would be generated automatically.
// For now, we manually define the host functions.
//
// Component model func_wrap signature: F: Fn(StoreContextMut<T>, Params) -> Result<Return>
// where Params is a tuple of the function arguments.
// host.log(level: log-level, message: string)
linker
.root()
.func_wrap(
"log",
|mut ctx: wasmtime::StoreContextMut<'_, StoreData>,
(level, message): (i32, String)| {
let log_level = match level {
0 => LogLevel::Trace,
1 => LogLevel::Debug,
2 => LogLevel::Info,
3 => LogLevel::Warn,
4 => LogLevel::Error,
_ => LogLevel::Info,
};
// Ignore errors from logging (rate limiting)
let _ = ctx.data_mut().host_state.log(log_level, message);
Ok(())
},
)
.map_err(|e| WasmError::ConfigError(format!("Failed to add log function: {}", e)))?;
// host.now-millis() -> u64
linker
.root()
.func_wrap(
"now-millis",
|ctx: wasmtime::StoreContextMut<'_, StoreData>, (): ()| -> anyhow::Result<(u64,)> {
Ok((ctx.data().host_state.now_millis(),))
},
)
.map_err(|e| {
WasmError::ConfigError(format!("Failed to add now-millis function: {}", e))
})?;
// host.workspace-read(path: string) -> option<string>
linker
.root()
.func_wrap(
"workspace-read",
|ctx: wasmtime::StoreContextMut<'_, StoreData>,
(path,): (String,)|
-> anyhow::Result<(Option<String>,)> {
let result = ctx.data().host_state.workspace_read(&path).ok().flatten();
Ok((result,))
},
)
.map_err(|e| {
WasmError::ConfigError(format!("Failed to add workspace-read function: {}", e))
})?;
Ok(())
}
}
/// Extract result and error from a WIT response record.
fn extract_response(response: &Val) -> Result<(Option<String>, Option<String>), WasmError> {
match response {
Val::Record(fields) => {
let mut result = None;
let mut error = None;
for (name, val) in fields {
match name.as_str() {
"output" => {
if let Val::Option(Some(inner)) = val {
if let Val::String(s) = inner.as_ref() {
result = Some(s.to_string());
}
}
}
"error" => {
if let Val::Option(Some(inner)) = val {
if let Val::String(s) = inner.as_ref() {
error = Some(s.to_string());
}
}
}
_ => {}
}
}
Ok((result, error))
}
_ => Err(WasmError::InvalidResponseJson(
"Expected record response".to_string(),
)),
Ok((response.output.unwrap_or_default(), logs))
}
}
@@ -315,6 +465,7 @@ impl Tool for WasmToolWrapper {
let capabilities = self.capabilities.clone();
let description = self.description.clone();
let schema = self.schema.clone();
let credentials = self.credentials.clone();
// Execute in blocking task with timeout
let result = tokio::time::timeout(timeout, async move {
@@ -324,6 +475,7 @@ impl Tool for WasmToolWrapper {
capabilities,
description,
schema,
credentials,
};
tokio::task::spawn_blocking(move || wrapper.execute_sync(params, context_json))
@@ -359,7 +511,7 @@ impl Tool for WasmToolWrapper {
}
fn requires_sanitization(&self) -> bool {
// WASM tools always require sanitization - they're untrusted by definition
// WASM tools always require sanitization, they're untrusted by definition
true
}
src/worker/api.rs
+400
View File
@@ -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);
}
}
src/worker/claude_bridge.rs
+644
View File
@@ -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), "");
}
}
src/worker/mod.rs
+35
View File
@@ -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;
src/worker/proxy_llm.rs
+95
View File
@@ -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);
}
}
src/worker/runtime.rs
+491
View File
@@ -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])
}
}
src/workspace/mod.rs
+91
View File
@@ -442,6 +442,97 @@ impl Workspace {
Ok(())
}
// ==================== Seeding ====================
/// Seed any missing core identity files in the workspace.
///
/// Called on every boot. Only creates files that don't already exist,
/// so user edits are never overwritten. Returns the number of files
/// created (0 if all core files already existed).
pub async fn seed_if_empty(&self) -> Result<usize, WorkspaceError> {
let seed_files: &[(&str, &str)] = &[
(
paths::README,
"# Workspace\n\n\
This is your agent's persistent memory. Files here are indexed for search\n\
and used to build the agent's context.\n\n\
## Structure\n\n\
- `MEMORY.md` - Long-term notes and facts worth remembering\n\
- `IDENTITY.md` - Agent name, nature, personality\n\
- `SOUL.md` - Core values and principles\n\
- `AGENTS.md` - Behavior instructions for the agent\n\
- `USER.md` - Information about you (the user)\n\
- `HEARTBEAT.md` - Periodic background task checklist\n\
- `daily/` - Automatic daily session logs\n\
- `context/` - Additional context documents\n\n\
Edit these files to shape how your agent thinks and acts.",
),
(
paths::MEMORY,
"# Memory\n\n\
Long-term notes, decisions, and facts worth remembering.\n\
The agent appends here during conversations.",
),
(
paths::IDENTITY,
"# Identity\n\n\
Name: IronClaw\n\
Nature: A secure personal AI assistant\n\n\
Edit this file to give your agent a custom name and personality.",
),
(
paths::SOUL,
"# Core Values\n\n\
- Protect user privacy and data security above all else\n\
- Be honest about limitations and uncertainty\n\
- Prefer action over lengthy deliberation\n\
- Ask for clarification rather than guessing on important decisions\n\
- Learn from mistakes and remember lessons",
),
(
paths::AGENTS,
"# Agent Instructions\n\n\
You are a personal AI assistant with access to tools and persistent memory.\n\n\
## Guidelines\n\n\
- Always search memory before answering questions about prior conversations\n\
- Write important facts and decisions to memory for future reference\n\
- Use the daily log for session-level notes\n\
- Be concise but thorough",
),
(
paths::USER,
"# User Context\n\n\
The agent will fill this in as it learns about you.\n\
You can also edit this directly to provide context upfront.",
),
(paths::HEARTBEAT, HEARTBEAT_SEED),
];
let mut count = 0;
for (path, content) in seed_files {
// Skip files that already exist (never overwrite user edits)
match self.read(path).await {
Ok(_) => continue,
Err(WorkspaceError::DocumentNotFound { .. }) => {}
Err(e) => {
tracing::warn!("Failed to check {}: {}", path, e);
continue;
}
}
if let Err(e) = self.write(path, content).await {
tracing::warn!("Failed to seed {}: {}", path, e);
} else {
count += 1;
}
}
if count > 0 {
tracing::info!("Seeded {} workspace files", count);
}
Ok(count)
}
/// Generate embeddings for chunks that don't have them yet.
///
/// This is useful for backfilling embeddings after enabling the provider.
tests/ws_gateway_integration.rs
+7 -1
View File
@@ -41,11 +41,13 @@ async fn start_test_server() -> (
msg_tx: tokio::sync::RwLock::new(Some(agent_tx)),
sse: SseManager::new(),
workspace: None,
context_manager: None,
session_manager: None,
log_broadcaster: None,
extension_manager: None,
tool_registry: None,
store: None,
job_manager: None,
prompt_queue: None,
user_id: "test-user".to_string(),
shutdown_tx: tokio::sync::RwLock::new(None),
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
@@ -162,6 +164,7 @@ async fn test_ws_thinking_event() {
state.sse.broadcast(SseEvent::Thinking {
message: "analyzing...".to_string(),
thread_id: None,
});
let text = recv_text(&mut ws).await;
@@ -286,13 +289,16 @@ async fn test_ws_multiple_events_in_sequence() {
// Broadcast multiple events rapidly
state.sse.broadcast(SseEvent::Thinking {
message: "step 1".to_string(),
thread_id: None,
});
state.sse.broadcast(SseEvent::ToolStarted {
name: "shell".to_string(),
thread_id: None,
});
state.sse.broadcast(SseEvent::ToolCompleted {
name: "shell".to_string(),
success: true,
thread_id: None,
});
state.sse.broadcast(SseEvent::Response {
content: "done".to_string(),
tools-src/gmail/src/api.rs
+2 -2
View File
@@ -26,7 +26,7 @@ fn api_call(method: &str, path: &str, body: Option<&str>) -> Result<String, Stri
&format!("Gmail API: {} {}", method, path),
);
let response = host::http_request(method, &url, headers, body_bytes.as_deref())?;
let response = host::http_request(method, &url, headers, body_bytes.as_deref(), None)?;
if response.status < 200 || response.status >= 300 {
let body_text = String::from_utf8_lossy(&response.body);
@@ -389,7 +389,7 @@ const BASE64URL_CHARS: &[u8; 64] =
/// Base64url-encode bytes (no padding, URL-safe alphabet).
fn base64url_encode(input: &[u8]) -> String {
let mut result = String::with_capacity((input.len() + 2) / 3 * 4);
let mut result = String::with_capacity(input.len().div_ceil(3) * 4);
for chunk in input.chunks(3) {
let b0 = chunk[0] as u32;
tools-src/google-calendar/src/api.rs
+57 -50
View File
@@ -26,7 +26,7 @@ fn api_call(method: &str, path: &str, body: Option<&str>) -> Result<String, Stri
&format!("Google Calendar API: {} {}", method, path),
);
let response = host::http_request(method, &url, headers, body_bytes.as_deref())?;
let response = host::http_request(method, &url, headers, body_bytes.as_deref(), None)?;
if response.status < 200 || response.status >= 300 {
let body_text = String::from_utf8_lossy(&response.body);
@@ -144,64 +144,68 @@ pub fn get_event(calendar_id: &str, event_id: &str) -> Result<EventResult, Strin
})
}
/// Parameters for creating a calendar event.
pub struct CreateEventParams<'a> {
pub calendar_id: &'a str,
pub summary: &'a str,
pub description: Option<&'a str>,
pub location: Option<&'a str>,
pub start_datetime: Option<&'a str>,
pub end_datetime: Option<&'a str>,
pub start_date: Option<&'a str>,
pub end_date: Option<&'a str>,
pub timezone: Option<&'a str>,
pub attendees: &'a [String],
}
/// Create a new event.
pub fn create_event(
calendar_id: &str,
summary: &str,
description: Option<&str>,
location: Option<&str>,
start_datetime: Option<&str>,
end_datetime: Option<&str>,
start_date: Option<&str>,
end_date: Option<&str>,
timezone: Option<&str>,
attendees: &[String],
) -> Result<EventResult, String> {
pub fn create_event(p: &CreateEventParams<'_>) -> Result<EventResult, String> {
let mut event = serde_json::json!({
"summary": summary,
"summary": p.summary,
});
if let Some(desc) = description {
if let Some(desc) = p.description {
event["description"] = serde_json::Value::String(desc.to_string());
}
if let Some(loc) = location {
if let Some(loc) = p.location {
event["location"] = serde_json::Value::String(loc.to_string());
}
// Build start/end, preferring datetime over date
if let Some(dt) = start_datetime {
if let Some(dt) = p.start_datetime {
let mut start = serde_json::json!({ "dateTime": dt });
if let Some(tz) = timezone {
if let Some(tz) = p.timezone {
start["timeZone"] = serde_json::Value::String(tz.to_string());
}
event["start"] = start;
} else if let Some(d) = start_date {
} else if let Some(d) = p.start_date {
event["start"] = serde_json::json!({ "date": d });
} else {
return Err("Either start_datetime or start_date is required".to_string());
}
if let Some(dt) = end_datetime {
if let Some(dt) = p.end_datetime {
let mut end = serde_json::json!({ "dateTime": dt });
if let Some(tz) = timezone {
if let Some(tz) = p.timezone {
end["timeZone"] = serde_json::Value::String(tz.to_string());
}
event["end"] = end;
} else if let Some(d) = end_date {
} else if let Some(d) = p.end_date {
event["end"] = serde_json::json!({ "date": d });
} else {
return Err("Either end_datetime or end_date is required".to_string());
}
if !attendees.is_empty() {
event["attendees"] = serde_json::json!(attendees
if !p.attendees.is_empty() {
event["attendees"] = serde_json::json!(p
.attendees
.iter()
.map(|e| serde_json::json!({ "email": e }))
.collect::<Vec<_>>());
}
let body = serde_json::to_string(&event).map_err(|e| e.to_string())?;
let path = format!("calendars/{}/events", url_encode(calendar_id));
let path = format!("calendars/{}/events", url_encode(p.calendar_id));
let response = api_call("POST", &path, Some(&body))?;
let parsed: serde_json::Value =
@@ -212,53 +216,56 @@ pub fn create_event(
})
}
/// Parameters for updating a calendar event.
pub struct UpdateEventParams<'a> {
pub calendar_id: &'a str,
pub event_id: &'a str,
pub summary: Option<&'a str>,
pub description: Option<&'a str>,
pub location: Option<&'a str>,
pub start_datetime: Option<&'a str>,
pub end_datetime: Option<&'a str>,
pub start_date: Option<&'a str>,
pub end_date: Option<&'a str>,
pub timezone: Option<&'a str>,
pub attendees: Option<&'a [String]>,
}
/// Update an existing event (PATCH for partial updates).
pub fn update_event(
calendar_id: &str,
event_id: &str,
summary: Option<&str>,
description: Option<&str>,
location: Option<&str>,
start_datetime: Option<&str>,
end_datetime: Option<&str>,
start_date: Option<&str>,
end_date: Option<&str>,
timezone: Option<&str>,
attendees: Option<&[String]>,
) -> Result<EventResult, String> {
pub fn update_event(p: &UpdateEventParams<'_>) -> Result<EventResult, String> {
let mut patch = serde_json::json!({});
if let Some(s) = summary {
if let Some(s) = p.summary {
patch["summary"] = serde_json::Value::String(s.to_string());
}
if let Some(d) = description {
if let Some(d) = p.description {
patch["description"] = serde_json::Value::String(d.to_string());
}
if let Some(l) = location {
if let Some(l) = p.location {
patch["location"] = serde_json::Value::String(l.to_string());
}
if let Some(dt) = start_datetime {
if let Some(dt) = p.start_datetime {
let mut start = serde_json::json!({ "dateTime": dt });
if let Some(tz) = timezone {
if let Some(tz) = p.timezone {
start["timeZone"] = serde_json::Value::String(tz.to_string());
}
patch["start"] = start;
} else if let Some(d) = start_date {
} else if let Some(d) = p.start_date {
patch["start"] = serde_json::json!({ "date": d });
}
if let Some(dt) = end_datetime {
if let Some(dt) = p.end_datetime {
let mut end = serde_json::json!({ "dateTime": dt });
if let Some(tz) = timezone {
if let Some(tz) = p.timezone {
end["timeZone"] = serde_json::Value::String(tz.to_string());
}
patch["end"] = end;
} else if let Some(d) = end_date {
} else if let Some(d) = p.end_date {
patch["end"] = serde_json::json!({ "date": d });
}
if let Some(att) = attendees {
if let Some(att) = p.attendees {
patch["attendees"] = serde_json::json!(att
.iter()
.map(|e| serde_json::json!({ "email": e }))
@@ -268,8 +275,8 @@ pub fn update_event(
let body = serde_json::to_string(&patch).map_err(|e| e.to_string())?;
let path = format!(
"calendars/{}/events/{}",
url_encode(calendar_id),
url_encode(event_id)
url_encode(p.calendar_id),
url_encode(p.event_id)
);
let response = api_call("PATCH", &path, Some(&body))?;
tools-src/google-calendar/src/lib.rs
+25 -25
View File
@@ -279,18 +279,18 @@ fn execute_inner(params: &str) -> Result<String, String> {
timezone,
attendees,
} => {
let result = api::create_event(
&calendar_id,
&summary,
description.as_deref(),
location.as_deref(),
start_datetime.as_deref(),
end_datetime.as_deref(),
start_date.as_deref(),
end_date.as_deref(),
timezone.as_deref(),
&attendees,
)?;
let result = api::create_event(&api::CreateEventParams {
calendar_id: &calendar_id,
summary: &summary,
description: description.as_deref(),
location: location.as_deref(),
start_datetime: start_datetime.as_deref(),
end_datetime: end_datetime.as_deref(),
start_date: start_date.as_deref(),
end_date: end_date.as_deref(),
timezone: timezone.as_deref(),
attendees: &attendees,
})?;
serde_json::to_string(&result).map_err(|e| e.to_string())?
}
@@ -307,19 +307,19 @@ fn execute_inner(params: &str) -> Result<String, String> {
timezone,
attendees,
} => {
let result = api::update_event(
&calendar_id,
&event_id,
summary.as_deref(),
description.as_deref(),
location.as_deref(),
start_datetime.as_deref(),
end_datetime.as_deref(),
start_date.as_deref(),
end_date.as_deref(),
timezone.as_deref(),
attendees.as_deref(),
)?;
let result = api::update_event(&api::UpdateEventParams {
calendar_id: &calendar_id,
event_id: &event_id,
summary: summary.as_deref(),
description: description.as_deref(),
location: location.as_deref(),
start_datetime: start_datetime.as_deref(),
end_datetime: end_datetime.as_deref(),
start_date: start_date.as_deref(),
end_date: end_date.as_deref(),
timezone: timezone.as_deref(),
attendees: attendees.as_deref(),
})?;
serde_json::to_string(&result).map_err(|e| e.to_string())?
}
tools-src/google-docs/src/api.rs
+1 -1
View File
@@ -30,7 +30,7 @@ fn api_call(method: &str, path: &str, body: Option<&str>) -> Result<String, Stri
&format!("Google Docs API: {} {}", method, url),
);
let response = host::http_request(method, &url, headers, body_bytes.as_deref())?;
let response = host::http_request(method, &url, headers, body_bytes.as_deref(), None)?;
if response.status < 200 || response.status >= 300 {
let body_text = String::from_utf8_lossy(&response.body);
tools-src/google-drive/src/api.rs
+3 -3
View File
@@ -32,7 +32,7 @@ fn api_call(method: &str, path: &str, body: Option<&str>) -> Result<String, Stri
&format!("Drive API: {} {}", method, path),
);
let response = host::http_request(method, &url, headers, body_bytes.as_deref())?;
let response = host::http_request(method, &url, headers, body_bytes.as_deref(), None)?;
if response.status < 200 || response.status >= 300 {
let body_text = String::from_utf8_lossy(&response.body);
@@ -56,7 +56,7 @@ fn api_call_raw(method: &str, url: &str) -> Result<Vec<u8>, String> {
&format!("Drive API raw: {} {}", method, url),
);
let response = host::http_request(method, url, "{}", None)?;
let response = host::http_request(method, url, "{}", None, None)?;
if response.status < 200 || response.status >= 300 {
let body_text = String::from_utf8_lossy(&response.body);
@@ -262,7 +262,7 @@ pub fn upload_file(
"Drive API: POST upload/files (multipart)",
);
let response = host::http_request("POST", &url, &headers, Some(body.as_bytes()))?;
let response = host::http_request("POST", &url, &headers, Some(body.as_bytes()), None)?;
if response.status < 200 || response.status >= 300 {
let body_text = String::from_utf8_lossy(&response.body);
tools-src/google-sheets/src/api.rs
+1 -1
View File
@@ -30,7 +30,7 @@ fn api_call(method: &str, path: &str, body: Option<&str>) -> Result<String, Stri
&format!("Google Sheets API: {} {}", method, url),
);
let response = host::http_request(method, &url, headers, body_bytes.as_deref())?;
let response = host::http_request(method, &url, headers, body_bytes.as_deref(), None)?;
if response.status < 200 || response.status >= 300 {
let body_text = String::from_utf8_lossy(&response.body);
tools-src/google-slides/src/api.rs
+1 -1
View File
@@ -30,7 +30,7 @@ fn api_call(method: &str, path: &str, body: Option<&str>) -> Result<String, Stri
&format!("Google Slides API: {} {}", method, url),
);
let response = host::http_request(method, &url, headers, body_bytes.as_deref())?;
let response = host::http_request(method, &url, headers, body_bytes.as_deref(), None)?;
if response.status < 200 || response.status >= 300 {
let body_text = String::from_utf8_lossy(&response.body);
tools-src/slack/src/api.rs
+1 -1
View File
@@ -45,7 +45,7 @@ fn slack_api_call(method: &str, endpoint: &str, body: Option<&str>) -> Result<St
&format!("Slack API: {} {}", method, endpoint),
);
let response = host::http_request(method, &url, headers, body_bytes.as_deref())?;
let response = host::http_request(method, &url, headers, body_bytes.as_deref(), None)?;
if response.status < 200 || response.status >= 300 {
return Err(format!(
tools-src/telegram/src/transport.rs
+1 -1
View File
@@ -89,7 +89,7 @@ pub fn post_encrypted(
/// HTTP POST with raw binary body via the WASM host's http-request capability.
fn http_post_binary(url: &str, body: &[u8]) -> Result<Vec<u8>, String> {
let resp = host::http_request("POST", url, "{}", Some(body))?;
let resp = host::http_request("POST", url, "{}", Some(body), None)?;
if resp.status < 200 || resp.status >= 300 {
let body_text = String::from_utf8_lossy(&resp.body);
wit/channel.wit
+7 -1
View File
@@ -90,11 +90,17 @@ interface channel-host {
/// - Credentials are injected by the host; WASM never sees them
/// - Response is scanned for leaked secrets before returning
/// - Rate-limited per channel
///
/// The optional timeout-ms parameter controls the HTTP client timeout
/// in milliseconds. Defaults to 30000 (30s) when not provided. Use a
/// longer timeout for long-polling requests (e.g., Telegram getUpdates).
/// Capped at the channel's callback_timeout to prevent hangs.
http-request: func(
method: string,
url: string,
headers-json: string,
body: option<list<u8>>
body: option<list<u8>>,
timeout-ms: option<u32>,
) -> result<http-response, string>;
/// Check if a secret exists (if capability granted).
wit/tool.wit
+6 -1
View File
@@ -67,11 +67,16 @@ interface host {
/// - Network error
/// - Timeout
/// - Secret leak detected in response
///
/// The optional timeout-ms parameter controls the HTTP client timeout
/// in milliseconds. Defaults to 30000 (30s) when not provided.
/// Capped at the callback timeout to prevent hangs.
http-request: func(
method: string,
url: string,
headers-json: string,
body: option<list<u8>>
body: option<list<u8>>,
timeout-ms: option<u32>,
) -> result<http-response, string>;
// ==================== Tool Invocation Capability ====================