Compare commits

...
Author SHA1 Message Date
Illia PolosukhinandClaude Opus 4.6 1aca76b1a7 style: fix rustfmt formatting in bootstrap test
Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-02-15 00:22:18 -08:00
Illia PolosukhinandClaude Opus 4.6 414c1b28b9 fix: status command shows libSQL backend and skips keychain probe
The status command only checked DATABASE_URL (postgres), showing
"not configured" for libSQL users. Now detects the DATABASE_BACKEND
env var and reports libSQL path and Turso sync status.

Also remove the keychain probe from status. get_generic_password()
triggers macOS unlock+authorization dialogs which is terrible UX
for a read-only diagnostic command.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-02-15 00:13:33 -08:00
Illia PolosukhinandClaude Opus 4.6 d771f99f9e fix: persist DATABASE_BACKEND to ~/.ironclaw/.env for libSQL startup
The wizard saved database_backend only to the database, but
Config::from_env() needs it BEFORE connecting to any database (to
decide which backend to use). Without it, the backend defaults to
Postgres and then fails with "Missing required setting database_url".

Now save all database bootstrap vars (DATABASE_BACKEND, DATABASE_URL,
LIBSQL_PATH, LIBSQL_URL) to ~/.ironclaw/.env via save_bootstrap_env().

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-02-15 00:12:01 -08:00
Illia PolosukhinandClaude Opus 4.6 1facda4a75 fix: cache keychain key eagerly to avoid redundant macOS password dialogs
Replace has_master_key() with get_master_key() in step_security() and
immediately build SecretsCrypto from the result. This eliminates redundant
keychain accesses later in init_secrets_context(), each of which triggers
macOS system dialogs (keychain unlock + app authorization).

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-02-15 00:07:02 -08:00
Illia PolosukhinandClaude Opus 4.6 e8caab1a12 fix: OAuth callback listener binds IPv4 first to match redirect URLs
The listener was binding to [::1] (IPv6) first, but NEAR AI and other
OAuth flows redirect to http://127.0.0.1:9876/... (IPv4 explicit).
On macOS and most systems, [::1] and 127.0.0.1 are separate addresses,
so the browser's connection to 127.0.0.1 was refused when the listener
was on [::1]. Reversed the bind order: try 127.0.0.1 first, fall back
to [::1] if IPv4 is unavailable.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-02-14 22:16:40 -08:00
Illia PolosukhinandClaude Opus 4.6 fac91aec3a fix: address latest PR review comments (SecretString, empty env, docs, embeddings)
- Change wizard llm_api_key from String to SecretString to prevent
  accidental logging of API keys
- Fix inject_llm_keys_from_secrets skipping when env var is set but
  empty, matching optional_env's treatment of empty as unset
- Fix inverted doc comment on INJECTED_VARS (env checked first, overlay
  is the fallback, not the other way around)
- Update stale "env vars" comments in main.rs to reflect overlay pattern
- Fix step_embeddings not seeing cached OpenAI key from wizard session

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-02-14 17:58:12 -08:00
Illia PolosukhinandClaude Opus 4.6 0aae66c9dc fix: address remaining PR review comments (clippy, TODO, secrets backend ordering)
- Fix empty line after doc comment (clippy: empty_line_after_doc_comments)
- Collapse nested if in optional_env overlay check (clippy: collapsible_if)
- Remove dangling TODO(#XX) placeholder issue ref in channels.rs
- Fix init_secrets_context to respect selected database_backend when both
  postgres and libsql features are compiled, preventing wrong-backend
  secrets storage when DATABASE_URL is set but libsql was chosen

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-02-14 17:19:30 -08:00
Illia PolosukhinandClaude Opus 4.6 aa808ca94e fix: remove unsafe set_var, use thread-safe overlay for injected secrets
Address PR #92 review comments:
- Replace all 5 unsafe `std::env::set_var()` calls with safe alternatives
- Add INJECTED_VARS OnceLock<HashMap> overlay in config.rs, checked by
  optional_env() before falling back to std::env::var()
- Cache wizard API key in SetupWizard.llm_api_key field instead of env
- Pass explicit key param to fetch_anthropic_models/fetch_openai_models
- Persist env-provided API keys to secrets store during onboarding

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-02-14 14:14:22 -08:00
Illia PolosukhinandClaude Opus 4.6 c7e6833d14 Merge remote-tracking branch 'origin/main' into fix/setup-audit-fixes
Resolve conflicts between main's simplified config (no bootstrap param,
env-only DatabaseConfig) and our branch's typed ChannelSetupError.

- config.rs: take main's simpler resolve() signatures (no bootstrap)
- main.rs: remove dead check_onboard_needed block and CACHED_KEYCHAIN_KEY ref
- channels.rs: keep ChannelSetupError types, restore settings params from main
- wizard.rs: pass &self.settings to setup_telegram, use ? with From impl
- settings.rs: fix test_llm_backend_round_trip (use std::fs::write, tempfile::tempdir)

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-02-14 14:05:54 -08:00
9fed8453c7 fix: shell destructive-command check bypassed by Value::Object arguments (#72)
Co-authored-by: Yi LIU <[email protected]>
Co-authored-by: Illia Polosukhin <[email protected]>
2026-02-14 21:54:13 +00:00
Illia PolosukhinandClaude Opus 4.6 1885d61d46 fix: replace unreachable!() with error return in setup wizard
The provider match in step_inference_provider was guarded by
is_known but used unreachable!() as the catch-all. If a new
provider is added to the is_known check without a corresponding
match arm, this would panic at runtime. Return a typed error
instead.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-02-14 13:51:37 -08:00
Illia PolosukhinandClaude Opus 4.6 da47903108 fix: harden setup module error handling and secret safety
- Introduce ChannelSetupError typed enum replacing raw String errors
  across all channel setup functions (setup_telegram, setup_http,
  setup_tunnel, setup_wasm_channel, validate_telegram_token)
- Add From<ChannelSetupError> for SetupError to simplify call sites
- Convert setup_telegram retry from recursion to loop (unbounded stack)
- Stop printing HTTP webhook secret plaintext to terminal
- Use secret_input() for Turso auth token (was visible input())
- Replace dirs::home_dir().unwrap_or_default() with proper error
- Fix UTF-8 panic in model name truncation (byte-index to chars-based)
- Log warning in secret_exists() instead of silently swallowing errors
- Deduplicate generate_webhook_secret() to delegate to shared helper

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-02-14 13:43:25 -08:00
eaef335db6 fix: propagate real tool_call_id instead of hardcoded placeholder (#73)
The worker (both agent/worker.rs and worker/runtime.rs) was passing the
literal string "tool_call_id" to ChatMessage::tool_result instead of
the actual tool call ID from the LLM response. This breaks
OpenAI-compatible providers that match tool results to their
corresponding calls by ID.

- Add tool_call_id field to ToolSelection struct
- Propagate ToolCall.id through select_tools() into ToolSelection
- Replace all hardcoded "tool_call_id" usages with selection.tool_call_id
- Generate unique IDs for plan-based synthetic selections
- Add test verifying tool_call_id is preserved

Co-authored-by: Yi LIU <[email protected]>
Co-authored-by: Illia Polosukhin <[email protected]>
2026-02-14 21:39:25 +00:00
Eric WinerandGitHub 225af29db2 Reformat architecture diagram in README (#64) 2026-02-14 21:22:58 +00:00
a53b2c10b5 fix: Fix wasm tool schemas and runtime (#42)
* feat: Move debug log truncation from agent loop to REPL channel

Full tool output now flows through StatusUpdate so the web gateway
gets untruncated content. The REPL channel truncates at display time
(200 chars for tool results, thinking, and status messages).

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

* fix: Flatten WASM tool schemas and fix host HTTP runtime contention

LLMs can't reliably follow oneOf + const discriminator patterns in JSON
Schema, causing tools like Google Calendar to receive malformed params
(e.g., {"operation":"list_events","data":{"calendarId":"primary"}} instead
of {"action":"list_events","calendar_id":"primary"}). Replace all 9 WASM
tool schemas with flat action enum + top-level properties. The serde
#[serde(tag = "action")] deserialization works identically.

Also fixes WASM host HTTP requests (channels and tools) stalling during
startup by replacing Handle::current().block_on() with a dedicated
single-threaded runtime per request, avoiding I/O driver contention.

Reduces verbose LLM debug logging (full request/response payloads) and
changes tower_http default from debug to warn.

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

* feat: Built-in OAuth credentials and combined Google scopes

Add infrastructure for shipping default OAuth credentials with the binary,
similar to how gcloud/rclone bake in their client_id. Credentials are set
at compile time via IRONCLAW_GOOGLE_CLIENT_ID / IRONCLAW_GOOGLE_CLIENT_SECRET
env vars, or can be hardcoded in src/cli/oauth_defaults.rs.

The fallback chain is: capabilities file > runtime env var > built-in defaults.

Also, when authing any Google tool, scopes from ALL installed Google tools
are now combined into a single OAuth request (they all share the same
google_oauth_token secret). One login covers Gmail, Calendar, Drive, etc.

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

* feat: Ship default Google OAuth credentials for zero-config auth

Google Desktop App credentials are not secret (per Google's own docs).
Hardcode them so `ironclaw tool auth <google-tool>` works out of the box
without requiring users to register their own OAuth app.

Credentials can still be overridden at compile time
(IRONCLAW_GOOGLE_CLIENT_ID) or runtime (GOOGLE_OAUTH_CLIENT_ID).

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

* fix: Consistent OAuth callback port and polished landing page

- Use fixed port 9876 instead of scanning 9876-9886 (one redirect URI
  to register in provider OAuth apps, deterministic behavior)
- Replace broken unicode checkmark with SVG icons (charset was missing,
  rendered as mojibake)
- Dark themed landing page with proper card layout for both success
  and error states
- Add charset=utf-8 to Content-Type headers

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

* refactor: Unify OAuth callback server across all auth flows

All three OAuth flows (WASM tool auth, MCP server auth, NEAR AI login)
now share the same code from cli::oauth_defaults:

- Fixed port 9876 (one redirect URI to register per provider)
- Shared landing page HTML (dark card with SVG icons, proper charset)
- Parameterized wait_for_callback(listener, path, param, display_name)

Removes ~120 lines of duplicated callback/HTML code.

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

* Support for oauth token refresh

* refactor: Replace bootstrap.json with ~/.ironclaw/.env for DATABASE_URL

Kill the 4-field BootstrapConfig JSON file. Only DATABASE_URL actually
needs disk persistence (chicken-and-egg before DB connect). The other
three fields are now derived: pool_size defaults to 10 via env var,
secrets master key is auto-detected (env then keychain probe), and
onboard_completed is inferred from DATABASE_URL presence.

The new format is a standard .env file loaded via dotenvy early in
main, so DATABASE_URL is available as a regular env var everywhere.

Handles three upgrade paths:
- Clean start: wizard writes .env, reload after wizard completes
- Returning user: .env loaded at startup, business as usual
- Legacy upgrade: bootstrap.json auto-migrated to .env on first run

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

* fix: Address PR review findings

- Fix UTF-8 panic in truncate_for_preview (byte-slice on char boundary)
- Cap WASM guest timeout_ms at 5 minutes to prevent resource exhaustion
- Fix localhost detection in requires_auth() to avoid substring matches
  (e.g. "notlocalhost.com" no longer matches)
- Fix query param injection to insert before URL fragment
- Fix extract_host_from_url for IPv6 bracket notation
- Remove misleading schema defaults: Slack limit, Slides insertion_index,
  Docs index (per-action defaults documented in descriptions instead)

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

* style: Fix cargo fmt formatting

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

* fix: IPv6 loopback support for OAuth listener and localhost detection

- bind_callback_listener: try [::1] first, fall back to 127.0.0.1,
  so OAuth redirects work on systems where localhost resolves to ::1
- is_localhost_url: replace manual string parsing with url::Url for
  correct handling of IPv6 brackets, ports, userinfo, etc.
- Add url crate as direct dependency (already a transitive dep)

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

* fix: Address PR review feedback on runtime reuse, onboard check, and OAuth binding

- Remove session file check from check_onboard_needed(); DATABASE_URL is sufficient
- Detect AddrInUse on IPv6 bind and fail immediately instead of falling through to IPv4
- Reuse dedicated tokio runtime across HTTP calls in both tool and channel WASM wrappers

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

* fix: HTML-escape provider name in OAuth landing page, simplify Slack limit description

- Add html_escape() to prevent XSS in landing_html() where provider_name
  was interpolated directly into HTML (defense-in-depth, source is trusted
  but escaping costs nothing)
- Remove per-action default numbers from Slack limit field description to
  avoid confusing LLMs with conflicting defaults

Addresses review feedback from zmanian on PR #42.

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

* fix: Save all bootstrap fields from wizard, fix config module comment

- Wizard now saves secrets_master_key_source and database_pool_size to
  bootstrap.json (was only saving database_url and onboard_completed,
  which broke secrets after fresh onboard since SecretsConfig::resolve
  reads key source from bootstrap)
- Update config.rs module doc to reflect bootstrap.json priority chain
  instead of the removed ~/.ironclaw/.env approach

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

* refactor: Replace BootstrapConfig with .env-based bootstrap

DATABASE_URL is the only setting that needs disk persistence before
the database is available. Instead of a custom bootstrap.json with 4
fields, use a standard ~/.ironclaw/.env file loaded via dotenvy.

- Remove BootstrapConfig struct entirely
- Restore ironclaw_env_path(), load_ironclaw_env(), save_database_url()
- SecretsConfig::resolve() now auto-detects (env var then keychain probe)
  instead of reading a saved source from bootstrap.json
- DatabaseConfig::resolve() reads DATABASE_URL from env only (dotenvy
  loads ~/.ironclaw/.env into the environment early in startup)
- check_onboard_needed() is now sync (just checks env vars)
- Wizard save_and_summarize() works for both postgres and libsql backends
- One-time migration from bootstrap.json to .env preserved

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

* fix: Ensure load_ironclaw_env() runs in all Config paths, fix .env priority

- Config::from_env() and Config::from_db() now call load_ironclaw_env()
  internally (after dotenvy::dotenv()), so CLI commands like `memory`
  and `config` correctly load DATABASE_URL from ~/.ironclaw/.env
- Fix load order: standard ./.env first (higher priority), then
  ~/.ironclaw/.env, matching the documented priority chain
- Collapse nested if/if-let into let-chains (clippy::collapsible_if)
  in oauth_defaults.rs, tool.rs, and secrets/store.rs
- Fix rename_to_migrated to take &Path instead of &PathBuf

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

* fix: Address PR review comments (quoting, SSRF, error mapping)

- Quote DATABASE_URL in .env writes so `#` in passwords isn't treated
  as a dotenv comment (e.g., `DATABASE_URL="postgres://..."`)
- Add SSRF defenses to refresh_oauth_token(): require HTTPS, reject
  private/loopback IPs (with DNS resolution), disable redirects.
  token_url comes from tool capabilities JSON, so a malicious tool
  could otherwise exfiltrate refresh tokens.
- Fix IPv4 bind error mapping: only map AddrInUse to PortInUse,
  use generic Io variant for other bind failures

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

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-02-14 21:21:22 +00:00
Illia PolosukhinandClaude Opus 4.6 85196cd527 fix: address second-round PR review feedback
- Validate custom model ID is non-empty (loop until valid input)
- Warn on unknown DATABASE_BACKEND env var before defaulting to Postgres
- Force re-selection when llm_backend contains unknown provider value
- Use ok_or_else for proper String error type in google-sheets

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-02-14 13:20:06 -08:00
Illia PolosukhinandClaude Opus 4.6 a0e01f04d3 fix: address critical/high audit findings across WASM sub-crates
- Telegram: remove .unwrap() panic on workspace_read (owner_id check)
- WhatsApp: use configured api_version instead of hardcoded v18.0
- WhatsApp: log config parse errors before falling back to defaults
- Slack: log serialization errors in emit_message and json_response
- Google Docs: safe array access for batch update replies
- Google Sheets: safe array access for add_sheet replies
- Google Calendar: fix doc comment secret name mismatch
- Gmail: avoid unnecessary String allocation in UNREAD check

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-02-14 12:54:20 -08:00
408ae8a29a feat: add multi-provider LLM failover with retry backoff (#28)
* feat: add multi-provider LLM failover

Add FailoverProvider that wraps multiple LlmProvider instances and
tries each in sequence on transient failures. Non-retryable errors
(auth, context length, model not available) propagate immediately.

- New `FailoverProvider` with generic `try_providers` helper
- `is_retryable()` classifies transient errors (request failed,
  rate limited, invalid response, session renewal, HTTP, IO)
- Configurable via `NEARAI_FALLBACK_MODEL` env var
- Returns `Result` from constructor (no panics in production)
- Updates FEATURE_PARITY.md: failover chains , cooldown 

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

* fix: track last-used provider for accurate cost/model reporting

After failover, model_name() and cost_per_token() now reflect the
provider that actually handled the request, not always the primary.
Also corrects is_retryable() docs to list ModelNotAvailable as retryable.

Addresses PR #28 review comments.

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

* feat: add retry with exponential backoff for LLM providers

Add retry logic with exponential backoff and jitter to both NearAiProvider
and NearAiChatProvider for transient errors (HTTP 429, 500, 502, 503, 504).

Extract shared retry helpers (is_retryable_status, retry_backoff_delay)
into src/llm/retry.rs so both providers reuse the same logic.

Configurable via NEARAI_MAX_RETRIES env var (default: 3).

* docs: clarify max_retries means N retries, not N total attempts

* warn when fallback model equals primary model

* fix: saturating_mul in backoff delay, dedupe to_lowercase allocation

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-02-14 15:43:38 +04:00
Illia PolosukhinandClaude Opus 4.6 e982699e09 fix: address PR review feedback (set_var safety, parse warnings, db_map efficiency)
1. Replace unsafe set_var keychain caching with OnceLock<String> in
   SecretsConfig::resolve(). Eliminates the env var write from main.rs
   entirely, using a process-wide OnceLock cache instead.

2. Log tracing::warn when database_backend or llm_backend settings
   fail to parse, instead of silently falling back to defaults.

3. Remove O(K*S) get() pre-check in from_db_map(). Instead, let set()
   run and match on "Path not found" errors to skip unknown keys,
   avoiding full Settings serialization per key.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-02-14 01:09:20 -08:00
Illia Polosukhin 5e73dbbdc8 Merge remote-tracking branch 'origin/main' into feat/onboarding-libsql-selection 2026-02-14 00:55:48 -08:00
Illia PolosukhinandClaude Opus 4.6 92863bf860 fix: resolve libSQL onboarding crash, keychain double-prompt, and setup audit findings
Three bugs fixed:

1. libSQL onboarding crash ("Missing required setting 'database_url'"):
   DatabaseConfig::resolve() only checked DATABASE_BACKEND env var, falling
   back to Postgres default. Now reads settings.database_backend, plus
   settings.libsql_path and settings.libsql_url as fallbacks.

2. OS keychain prompts twice during startup: Config::from_env() and
   Config::from_db() both called get_master_key(). Now caches the key in
   SECRETS_MASTER_KEY env var after first read so from_db() skips keychain.

3. "Path not found: nearai.session" warning: from_db_map() tried to apply
   app-specific DB keys (nearai.session_token) to the Settings struct.
   Now skips keys that don't map to known Settings fields. Also fixed
   bootstrap migration key mismatch (nearai.session -> nearai.session_token).

Setup module audit fixes (14 findings):
- Replace unreachable!() with proper error in provider match
- Extract setup_api_key_provider() to deduplicate setup_anthropic/setup_openai
- Add SAFETY comments to all unsafe std::env::set_var blocks
- Fix .unwrap() calls with proper error handling
- Remove incorrect #[allow(dead_code)] on used TelegramUpdate::update_id
- Log warnings instead of silently discarding HTTP errors in Telegram binding
- Guard select_many against empty options, fix mask_api_key for non-ASCII
- Update stale doc comment in mod.rs, rename misleading variable
- Add 7 new tests (model fetcher fallbacks, channel discovery, secret gen)

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-02-13 21:46:19 -08:00
Zaki ManianGitHubClaude Opus 4.6Illia Polosukhingemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
d9ff86d7e0 docs: Add review discipline guidelines to CLAUDE.md (#68)
* docs: Add review discipline guidelines to CLAUDE.md

Codifies lessons learned from Illia's review fixes on the libSQL
backend PR -- patterns we missed that should be caught systematically
going forward.

- Ban .expect() alongside .unwrap() in production code
- Add mechanical grep checks before committing
- New "Review & Fix Discipline" section covering:
  - Fix all instances of a pattern, not just the one flagged
  - Propagate architectural changes to satellite types
  - Schema translation must include indexes and seed data
  - Feature flag testing with each feature in isolation

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

* Apply suggestions from code review

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
Co-authored-by: Illia Polosukhin <[email protected]>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-02-14 04:25:53 +00:00
Illia PolosukhinandClaude Opus 4.6 46c1daca5e feat: add interactive database backend selection during onboarding
Previously the onboarding wizard silently defaulted to PostgreSQL because
libsql wasn't in the default feature set. Now both backends ship by default
and the wizard presents a selection prompt when both are available.

DATABASE_BACKEND env var still bypasses the prompt for headless/CI use.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-02-13 18:27:53 -08:00
e843c18141 feat: add libSQL/Turso embedded database backend (#47)
* feat: add libSQL/Turso database backend with full feature parity

Introduce a Database trait abstraction (~60 async methods) enabling
compile-time backend selection between PostgreSQL and libSQL/Turso.
Convert all modules from concrete Store to Arc<dyn Database>, add
LibSqlSecretsStore and LibSqlWasmToolStore implementations, wire
libsql stores throughout CLI and main entry points, and make the
setup wizard backend-agnostic.

Key changes:
- src/db/: Database trait, PostgresDatabase adapter, LibSqlBackend
  with native SQLite-dialect SQL, and idempotent migration system
- src/secrets/store.rs: LibSqlSecretsStore (all 8 trait methods)
- src/tools/wasm/storage.rs: LibSqlWasmToolStore (all 7 trait methods)
- src/main.rs, cli/tool.rs, cli/mcp.rs: backend-conditional wiring
- src/setup/channels.rs: SecretsContext uses Arc<dyn SecretsStore>
- Feature-gate postgres-only tests and examples

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

* feat: enable onboarding wizard for libSQL builds

Refactor the setup wizard to work with both postgres and libsql feature
flags. Previously the wizard was gated behind #[cfg(feature = "postgres")]
only, so libsql-only builds would print an error on `ironclaw onboard`.

- Add libsql fields to Settings (database_backend, libsql_path, libsql_url)
- Split wizard database/migration/secrets methods into feature-gated variants
- Add step_database_libsql() with local path and Turso remote replica prompts
- Update setup/mod.rs and main.rs feature gates to any(postgres, libsql)
- Extend check_onboard_needed() to detect libsql database presence

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

* fix: address PR review feedback for libSQL backend

- P0: Switch libsql_backend to connection-per-operation pattern to fix
  shared Connection concurrency issue across tokio tasks
- P0: Wrap secrets store INSERT+SELECT in transaction to fix TOCTOU race
- P0: Document encryption-at-rest limitations and json_patch divergence
- P1: Fix get_opt_text removing .filter(|s| !s.is_empty()) that conflated
  empty strings with NULL
- P1: Replace datetime('now') with fmt_ts(&Utc::now()) for consistent
  RFC 3339 timestamps across all queries
- P2: Use explicit _rowid column in FTS5 triggers and joins for stability
  across VACUUM operations
- P2: Add tracing::warn when embedding provided but vector search disabled
  in hybrid_search
- Extract shared connect_from_config() helper to deduplicate DB connection
  logic across main.rs, cli/config.rs, and cli/mcp.rs

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

* fix: add missing JobContext fields and resolve fmt/clippy warnings

Add total_tokens_used and max_tokens fields to JobContext in
libsql_backend.rs, apply cargo fmt, and fix clippy warnings.

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

* fix: review fixes for libSQL backend (shared connections, panics, indexes)

- Replace .expect() with proper error propagation in 3 call sites
- Share Arc<Database> between backend and stores instead of single Connection
- Add connect-per-operation pattern to LibSqlSecretsStore and LibSqlWasmToolStore
- Wrap store() INSERT + SELECT-back in a transaction
- Add ~22 missing indexes for parity with PostgreSQL schema
- Add 18 leak_detection_patterns seed rows matching PostgreSQL V2 migration
- Fix super:: import to use crate:: style
- Gate mask_password_in_url behind #[cfg(feature = "postgres")]
- Rewrite secrets store init with or_else chain for runtime backend selection

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

* fix: Resolve clippy lints (collapsible_if, too_many_arguments)

Collapse nested if blocks into let_chains to satisfy clippy's
collapsible_if lint (CI uses -D warnings). Suppress too_many_arguments
on libsql_row_to_tool_at since refactoring the positional index
pattern would be a larger change.

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

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
Co-authored-by: Illia Polosukhin <[email protected]>
2026-02-14 02:05:05 +00:00
54e9206f0b feat: Move debug log truncation from agent loop to REPL channel (#65)
* feat: Move debug log truncation from agent loop to REPL channel

Full tool output now flows through StatusUpdate so the web gateway
gets untruncated content. The REPL channel truncates at display time
(200 chars for tool results, thinking, and status messages).

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

* feat: truncating fmt layer for terminal, full logs for web gateway

Instead of truncating debug output at each LLM call site (fragile),
use a custom MakeWriter on the fmt layer that caps each tracing event
at 500 bytes before flushing to stderr. The web gateway WebLogLayer
still receives full untruncated content for /api/logs/events SSE.

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

* fix: UTF-8 safe truncation in truncate_for_preview, remove double truncation

- Use char_indices() instead of byte-based slicing to find the cut
  point, preventing panics on multi-byte characters (emoji, CJK, etc.)
- Remove redundant truncation in REPL channel (agent loop already
  truncates ToolResult previews to 200 chars)
- Add 9 unit tests covering edge cases: empty, exact length, multi-byte
  UTF-8 (emoji, CJK), mixed scripts, newline collapsing, whitespace

Addresses PR #65 review comments.

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

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-02-13 23:24:07 +00:00
5df0d13b59 Bump MSRV to 1.92, add GCP deployment files (#40)
* Bump MSRV to 1.92 and add GCP deployment files

rig-core 0.30 uses let_chains (stabilized post-1.87), which breaks
builds on Rust 1.85. Bump rust-version in Cargo.toml and both
Dockerfiles to 1.92 (verified working).

Add cloud deployment scaffolding:
- Dockerfile: multi-stage build for the main agent container
- deploy/cloud-sql-proxy.service: systemd unit for Cloud SQL Auth Proxy
- deploy/ironclaw.service: systemd unit for the IronClaw container
- deploy/setup.sh: VM bootstrap script (Docker, proxy, services)
- deploy/env.example: reference environment configuration

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

* Address review feedback: harden deploy scaffolding

- Add comment explaining GATEWAY_HOST=0.0.0.0 and when to use 127.0.0.1
- Document /opt/ironclaw ownership model (root-owned, Docker reads as root)
- Switch cloud-sql-proxy service from User=root to DynamicUser=yes

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

* fix: Resolve clippy lints (Rust 1.93) and fix CI test workflow

- Fix 97 collapsible_if warnings using let-chains syntax (auto-fixed)
- Fix ptr_arg: change &PathBuf to &Path in pairing store functions
- Fix suspicious_open_options: add .truncate(false) to OpenOptions
- Fix too_many_arguments: add clippy allow on execute_status
- Fix unnecessary_unwrap: use if-let in repository.rs hybrid_search
- Gate unused EchoTool with #[cfg(test)]
- Add PairingStore argument to ChannelStoreData::new() test call sites
- Add skip guard for bundled channel test when WASM artifacts unavailable
- Split CI test workflow to exclude PostgreSQL-dependent integration tests

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

* fix: Address review feedback from ilblackdragon

- Add root check to setup.sh (exits with error if not root)
- Add warning comment to env.example about placeholder passwords
- Dockerfile.worker already uses rust:1.92 (no change needed)
- PR #41 overlap noted; will rebase after #41 merges

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

* fix: resolve 47 collapsible_if clippy warnings

Collapse nested if statements across the codebase to satisfy
clippy::collapsible_if on Rust 1.93.

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

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-02-13 22:21:50 +04:00
bbb68f7490 Add OpenAI-compatible HTTP API (/v1/chat/completions, /v1/models) (#31)
* feat: add OpenAI-compatible HTTP API (/v1/chat/completions, /v1/models)

* - Reject model mismatches: validate req.model against the active model
    and return 404 model_not_found instead of silently ignoring it
  - Add x-ironclaw-streaming: simulated response header so clients know
    streaming is not true token-by-token delivery
  - Use SSE event type "error" for mid-stream LLM failures so clients can
    distinguish errors from content chunks
  - Mark docker-compose credentials as dev-only
  - Add integration tests for model mismatch, streaming header, and body
    size limit (axum's default 2MB)

* fix: address Copilot review feedback on OpenAI-compat API

- Wire chat_rate_limiter into /v1/chat/completions handler
- Execute LLM before starting SSE stream so failures return proper HTTP
  errors instead of SSE error events
- Validate tool-role messages require tool_call_id and name fields
- Surface list_models() errors in models_handler via map_llm_error
- Reject unknown roles with 400 instead of defaulting to User

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

---------

Co-authored-by: firat.sertgoz <[email protected]>
Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-02-13 18:26:49 +04:00
b3dee13954 fix: flatten tool messages for NEAR AI cloud-api compatibility (#41)
* fix: flatten tool messages for NEAR AI cloud-api compatibility

NEAR AI cloud-api does not support the OpenAI multi-turn tool-calling
protocol (role:"tool" messages cause HTTP 400). This adds a
flatten_tool_messages() pass in NearAiChatProvider that rewrites
assistant tool_call messages and tool result messages into plain
assistant/user text before sending to the API. The model still sees
the tool execution history, just in a text format it can process.

Also includes a minor fix to telegram channel send_pairing_reply
for updated WASM host function signature.

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

* fix: resolve CI failures in fmt, rate limiting, and test configuration

- Apply cargo fmt to nearai_chat.rs formatting violations
- Fix truncate(true) bug in record_failed_approve that cleared the
  attempts file before reading, preventing rate limit from ever
  triggering
- Skip bundled channel test when WASM build artifacts are unavailable
  (CI lacks wasm32-wasip2 target)
- Split CI test workflow to exclude workspace_integration tests that
  require PostgreSQL

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

* fix: resolve clippy unnecessary_unwrap lint (Rust 1.93)

Replace is_some() + unwrap() pattern with if-let binding to satisfy
clippy::unnecessary_unwrap which is now deny-by-default.

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

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
Co-authored-by: firat.sertgoz <[email protected]>
Co-authored-by: Illia Polosukhin <[email protected]>
2026-02-13 06:18:43 +00:00
33ef0a6ea5 fix: security hardening across all layers (#35)
* fix: comprehensive security hardening across all layers

Critical:
- Replace --dangerously-skip-permissions with explicit tool allowlist
  via settings.json (Claude Code bridge)
- Constant-time token comparison (subtle crate) in web auth and
  orchestrator auth to prevent timing attacks

High:
- Revoke tokens and clean up handles on container creation failure
- Drop SETUID/SETGID capabilities from containers (keep only CHOWN)
- Disable redirect following in HTTP tool and WASM wrapper (SSRF)
- Reject URL userinfo (@) in WASM allowlist parser (host confusion)
- Fix binary body bypassing leak detection (from_utf8 -> from_utf8_lossy)
- Protect identity files from LLM overwrites (prompt injection defense)
- Prevent tool shadowing: built-in tools cannot be replaced dynamically
- User-scoped job APIs: list/detail/cancel/restart/prompt/events/files
- CORS restricted to localhost origins, WebSocket origin validation
- Sandbox shell fail-closed: no silent fallback to unsandboxed execution
- Scrub secrets from log broadcaster before SSE broadcast
- XSS sanitization on rendered markdown in web UI
- WASM epoch ticker thread so timeout deadlines actually fire

Medium:
- Cap state transition history at 200 entries
- SSE/WebSocket connection limit (100 max)
- Request body size limit (1MB)
- Response body size limit enforcement in WASM HTTP
- UTF-8 safe string truncation (routine engine, shell tool)
- Fix PolicyAction::Sanitize to actually run the sanitizer
- TOCTOU fix in scheduler and context manager (hold write lock)
- Project file serving moved behind auth
- Path traversal guard on project_id
- Session file permissions set to 0600 on unix
- AtomicUsize for routine running_count (panic-safe)
- Completion detection hardened against false positives and tool injection
- Tool output no longer drives job completion (only LLM response)

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

* fix: address security review findings across all layers

- Fix path traversal sandbox bypass via lexical normalization (file.rs)
- Fix SSRF via DNS rebinding with pre-request hostname resolution (http.rs)
- Add token budget enforcement on LLM calls (reasoning.rs, state.rs)
- Fix cross-user chat history leak with ownership verification (store.rs, server.rs)
- Add sliding-window rate limiter on gateway chat endpoint (server.rs)
- Harden extension install: HTTPS-only, 50MB cap, WASM magic validation (manager.rs)
- Add destructive command blocklist that overrides shell auto-approval (shell.rs)
- Add 5MB response body size cap to HTTP tool (http.rs)

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

* refactor: deduplicate shared helpers and remove dead code

Extract floor_char_boundary and llm_signals_completion into src/util.rs,
unifying diverging phrase lists from agent/worker.rs and worker/runtime.rs.
Remove dead RespondResult::usage(), duplicate PROTECTED_IDENTITY_FILES
constant, double LeakDetector scanning in WebLogLayer, and invalid
0.0.0.0 origin from WebSocket allow list.

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

* fix: address PR review findings and CI test failures

- Fix record_failed_approve: .truncate(true) wiped the attempts file
  before reading, so failed pairing attempts never accumulated and
  rate limiting never triggered.
- Guard wizard WASM test: skip gracefully when channel build artifacts
  are absent (CI doesn't compile wasm32-wasip2 targets).
- Fix DNS rebinding check: use port 0 instead of hardcoded 443, since
  the port is irrelevant for hostname resolution.
- Remove hardcoded CORS port 3001: the dynamic addr.port() entries
  already cover the actual server port.
- Require WebSocket Origin header: reject connections that omit it
  entirely, since browsers always send Origin for WS upgrades and a
  missing header indicates a non-browser client bypassing the check.

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

* fix: address second round of PR review findings

- store.rs: reintroduce file locking around read-modify-write in
  record_failed_approve (concurrent callers could clobber each other).
- sse.rs: replace load+check+fetch_add with atomic fetch_update in both
  subscribe_raw() and subscribe() to prevent overshooting max_connections.
- ws.rs: decrement WS tracker before early return when subscribe_raw()
  returns None (connection limit reached), fixing a counter leak.
- server.rs: parse WS Origin host exactly instead of prefix matching,
  preventing bypass via crafted origins like http://localhost.evil.com.
- workspace_integration.rs: skip tests gracefully when Postgres is
  unreachable instead of panicking (fixes 10 CI failures).

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

* fix: add Origin header to WS integration tests

The Origin header requirement added in a3b0190 broke the WS gateway
integration tests. Test clients now send Origin: http://127.0.0.1:{port}
to match the server's localhost validation.

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

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-02-13 05:25:20 +00:00
github-actions[bot]GitHubgithub-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
e0a43c81f9 chore: release v0.1.3 (#56)
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
2026-02-13 00:28:13 +01:00
Vlad Frolov bada79ba4a ci: Enabled builds caching during CI/CD 2026-02-13 00:17:55 +01:00
Vlad Frolov a70c89d9e3 ci: Disabled npm publishing as the name is already taken 2026-02-13 00:17:55 +01:00
github-actions[bot]GitHubgithub-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
247445f819 chore: release v0.1.2 (#55)
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
2026-02-12 23:47:16 +01:00
Vlad Frolov 14254a699f docs: Added Installation instructions for the pre-built binaries 2026-02-12 23:42:46 +01:00
Vlad Frolov 2039442885 ci: Disabled Windows ARM64 builds as auto-updater [provided by cargo-dist] does not support this platform yet and it is not a common platform for us to support 2026-02-12 23:41:27 +01:00
Vlad Frolov e796b838fa ci: Skip creating GitHub Release with release-plz [cargo-dist will do it in the release workflow] 2026-02-12 23:11:30 +01:00
github-actions[bot]GitHubgithub-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
54ce7434fb chore: release v0.1.1 (#54)
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
2026-02-12 23:09:48 +01:00
Vlad Frolov e9e0374c85 ci: Renamed the secrets in release-plz.yml to match the configuration 2026-02-12 23:05:30 +01:00
Vlad Frolov 5582a0dfbf ci: Make sure that the binaries release CD it kicking in after release-plz 2026-02-12 22:55:08 +01:00
github-actions[bot]GitHubgithub-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
1cf08a4b42 chore: release v0.1.0 (#46)
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
2026-02-12 22:14:28 +01:00
130 changed files with 17529 additions and 4107 deletions
+2 -1
View File
@@ -8,12 +8,13 @@ jobs:
runs-on: ubuntu-latest
steps:
- name: Checkout repository
uses: actions/checkout@v4
uses: actions/checkout@v6
- name: Install Rust
uses: dtolnay/rust-toolchain@stable
with:
profile: minimal
components: rustfmt, clippy
- uses: Swatinem/rust-cache@v2
- name: Check formatting
run: |
cargo fmt --all -- --check
+13 -1
View File
@@ -24,12 +24,23 @@ jobs:
- &install-rust
name: Install Rust toolchain
uses: dtolnay/rust-toolchain@stable
- uses: Swatinem/rust-cache@v2
# Generating a GitHub token, so that PRs and tags created by
# the release-plz-action can trigger actions workflows.
- name: Generate GitHub token
uses: actions/create-github-app-token@v2
id: generate-token
with:
# GitHub App ID secret name
app-id: ${{ secrets.GH_RELEASES_MANAGER_APP_ID }}
# GitHub App private key secret name
private-key: ${{ secrets.GH_RELEASES_MANAGER_APP_PRIVATE_KEY }}
- name: Run release-plz
uses: release-plz/[email protected]
with:
command: release
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
GITHUB_TOKEN: ${{ steps.generate-token.outputs.token }}
CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }}
# Create a PR with the new versions and changelog, preparing the next release.
@@ -46,6 +57,7 @@ jobs:
steps:
- *checkout
- *install-rust
- uses: Swatinem/rust-cache@v2
- name: Run release-plz
uses: release-plz/[email protected]
with:
+9 -33
View File
@@ -1,4 +1,4 @@
# This file was autogenerated by dist: https://opensource.axo.dev/cargo-dist/
# This file was autogenerated by dist: https://axodotdev.github.io/cargo-dist
#
# Copyright 2022-2024, axodotdev
# SPDX-License-Identifier: MIT or Apache-2.0
@@ -58,6 +58,7 @@ jobs:
steps:
- uses: actions/checkout@v4
with:
persist-credentials: false
submodules: recursive
- name: Install dist
# we specify bash to get pipefail; it guards against the `curl` command
@@ -117,6 +118,7 @@ jobs:
git config --global core.longpaths true
- uses: actions/checkout@v4
with:
persist-credentials: false
submodules: recursive
- name: Install Rust non-interactively if not already installed
if: ${{ matrix.container }}
@@ -179,6 +181,7 @@ jobs:
steps:
- uses: actions/checkout@v4
with:
persist-credentials: false
submodules: recursive
- name: Install cached dist
uses: actions/download-artifact@v4
@@ -218,8 +221,8 @@ jobs:
- plan
- build-local-artifacts
- build-global-artifacts
# Only run if we're "publishing", and only if local and global didn't fail (skipped is fine)
if: ${{ always() && needs.plan.outputs.publishing == 'true' && (needs.build-global-artifacts.result == 'skipped' || needs.build-global-artifacts.result == 'success') && (needs.build-local-artifacts.result == 'skipped' || needs.build-local-artifacts.result == 'success') }}
# Only run if we're "publishing", and only if plan, local and global didn't fail (skipped is fine)
if: ${{ always() && needs.plan.result == 'success' && needs.plan.outputs.publishing == 'true' && (needs.build-global-artifacts.result == 'skipped' || needs.build-global-artifacts.result == 'success') && (needs.build-local-artifacts.result == 'skipped' || needs.build-local-artifacts.result == 'success') }}
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
runs-on: "ubuntu-22.04"
@@ -228,6 +231,7 @@ jobs:
steps:
- uses: actions/checkout@v4
with:
persist-credentials: false
submodules: recursive
- name: Install cached dist
uses: actions/download-artifact@v4
@@ -278,47 +282,19 @@ jobs:
gh release create "${{ needs.plan.outputs.tag }}" --target "$RELEASE_COMMIT" $PRERELEASE_FLAG --title "$ANNOUNCEMENT_TITLE" --notes-file "$RUNNER_TEMP/notes.txt" artifacts/*
publish-npm:
needs:
- plan
- host
runs-on: "ubuntu-22.04"
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
PLAN: ${{ needs.plan.outputs.val }}
if: ${{ !fromJson(needs.plan.outputs.val).announcement_is_prerelease || fromJson(needs.plan.outputs.val).publish_prereleases }}
steps:
- name: Fetch npm packages
uses: actions/download-artifact@v4
with:
pattern: artifacts-*
path: npm/
merge-multiple: true
- uses: actions/setup-node@v4
with:
node-version: '20.x'
registry-url: 'https://registry.npmjs.org'
- run: |
for release in $(echo "$PLAN" | jq --compact-output '.releases[] | select([.artifacts[] | endswith("-npm-package.tar.gz")] | any)'); do
pkg=$(echo "$release" | jq '.artifacts[] | select(endswith("-npm-package.tar.gz"))' --raw-output)
npm publish --access public "./npm/${pkg}"
done
env:
NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
announce:
needs:
- plan
- host
- publish-npm
# use "always() && ..." to allow us to wait for all publish jobs while
# still allowing individual publish jobs to skip themselves (for prereleases).
# "host" however must run to completion, no skipping allowed!
if: ${{ always() && needs.host.result == 'success' && (needs.publish-npm.result == 'skipped' || needs.publish-npm.result == 'success') }}
if: ${{ always() && needs.host.result == 'success' }}
runs-on: "ubuntu-22.04"
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
steps:
- uses: actions/checkout@v4
with:
persist-credentials: false
submodules: recursive
+2 -1
View File
@@ -11,10 +11,11 @@ jobs:
runs-on: ubuntu-latest
steps:
- name: Checkout repository
uses: actions/checkout@v4
uses: actions/checkout@v6
- name: Install Rust
uses: dtolnay/rust-toolchain@stable
with:
profile: minimal
- uses: Swatinem/rust-cache@v2
- name: Run Tests
run: cargo test --all-features -- --nocapture
+1
View File
@@ -1,6 +1,7 @@
.env
.env.local
.env.*
target/
+117
View File
@@ -0,0 +1,117 @@
# Changelog
All notable changes to this project will be documented in this file.
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
## [Unreleased]
## [0.1.3](https://github.com/nearai/ironclaw/compare/v0.1.2...v0.1.3) - 2026-02-12
### Other
- Enabled builds caching during CI/CD
- Disabled npm publishing as the name is already taken
## [0.1.2](https://github.com/nearai/ironclaw/compare/v0.1.1...v0.1.2) - 2026-02-12
### Other
- Added Installation instructions for the pre-built binaries
- Disabled Windows ARM64 builds as auto-updater [provided by cargo-dist] does not support this platform yet and it is not a common platform for us to support
## [0.1.1](https://github.com/nearai/ironclaw/compare/v0.1.0...v0.1.1) - 2026-02-12
### Other
- Renamed the secrets in release-plz.yml to match the configuration
- Make sure that the binaries release CD it kicking in after release-plz
## [0.1.0](https://github.com/nearai/ironclaw/releases/tag/v0.1.0) - 2026-02-12
### Added
- Add multi-provider LLM support via rig-core adapter ([#36](https://github.com/nearai/ironclaw/pull/36))
- Sandbox jobs ([#4](https://github.com/nearai/ironclaw/pull/4))
- Add Google Suite & Telegram WASM tools ([#9](https://github.com/nearai/ironclaw/pull/9))
- Improve CLI ([#5](https://github.com/nearai/ironclaw/pull/5))
### Fixed
- resolve runtime panic in Linux keychain integration ([#32](https://github.com/nearai/ironclaw/pull/32))
### Other
- Skip release-plz on forks
- Upgraded release-plz CD pipeline
- Added CI/CD and release pipelines ([#45](https://github.com/nearai/ironclaw/pull/45))
- DM pairing + Telegram channel improvements ([#17](https://github.com/nearai/ironclaw/pull/17))
- Fixes build, adds missing sse event and correct command ([#11](https://github.com/nearai/ironclaw/pull/11))
- Codex/feature parity pr hook ([#6](https://github.com/nearai/ironclaw/pull/6))
- Add WebSocket gateway and control plane ([#8](https://github.com/nearai/ironclaw/pull/8))
- select bundled Telegram channel and auto-install ([#3](https://github.com/nearai/ironclaw/pull/3))
- Adding skills for reusable work
- Fix MCP tool calls, approval loop, shutdown, and improve web UI
- Add auth mode, fix MCP token handling, and parallelize startup loading
- Merge remote-tracking branch 'origin/main' into ui
- Adding web UI
- Rename `setup` CLI command to `onboard` for compatibility
- Add in-chat extension discovery, auth, and activation system
- Add Telegram typing indicator via WIT on-status callback
- Add proactivity features: memory CLI, session pruning, self-repair notifications, slash commands, status diagnostics, context warnings
- Add hosted MCP server support with OAuth 2.1 and token refresh
- Add interactive setup wizard and persistent settings
- Rebrand to IronClaw with security-first mission
- Fix build_software tool stuck in planning mode loop
- Enable sandbox by default
- Fix Telegram Markdown formatting and clarify tool/memory distinctions
- Simplify Telegram channel config with host-injected tunnel/webhook settings
- Apply Telegram channel learnings to WhatsApp implementation
- Merge remote-tracking branch 'origin/main'
- Docker file for sandbox
- Replace hardcoded intent patterns with job tools
- Fix router test to match intentional job creation patterns
- Add Docker execution sandbox for secure shell command isolation
- Move setup wizard credentials to database storage
- Add interactive setup wizard for first-run configuration
- Add Telegram Bot API channel as WASM module
- Add OpenClaw feature parity tracking matrix
- Add Chat Completions API support and expand REPL debugging
- Implementing channels to be handled in wasm
- Support non interactive mode and model selection
- Implement tool approval, fix tool definition refresh, and wire embeddings
- Tool use
- Wiring more
- Add heartbeat integration, planning phase, and auto-repair
- Login flow
- Extend support for session management
- Adding builder capability
- Load tools at launch
- Fix multiline message rendering in TUI
- Parse NEAR AI alternative response format with output field
- Handle NEAR AI plain text responses
- Disable mouse capture to allow text selection in TUI
- Add verbose logging to debug empty NEAR AI responses
- Improve NEAR AI response parsing for varying response formats
- Show status/thinking messages in chat window, debug empty responses
- Add timeout and logging to NEAR AI provider
- Add status updates to show agent thinking/processing state
- Add CLI subcommands for WASM tool management
- Fix TUI shutdown: send /shutdown message and handle in agent loop
- Remove SimpleCliChannel, add Ctrl+D twice quit, redirect logs to TUI
- Fix TuiChannel integration and enable in main.rs
- Integrate Codex patterns: task scheduler, TUI, sessions, compaction
- Adding LICENSE
- Add README with IronClaw branding
- Add WASM sandbox secure API extension
- Wire database Store into agent loop
- Implementing WASM runtime
- Add workspace integration tests
- Compact memory_tree output format
- Replace memory_list with memory_tree tool
- Simplify workspace to path-based storage, remove legacy code
- Add NEAR AI chat-api as default LLM provider
- Add CLAUDE.md project documentation
- Add workspace and memory system (OpenClaw-inspired)
- Initial implementation of the agent framework
+128 -6
View File
@@ -151,6 +151,12 @@ src/
│ ├── rate_limiter.rs # Per-tool rate limiting
│ └── storage.rs # Linear memory persistence
├── db/ # Database abstraction layer
│ ├── mod.rs # Database trait (~60 async methods)
│ ├── postgres.rs # PostgreSQL backend (delegates to Store + Repository)
│ ├── libsql_backend.rs # libSQL/Turso backend (embedded SQLite)
│ └── libsql_migrations.rs # SQLite-dialect schema (idempotent)
├── workspace/ # Persistent memory system (OpenClaw-inspired)
│ ├── mod.rs # Workspace struct, memory operations
│ ├── document.rs # MemoryDocument, MemoryChunk, WorkspaceEntry
@@ -192,8 +198,9 @@ When designing new features or systems, always prefer generic/extensible archite
### Error Handling
- Use `thiserror` for error types in `error.rs`
- Never use `.unwrap()` in production code (tests are fine)
- Never use `.unwrap()` or `.expect()` in production code (tests are fine)
- Map errors with context: `.map_err(|e| SomeError::Variant { reason: e.to_string() })?`
- Before committing, grep for `.unwrap()` and `.expect(` in changed files to catch violations mechanically
### Async
- All I/O is async with tokio
@@ -201,6 +208,7 @@ When designing new features or systems, always prefer generic/extensible archite
- Use `RwLock` for concurrent read/write access
### Traits for Extensibility
- `Database` - Add new database backends (must implement all ~60 methods)
- `Channel` - Add new input sources
- `Tool` - Add new capabilities
- `LlmProvider` - Add new LLM backends
@@ -248,7 +256,12 @@ Pending -> InProgress -> Completed -> Submitted -> Accepted
Environment variables (see `.env.example`):
```bash
# Database backend (default: postgres)
DATABASE_BACKEND=postgres # or "libsql" / "turso"
DATABASE_URL=postgres://user:pass@localhost/ironclaw
LIBSQL_PATH=~/.ironclaw/ironclaw.db # libSQL local path (default)
# LIBSQL_URL=libsql://xxx.turso.io # Turso cloud (optional)
# LIBSQL_AUTH_TOKEN=xxx # Required with LIBSQL_URL
# NEAR AI (required)
NEARAI_SESSION_TOKEN=sess_...
@@ -308,7 +321,51 @@ Session tokens have the format `sess_xxx` (37 characters). They are authenticate
## Database
Single migration in `migrations/V1__initial.sql`. Tables:
IronClaw supports two database backends, selected at compile time via Cargo feature flags and at runtime via the `DATABASE_BACKEND` environment variable.
**IMPORTANT: All new features that touch persistence MUST support both backends.** Implement the operation as a method on the `Database` trait in `src/db/mod.rs`, then add the implementation in both `src/db/postgres.rs` (delegate to Store/Repository) and `src/db/libsql_backend.rs` (native SQL).
### Backends
| Backend | Feature Flag | Default | Use Case |
|---------|-------------|---------|----------|
| PostgreSQL | `postgres` (default) | Yes | Production, existing deployments |
| libSQL/Turso | `libsql` | No | Zero-dependency local mode, edge, Turso cloud |
```bash
# Build with PostgreSQL only (default)
cargo build
# Build with libSQL only
cargo build --no-default-features --features libsql
# Build with both backends available
cargo build --features "postgres,libsql"
```
### Database Trait
The `Database` trait (`src/db/mod.rs`) defines ~60 async methods covering all persistence:
- Conversations, messages, metadata
- Jobs, actions, LLM calls, estimation snapshots
- Sandbox jobs, job events
- Routines, routine runs
- Tool failures, settings
- Workspace: documents, chunks, hybrid search
Both backends implement this trait. PostgreSQL delegates to the existing `Store` + `Repository`. libSQL implements native SQLite-dialect SQL.
### Schema
**PostgreSQL:** `migrations/V1__initial.sql` (351 lines). Uses pgvector for embeddings, tsvector for FTS, PL/pgSQL functions. Managed by `refinery`.
**libSQL:** `src/db/libsql_migrations.rs` (consolidated schema, ~480 lines). Translates PG types:
- `UUID` -> `TEXT`, `TIMESTAMPTZ` -> `TEXT` (ISO-8601), `JSONB` -> `TEXT`
- `VECTOR(1536)` -> `F32_BLOB(1536)` with `libsql_vector_idx`
- `tsvector`/`ts_rank_cd` -> FTS5 virtual table with sync triggers
- PL/pgSQL functions -> SQLite triggers
**Tables (both backends):**
**Core:**
- `conversations` - Multi-channel conversation tracking
@@ -320,12 +377,41 @@ Single migration in `migrations/V1__initial.sql`. Tables:
**Workspace/Memory:**
- `memory_documents` - Flexible path-based files (e.g., "context/vision.md", "daily/2024-01-15.md")
- `memory_chunks` - Chunked content with FTS (tsvector) and vector (pgvector) indexes
- `memory_chunks` - Chunked content with FTS and vector indexes
- `heartbeat_state` - Periodic execution tracking
Requires pgvector extension: `CREATE EXTENSION IF NOT EXISTS vector;`
**Other:**
- `routines`, `routine_runs` - Scheduled/reactive execution
- `settings` - Per-user key-value settings
- `tool_failures` - Self-repair tracking
- `secrets`, `wasm_tools`, `tool_capabilities` - Extension infrastructure
Run migrations: `refinery migrate -c refinery.toml`
### Configuration
```bash
# Backend selection (default: postgres)
DATABASE_BACKEND=libsql
# PostgreSQL
DATABASE_URL=postgres://user:pass@localhost/ironclaw
# libSQL (embedded)
LIBSQL_PATH=~/.ironclaw/ironclaw.db # Default path
# libSQL (Turso cloud sync)
LIBSQL_URL=libsql://your-db.turso.io
LIBSQL_AUTH_TOKEN=your-token # Required when LIBSQL_URL is set
```
### Current Limitations (libSQL backend)
- **Workspace/memory system** not yet wired through Database trait (requires Store migration)
- **Secrets store** not yet available (still requires PostgresSecretsStore)
- **Hybrid search** uses FTS5 only (vector search via libsql_vector_idx not yet implemented)
- **Settings reload from DB** skipped (Config::from_db requires Store)
- No incremental migration versioning (schema is CREATE IF NOT EXISTS, no ALTER TABLE support yet)
- **No encryption at rest** -- The local SQLite database file stores conversation content, job data, workspace memory, and other application data in plaintext. Only secrets (API tokens, credentials) are encrypted via AES-256-GCM before storage. Users handling sensitive data should use full-disk encryption (FileVault, LUKS, BitLocker) or consider the PostgreSQL backend with TDE/encrypted storage.
- **JSON merge patch vs path-targeted update** -- The libSQL backend uses RFC 7396 JSON Merge Patch (`json_patch`) for metadata updates, while PostgreSQL uses path-targeted `jsonb_set`. Merge patch replaces top-level keys entirely, which may drop nested keys not present in the patch. Callers should avoid relying on partial nested object updates in metadata fields.
## Safety Layer
@@ -387,6 +473,7 @@ Key test patterns:
-**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
-**libSQL/Turso backend** - Database trait abstraction (`src/db/`), feature-gated dual backend support (postgres/libsql), embedded SQLite for zero-dependency local mode
## Adding a New Tool
@@ -551,6 +638,37 @@ RUST_LOG=ironclaw=debug,tower_http=debug cargo run
- Keep functions focused, extract helpers when logic is reused
- Comments for non-obvious logic only
## Review & Fix Discipline
Hard-won lessons from code review -- follow these when fixing bugs or addressing review feedback.
### Fix the pattern, not just the instance
When a reviewer flags a bug (e.g., TOCTOU race in INSERT + SELECT-back), search the entire codebase for all instances of that same pattern. A fix in `SecretsStore::create()` that doesn't also fix `WasmToolStore::store()` is half a fix.
### Propagate architectural fixes to satellite types
If a core type changes its concurrency model (e.g., `LibSqlBackend` switches to connection-per-operation), every type that was handed a resource from the old model (e.g., `LibSqlSecretsStore`, `LibSqlWasmToolStore` holding a single `Connection`) must also be updated. Grep for the old type across the codebase.
### Schema translation is more than DDL
When translating a database schema between backends (PostgreSQL to libSQL, etc.), check for:
- **Indexes** -- diff `CREATE INDEX` statements between the two schemas
- **Seed data** -- check for `INSERT INTO` in migrations (e.g., `leak_detection_patterns`)
- **Semantic differences** -- document where SQL functions behave differently (e.g., `json_patch` vs `jsonb_set`)
### Feature flag testing
When adding feature-gated code, test compilation with each feature in isolation:
```bash
cargo check # default features
cargo check --no-default-features --features libsql # libsql only
cargo check --all-features # all features
```
Dead code behind the wrong `#[cfg]` gate will only show up when building with a single feature.
### Mechanical verification before committing
Run these checks on changed files before committing:
- `grep -rnE '\.unwrap\(|\.expect\(' <files>` -- no panics in production
- `grep -rn 'super::' <files>` -- use `crate::` imports
- If you fixed a pattern bug, `grep` for other instances of that pattern across `src/`
## Workspace & Memory System
Inspired by [OpenClaw](https://github.com/openclaw/openclaw), the workspace provides persistent memory for agents with a flexible filesystem-like structure.
@@ -625,7 +743,7 @@ Four tools for LLM use:
### Hybrid Search (RRF)
Combines full-text search (PostgreSQL `ts_rank_cd`) and vector similarity (pgvector cosine) using Reciprocal Rank Fusion:
Combines full-text search and vector similarity using Reciprocal Rank Fusion:
```
score(d) = Σ 1/(k + rank(d)) for each method where d appears
@@ -633,6 +751,10 @@ score(d) = Σ 1/(k + rank(d)) for each method where d appears
Default k=60. Results from both methods are combined, with documents appearing in both getting boosted scores.
**Backend differences:**
- **PostgreSQL:** `ts_rank_cd` for FTS, pgvector cosine distance for vectors, full RRF
- **libSQL:** FTS5 for keyword search only (vector search via `libsql_vector_idx` not yet wired)
### Heartbeat System
Proactive periodic execution (default: 30 minutes):
Generated
+635 -66
View File
File diff suppressed because it is too large Load Diff
+32 -17
View File
@@ -1,8 +1,8 @@
[package]
name = "ironclaw"
version = "0.1.0"
version = "0.1.3"
edition = "2024"
rust-version = "1.85"
rust-version = "1.92"
description = "Secure personal AI assistant that protects your data and expands its capabilities on the fly"
authors = ["NEAR AI <[email protected]>"]
license = "MIT OR Apache-2.0"
@@ -22,17 +22,20 @@ tokio-stream = { version = "0.1", features = ["sync"] }
futures = "0.3"
# HTTP client
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls", "stream"] }
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls-native-roots", "stream"] }
# Serialization
serde = { version = "1", features = ["derive"] }
serde_json = "1"
# Database
deadpool-postgres = "0.14"
tokio-postgres = { version = "0.7", features = ["with-uuid-1", "with-chrono-0_4", "with-serde_json-1"] }
postgres-types = { version = "0.2", features = ["with-serde_json-1"] }
refinery = { version = "0.8", features = ["tokio-postgres"] }
# Database - PostgreSQL (default, feature-gated)
deadpool-postgres = { version = "0.14", optional = true }
tokio-postgres = { version = "0.7", features = ["with-uuid-1", "with-chrono-0_4", "with-serde_json-1"], optional = true }
postgres-types = { version = "0.2", features = ["with-serde_json-1"], optional = true }
refinery = { version = "0.8", features = ["tokio-postgres"], optional = true }
# Database - libSQL/Turso (optional embedded database)
libsql = { version = "0.6", optional = true, default-features = false, features = ["core", "replication"] }
# Error handling
thiserror = "2"
@@ -48,7 +51,7 @@ dotenvy = "0.15"
# Core types
uuid = { version = "1", features = ["v4", "serde"] }
chrono = { version = "0.4", features = ["serde"] }
rust_decimal = { version = "1", features = ["serde", "serde-with-str", "db-tokio-postgres", "maths"] }
rust_decimal = { version = "1", features = ["serde", "serde-with-str", "maths"] }
rust_decimal_macros = "1"
# Async traits
@@ -81,7 +84,8 @@ fs4 = "0.6"
# Secrecy for sensitive values
secrecy = { version = "0.10", features = ["serde"] }
# URL encoding for OAuth flow
# URL parsing and encoding
url = "2"
urlencoding = "2"
# Open URLs in browser
@@ -89,7 +93,7 @@ open = "5"
# Vector embeddings for semantic search
# The postgres feature provides ToSql/FromSql for postgres-types (shared by tokio-postgres)
pgvector = { version = "0.4", features = ["postgres"] }
pgvector = { version = "0.4", features = ["postgres"], optional = true }
# WASM sandbox for untrusted tool execution
wasmtime = { version = "28", features = ["component-model"] }
@@ -102,6 +106,7 @@ hkdf = "0.12"
sha2 = "0.10"
blake3 = "1"
rand = "0.8"
subtle = "2" # Constant-time comparisons for token validation
# Multi-provider LLM support
rig-core = "0.30"
@@ -134,9 +139,22 @@ pretty_assertions = "1"
tempfile = "3"
[features]
default = []
default = ["postgres", "libsql"]
postgres = [
"dep:deadpool-postgres",
"dep:tokio-postgres",
"dep:postgres-types",
"dep:refinery",
"dep:pgvector",
"rust_decimal/db-tokio-postgres",
]
libsql = ["dep:libsql"]
integration = []
[[example]]
name = "test_heartbeat"
required-features = ["postgres"]
# The profile that 'cargo dist' will build with
[profile.dist]
inherits = "release"
@@ -146,18 +164,16 @@ lto = "thin"
[workspace.metadata.dist]
# The preferred dist version to use in CI (Cargo.toml SemVer syntax)
cargo-dist-version = "0.30.3"
allow-dirty = ["ci"]
# CI backends to support
ci = "github"
# The installers to generate for each app
installers = ["shell", "powershell", "npm", "msi"]
# Publish jobs to run in CI
publish-jobs = ["npm"]
publish-jobs = []
# Target platforms to build apps for (Rust target-triple syntax)
targets = [
"aarch64-apple-darwin",
"aarch64-unknown-linux-gnu",
"aarch64-pc-windows-msvc",
"x86_64-apple-darwin",
"x86_64-unknown-linux-gnu",
"x86_64-pc-windows-msvc",
@@ -171,12 +187,11 @@ pr-run-mode = "upload"
# Path that installers should place binaries in
install-path = "CARGO_HOME"
# Whether to install an updater program
install-updater = false
install-updater = true
[workspace.metadata.dist.github-custom-runners]
aarch64-unknown-linux-gnu = "ubuntu-24.04-arm"
x86_64-unknown-linux-gnu = "ubuntu-22.04"
x86_64-pc-windows-msvc = "windows-2022"
aarch64-pc-windows-msvc = "windows-2025"
x86_64-apple-darwin = "macos-15-intel"
aarch64-apple-darwin = "macos-14"
+46
View File
@@ -0,0 +1,46 @@
# Multi-stage Dockerfile for the IronClaw agent (cloud deployment).
#
# Build:
# docker build --platform linux/amd64 -t ironclaw:latest .
#
# Run:
# docker run --env-file .env -p 3000:3000 ironclaw:latest
# Stage 1: Build
FROM rust:1.92-slim-bookworm AS builder
RUN apt-get update && apt-get install -y --no-install-recommends \
pkg-config libssl-dev cmake gcc g++ \
&& rm -rf /var/lib/apt/lists/*
WORKDIR /app
# Copy manifests first for layer caching
COPY Cargo.toml Cargo.lock ./
# Copy source and build artifacts
COPY src/ src/
COPY migrations/ migrations/
COPY wit/ wit/
RUN cargo build --release --bin ironclaw
# Stage 2: Runtime
FROM debian:bookworm-slim
RUN apt-get update && apt-get install -y --no-install-recommends \
ca-certificates libssl3 \
&& rm -rf /var/lib/apt/lists/*
COPY --from=builder /app/target/release/ironclaw /usr/local/bin/ironclaw
COPY --from=builder /app/migrations /app/migrations
# Non-root user
RUN useradd -m -u 1000 -s /bin/bash ironclaw
USER ironclaw
EXPOSE 3000
ENV RUST_LOG=ironclaw=info
ENTRYPOINT ["ironclaw"]
+2 -2
View File
@@ -9,7 +9,7 @@
# The image includes common development tools so workers can build software,
# run tests, and execute shell commands.
FROM rust:1.85-bookworm AS builder
FROM rust:1.92-bookworm AS builder
WORKDIR /build
COPY . .
@@ -40,7 +40,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
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 \
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain 1.92.0 \
&& chmod -R a+r /usr/local/rustup /usr/local/cargo
# Install Claude Code CLI (for claude-bridge mode)
+4 -4
View File
@@ -37,7 +37,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
| Session management/routing | ✅ | ✅ | SessionManager exists |
| Configuration hot-reload | ✅ | ❌ | |
| Network modes (loopback/LAN/remote) | ✅ | 🚧 | HTTP only |
| OpenAI-compatible HTTP API | ✅ | | /v1/chat/completions |
| OpenAI-compatible HTTP API | ✅ | | /v1/chat/completions |
| Canvas hosting | ✅ | ❌ | Agent-driven UI |
| Gateway lock (PID-based) | ✅ | ❌ | |
| launchd/systemd integration | ✅ | ❌ | |
@@ -133,7 +133,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|---------|----------|----------|-------|
| Pi agent runtime | ✅ | | IronClaw uses custom runtime |
| RPC-based execution | ✅ | ✅ | Orchestrator/worker pattern |
| Multi-provider failover | ✅ | | Provider fallback chains |
| Multi-provider failover | ✅ | | `FailoverProvider` tries providers sequentially on retryable errors |
| Per-sender sessions | ✅ | ✅ | |
| Global sessions | ✅ | ❌ | Optional shared context |
| Session pruning | ✅ | ❌ | Auto cleanup old sessions |
@@ -173,7 +173,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
| Feature | OpenClaw | IronClaw | Notes |
|---------|----------|----------|-------|
| Auto-discovery | ✅ | ❌ | |
| Failover chains | ✅ | | Provider fallback |
| Failover chains | ✅ | | `FailoverProvider` with configurable `fallback_model` |
| Cooldown management | ✅ | ❌ | Skip failed providers |
| Per-session model override | ✅ | ✅ | Model selector in TUI |
| Model selection UI | ✅ | ✅ | TUI keyboard shortcut |
@@ -419,7 +419,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
- ❌ Slack channel (real implementation)
- ✅ Telegram channel (WASM, DM pairing, caption, /start)
- ❌ WhatsApp channel
- Multi-provider failover
- Multi-provider failover (`FailoverProvider` with retryable error classification)
- ❌ Hooks system (beforeInbound, beforeToolCall, etc.)
### P2 - Medium Priority
+70 -37
View File
@@ -71,7 +71,38 @@ IronClaw is the AI assistant you can actually trust with your personal and profe
- PostgreSQL 15+ with [pgvector](https://github.com/pgvector/pgvector) extension
- NEAR AI account (authentication handled via setup wizard)
### Build
## Download or Build
Visit [Releases page](https://github.com/nearai/ironclaw/releases/) to see the latest updates.
<details>
<summary>Install via Windows Installer (Windows)</summary>
Download the [Windows Installer](https://github.com/nearai/ironclaw/releases/latest/download/ironclaw-x86_64-pc-windows-msvc.msi) and run it.
</details>
<details>
<summary>Install via powershell script (Windows)</summary>
```sh
irm https://github.com/nearai/ironclaw/releases/latest/download/ironclaw-installer.ps1 | iex
```
</details>
<details>
<summary>Install via shell script (macOS, Linux, Windows/WSL)</summary>
```sh
curl --proto '=https' --tlsv1.2 -LsSf https://github.com/nearai/ironclaw/releases/latest/download/ironclaw-installer.sh | sh
```
</details>
<details>
<summary>Compile the source code (Cargo on Windows, Linux, macOS)</summary>
Install it with `cargo`, just make sure you have [Rust](https://rustup.rs) installed on your computer.
```bash
# Clone the repository
@@ -87,6 +118,8 @@ cargo test
For **full release** (after modifying channel sources), run `./scripts/build-all.sh` to rebuild channels first.
</details>
### Database Setup
```bash
@@ -148,42 +181,42 @@ External content passes through multiple security layers:
## Architecture
```
┌────────────────────────────────────────────────────────────────────
│ 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 │
│ └──────────────────────┘
└────────────────────────────────────────────────────────────────────
┌────────────────────────────────────────────────────────────────┐
│ 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
+14 -2
View File
@@ -338,7 +338,13 @@ 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(|e| {
channel_host::log(
channel_host::LogLevel::Error,
&format!("Failed to serialize Slack metadata: {}", e),
);
"{}".to_string()
});
// Strip @ mentions of the bot from the text for cleaner messages
let cleaned_text = strip_bot_mention(&text);
@@ -366,7 +372,13 @@ fn strip_bot_mention(text: &str) -> String {
/// Create a JSON HTTP response.
fn json_response(status: u16, value: serde_json::Value) -> OutgoingHttpResponse {
let body = serde_json::to_vec(&value).unwrap_or_default();
let body = serde_json::to_vec(&value).unwrap_or_else(|e| {
channel_host::log(
channel_host::LogLevel::Error,
&format!("Failed to serialize JSON response: {}", e),
);
Vec::new()
});
let headers = serde_json::json!({"Content-Type": "application/json"});
OutgoingHttpResponse {
+10 -19
View File
@@ -285,11 +285,7 @@ impl Guest for TelegramChannel {
}
// Persist dm_policy and allow_from for DM pairing in handle_message
let dm_policy = config
.dm_policy
.as_deref()
.unwrap_or("pairing")
.to_string();
let dm_policy = config.dm_policy.as_deref().unwrap_or("pairing").to_string();
let _ = channel_host::workspace_write(DM_POLICY_PATH, &dm_policy);
let allow_from_json = serde_json::to_string(&config.allow_from.unwrap_or_default())
@@ -844,8 +840,8 @@ fn send_pairing_reply(chat_id: i64, code: &str) -> Result<(), String> {
"parse_mode": "Markdown",
});
let payload_bytes = serde_json::to_vec(&payload)
.map_err(|e| format!("Failed to serialize payload: {}", e))?;
let payload_bytes =
serde_json::to_vec(&payload).map_err(|e| format!("Failed to serialize payload: {}", e))?;
let headers = serde_json::json!({
"Content-Type": "application/json"
@@ -856,6 +852,7 @@ fn send_pairing_reply(chat_id: i64, code: &str) -> Result<(), String> {
"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/sendMessage",
&headers.to_string(),
Some(&payload_bytes),
None,
);
match result {
@@ -914,15 +911,10 @@ fn handle_message(message: TelegramMessage) {
let is_private = message.chat.chat_type == "private";
// Owner validation: when owner_id is set, only that user can message
let owner_configured = channel_host::workspace_read(OWNER_ID_PATH)
.map(|s| !s.is_empty())
.unwrap_or(false);
let owner_id_str = channel_host::workspace_read(OWNER_ID_PATH).filter(|s| !s.is_empty());
if owner_configured {
if let Ok(owner_id) = channel_host::workspace_read(OWNER_ID_PATH)
.unwrap()
.parse::<i64>()
{
if let Some(ref id_str) = owner_id_str {
if let Ok(owner_id) = id_str.parse::<i64>() {
if from.id != owner_id {
channel_host::log(
channel_host::LogLevel::Debug,
@@ -936,8 +928,8 @@ fn handle_message(message: TelegramMessage) {
}
} else if is_private {
// No owner_id: apply dm_policy for private chats
let dm_policy = channel_host::workspace_read(DM_POLICY_PATH)
.unwrap_or_else(|| "pairing".to_string());
let dm_policy =
channel_host::workspace_read(DM_POLICY_PATH).unwrap_or_else(|| "pairing".to_string());
if dm_policy != "open" {
// Build effective allow list: config allow_from + pairing store
@@ -1000,8 +992,7 @@ fn handle_message(message: TelegramMessage) {
if !respond_to_all {
let has_command = content.starts_with('/');
let bot_username = channel_host::workspace_read(BOT_USERNAME_PATH)
.unwrap_or_default();
let bot_username = channel_host::workspace_read(BOT_USERNAME_PATH).unwrap_or_default();
let has_bot_mention = if bot_username.is_empty() {
content.contains('@')
} else {
+23 -6
View File
@@ -254,10 +254,19 @@ struct WhatsAppChannel;
impl Guest for WhatsAppChannel {
fn on_start(config_json: String) -> Result<ChannelConfig, String> {
let config: WhatsAppConfig = serde_json::from_str(&config_json).unwrap_or(WhatsAppConfig {
api_version: default_api_version(),
reply_to_message: default_reply_to_message(),
});
let config: WhatsAppConfig = match serde_json::from_str(&config_json) {
Ok(c) => c,
Err(e) => {
channel_host::log(
channel_host::LogLevel::Warn,
&format!("Failed to parse WhatsApp config, using defaults: {}", e),
);
WhatsAppConfig {
api_version: default_api_version(),
reply_to_message: default_reply_to_message(),
}
}
};
channel_host::log(
channel_host::LogLevel::Info,
@@ -267,6 +276,9 @@ impl Guest for WhatsAppChannel {
),
);
// Persist api_version in workspace so on_respond() can read it
let _ = channel_host::workspace_write("channels/whatsapp/api_version", &config.api_version);
// WhatsApp Cloud API is webhook-only, no polling available
Ok(ChannelConfig {
display_name: "WhatsApp".to_string(),
@@ -327,11 +339,16 @@ impl Guest for WhatsAppChannel {
let metadata: WhatsAppMessageMetadata = serde_json::from_str(&response.metadata_json)
.map_err(|e| format!("Failed to parse metadata: {}", e))?;
// Read api_version from workspace (set during on_start), fallback to default
let api_version = channel_host::workspace_read("channels/whatsapp/api_version")
.filter(|s| !s.is_empty())
.unwrap_or_else(|| "v18.0".to_string());
// Build WhatsApp API URL with token placeholder
// Host will replace {WHATSAPP_ACCESS_TOKEN} with actual token in Authorization header
let api_url = format!(
"https://graph.facebook.com/v18.0/{}/messages",
metadata.phone_number_id
"https://graph.facebook.com/{}/{}/messages",
api_version, metadata.phone_number_id
);
// Build sendMessage payload
+13
View File
@@ -0,0 +1,13 @@
[Unit]
Description=Cloud SQL Auth Proxy
After=network.target
[Service]
Type=simple
DynamicUser=yes
ExecStart=/usr/local/bin/cloud-sql-proxy ironclaw-prod:us-central1:ironclaw-db --port=5432
Restart=always
RestartSec=5
[Install]
WantedBy=multi-user.target
+27
View File
@@ -0,0 +1,27 @@
# WARNING: Replace all CHANGE_ME values before deploying.
# Do not use placeholder passwords in production.
DATABASE_URL=postgres://ironclaw:CHANGE_ME@localhost:5432/ironclaw
# NEAR AI
NEARAI_SESSION_TOKEN=CHANGE_ME
NEARAI_MODEL=claude-3-5-sonnet-20241022
NEARAI_BASE_URL=https://cloud-api.near.ai
NEARAI_AUTH_URL=https://private.near.ai
NEARAI_API_MODE=chat_completions
# Agent
AGENT_NAME=ironclaw
CLI_ENABLED=false
# Web Gateway
GATEWAY_ENABLED=true
# 0.0.0.0 binds to all interfaces (required for Docker --network=host).
# Use 127.0.0.1 if running outside Docker or for local-only access.
GATEWAY_HOST=0.0.0.0
GATEWAY_PORT=3000
GATEWAY_AUTH_TOKEN=CHANGE_ME
# Disabled for initial deploy
SANDBOX_ENABLED=false
HEARTBEAT_ENABLED=false
EMBEDDING_ENABLED=false
+20
View File
@@ -0,0 +1,20 @@
[Unit]
Description=IronClaw AI Assistant
After=cloud-sql-proxy.service docker.service
Requires=cloud-sql-proxy.service
[Service]
Type=simple
ExecStartPre=/usr/bin/docker pull us-central1-docker.pkg.dev/ironclaw-prod/ironclaw/agent:latest
ExecStart=/usr/bin/docker run --rm \
--name ironclaw \
--env-file /opt/ironclaw/.env \
--network=host \
us-central1-docker.pkg.dev/ironclaw-prod/ironclaw/agent:latest \
--no-onboard
ExecStop=/usr/bin/docker stop ironclaw
Restart=always
RestartSec=10
[Install]
WantedBy=multi-user.target
+68
View File
@@ -0,0 +1,68 @@
#!/usr/bin/env bash
# VM bootstrap script for IronClaw on GCP Compute Engine.
#
# Run on a fresh Debian 12 VM after SSH:
# sudo bash setup.sh
#
# Prerequisites:
# - VM has the ironclaw-vm service account attached
# - Cloud SQL Auth Proxy accessible via IAM
# - Artifact Registry image pushed
set -euo pipefail
# Must run as root
if [ "$(id -u)" -ne 0 ]; then
echo "ERROR: This script must be run as root (sudo bash setup.sh)"
exit 1
fi
echo "==> Installing Docker"
apt-get update
apt-get install -y docker.io
systemctl enable docker
systemctl start docker
echo "==> Installing Cloud SQL Auth Proxy"
curl -fsSL -o /usr/local/bin/cloud-sql-proxy \
https://storage.googleapis.com/cloud-sql-connectors/cloud-sql-proxy/v2.14.3/cloud-sql-proxy.linux.amd64
chmod +x /usr/local/bin/cloud-sql-proxy
echo "==> Installing systemd services"
cp /tmp/deploy/cloud-sql-proxy.service /etc/systemd/system/
cp /tmp/deploy/ironclaw.service /etc/systemd/system/
systemctl daemon-reload
echo "==> Starting Cloud SQL Auth Proxy"
systemctl enable cloud-sql-proxy
systemctl start cloud-sql-proxy
echo "==> Configuring Docker registry auth"
# The VM service account provides Artifact Registry access
gcloud auth configure-docker us-central1-docker.pkg.dev --quiet
echo "==> Creating config directory"
# Owned by root, readable only by root. Docker reads --env-file as root
# before dropping to uid 1000 (ironclaw) inside the container.
mkdir -p /opt/ironclaw
chmod 700 /opt/ironclaw
if [ ! -f /opt/ironclaw/.env ]; then
echo "WARNING: /opt/ironclaw/.env does not exist."
echo "Create it with your configuration before starting IronClaw."
echo "See deploy/env.example for the required variables."
echo ""
echo "Then run: systemctl enable ironclaw && systemctl start ironclaw"
else
chmod 600 /opt/ironclaw/.env
echo "==> Starting IronClaw"
systemctl enable ironclaw
systemctl start ironclaw
fi
echo "==> Setup complete"
echo ""
echo "Verify with:"
echo " systemctl status cloud-sql-proxy"
echo " systemctl status ironclaw"
echo " docker logs ironclaw"
+20
View File
@@ -0,0 +1,20 @@
# Local development only — do NOT use these credentials in production.
services:
postgres:
image: pgvector/pgvector:pg16
ports:
- "5432:5432"
environment:
POSTGRES_DB: ironclaw
POSTGRES_USER: ironclaw
POSTGRES_PASSWORD: ironclaw # dev-only, change for any non-local deployment
volumes:
- pgdata:/var/lib/postgresql/data
healthcheck:
test: ["CMD-SHELL", "pg_isready -U ironclaw"]
interval: 5s
timeout: 3s
retries: 5
volumes:
pgdata:
-1
View File
@@ -81,7 +81,6 @@ async fn main() -> anyhow::Result<()> {
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)?;
+2
View File
@@ -0,0 +1,2 @@
[workspace]
git_release_enable = false
+231 -126
View File
@@ -19,16 +19,16 @@ use crate::channels::{ChannelManager, IncomingMessage, OutgoingResponse, StatusU
use crate::config::{AgentConfig, HeartbeatConfig, RoutineConfig};
use crate::context::ContextManager;
use crate::context::JobContext;
use crate::db::Database;
use crate::error::Error;
use crate::extensions::ExtensionManager;
use crate::history::Store;
use crate::llm::{ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult};
use crate::safety::SafetyLayer;
use crate::tools::ToolRegistry;
use crate::workspace::Workspace;
/// Collapse a tool output string into a single-line preview for display.
fn truncate_for_preview(output: &str, max_chars: usize) -> String {
pub(crate) fn truncate_for_preview(output: &str, max_chars: usize) -> String {
let collapsed: String = output
.chars()
.take(max_chars + 50)
@@ -37,8 +37,14 @@ fn truncate_for_preview(output: &str, max_chars: usize) -> String {
.split_whitespace()
.collect::<Vec<_>>()
.join(" ");
if collapsed.len() > max_chars {
format!("{}...", &collapsed[..max_chars])
// char_indices gives us byte offsets at char boundaries, so the slice is always valid UTF-8.
if collapsed.chars().count() > max_chars {
let byte_offset = collapsed
.char_indices()
.nth(max_chars)
.map(|(i, _)| i)
.unwrap_or(collapsed.len());
format!("{}...", &collapsed[..byte_offset])
} else {
collapsed
}
@@ -59,7 +65,7 @@ enum AgenticLoopResult {
///
/// Bundles the shared components to reduce argument count.
pub struct AgentDeps {
pub store: Option<Arc<Store>>,
pub store: Option<Arc<dyn Database>>,
pub llm: Arc<dyn LlmProvider>,
pub safety: Arc<SafetyLayer>,
pub tools: Arc<ToolRegistry>,
@@ -124,7 +130,7 @@ impl Agent {
}
// Convenience accessors
fn store(&self) -> Option<&Arc<Store>> {
fn store(&self) -> Option<&Arc<dyn Database>> {
self.deps.store.as_ref()
}
@@ -654,19 +660,17 @@ impl Agent {
}
// 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);
}
}
if let Some(store) = self.store()
&& let Ok(Some(metadata)) = store.get_conversation_metadata(thread_uuid).await
&& 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
@@ -954,13 +958,12 @@ impl Agent {
return;
}
if let Some(ref resp) = response {
if let Err(e) = store
if let Some(ref resp) = response
&& let Err(e) = store
.add_conversation_message(thread_id, "assistant", resp)
.await
{
tracing::warn!("Failed to persist assistant message: {}", e);
}
{
tracing::warn!("Failed to persist assistant message: {}", e);
}
});
}
@@ -1058,14 +1061,14 @@ impl Agent {
// Check if interrupted
{
let sess = session.lock().await;
if let Some(thread) = sess.threads.get(&thread_id) {
if thread.state == ThreadState::Interrupted {
return Err(crate::error::JobError::ContextError {
id: thread_id,
reason: "Interrupted".to_string(),
}
.into());
if let Some(thread) = sess.threads.get(&thread_id)
&& thread.state == ThreadState::Interrupted
{
return Err(crate::error::JobError::ContextError {
id: thread_id,
reason: "Interrupted".to_string(),
}
.into());
}
}
@@ -1082,9 +1085,16 @@ impl Agent {
m
});
let result = reasoning.respond_with_tools(&context).await?;
let output = reasoning.respond_with_tools(&context).await?;
match result {
// Track token usage for budget enforcement
tracing::debug!(
"LLM call used {} input + {} output tokens",
output.usage.input_tokens,
output.usage.output_tokens
);
match output.result {
RespondResult::Text(text) => {
// If no tools have been executed yet, prompt the LLM to use tools
// This handles the case where the model explains what it will do
@@ -1133,11 +1143,11 @@ impl Agent {
// Record tool calls in the thread
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) {
if let Some(turn) = thread.last_turn_mut() {
for tc in &tool_calls {
turn.record_tool_call(&tc.name, tc.arguments.clone());
}
if let Some(thread) = sess.threads.get_mut(&thread_id)
&& let Some(turn) = thread.last_turn_mut()
{
for tc in &tool_calls {
turn.record_tool_call(&tc.name, tc.arguments.clone());
}
}
}
@@ -1145,27 +1155,56 @@ impl Agent {
// Execute each tool (with approval checking)
for tc in tool_calls {
// Check if tool requires approval
if let Some(tool) = self.tools().get(&tc.name).await {
if tool.requires_approval() {
// Check if auto-approved for this session
let is_auto_approved = {
let sess = session.lock().await;
sess.is_tool_auto_approved(&tc.name)
if let Some(tool) = self.tools().get(&tc.name).await
&& tool.requires_approval()
{
// Check if auto-approved for this session
let mut is_auto_approved = {
let sess = session.lock().await;
sess.is_tool_auto_approved(&tc.name)
};
// For shell commands, override auto-approval for
// destructive patterns that should always require
// explicit per-invocation approval.
if is_auto_approved
&& tc.name == "shell"
&& let Some(cmd) = tc
.arguments
.get("command")
.and_then(|c| c.as_str().map(String::from))
.or_else(|| {
tc.arguments
.as_str()
.and_then(|s| {
serde_json::from_str::<serde_json::Value>(s).ok()
})
.and_then(|v| {
v.get("command")
.and_then(|c| c.as_str().map(String::from))
})
})
&& crate::tools::builtin::shell::requires_explicit_approval(&cmd)
{
tracing::info!(
"Shell command '{}' requires explicit approval despite auto-approve",
cmd.chars().take(80).collect::<String>()
);
is_auto_approved = false;
}
if !is_auto_approved {
// Need approval - store pending request and return
let pending = PendingApproval {
request_id: Uuid::new_v4(),
tool_name: tc.name.clone(),
parameters: tc.arguments.clone(),
description: tool.description().to_string(),
tool_call_id: tc.id.clone(),
context_messages: context_messages.clone(),
};
if !is_auto_approved {
// Need approval - store pending request and return
let pending = PendingApproval {
request_id: Uuid::new_v4(),
tool_name: tc.name.clone(),
parameters: tc.arguments.clone(),
description: tool.description().to_string(),
tool_call_id: tc.id.clone(),
context_messages: context_messages.clone(),
};
return Ok(AgenticLoopResult::NeedApproval { pending });
}
return Ok(AgenticLoopResult::NeedApproval { pending });
}
}
@@ -1196,34 +1235,34 @@ impl Agent {
)
.await;
if let Ok(ref output) = tool_result {
if !output.is_empty() {
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::ToolResult {
name: tc.name.clone(),
preview: truncate_for_preview(output, 200),
},
&message.metadata,
)
.await;
}
if let Ok(ref output) = tool_result
&& !output.is_empty()
{
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::ToolResult {
name: tc.name.clone(),
preview: output.clone(),
},
&message.metadata,
)
.await;
}
// Record result in thread
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) {
if let Some(turn) = thread.last_turn_mut() {
match &tool_result {
Ok(output) => {
turn.record_tool_result(serde_json::json!(output));
}
Err(e) => {
turn.record_tool_error(e.to_string());
}
if let Some(thread) = sess.threads.get_mut(&thread_id)
&& let Some(turn) = thread.last_turn_mut()
{
match &tool_result {
Ok(output) => {
turn.record_tool_result(serde_json::json!(output));
}
Err(e) => {
turn.record_tool_error(e.to_string());
}
}
}
@@ -1606,17 +1645,17 @@ impl Agent {
};
// Verify request ID if provided
if let Some(req_id) = request_id {
if req_id != pending.request_id {
// Put it back and return error
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) {
thread.await_approval(pending);
}
return Ok(SubmissionResult::error(
"Request ID mismatch. Use the correct request ID.",
));
if let Some(req_id) = request_id
&& req_id != pending.request_id
{
// Put it back and return error
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) {
thread.await_approval(pending);
}
return Ok(SubmissionResult::error(
"Request ID mismatch. Use the correct request ID.",
));
}
if approved {
@@ -1670,20 +1709,20 @@ impl Agent {
)
.await;
if let Ok(ref output) = tool_result {
if !output.is_empty() {
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::ToolResult {
name: pending.tool_name.clone(),
preview: truncate_for_preview(output, 200),
},
&message.metadata,
)
.await;
}
if let Ok(ref output) = tool_result
&& !output.is_empty()
{
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::ToolResult {
name: pending.tool_name.clone(),
preview: output.clone(),
},
&message.metadata,
)
.await;
}
// Build context including the tool result
@@ -1692,15 +1731,15 @@ impl Agent {
// Record result in thread
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) {
if let Some(turn) = thread.last_turn_mut() {
match &tool_result {
Ok(output) => {
turn.record_tool_result(serde_json::json!(output));
}
Err(e) => {
turn.record_tool_error(e.to_string());
}
if let Some(thread) = sess.threads.get_mut(&thread_id)
&& let Some(turn) = thread.last_turn_mut()
{
match &tool_result {
Ok(output) => {
turn.record_tool_result(serde_json::json!(output));
}
Err(e) => {
turn.record_tool_error(e.to_string());
}
}
}
@@ -2060,15 +2099,15 @@ impl Agent {
}
// Persist new job to database (fire-and-forget)
if let Some(store) = self.store() {
if let Ok(ctx) = self.context_manager.get_context(job_id).await {
let store = store.clone();
tokio::spawn(async move {
if let Err(e) = store.save_job(&ctx).await {
tracing::warn!("Failed to persist new job {}: {}", job_id, e);
}
});
}
if let Some(store) = self.store()
&& let Ok(ctx) = self.context_manager.get_context(job_id).await
{
let store = store.clone();
tokio::spawn(async move {
if let Err(e) = store.save_job(&ctx).await {
tracing::warn!("Failed to persist new job {}: {}", job_id, e);
}
});
}
// Schedule for execution
@@ -2148,10 +2187,10 @@ impl Agent {
let mut output = String::from("Jobs:\n");
for job_id in jobs {
if let Ok(ctx) = self.context_manager.get_context(job_id).await {
if ctx.user_id == user_id {
output.push_str(&format!(" {} - {} ({:?})\n", job_id, ctx.title, ctx.state));
}
if let Ok(ctx) = self.context_manager.get_context(job_id).await
&& ctx.user_id == user_id
{
output.push_str(&format!(" {} - {} ({:?})\n", job_id, ctx.title, ctx.state));
}
}
@@ -2602,4 +2641,70 @@ mod tests {
assert!(detect_auth_awaiting("tool_activate", &result).is_none());
}
// --- truncate_for_preview tests ---
use super::truncate_for_preview;
#[test]
fn test_truncate_short_input() {
assert_eq!(truncate_for_preview("hello", 10), "hello");
}
#[test]
fn test_truncate_empty_input() {
assert_eq!(truncate_for_preview("", 10), "");
}
#[test]
fn test_truncate_exact_length() {
assert_eq!(truncate_for_preview("hello", 5), "hello");
}
#[test]
fn test_truncate_over_limit() {
let result = truncate_for_preview("hello world, this is long", 10);
assert!(result.ends_with("..."));
// "hello worl" = 10 chars + "..."
assert_eq!(result, "hello worl...");
}
#[test]
fn test_truncate_collapses_newlines() {
let result = truncate_for_preview("line1\nline2\nline3", 100);
assert!(!result.contains('\n'));
assert_eq!(result, "line1 line2 line3");
}
#[test]
fn test_truncate_collapses_whitespace() {
let result = truncate_for_preview("hello world", 100);
assert_eq!(result, "hello world");
}
#[test]
fn test_truncate_multibyte_utf8() {
// Each emoji is 4 bytes. Truncating at char boundary must not panic.
let input = "😀😁😂🤣😃😄😅😆😉😊";
let result = truncate_for_preview(input, 5);
assert!(result.ends_with("..."));
// First 5 chars = 5 emoji
assert_eq!(result, "😀😁😂🤣😃...");
}
#[test]
fn test_truncate_cjk_characters() {
// CJK chars are 3 bytes each in UTF-8.
let input = "你好世界测试数据很长的字符串";
let result = truncate_for_preview(input, 4);
assert_eq!(result, "你好世界...");
}
#[test]
fn test_truncate_mixed_multibyte_and_ascii() {
let input = "hello 世界 foo";
let result = truncate_for_preview(input, 8);
// 'h','e','l','l','o',' ','世','界' = 8 chars
assert_eq!(result, "hello 世界...");
}
}
+1
View File
@@ -26,6 +26,7 @@ pub mod task;
pub mod undo;
pub mod worker;
pub(crate) use agent_loop::truncate_for_preview;
pub use agent_loop::{Agent, AgentDeps};
pub use compaction::{CompactionResult, ContextCompactor};
pub use context_monitor::{CompactionStrategy, ContextBreakdown, ContextMonitor};
+17 -22
View File
@@ -11,6 +11,7 @@
//! Full-job routines are delegated to the existing `Scheduler`.
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use chrono::Utc;
@@ -23,20 +24,20 @@ use crate::agent::routine::{
};
use crate::channels::{IncomingMessage, OutgoingResponse};
use crate::config::RoutineConfig;
use crate::history::Store;
use crate::db::Database;
use crate::llm::{ChatMessage, CompletionRequest, FinishReason, LlmProvider};
use crate::workspace::Workspace;
/// The routine execution engine.
pub struct RoutineEngine {
config: RoutineConfig,
store: Arc<Store>,
store: Arc<dyn Database>,
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>>,
running_count: Arc<AtomicUsize>,
/// Compiled event regex cache: routine_id -> compiled regex.
event_cache: Arc<RwLock<Vec<(Uuid, Routine, Regex)>>>,
}
@@ -44,7 +45,7 @@ pub struct RoutineEngine {
impl RoutineEngine {
pub fn new(
config: RoutineConfig,
store: Arc<Store>,
store: Arc<dyn Database>,
llm: Arc<dyn LlmProvider>,
workspace: Arc<Workspace>,
notify_tx: mpsc::Sender<OutgoingResponse>,
@@ -55,7 +56,7 @@ impl RoutineEngine {
llm,
workspace,
notify_tx,
running_count: Arc::new(RwLock::new(0)),
running_count: Arc::new(AtomicUsize::new(0)),
event_cache: Arc::new(RwLock::new(Vec::new())),
}
}
@@ -102,10 +103,9 @@ impl RoutineEngine {
if let Trigger::Event {
channel: Some(ch), ..
} = &routine.trigger
&& ch != &message.channel
{
if ch != &message.channel {
continue;
}
continue;
}
// Regex match
@@ -126,7 +126,7 @@ impl RoutineEngine {
}
// Global capacity check
if *self.running_count.read().await >= self.config.max_concurrent_routines {
if self.running_count.load(Ordering::Relaxed) >= self.config.max_concurrent_routines {
tracing::warn!(routine = %routine.name, "Skipped: global max concurrent reached");
continue;
}
@@ -150,7 +150,7 @@ impl RoutineEngine {
};
for routine in routines {
if *self.running_count.read().await >= self.config.max_concurrent_routines {
if self.running_count.load(Ordering::Relaxed) >= self.config.max_concurrent_routines {
tracing::warn!("Global max concurrent routines reached, skipping remaining");
break;
}
@@ -293,21 +293,18 @@ impl RoutineEngine {
/// Shared context passed to the execution function.
struct EngineContext {
store: Arc<Store>,
store: Arc<dyn Database>,
llm: Arc<dyn LlmProvider>,
workspace: Arc<Workspace>,
notify_tx: mpsc::Sender<OutgoingResponse>,
running_count: Arc<RwLock<usize>>,
running_count: Arc<AtomicUsize>,
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;
}
// Increment running count (atomic: survives panics in the execution below)
ctx.running_count.fetch_add(1, Ordering::Relaxed);
let result = match &routine.action {
RoutineAction::Lightweight {
@@ -327,10 +324,7 @@ async fn execute_routine(ctx: EngineContext, routine: Routine, run: RoutineRun)
};
// Decrement running count
{
let mut count = ctx.running_count.write().await;
*count = count.saturating_sub(1);
}
ctx.running_count.fetch_sub(1, Ordering::Relaxed);
// Process result
let (status, summary, tokens) = match result {
@@ -568,7 +562,8 @@ fn truncate(s: &str, max: usize) -> String {
if s.len() <= max {
s.to_string()
} else {
format!("{}...", &s[..max])
let end = crate::util::floor_char_boundary(s, max);
format!("{}...", &s[..end])
}
}
+56 -56
View File
@@ -12,8 +12,8 @@ use crate::agent::task::{Task, TaskContext, TaskOutput};
use crate::agent::worker::{Worker, WorkerDeps};
use crate::config::AgentConfig;
use crate::context::{ContextManager, JobContext, JobState};
use crate::db::Database;
use crate::error::{Error, JobError};
use crate::history::Store;
use crate::llm::LlmProvider;
use crate::safety::SafetyLayer;
use crate::tools::ToolRegistry;
@@ -48,7 +48,7 @@ pub struct Scheduler {
llm: Arc<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
tools: Arc<ToolRegistry>,
store: Option<Arc<Store>>,
store: Option<Arc<dyn Database>>,
/// Running jobs (main LLM-driven jobs).
jobs: Arc<RwLock<HashMap<Uuid, ScheduledJob>>>,
/// Running sub-tasks (tool executions, background tasks).
@@ -63,7 +63,7 @@ impl Scheduler {
llm: Arc<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
tools: Arc<ToolRegistry>,
store: Option<Arc<Store>>,
store: Option<Arc<dyn Database>>,
) -> Self {
Self {
config,
@@ -79,63 +79,63 @@ impl Scheduler {
/// Schedule a job for execution.
pub async fn schedule(&self, job_id: Uuid) -> Result<(), JobError> {
// Check if already scheduled
if self.jobs.read().await.contains_key(&job_id) {
return Ok(());
}
// Hold write lock for the entire check-insert sequence to prevent
// TOCTOU races where two concurrent calls both pass the checks.
{
let mut jobs = self.jobs.write().await;
// Check capacity
let current_count = self.jobs.read().await.len();
if current_count >= self.config.max_parallel_jobs {
return Err(JobError::MaxJobsExceeded {
max: self.config.max_parallel_jobs,
});
}
// Transition job to in_progress
self.context_manager
.update_context(job_id, |ctx| {
ctx.transition_to(
JobState::InProgress,
Some("Scheduled for execution".to_string()),
)
})
.await?
.map_err(|s| JobError::ContextError {
id: job_id,
reason: s,
})?;
// Create worker channel
let (tx, rx) = mpsc::channel(16);
// Create worker with shared dependencies
let deps = WorkerDeps {
context_manager: self.context_manager.clone(),
llm: self.llm.clone(),
safety: self.safety.clone(),
tools: self.tools.clone(),
store: self.store.clone(),
timeout: self.config.job_timeout,
use_planning: self.config.use_planning,
};
let worker = Worker::new(job_id, deps);
// Spawn worker task
let handle = tokio::spawn(async move {
if let Err(e) = worker.run(rx).await {
tracing::error!("Worker for job {} failed: {}", job_id, e);
if jobs.contains_key(&job_id) {
return Ok(());
}
});
// Start the worker
let _ = tx.send(WorkerMessage::Start).await;
if jobs.len() >= self.config.max_parallel_jobs {
return Err(JobError::MaxJobsExceeded {
max: self.config.max_parallel_jobs,
});
}
// Store the scheduled job
self.jobs
.write()
.await
.insert(job_id, ScheduledJob { handle, tx });
// Transition job to in_progress
self.context_manager
.update_context(job_id, |ctx| {
ctx.transition_to(
JobState::InProgress,
Some("Scheduled for execution".to_string()),
)
})
.await?
.map_err(|s| JobError::ContextError {
id: job_id,
reason: s,
})?;
// Create worker channel
let (tx, rx) = mpsc::channel(16);
// Create worker with shared dependencies
let deps = WorkerDeps {
context_manager: self.context_manager.clone(),
llm: self.llm.clone(),
safety: self.safety.clone(),
tools: self.tools.clone(),
store: self.store.clone(),
timeout: self.config.job_timeout,
use_planning: self.config.use_planning,
};
let worker = Worker::new(job_id, deps);
// Spawn worker task
let handle = tokio::spawn(async move {
if let Err(e) = worker.run(rx).await {
tracing::error!("Worker for job {} failed: {}", job_id, e);
}
});
// Start the worker
let _ = tx.send(WorkerMessage::Start).await;
// Insert while still holding the write lock
jobs.insert(job_id, ScheduledJob { handle, tx });
}
// Cleanup task for this job to avoid capacity leaks
let jobs = Arc::clone(&self.jobs);
+21 -21
View File
@@ -8,8 +8,8 @@ use chrono::{DateTime, Utc};
use uuid::Uuid;
use crate::context::{ContextManager, JobState};
use crate::db::Database;
use crate::error::RepairError;
use crate::history::Store;
use crate::tools::{BuildRequirement, Language, SoftwareBuilder, SoftwareType, ToolRegistry};
/// A job that has been detected as stuck.
@@ -69,7 +69,7 @@ pub struct DefaultSelfRepair {
#[allow(dead_code)] // Will be used for time-based stuck detection
stuck_threshold: Duration,
max_repair_attempts: u32,
store: Option<Arc<Store>>,
store: Option<Arc<dyn Database>>,
builder: Option<Arc<dyn SoftwareBuilder>>,
#[allow(dead_code)] // Will be used for tool hot-reload after repair
tools: Option<Arc<ToolRegistry>>,
@@ -94,7 +94,7 @@ impl DefaultSelfRepair {
/// Add a Store for tool failure tracking.
#[allow(dead_code)] // Public API for configuring repair with persistence
pub fn with_store(mut self, store: Arc<Store>) -> Self {
pub fn with_store(mut self, store: Arc<dyn Database>) -> Self {
self.store = Some(store);
self
}
@@ -119,25 +119,25 @@ impl SelfRepair for DefaultSelfRepair {
let mut stuck_jobs = Vec::new();
for job_id in stuck_ids {
if let Ok(ctx) = self.context_manager.get_context(job_id).await {
if ctx.state == JobState::Stuck {
let stuck_duration = ctx
.started_at
.map(|start| {
let now = Utc::now();
let duration = now.signed_duration_since(start);
Duration::from_secs(duration.num_seconds().max(0) as u64)
})
.unwrap_or_default();
if let Ok(ctx) = self.context_manager.get_context(job_id).await
&& ctx.state == JobState::Stuck
{
let stuck_duration = ctx
.started_at
.map(|start| {
let now = Utc::now();
let duration = now.signed_duration_since(start);
Duration::from_secs(duration.num_seconds().max(0) as u64)
})
.unwrap_or_default();
stuck_jobs.push(StuckJob {
job_id,
last_activity: ctx.started_at.unwrap_or(ctx.created_at),
stuck_duration,
last_error: None,
repair_attempts: ctx.repair_attempts,
});
}
stuck_jobs.push(StuckJob {
job_id,
last_activity: ctx.started_at.unwrap_or(ctx.created_at),
stuck_duration,
last_error: None,
repair_attempts: ctx.repair_attempts,
});
}
}
+5 -5
View File
@@ -346,11 +346,11 @@ impl Thread {
let mut turn = Turn::new(turn_number, &msg.content);
// Check if next is assistant response
if let Some(next) = iter.peek() {
if next.role == crate::llm::Role::Assistant {
let response = iter.next().expect("peeked");
turn.complete(&response.content);
}
if let Some(next) = iter.peek()
&& next.role == crate::llm::Role::Assistant
{
let response = iter.next().expect("peeked");
turn.complete(&response.content);
}
self.turns.push(turn);
+4 -4
View File
@@ -199,10 +199,10 @@ impl SessionManager {
{
let sessions = self.sessions.read().await;
for user_id in &stale_users {
if let Some(session) = sessions.get(user_id) {
if let Ok(sess) = session.try_lock() {
stale_thread_ids.extend(sess.threads.keys());
}
if let Some(session) = sessions.get(user_id)
&& let Ok(sess) = session.try_lock()
{
stale_thread_ids.extend(sess.threads.keys());
}
}
}
+13 -14
View File
@@ -93,27 +93,26 @@ impl SubmissionParser {
// /thread <uuid> - switch thread
if let Some(rest) = lower.strip_prefix("/thread ") {
let rest = rest.trim();
if rest != "new" {
if let Ok(id) = Uuid::parse_str(rest) {
return Submission::SwitchThread { thread_id: id };
}
if rest != "new"
&& let Ok(id) = Uuid::parse_str(rest)
{
return Submission::SwitchThread { thread_id: id };
}
}
// /resume <uuid> - resume from checkpoint
if let Some(rest) = lower.strip_prefix("/resume ") {
if let Ok(id) = Uuid::parse_str(rest.trim()) {
return Submission::Resume { checkpoint_id: id };
}
if let Some(rest) = lower.strip_prefix("/resume ")
&& let Ok(id) = Uuid::parse_str(rest.trim())
{
return Submission::Resume { checkpoint_id: id };
}
// Try structured JSON approval (from web gateway's /api/chat/approval endpoint)
if trimmed.starts_with('{') {
if let Ok(submission) = serde_json::from_str::<Submission>(trimmed) {
if matches!(submission, Submission::ExecApproval { .. }) {
return submission;
}
}
if trimmed.starts_with('{')
&& let Ok(submission) = serde_json::from_str::<Submission>(trimmed)
&& matches!(submission, Submission::ExecApproval { .. })
{
return submission;
}
// Approval responses (simple yes/no/always for pending approvals)
+110 -31
View File
@@ -10,8 +10,8 @@ use uuid::Uuid;
use crate::agent::scheduler::WorkerMessage;
use crate::agent::task::TaskOutput;
use crate::context::{ContextManager, JobState};
use crate::db::Database;
use crate::error::Error;
use crate::history::Store;
use crate::llm::{
ActionPlan, ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult, ToolSelection,
};
@@ -28,7 +28,7 @@ pub struct WorkerDeps {
pub llm: Arc<dyn LlmProvider>,
pub safety: Arc<SafetyLayer>,
pub tools: Arc<ToolRegistry>,
pub store: Option<Arc<Store>>,
pub store: Option<Arc<dyn Database>>,
pub timeout: Duration,
pub use_planning: bool,
}
@@ -67,7 +67,7 @@ impl Worker {
&self.deps.tools
}
fn store(&self) -> Option<&Arc<Store>> {
fn store(&self) -> Option<&Arc<dyn Database>> {
self.deps.store.as_ref()
}
@@ -227,11 +227,11 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
}
// Check for cancellation
if let Ok(ctx) = self.context_manager().get_context(self.job_id).await {
if ctx.state == JobState::Cancelled {
tracing::info!("Worker for job {} detected cancellation", self.job_id);
return Ok(());
}
if let Ok(ctx) = self.context_manager().get_context(self.job_id).await
&& ctx.state == JobState::Cancelled
{
tracing::info!("Worker for job {} detected cancellation", self.job_id);
return Ok(());
}
iteration += 1;
@@ -248,16 +248,15 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
if selections.is_empty() {
// No tools from select_tools, ask LLM directly (may still return tool calls)
let respond_result = reasoning.respond_with_tools(reason_ctx).await?;
let respond_output = reasoning.respond_with_tools(reason_ctx).await?;
match respond_result {
match respond_output.result {
RespondResult::Text(response) => {
// Check for completion keywords
let response_lower = response.to_lowercase();
if response_lower.contains("complete")
|| response_lower.contains("finished")
|| response_lower.contains("done")
{
// Check for explicit completion phrases. Use word-boundary
// aware checks to avoid false positives like "incomplete",
// "not done", or "unfinished". Only the LLM's own response
// (not tool output) can trigger this.
if crate::util::llm_signals_completion(&response) {
self.mark_completed().await?;
return Ok(());
}
@@ -300,6 +299,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
parameters: tc.arguments.clone(),
reasoning: String::new(),
alternatives: vec![],
tool_call_id: tc.id.clone(),
};
self.process_tool_result(reason_ctx, &selection, result)
@@ -382,7 +382,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
tools: Arc<ToolRegistry>,
context_manager: Arc<ContextManager>,
safety: Arc<SafetyLayer>,
store: Option<Arc<Store>>,
store: Option<Arc<dyn Database>>,
job_id: Uuid,
tool_name: &str,
params: &serde_json::Value,
@@ -566,17 +566,14 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
);
reason_ctx.messages.push(ChatMessage::tool_result(
"tool_call_id",
&selection.tool_call_id,
&selection.tool_name,
wrapped,
));
// Check if job is complete
if output.contains("TASK_COMPLETE") || output.contains("JOB_DONE") {
self.mark_completed().await?;
return Ok(true);
}
// Tool output never drives job completion. A malicious tool could
// emit "TASK_COMPLETE" to force premature completion. Only the LLM's
// own structured response (in execution_loop) can mark a job done.
Ok(false)
}
Err(e) => {
@@ -601,7 +598,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
}
reason_ctx.messages.push(ChatMessage::tool_result(
"tool_call_id",
&selection.tool_call_id,
&selection.tool_name,
format!("Error: {}", e),
));
@@ -651,12 +648,15 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
.execute_tool(&action.tool_name, &action.parameters)
.await;
// Create a synthetic ToolSelection for process_tool_result
// Create a synthetic ToolSelection for process_tool_result.
// Plan actions don't originate from an LLM tool_call response so
// there is no real tool_call_id; generate a unique one.
let selection = ToolSelection {
tool_name: action.tool_name.clone(),
parameters: action.parameters.clone(),
reasoning: action.reasoning.clone(),
alternatives: vec![],
tool_call_id: format!("plan_{}_{}", self.job_id, i),
};
// Process the result
@@ -680,11 +680,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
let response = reasoning.respond(reason_ctx).await?;
reason_ctx.messages.push(ChatMessage::assistant(&response));
let response_lower = response.to_lowercase();
if response_lower.contains("complete")
|| response_lower.contains("finished")
|| response_lower.contains("done")
{
if crate::util::llm_signals_completion(&response) {
self.mark_completed().await?;
} else {
// Job not complete, could re-plan or fall back to direct selection
@@ -779,3 +775,86 @@ impl From<TaskOutput> for Result<String, Error> {
})
}
}
#[cfg(test)]
mod tests {
use crate::llm::ToolSelection;
use crate::util::llm_signals_completion;
#[test]
fn test_tool_selection_preserves_call_id() {
let selection = ToolSelection {
tool_name: "memory_search".to_string(),
parameters: serde_json::json!({"query": "test"}),
reasoning: "Need to search memory".to_string(),
alternatives: vec![],
tool_call_id: "call_abc123".to_string(),
};
assert_eq!(selection.tool_call_id, "call_abc123");
assert_ne!(
selection.tool_call_id, "tool_call_id",
"tool_call_id must not be the hardcoded placeholder string"
);
}
#[test]
fn test_completion_positive_signals() {
assert!(llm_signals_completion("The job is complete."));
assert!(llm_signals_completion(
"I have completed the task successfully."
));
assert!(llm_signals_completion("The task is done."));
assert!(llm_signals_completion("The task is finished."));
assert!(llm_signals_completion(
"All steps are complete and verified."
));
assert!(llm_signals_completion(
"I've done all the work. The work is done."
));
assert!(llm_signals_completion(
"Successfully completed the migration."
));
}
#[test]
fn test_completion_negative_signals_block_false_positives() {
// These contain completion keywords but also negation, should NOT trigger.
assert!(!llm_signals_completion("The task is not complete yet."));
assert!(!llm_signals_completion("This is not done."));
assert!(!llm_signals_completion("The work is incomplete."));
assert!(!llm_signals_completion(
"The migration is not yet finished."
));
assert!(!llm_signals_completion("The job isn't done yet."));
assert!(!llm_signals_completion("This remains unfinished."));
}
#[test]
fn test_completion_does_not_match_bare_substrings() {
// Bare words embedded in other text should NOT trigger completion.
assert!(!llm_signals_completion(
"I need to complete more work first."
));
assert!(!llm_signals_completion(
"Let me finish the remaining steps."
));
assert!(!llm_signals_completion(
"I'm done analyzing, now let me fix it."
));
assert!(!llm_signals_completion(
"I completed step 1 but step 2 remains."
));
}
#[test]
fn test_completion_tool_output_injection() {
// A malicious tool output echoed by the LLM should not trigger
// completion unless it forms a genuine completion phrase.
assert!(!llm_signals_completion("TASK_COMPLETE"));
assert!(!llm_signals_completion("JOB_DONE"));
assert!(!llm_signals_completion(
"The tool returned: TASK_COMPLETE signal"
));
}
}
+306 -167
View File
@@ -1,147 +1,145 @@
//! Bootstrap configuration for IronClaw.
//! Bootstrap helpers 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.
//! The only setting that truly needs disk persistence before the database is
//! available is `DATABASE_URL` (chicken-and-egg: can't connect to DB without
//! it). Everything else is auto-detected or read from env vars.
//!
//! File: `~/.ironclaw/bootstrap.json`
//! File: `~/.ironclaw/.env` (standard dotenvy format)
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,
/// Path to the IronClaw-specific `.env` file: `~/.ironclaw/.env`.
pub fn ironclaw_env_path() -> PathBuf {
dirs::home_dir()
.unwrap_or_else(|| PathBuf::from("."))
.join(".ironclaw")
.join(".env")
}
impl Default for BootstrapConfig {
fn default() -> Self {
Self {
database_url: None,
database_pool_size: None,
secrets_master_key_source: KeySource::None,
onboard_completed: false,
}
/// Load env vars from `~/.ironclaw/.env` (in addition to the standard `.env`).
///
/// Call this **after** `dotenvy::dotenv()` so that the standard `./.env`
/// takes priority over `~/.ironclaw/.env`. dotenvy never overwrites
/// existing env vars, so the effective priority is:
///
/// explicit env vars > `./.env` > `~/.ironclaw/.env`
///
/// If `~/.ironclaw/.env` doesn't exist but the legacy `bootstrap.json` does,
/// extracts `DATABASE_URL` from it and writes the `.env` file (one-time
/// upgrade from the old config format).
pub fn load_ironclaw_env() {
let path = ironclaw_env_path();
if !path.exists() {
// One-time upgrade: extract DATABASE_URL from legacy bootstrap.json
migrate_bootstrap_json_to_env(&path);
}
if path.exists() {
let _ = dotenvy::from_path(&path);
}
}
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")
/// If `bootstrap.json` exists, pull `database_url` out of it and write `.env`.
fn migrate_bootstrap_json_to_env(env_path: &std::path::Path) {
let ironclaw_dir = env_path
.parent()
.unwrap_or_else(|| std::path::Path::new("."));
let bootstrap_path = ironclaw_dir.join("bootstrap.json");
if !bootstrap_path.exists() {
return;
}
/// 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")
}
let content = match std::fs::read_to_string(&bootstrap_path) {
Ok(c) => c,
Err(_) => return,
};
/// 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);
// Minimal parse: just grab database_url from the JSON
let parsed: serde_json::Value = match serde_json::from_str(&content) {
Ok(v) => v,
Err(_) => return,
};
if let Some(url) = parsed.get("database_url").and_then(|v| v.as_str()) {
if let Some(parent) = env_path.parent()
&& let Err(e) = std::fs::create_dir_all(parent)
{
eprintln!("Warning: failed to create {}: {}", parent.display(), e);
return;
}
// 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);
if let Err(e) = std::fs::write(env_path, format!("DATABASE_URL=\"{}\"\n", url)) {
eprintln!("Warning: failed to migrate bootstrap.json to .env: {}", e);
return;
}
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)
rename_to_migrated(&bootstrap_path);
eprintln!(
"Migrated DATABASE_URL from bootstrap.json to {}",
env_path.display()
);
}
}
/// One-time migration from disk config files to the database settings table.
/// Write database bootstrap vars to `~/.ironclaw/.env`.
///
/// On first boot after upgrade, checks if:
/// 1. `~/.ironclaw/settings.json` exists
/// 2. The DB settings table is empty for this user
/// These settings form the chicken-and-egg layer: they must be available
/// from the filesystem (env vars) BEFORE any database connection, because
/// they determine which database to connect to. Everything else is stored
/// in the database itself.
///
/// If both conditions hold, migrates settings, MCP servers, and session data
/// to the database, writes `bootstrap.json`, and renames old files to `.migrated`.
/// Creates the parent directory if it doesn't exist.
/// Values are double-quoted so that `#` (common in URL-encoded passwords)
/// and other shell-special characters are preserved by dotenvy.
pub fn save_bootstrap_env(vars: &[(&str, &str)]) -> std::io::Result<()> {
let path = ironclaw_env_path();
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let mut content = String::new();
for (key, value) in vars {
content.push_str(&format!("{}=\"{}\"\n", key, value));
}
std::fs::write(&path, content)
}
/// Write `DATABASE_URL` to `~/.ironclaw/.env`.
///
/// Convenience wrapper around `save_bootstrap_env` for single-value migration
/// paths. Prefer `save_bootstrap_env` for new code.
pub fn save_database_url(url: &str) -> std::io::Result<()> {
save_bootstrap_env(&[("DATABASE_URL", url)])
}
/// One-time migration of legacy `~/.ironclaw/settings.json` into the database.
///
/// Only runs when a `settings.json` exists on disk AND the DB has no settings
/// yet. After the wizard writes directly to the DB, this path is only hit by
/// users upgrading from the old disk-only configuration.
///
/// After syncing, renames `settings.json` to `.migrated` so it won't trigger again.
pub async fn migrate_disk_to_db(
store: &crate::history::Store,
store: &dyn crate::db::Database,
user_id: &str,
) -> Result<(), MigrationError> {
let legacy_settings_path = BootstrapConfig::legacy_settings_path();
let ironclaw_dir = dirs::home_dir()
.unwrap_or_else(|| PathBuf::from("."))
.join(".ironclaw");
let legacy_settings_path = ironclaw_dir.join("settings.json");
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
// If DB already has settings, this is not a first boot, the wizard already
// wrote directly to the DB. Just clean up the stale file.
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
);
tracing::info!("DB already has settings, renaming stale settings.json");
rename_to_migrated(&legacy_settings_path);
return Ok(());
}
@@ -160,22 +158,14 @@ pub async fn migrate_disk_to_db(
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");
// 2. Write DATABASE_URL to ~/.ironclaw/.env
if let Some(ref url) = settings.database_url {
save_database_url(url)
.map_err(|e| MigrationError::Io(format!("Failed to write .env: {}", e)))?;
tracing::info!("Wrote DATABASE_URL to {}", ironclaw_env_path().display());
}
// 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) {
@@ -211,7 +201,7 @@ pub async fn migrate_disk_to_db(
Ok(content) => match serde_json::from_str::<serde_json::Value>(&content) {
Ok(value) => {
store
.set_setting(user_id, "nearai.session", &value)
.set_setting(user_id, "nearai.session_token", &value)
.await
.map_err(|e| {
MigrationError::Database(format!(
@@ -236,12 +226,19 @@ pub async fn migrate_disk_to_db(
// 5. Rename settings.json to .migrated (don't delete, safety net)
rename_to_migrated(&legacy_settings_path);
// 6. Clean up old bootstrap.json if it exists (superseded by .env)
let old_bootstrap = ironclaw_dir.join("bootstrap.json");
if old_bootstrap.exists() {
rename_to_migrated(&old_bootstrap);
tracing::info!("Renamed old bootstrap.json to .migrated");
}
tracing::info!("Disk-to-DB migration complete");
Ok(())
}
/// Rename a file to `<name>.migrated` as a safety net.
fn rename_to_migrated(path: &PathBuf) {
fn rename_to_migrated(path: &std::path::Path) {
let mut migrated = path.as_os_str().to_owned();
migrated.push(".migrated");
if let Err(e) = std::fs::rename(path, &migrated) {
@@ -264,62 +261,204 @@ mod tests {
use tempfile::tempdir;
#[test]
fn test_bootstrap_save_load() {
fn test_save_and_load_database_url() {
let dir = tempdir().unwrap();
let path = dir.path().join("bootstrap.json");
let env_path = dir.path().join(".env");
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,
};
// Write in the quoted format that save_database_url uses
let url = "postgres://localhost:5432/ironclaw_test";
std::fs::write(&env_path, format!("DATABASE_URL=\"{}\"\n", url)).unwrap();
config.save_to(&path).unwrap();
let loaded = BootstrapConfig::load_from(&path);
// Verify the content is a valid dotenv line (quoted)
let content = std::fs::read_to_string(&env_path).unwrap();
assert_eq!(
loaded.database_url,
Some("postgres://localhost/test".to_string())
content,
"DATABASE_URL=\"postgres://localhost:5432/ironclaw_test\"\n"
);
assert_eq!(loaded.database_pool_size, Some(5));
assert_eq!(loaded.secrets_master_key_source, KeySource::Keychain);
assert!(loaded.onboard_completed);
// Verify dotenvy can parse it (strips quotes automatically)
let parsed: Vec<(String, String)> = dotenvy::from_path_iter(&env_path)
.unwrap()
.filter_map(|r| r.ok())
.collect();
assert_eq!(parsed.len(), 1);
assert_eq!(parsed[0].0, "DATABASE_URL");
assert_eq!(parsed[0].1, url);
}
#[test]
fn test_bootstrap_from_legacy_settings() {
fn test_save_database_url_with_hash_in_password() {
let dir = tempdir().unwrap();
let path = dir.path().join("settings.json");
let env_path = dir.path().join(".env");
// Write a legacy settings.json with many extra fields
let legacy = serde_json::json!({
"database_url": "postgres://localhost/ironclaw",
"database_pool_size": 10,
// URLs with # in the password are common (URL-encoded special chars).
// Without quoting, dotenvy treats # as a comment delimiter.
let url = "postgres://user:p%23ss@localhost:5432/ironclaw";
std::fs::write(&env_path, format!("DATABASE_URL=\"{}\"\n", url)).unwrap();
let parsed: Vec<(String, String)> = dotenvy::from_path_iter(&env_path)
.unwrap()
.filter_map(|r| r.ok())
.collect();
assert_eq!(parsed.len(), 1);
assert_eq!(parsed[0].0, "DATABASE_URL");
assert_eq!(parsed[0].1, url);
}
#[test]
fn test_save_database_url_creates_parent_dirs() {
let dir = tempdir().unwrap();
let nested = dir.path().join("deep").join("nested");
let env_path = nested.join(".env");
// Parent doesn't exist yet
assert!(!nested.exists());
// The global function uses a fixed path, so we test the logic directly
std::fs::create_dir_all(&nested).unwrap();
std::fs::write(&env_path, "DATABASE_URL=postgres://test\n").unwrap();
assert!(env_path.exists());
let content = std::fs::read_to_string(&env_path).unwrap();
assert!(content.contains("DATABASE_URL=postgres://test"));
}
#[test]
fn test_ironclaw_env_path() {
let path = ironclaw_env_path();
assert!(path.ends_with(".ironclaw/.env"));
}
#[test]
fn test_migrate_bootstrap_json_to_env() {
let dir = tempdir().unwrap();
let env_path = dir.path().join(".env");
let bootstrap_path = dir.path().join("bootstrap.json");
// Write a legacy bootstrap.json
let bootstrap_json = serde_json::json!({
"database_url": "postgres://localhost/ironclaw_upgrade",
"database_pool_size": 5,
"secrets_master_key_source": "keychain",
"onboard_completed": true,
"selected_model": "claude-3-5-sonnet",
"agent": { "name": "testbot", "max_parallel_jobs": 3 },
"heartbeat": { "enabled": true }
"onboard_completed": true
});
std::fs::write(&path, serde_json::to_string_pretty(&legacy).unwrap()).unwrap();
std::fs::write(
&bootstrap_path,
serde_json::to_string_pretty(&bootstrap_json).unwrap(),
)
.unwrap();
let config = BootstrapConfig::load_from_legacy(&path);
assert!(!env_path.exists());
assert!(bootstrap_path.exists());
// Run the migration
migrate_bootstrap_json_to_env(&env_path);
// .env should now exist with DATABASE_URL
assert!(env_path.exists());
let content = std::fs::read_to_string(&env_path).unwrap();
assert_eq!(
config.database_url,
Some("postgres://localhost/ironclaw".to_string())
content,
"DATABASE_URL=\"postgres://localhost/ironclaw_upgrade\"\n"
);
assert_eq!(config.database_pool_size, Some(10));
assert_eq!(config.secrets_master_key_source, KeySource::Keychain);
assert!(config.onboard_completed);
// bootstrap.json should be renamed to .migrated
assert!(!bootstrap_path.exists());
assert!(dir.path().join("bootstrap.json.migrated").exists());
}
#[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);
fn test_migrate_bootstrap_json_no_database_url() {
let dir = tempdir().unwrap();
let env_path = dir.path().join(".env");
let bootstrap_path = dir.path().join("bootstrap.json");
// bootstrap.json with no database_url
let bootstrap_json = serde_json::json!({
"onboard_completed": false
});
std::fs::write(
&bootstrap_path,
serde_json::to_string_pretty(&bootstrap_json).unwrap(),
)
.unwrap();
migrate_bootstrap_json_to_env(&env_path);
// .env should NOT be created
assert!(!env_path.exists());
// bootstrap.json should remain (no migration happened)
assert!(bootstrap_path.exists());
}
#[test]
fn test_migrate_bootstrap_json_missing() {
let dir = tempdir().unwrap();
let env_path = dir.path().join(".env");
// No bootstrap.json at all
migrate_bootstrap_json_to_env(&env_path);
// Nothing should happen
assert!(!env_path.exists());
}
#[test]
fn test_save_bootstrap_env_multiple_vars() {
let dir = tempdir().unwrap();
let env_path = dir.path().join("nested").join(".env");
std::fs::create_dir_all(env_path.parent().unwrap()).unwrap();
let vars = [
("DATABASE_BACKEND", "libsql"),
("LIBSQL_PATH", "/home/user/.ironclaw/ironclaw.db"),
];
// Write manually to the temp path (save_bootstrap_env uses the global path)
let mut content = String::new();
for (key, value) in &vars {
content.push_str(&format!("{}=\"{}\"\n", key, value));
}
std::fs::write(&env_path, &content).unwrap();
// Verify dotenvy can parse all entries
let parsed: Vec<(String, String)> = dotenvy::from_path_iter(&env_path)
.unwrap()
.filter_map(|r| r.ok())
.collect();
assert_eq!(parsed.len(), 2);
assert_eq!(
parsed[0],
("DATABASE_BACKEND".to_string(), "libsql".to_string())
);
assert_eq!(
parsed[1],
(
"LIBSQL_PATH".to_string(),
"/home/user/.ironclaw/ironclaw.db".to_string()
)
);
}
#[test]
fn test_save_bootstrap_env_overwrites_previous() {
let dir = tempdir().unwrap();
let env_path = dir.path().join(".env");
// Write initial content
std::fs::write(&env_path, "DATABASE_URL=\"postgres://old\"\n").unwrap();
// Overwrite with new vars (simulating save_bootstrap_env behavior)
let content = "DATABASE_BACKEND=\"libsql\"\nLIBSQL_PATH=\"/new/path.db\"\n";
std::fs::write(&env_path, content).unwrap();
let parsed: Vec<(String, String)> = dotenvy::from_path_iter(&env_path)
.unwrap()
.filter_map(|r| r.ok())
.collect();
// Old DATABASE_URL should be gone
assert_eq!(parsed.len(), 2);
assert!(parsed.iter().all(|(k, _)| k != "DATABASE_URL"));
}
}
+17 -7
View File
@@ -33,9 +33,16 @@ use termimad::MadSkin;
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use crate::agent::truncate_for_preview;
use crate::channels::{Channel, IncomingMessage, MessageStream, OutgoingResponse, StatusUpdate};
use crate::error::ChannelError;
/// Max characters for tool result previews in the terminal.
const CLI_TOOL_RESULT_MAX: usize = 200;
/// Max characters for thinking/status messages in the terminal.
const CLI_STATUS_MAX: usize = 200;
/// Slash commands available in the REPL.
const SLASH_COMMANDS: &[&str] = &[
"/help",
@@ -261,7 +268,7 @@ impl Channel for ReplChannel {
std::thread::spawn(move || {
// Single message mode: send it and return
if let Some(msg) = single_message {
let incoming = IncomingMessage::new("repl", "user", &msg);
let incoming = IncomingMessage::new("repl", "default", &msg);
let _ = tx.blocking_send(incoming);
return;
}
@@ -329,21 +336,21 @@ impl Channel for ReplChannel {
_ => {}
}
let msg = IncomingMessage::new("repl", "user", line);
let msg = IncomingMessage::new("repl", "default", line);
if tx.blocking_send(msg).is_err() {
break;
}
}
Err(ReadlineError::Interrupted) => {
// Ctrl+C: send /interrupt
let msg = IncomingMessage::new("repl", "user", "/interrupt");
let msg = IncomingMessage::new("repl", "default", "/interrupt");
if tx.blocking_send(msg).is_err() {
break;
}
}
Err(ReadlineError::Eof) => {
// Ctrl+D: send /quit so the agent loop runs graceful shutdown
let msg = IncomingMessage::new("repl", "user", "/quit");
let msg = IncomingMessage::new("repl", "default", "/quit");
let _ = tx.blocking_send(msg);
break;
}
@@ -400,7 +407,8 @@ impl Channel for ReplChannel {
match status {
StatusUpdate::Thinking(msg) => {
eprintln!(" \x1b[90m\u{25CB} {msg}\x1b[0m");
let display = truncate_for_preview(&msg, CLI_STATUS_MAX);
eprintln!(" \x1b[90m\u{25CB} {display}\x1b[0m");
}
StatusUpdate::ToolStarted { name } => {
eprintln!(" \x1b[33m\u{25CB} {name}\x1b[0m");
@@ -413,7 +421,8 @@ impl Channel for ReplChannel {
}
}
StatusUpdate::ToolResult { name: _, preview } => {
eprintln!(" \x1b[90m{preview}\x1b[0m");
let display = truncate_for_preview(&preview, CLI_TOOL_RESULT_MAX);
eprintln!(" \x1b[90m{display}\x1b[0m");
}
StatusUpdate::StreamChunk(chunk) => {
// Print separator on the false-to-true transition
@@ -438,7 +447,8 @@ impl Channel for ReplChannel {
}
StatusUpdate::Status(msg) => {
if debug || msg.contains("approval") || msg.contains("Approval") {
eprintln!(" \x1b[90m{msg}\x1b[0m");
let display = truncate_for_preview(&msg, CLI_STATUS_MAX);
eprintln!(" \x1b[90m{display}\x1b[0m");
}
}
StatusUpdate::ApprovalNeeded {
+135 -38
View File
@@ -76,6 +76,9 @@ struct ChannelStoreData {
credentials: HashMap<String, String>,
/// Pairing store for DM pairing (guest access control).
pairing_store: Arc<PairingStore>,
/// Dedicated tokio runtime for HTTP requests, lazily initialized.
/// Reused across multiple `http_request` calls within one execution.
http_runtime: Option<tokio::runtime::Runtime>,
}
impl ChannelStoreData {
@@ -96,6 +99,7 @@ impl ChannelStoreData {
table: ResourceTable::new(),
credentials,
pairing_store,
http_runtime: None,
}
}
@@ -134,13 +138,13 @@ impl ChannelStoreData {
if result.contains('{') && result.contains('}') {
// Only warn if it looks like an unresolved placeholder (not JSON braces)
let brace_pattern = regex::Regex::new(r"\{[A-Z_]+\}").ok();
if let Some(re) = brace_pattern {
if re.is_match(&result) {
tracing::warn!(
context = %context,
"String may contain unresolved credential placeholders"
);
}
if let Some(re) = brace_pattern
&& re.is_match(&result)
{
tracing::warn!(
context = %context,
"String may contain unresolved credential placeholders"
);
}
}
@@ -273,10 +277,35 @@ impl near::agent::channel_host::Host for ChannelStoreData {
.scan_http_request(&url, &header_vec, body.as_deref())
.map_err(|e| format!("Potential secret leak blocked: {}", e))?;
// Make the HTTP request using blocking I/O
// We're already in a spawn_blocking context, so we can use block_on
let result = tokio::runtime::Handle::current().block_on(async {
let client = reqwest::Client::new();
// Get the max response size from capabilities (default 10MB).
let max_response_bytes = self
.host_state
.capabilities()
.tool_capabilities
.http
.as_ref()
.map(|h| h.max_response_bytes)
.unwrap_or(10 * 1024 * 1024);
// Make the HTTP request using a dedicated single-threaded runtime.
// We're inside spawn_blocking, so we can't rely on the main runtime's
// I/O driver (it may be busy with WASM compilation or other startup work).
// A dedicated runtime gives us our own I/O driver and avoids contention.
// The runtime is lazily created and reused across calls within one execution.
if self.http_runtime.is_none() {
self.http_runtime = Some(
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.map_err(|e| format!("Failed to create HTTP runtime: {e}"))?,
);
}
let rt = self.http_runtime.as_ref().expect("just initialized");
let result = rt.block_on(async {
let client = reqwest::Client::builder()
.connect_timeout(std::time::Duration::from_secs(10))
.build()
.map_err(|e| format!("Failed to build HTTP client: {e}"))?;
let mut request = match method.to_uppercase().as_str() {
"GET" => client.get(&url),
@@ -298,9 +327,9 @@ impl near::agent::channel_host::Host for ChannelStoreData {
request = request.body(body_bytes);
}
// 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);
// Send request with caller-specified timeout (default 30s, max 5min).
let timeout_ms = timeout_ms.unwrap_or(30_000).min(300_000) as u64;
let timeout = std::time::Duration::from_millis(timeout_ms);
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
@@ -325,11 +354,29 @@ impl near::agent::channel_host::Host for ChannelStoreData {
})
.collect();
let headers_json = serde_json::to_string(&response_headers).unwrap_or_default();
// Enforce max response body size to prevent memory exhaustion.
let max_response = max_response_bytes;
if let Some(cl) = response.content_length()
&& cl as usize > max_response
{
return Err(format!(
"Response body too large: {} bytes exceeds limit of {} bytes",
cl, max_response
));
}
let body = response
.bytes()
.await
.map_err(|e| format!("Failed to read response body: {}", e))?
.to_vec();
.map_err(|e| format!("Failed to read response body: {}", e))?;
if body.len() > max_response {
return Err(format!(
"Response body too large: {} bytes exceeds limit of {} bytes",
body.len(),
max_response
));
}
let body = body.to_vec();
tracing::info!(
status = status,
@@ -767,7 +814,21 @@ impl WasmChannel {
.await;
match result {
Ok(Ok((config, _host_state))) => {
Ok(Ok((config, mut host_state))) => {
// Surface WASM guest logs (errors/warnings from webhook setup, etc.)
for entry in host_state.take_logs() {
match entry.level {
crate::tools::wasm::LogLevel::Error => {
tracing::error!(channel = %self.name, "{}", entry.message);
}
crate::tools::wasm::LogLevel::Warn => {
tracing::warn!(channel = %self.name, "{}", entry.message);
}
_ => {
tracing::debug!(channel = %self.name, "{}", entry.message);
}
}
}
tracing::info!(
channel = %self.name,
display_name = %config.display_name,
@@ -1467,8 +1528,8 @@ impl WasmChannel {
match result {
Ok(emitted_messages) => {
// Process any emitted messages
if !emitted_messages.is_empty() {
if let Err(e) = Self::dispatch_emitted_messages(
if !emitted_messages.is_empty()
&& let Err(e) = Self::dispatch_emitted_messages(
&channel_name,
emitted_messages,
&message_tx,
@@ -1480,7 +1541,6 @@ impl WasmChannel {
"Failed to dispatch emitted messages from poll"
);
}
}
}
Err(e) => {
tracing::warn!(
@@ -1710,22 +1770,22 @@ impl Channel for WasmChannel {
*self.endpoints.write().await = endpoints;
// Start polling if configured
if let Some(poll_config) = &config.poll {
if poll_config.enabled {
let interval = self
.capabilities
.validate_poll_interval(poll_config.interval_ms)
.map_err(|e| ChannelError::StartupFailed {
name: self.name.clone(),
reason: e,
})?;
if let Some(poll_config) = &config.poll
&& poll_config.enabled
{
let interval = self
.capabilities
.validate_poll_interval(poll_config.interval_ms)
.map_err(|e| ChannelError::StartupFailed {
name: self.name.clone(),
reason: e,
})?;
// Create shutdown channel for polling and store the sender to keep it alive
let (poll_shutdown_tx, poll_shutdown_rx) = oneshot::channel();
*self.poll_shutdown_tx.write().await = Some(poll_shutdown_tx);
// Create shutdown channel for polling and store the sender to keep it alive
let (poll_shutdown_tx, poll_shutdown_rx) = oneshot::channel();
*self.poll_shutdown_tx.write().await = Some(poll_shutdown_tx);
self.start_polling(Duration::from_millis(interval as u64), poll_shutdown_rx);
}
self.start_polling(Duration::from_millis(interval as u64), poll_shutdown_rx);
}
tracing::info!(
@@ -2588,15 +2648,52 @@ mod tests {
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.
/// Verify that WASM HTTP host functions work using a dedicated
/// current-thread runtime inside spawn_blocking.
#[tokio::test]
async fn test_block_on_inside_spawn_blocking_does_not_deadlock() {
async fn test_dedicated_runtime_inside_spawn_blocking() {
let result = tokio::task::spawn_blocking(|| {
tokio::runtime::Handle::current().block_on(async { 42 })
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("failed to build runtime");
rt.block_on(async { 42 })
})
.await
.expect("spawn_blocking panicked");
assert_eq!(result, 42);
}
/// Verify a real HTTP request works using the dedicated-runtime pattern.
/// This catches DNS, TLS, and I/O driver issues that trivial tests miss.
#[tokio::test]
#[ignore] // requires network
async fn test_dedicated_runtime_real_http() {
let result = tokio::task::spawn_blocking(|| {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("failed to build runtime");
rt.block_on(async {
let client = reqwest::Client::builder()
.connect_timeout(std::time::Duration::from_secs(10))
.build()
.expect("failed to build client");
let resp = client
.get("https://api.telegram.org/bot000/getMe")
.timeout(std::time::Duration::from_secs(10))
.send()
.await;
match resp {
Ok(r) => r.status().as_u16(),
Err(e) if e.is_timeout() => panic!("request timed out: {e}"),
Err(e) => panic!("unexpected error: {e}"),
}
})
})
.await
.expect("spawn_blocking panicked");
// 404 because "000" is not a valid bot token
assert_eq!(result, 404);
}
}
+13 -14
View File
@@ -6,6 +6,7 @@ use axum::{
middleware::Next,
response::{IntoResponse, Response},
};
use subtle::ConstantTimeEq;
/// Shared auth state injected via axum middleware state.
#[derive(Clone)]
@@ -23,24 +24,22 @@ pub async fn auth_middleware(
request: Request,
next: Next,
) -> Response {
// Try Authorization header first
if let Some(auth_header) = headers.get("authorization") {
if let Ok(value) = auth_header.to_str() {
if let Some(token) = value.strip_prefix("Bearer ") {
if token == auth.token {
return next.run(request).await;
}
}
}
// Try Authorization header first (constant-time comparison)
if let Some(auth_header) = headers.get("authorization")
&& let Ok(value) = auth_header.to_str()
&& let Some(token) = value.strip_prefix("Bearer ")
&& bool::from(token.as_bytes().ct_eq(auth.token.as_bytes()))
{
return next.run(request).await;
}
// Fall back to query parameter (for SSE EventSource)
// Fall back to query parameter for SSE EventSource (constant-time comparison)
if let Some(query) = request.uri().query() {
for pair in query.split('&') {
if let Some(token) = pair.strip_prefix("token=") {
if token == auth.token {
return next.run(request).await;
}
if let Some(token) = pair.strip_prefix("token=")
&& bool::from(token.as_bytes().ct_eq(auth.token.as_bytes()))
{
return next.run(request).await;
}
}
}
+43 -1
View File
@@ -24,6 +24,8 @@ use tokio::sync::broadcast;
use tracing::field::{Field, Visit};
use tracing_subscriber::Layer;
use crate::safety::LeakDetector;
/// Maximum number of recent log entries kept for late-joining SSE subscribers.
const HISTORY_CAP: usize = 500;
@@ -46,6 +48,8 @@ pub struct LogEntry {
pub struct LogBroadcaster {
tx: broadcast::Sender<LogEntry>,
recent: Mutex<VecDeque<LogEntry>>,
/// Scrubs secrets from log messages before broadcasting to SSE clients.
leak_detector: LeakDetector,
}
impl LogBroadcaster {
@@ -54,10 +58,19 @@ impl LogBroadcaster {
Self {
tx,
recent: Mutex::new(VecDeque::with_capacity(HISTORY_CAP)),
leak_detector: LeakDetector::new(),
}
}
pub fn send(&self, entry: LogEntry) {
pub fn send(&self, mut entry: LogEntry) {
// Scrub secrets from the message before it reaches any subscriber.
// This is defense-in-depth: even if code elsewhere accidentally logs
// a secret, it won't be broadcast to SSE clients.
entry.message = self
.leak_detector
.scan_and_clean(&entry.message)
.unwrap_or_else(|_| "[log message redacted: contained blocked secret]".to_string());
// Stash in ring buffer (for late joiners)
if let Ok(mut buf) = self.recent.lock() {
if buf.len() >= HISTORY_CAP {
@@ -145,6 +158,9 @@ impl Visit for MessageVisitor {
///
/// Only forwards DEBUG and above. Attach to the tracing subscriber
/// alongside the existing fmt layer.
///
/// Log messages are scrubbed through `LeakDetector` in `LogBroadcaster::send()`
/// (the single funnel point for all log output, including late-joiner history).
pub struct WebLogLayer {
broadcaster: Arc<LogBroadcaster>,
}
@@ -178,6 +194,7 @@ impl<S: tracing::Subscriber> Layer<S> for WebLogLayer {
timestamp: chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
};
// LeakDetector scrubbing happens inside broadcaster.send()
self.broadcaster.send(entry);
}
}
@@ -313,4 +330,29 @@ mod tests {
let v = MessageVisitor::new();
assert_eq!(v.finish(), "");
}
#[test]
fn test_broadcaster_has_leak_detector() {
let broadcaster = LogBroadcaster::new();
// Verify the leak detector is initialized with default patterns
assert!(broadcaster.leak_detector.pattern_count() > 0);
}
#[test]
fn test_leak_detector_scrubs_api_key_in_log() {
let detector = crate::safety::LeakDetector::new();
let msg = "Connecting with token sk-proj-test1234567890abcdefghij";
let result = detector.scan_and_clean(msg);
// Should be blocked (OpenAI key pattern)
assert!(result.is_err());
}
#[test]
fn test_leak_detector_passes_clean_log() {
let detector = crate::safety::LeakDetector::new();
let msg = "Request completed status=200 url=https://api.example.com/data";
let result = detector.scan_and_clean(msg);
assert!(result.is_ok());
assert_eq!(result.unwrap(), msg);
}
}
+13 -2
View File
@@ -16,6 +16,7 @@
pub mod auth;
pub mod log_layer;
pub mod openai_compat;
pub mod server;
pub mod sse;
pub mod types;
@@ -31,9 +32,9 @@ use tokio_stream::wrappers::ReceiverStream;
use crate::agent::SessionManager;
use crate::channels::{Channel, IncomingMessage, MessageStream, OutgoingResponse, StatusUpdate};
use crate::config::GatewayConfig;
use crate::db::Database;
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;
@@ -81,6 +82,8 @@ impl GatewayChannel {
user_id: config.user_id.clone(),
shutdown_tx: tokio::sync::RwLock::new(None),
ws_tracker: Some(Arc::new(ws::WsConnectionTracker::new())),
llm_provider: None,
chat_rate_limiter: server::RateLimiter::new(30, 60),
});
Self {
@@ -106,6 +109,8 @@ impl GatewayChannel {
user_id: self.state.user_id.clone(),
shutdown_tx: tokio::sync::RwLock::new(None),
ws_tracker: self.state.ws_tracker.clone(),
llm_provider: self.state.llm_provider.clone(),
chat_rate_limiter: server::RateLimiter::new(30, 60),
};
mutate(&mut new_state);
self.state = Arc::new(new_state);
@@ -142,7 +147,7 @@ impl GatewayChannel {
}
/// Inject the database store for sandbox job persistence.
pub fn with_store(mut self, store: Arc<Store>) -> Self {
pub fn with_store(mut self, store: Arc<dyn Database>) -> Self {
self.rebuild_state(|s| s.store = Some(store));
self
}
@@ -169,6 +174,12 @@ impl GatewayChannel {
self
}
/// Inject the LLM provider for OpenAI-compatible API proxy.
pub fn with_llm_provider(mut self, llm: Arc<dyn crate::llm::LlmProvider>) -> Self {
self.rebuild_state(|s| s.llm_provider = Some(llm));
self
}
/// Get the auth token (for printing to console on startup).
pub fn auth_token(&self) -> &str {
&self.auth_token
File diff suppressed because it is too large Load Diff
+367 -153
View File
@@ -5,10 +5,11 @@
use std::convert::Infallible;
use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use axum::{
Json, Router,
extract::{Path, Query, State, WebSocketUpgrade},
extract::{DefaultBodyLimit, Path, Query, State, WebSocketUpgrade},
http::{StatusCode, header},
middleware,
response::{
@@ -20,6 +21,7 @@ use axum::{
use serde::Deserialize;
use tokio::sync::{mpsc, oneshot};
use tokio_stream::StreamExt;
use tower_http::cors::{AllowHeaders, CorsLayer};
use uuid::Uuid;
use crate::agent::SessionManager;
@@ -28,8 +30,8 @@ use crate::channels::web::auth::{AuthState, auth_middleware};
use crate::channels::web::log_layer::LogBroadcaster;
use crate::channels::web::sse::SseManager;
use crate::channels::web::types::*;
use crate::db::Database;
use crate::extensions::ExtensionManager;
use crate::history::Store;
use crate::orchestrator::job_manager::ContainerJobManager;
use crate::tools::ToolRegistry;
use crate::workspace::Workspace;
@@ -44,6 +46,69 @@ pub type PromptQueue = Arc<
>,
>;
/// Simple sliding-window rate limiter.
///
/// Tracks the number of requests in the current window. Resets when the window expires.
/// Not per-IP (since this is a single-user gateway with auth), but prevents flooding.
pub struct RateLimiter {
/// Requests remaining in the current window.
remaining: AtomicU64,
/// Epoch second when the current window started.
window_start: AtomicU64,
/// Maximum requests per window.
max_requests: u64,
/// Window duration in seconds.
window_secs: u64,
}
impl RateLimiter {
pub fn new(max_requests: u64, window_secs: u64) -> Self {
Self {
remaining: AtomicU64::new(max_requests),
window_start: AtomicU64::new(
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs(),
),
max_requests,
window_secs,
}
}
/// Try to consume one request. Returns `true` if allowed, `false` if rate limited.
pub fn check(&self) -> bool {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
let window = self.window_start.load(Ordering::Relaxed);
if now.saturating_sub(window) >= self.window_secs {
// Window expired, reset
self.window_start.store(now, Ordering::Relaxed);
self.remaining
.store(self.max_requests - 1, Ordering::Relaxed);
return true;
}
// Try to decrement remaining
loop {
let current = self.remaining.load(Ordering::Relaxed);
if current == 0 {
return false;
}
if self
.remaining
.compare_exchange_weak(current, current - 1, Ordering::Relaxed, Ordering::Relaxed)
.is_ok()
{
return true;
}
}
}
}
/// Shared state for all gateway handlers.
pub struct GatewayState {
/// Channel to send messages to the agent loop.
@@ -61,7 +126,7 @@ pub struct GatewayState {
/// Tool registry for listing registered tools.
pub tool_registry: Option<Arc<ToolRegistry>>,
/// Database store for sandbox job persistence.
pub store: Option<Arc<Store>>,
pub store: Option<Arc<dyn Database>>,
/// Container job manager for sandbox operations.
pub job_manager: Option<Arc<ContainerJobManager>>,
/// Prompt queue for Claude Code follow-up prompts.
@@ -72,6 +137,10 @@ pub struct GatewayState {
pub shutdown_tx: tokio::sync::RwLock<Option<oneshot::Sender<()>>>,
/// WebSocket connection tracker.
pub ws_tracker: Option<Arc<crate::channels::web::ws::WsConnectionTracker>>,
/// LLM provider for OpenAI-compatible API proxy.
pub llm_provider: Option<Arc<dyn crate::llm::LlmProvider>>,
/// Rate limiter for chat endpoints (30 messages per 60 seconds).
pub chat_rate_limiter: RateLimiter,
}
/// Start the gateway HTTP server.
@@ -168,7 +237,16 @@ pub async fn start_server(
)
// Gateway control plane
.route("/api/gateway/status", get(gateway_status_handler))
.route_layer(middleware::from_fn_with_state(auth_state, auth_middleware));
// OpenAI-compatible API
.route(
"/v1/chat/completions",
post(super::openai_compat::chat_completions_handler),
)
.route("/v1/models", get(super::openai_compat::models_handler))
.route_layer(middleware::from_fn_with_state(
auth_state.clone(),
auth_middleware,
));
// Static file routes (no auth, served from embedded strings)
let statics = Router::new()
@@ -176,19 +254,46 @@ pub async fn start_server(
.route("/style.css", get(css_handler))
.route("/app.js", get(js_handler));
// Project file serving (no auth, local browsing of sandbox outputs).
// The trailing-slash route serves index.html; the bare route redirects so
// relative paths in the HTML (e.g. href="style.css") resolve correctly.
// Project file serving (behind auth to prevent unauthorized file access).
let projects = Router::new()
.route("/projects/{project_id}", get(project_redirect_handler))
.route("/projects/{project_id}/", get(project_index_handler))
.route("/projects/{project_id}/{*path}", get(project_file_handler));
.route("/projects/{project_id}/{*path}", get(project_file_handler))
.route_layer(middleware::from_fn_with_state(
auth_state.clone(),
auth_middleware,
));
// CORS: restrict to same-origin by default. Only localhost/127.0.0.1
// origins are allowed, since the gateway is a local-first service.
let cors = CorsLayer::new()
.allow_origin([
format!("http://{}:{}", addr.ip(), addr.port())
.parse()
.expect("valid origin"),
format!("http://localhost:{}", addr.port())
.parse()
.expect("valid origin"),
])
.allow_methods([
axum::http::Method::GET,
axum::http::Method::POST,
axum::http::Method::PUT,
axum::http::Method::DELETE,
])
.allow_headers(AllowHeaders::list([
header::CONTENT_TYPE,
header::AUTHORIZATION,
]))
.allow_credentials(true);
let app = Router::new()
.merge(public)
.merge(statics)
.merge(projects)
.merge(protected)
.layer(cors)
.layer(DefaultBodyLimit::max(1024 * 1024)) // 1 MB max request body
.with_state(state.clone());
let (shutdown_tx, shutdown_rx) = oneshot::channel();
@@ -244,6 +349,13 @@ async fn chat_send_handler(
State(state): State<Arc<GatewayState>>,
Json(req): Json<SendMessageRequest>,
) -> Result<(StatusCode, Json<SendMessageResponse>), (StatusCode, String)> {
if !state.chat_rate_limiter.check() {
return Err((
StatusCode::TOO_MANY_REQUESTS,
"Rate limit exceeded. Try again shortly.".to_string(),
));
}
let mut msg = IncomingMessage::new("gateway", &state.user_id, &req.content);
if let Some(ref thread_id) = req.thread_id {
@@ -413,24 +525,58 @@ pub async fn clear_auth_mode(state: &GatewayState) {
if let Some(ref sm) = state.session_manager {
let session = sm.get_or_create_session(&state.user_id).await;
let mut sess = session.lock().await;
if let Some(thread_id) = sess.active_thread {
if let Some(thread) = sess.threads.get_mut(&thread_id) {
thread.pending_auth = None;
}
if let Some(thread_id) = sess.active_thread
&& let Some(thread) = sess.threads.get_mut(&thread_id)
{
thread.pending_auth = None;
}
}
}
async fn chat_events_handler(State(state): State<Arc<GatewayState>>) -> impl IntoResponse {
// subscribe() returns Sse<impl Stream + 'static + use<>> so no lifetime issues
state.sse.subscribe()
async fn chat_events_handler(
State(state): State<Arc<GatewayState>>,
) -> Result<impl IntoResponse, (StatusCode, String)> {
state.sse.subscribe().ok_or((
StatusCode::SERVICE_UNAVAILABLE,
"Too many connections".to_string(),
))
}
async fn chat_ws_handler(
headers: axum::http::HeaderMap,
ws: WebSocketUpgrade,
State(state): State<Arc<GatewayState>>,
) -> impl IntoResponse {
ws.on_upgrade(move |socket| crate::channels::web::ws::handle_ws_connection(socket, state))
) -> Result<impl IntoResponse, (StatusCode, String)> {
// Validate Origin header to prevent cross-site WebSocket hijacking.
// Require the header outright; browsers always send it for WS upgrades,
// so a missing Origin means a non-browser client trying to bypass the check.
let origin = headers
.get("origin")
.and_then(|v| v.to_str().ok())
.ok_or_else(|| {
(
StatusCode::FORBIDDEN,
"WebSocket Origin header required".to_string(),
)
})?;
// Extract the host from the origin and compare exactly, so that
// crafted origins like "http://localhost.evil.com" are rejected.
// Origin format is "scheme://host[:port]".
let host = origin
.strip_prefix("http://")
.or_else(|| origin.strip_prefix("https://"))
.and_then(|rest| rest.split(':').next()?.split('/').next())
.unwrap_or("");
let is_local = matches!(host, "localhost" | "127.0.0.1" | "[::1]");
if !is_local {
return Err((
StatusCode::FORBIDDEN,
"WebSocket origin not allowed".to_string(),
));
}
Ok(ws.on_upgrade(move |socket| crate::channels::web::ws::handle_ws_connection(socket, state)))
}
#[derive(Deserialize)]
@@ -477,57 +623,72 @@ async fn chat_history_handler(
.ok_or((StatusCode::NOT_FOUND, "No active thread".to_string()))?
};
// For paginated requests (before cursor set), always go to DB
if before_cursor.is_some() {
if let Some(ref store) = state.store {
let (messages, has_more) = store
.list_conversation_messages_paginated(thread_id, before_cursor, limit as i64)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
let oldest_timestamp = messages.first().map(|m| m.created_at.to_rfc3339());
let turns = build_turns_from_db_messages(&messages);
return Ok(Json(HistoryResponse {
thread_id,
turns,
has_more,
oldest_timestamp,
}));
// Verify the thread belongs to the authenticated user before returning any data.
// In-memory threads are already scoped by user via session_manager, but DB
// lookups could expose another user's conversation if the UUID is guessed.
if query.thread_id.is_some()
&& let Some(ref store) = state.store
{
let owned = store
.conversation_belongs_to_user(thread_id, &state.user_id)
.await
.unwrap_or(false);
if !owned && !sess.threads.contains_key(&thread_id) {
return Err((StatusCode::NOT_FOUND, "Thread not found".to_string()));
}
}
// Try in-memory first (freshest data for active threads)
if let Some(thread) = sess.threads.get(&thread_id) {
if !thread.turns.is_empty() {
let turns: Vec<TurnInfo> = thread
.turns
.iter()
.map(|t| TurnInfo {
turn_number: t.turn_number,
user_input: t.user_input.clone(),
response: t.response.clone(),
state: format!("{:?}", t.state),
started_at: t.started_at.to_rfc3339(),
completed_at: t.completed_at.map(|dt| dt.to_rfc3339()),
tool_calls: t
.tool_calls
.iter()
.map(|tc| ToolCallInfo {
name: tc.name.clone(),
has_result: tc.result.is_some(),
has_error: tc.error.is_some(),
})
.collect(),
})
.collect();
// For paginated requests (before cursor set), always go to DB
if before_cursor.is_some()
&& let Some(ref store) = state.store
{
let (messages, has_more) = store
.list_conversation_messages_paginated(thread_id, before_cursor, limit as i64)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
return Ok(Json(HistoryResponse {
thread_id,
turns,
has_more: false,
oldest_timestamp: None,
}));
}
let oldest_timestamp = messages.first().map(|m| m.created_at.to_rfc3339());
let turns = build_turns_from_db_messages(&messages);
return Ok(Json(HistoryResponse {
thread_id,
turns,
has_more,
oldest_timestamp,
}));
}
// Try in-memory first (freshest data for active threads)
if let Some(thread) = sess.threads.get(&thread_id)
&& !thread.turns.is_empty()
{
let turns: Vec<TurnInfo> = thread
.turns
.iter()
.map(|t| TurnInfo {
turn_number: t.turn_number,
user_input: t.user_input.clone(),
response: t.response.clone(),
state: format!("{:?}", t.state),
started_at: t.started_at.to_rfc3339(),
completed_at: t.completed_at.map(|dt| dt.to_rfc3339()),
tool_calls: t
.tool_calls
.iter()
.map(|tc| ToolCallInfo {
name: tc.name.clone(),
has_result: tc.result.is_some(),
has_error: tc.error.is_some(),
})
.collect(),
})
.collect();
return Ok(Json(HistoryResponse {
thread_id,
turns,
has_more: false,
oldest_timestamp: None,
}));
}
// Fall back to DB for historical threads not in memory (paginated)
@@ -577,12 +738,12 @@ fn build_turns_from_db_messages(messages: &[crate::history::ConversationMessage]
};
// Check if next message is an assistant response
if let Some(next) = iter.peek() {
if next.role == "assistant" {
let assistant_msg = iter.next().expect("peeked");
turn.response = Some(assistant_msg.content.clone());
turn.completed_at = Some(assistant_msg.created_at.to_rfc3339());
}
if let Some(next) = iter.peek()
&& next.role == "assistant"
{
let assistant_msg = iter.next().expect("peeked");
turn.response = Some(assistant_msg.content.clone());
turn.completed_at = Some(assistant_msg.created_at.to_rfc3339());
}
// Incomplete turn (user message without response)
@@ -901,14 +1062,16 @@ async fn jobs_list_handler(
"Database not available".to_string(),
))?;
// Fetch sandbox jobs from the DB.
// Fetch sandbox jobs scoped to the authenticated user.
let sandbox_jobs = store
.list_sandbox_jobs()
.list_sandbox_jobs_for_user(&state.user_id)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
// Scope jobs to the authenticated user.
let mut jobs: Vec<JobInfo> = sandbox_jobs
.iter()
.filter(|j| j.user_id == state.user_id)
.map(|j| {
let ui_state = match j.status.as_str() {
"creating" => "pending",
@@ -941,7 +1104,7 @@ async fn jobs_summary_handler(
))?;
let s = store
.sandbox_job_summary()
.sandbox_job_summary_for_user(&state.user_id)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
@@ -962,63 +1125,66 @@ async fn jobs_detail_handler(
let job_id = Uuid::parse_str(&id)
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?;
// Try sandbox job from DB first.
if let Some(ref store) = state.store {
if let Ok(Some(job)) = store.get_sandbox_job(job_id).await {
let browse_id = std::path::Path::new(&job.project_dir)
.file_name()
.map(|n| n.to_string_lossy().to_string())
.unwrap_or_else(|| job.id.to_string());
let ui_state = match job.status.as_str() {
"creating" => "pending",
"running" => "in_progress",
s => s,
};
let elapsed_secs = job.started_at.map(|start| {
let end = job.completed_at.unwrap_or_else(chrono::Utc::now);
(end - start).num_seconds().max(0) as u64
});
// Synthesize transitions from timestamps.
let mut transitions = Vec::new();
if let Some(started) = job.started_at {
transitions.push(TransitionInfo {
from: "creating".to_string(),
to: "running".to_string(),
timestamp: started.to_rfc3339(),
reason: None,
});
}
if let Some(completed) = job.completed_at {
transitions.push(TransitionInfo {
from: "running".to_string(),
to: job.status.clone(),
timestamp: completed.to_rfc3339(),
reason: job.failure_reason.clone(),
});
}
return Ok(Json(JobDetailResponse {
id: job.id,
title: job.task.clone(),
description: String::new(),
state: ui_state.to_string(),
user_id: job.user_id.clone(),
created_at: job.created_at.to_rfc3339(),
started_at: job.started_at.map(|dt| dt.to_rfc3339()),
completed_at: job.completed_at.map(|dt| dt.to_rfc3339()),
elapsed_secs,
project_dir: Some(job.project_dir.clone()),
browse_url: Some(format!("/projects/{}/", browse_id)),
job_mode: {
let mode = store.get_sandbox_job_mode(job.id).await.ok().flatten();
mode.filter(|m| m != "worker")
},
transitions,
}));
// Try sandbox job from DB first, scoped to the authenticated user.
if let Some(ref store) = state.store
&& let Ok(Some(job)) = store.get_sandbox_job(job_id).await
{
if job.user_id != state.user_id {
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
}
let browse_id = std::path::Path::new(&job.project_dir)
.file_name()
.map(|n| n.to_string_lossy().to_string())
.unwrap_or_else(|| job.id.to_string());
let ui_state = match job.status.as_str() {
"creating" => "pending",
"running" => "in_progress",
s => s,
};
let elapsed_secs = job.started_at.map(|start| {
let end = job.completed_at.unwrap_or_else(chrono::Utc::now);
(end - start).num_seconds().max(0) as u64
});
// Synthesize transitions from timestamps.
let mut transitions = Vec::new();
if let Some(started) = job.started_at {
transitions.push(TransitionInfo {
from: "creating".to_string(),
to: "running".to_string(),
timestamp: started.to_rfc3339(),
reason: None,
});
}
if let Some(completed) = job.completed_at {
transitions.push(TransitionInfo {
from: "running".to_string(),
to: job.status.clone(),
timestamp: completed.to_rfc3339(),
reason: job.failure_reason.clone(),
});
}
return Ok(Json(JobDetailResponse {
id: job.id,
title: job.task.clone(),
description: String::new(),
state: ui_state.to_string(),
user_id: job.user_id.clone(),
created_at: job.created_at.to_rfc3339(),
started_at: job.started_at.map(|dt| dt.to_rfc3339()),
completed_at: job.completed_at.map(|dt| dt.to_rfc3339()),
elapsed_secs,
project_dir: Some(job.project_dir.clone()),
browse_url: Some(format!("/projects/{}/", browse_id)),
job_mode: {
let mode = store.get_sandbox_job_mode(job.id).await.ok().flatten();
mode.filter(|m| m != "worker")
},
transitions,
}));
}
Err((StatusCode::NOT_FOUND, "Job not found".to_string()))
@@ -1031,31 +1197,36 @@ async fn jobs_cancel_handler(
let job_id = Uuid::parse_str(&id)
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?;
// Try sandbox job cancellation.
if let Some(ref store) = state.store {
if let Ok(Some(job)) = store.get_sandbox_job(job_id).await {
if job.status == "running" || job.status == "creating" {
// Stop the container if we have a job manager.
if let Some(ref jm) = state.job_manager {
let _ = jm.stop_job(job_id).await;
}
store
.update_sandbox_job_status(
job_id,
"failed",
Some(false),
Some("Cancelled by user"),
None,
Some(chrono::Utc::now()),
)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
}
return Ok(Json(serde_json::json!({
"status": "cancelled",
"job_id": job_id,
})));
// Try sandbox job cancellation, scoped to the authenticated user.
if let Some(ref store) = state.store
&& let Ok(Some(job)) = store.get_sandbox_job(job_id).await
{
if job.user_id != state.user_id {
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
}
if job.status == "running" || job.status == "creating" {
// Stop the container if we have a job manager.
if let Some(ref jm) = state.job_manager
&& let Err(e) = jm.stop_job(job_id).await
{
tracing::warn!(job_id = %job_id, error = %e, "Failed to stop container during cancellation");
}
store
.update_sandbox_job_status(
job_id,
"failed",
Some(false),
Some("Cancelled by user"),
None,
Some(chrono::Utc::now()),
)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
}
return Ok(Json(serde_json::json!({
"status": "cancelled",
"job_id": job_id,
})));
}
Err((StatusCode::NOT_FOUND, "Job not found".to_string()))
@@ -1083,6 +1254,11 @@ async fn jobs_restart_handler(
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
.ok_or((StatusCode::NOT_FOUND, "Job not found".to_string()))?;
// Scope to the authenticated user.
if old_job.user_id != state.user_id {
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
}
if old_job.status != "interrupted" && old_job.status != "failed" {
return Err((
StatusCode::CONFLICT,
@@ -1157,6 +1333,16 @@ async fn jobs_prompt_handler(
.parse()
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?;
// Verify user owns this job.
if let Some(ref store) = state.store
&& !store
.sandbox_job_belongs_to_user(job_id, &state.user_id)
.await
.unwrap_or(false)
{
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
}
let content = body
.get("content")
.and_then(|v| v.as_str())
@@ -1195,6 +1381,15 @@ async fn jobs_events_handler(
.parse()
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?;
// Verify user owns this job.
if !store
.sandbox_job_belongs_to_user(job_id, &state.user_id)
.await
.unwrap_or(false)
{
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
}
let events = store
.list_job_events(job_id)
.await
@@ -1244,6 +1439,11 @@ async fn job_files_list_handler(
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
.ok_or((StatusCode::NOT_FOUND, "Job not found".to_string()))?;
// Verify user owns this job.
if job.user_id != state.user_id {
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
}
let base = std::path::PathBuf::from(&job.project_dir);
let rel_path = query.path.as_deref().unwrap_or("");
let target = base.join(rel_path);
@@ -1307,6 +1507,11 @@ async fn job_files_read_handler(
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
.ok_or((StatusCode::NOT_FOUND, "Job not found".to_string()))?;
// Verify user owns this job.
if job.user_id != state.user_id {
return Err((StatusCode::NOT_FOUND, "Job not found".to_string()));
}
let path = query.path.as_deref().ok_or((
StatusCode::BAD_REQUEST,
"path parameter required".to_string(),
@@ -1525,6 +1730,15 @@ async fn project_file_handler(
/// Shared logic: resolve the file inside `~/.ironclaw/projects/{project_id}/`,
/// guard against path traversal, and stream the content with the right MIME type.
async fn serve_project_file(project_id: &str, path: &str) -> axum::response::Response {
// Reject project_id values that could escape the projects directory.
if project_id.contains('/')
|| project_id.contains('\\')
|| project_id.contains("..")
|| project_id.is_empty()
{
return (StatusCode::BAD_REQUEST, "Invalid project ID").into_response();
}
let base = dirs::home_dir()
.unwrap_or_else(|| std::path::PathBuf::from("."))
.join(".ironclaw")
+59 -12
View File
@@ -13,10 +13,15 @@ use tokio_stream::wrappers::BroadcastStream;
use crate::channels::web::types::SseEvent;
/// Maximum number of concurrent SSE/WebSocket connections.
/// Prevents resource exhaustion from connection flooding.
const MAX_CONNECTIONS: u64 = 100;
/// Manages SSE broadcast to all connected browser tabs.
pub struct SseManager {
tx: broadcast::Sender<SseEvent>,
connection_count: Arc<AtomicU64>,
max_connections: u64,
}
impl SseManager {
@@ -27,6 +32,7 @@ impl SseManager {
Self {
tx,
connection_count: Arc::new(AtomicU64::new(0)),
max_connections: MAX_CONNECTIONS,
}
}
@@ -45,25 +51,50 @@ impl SseManager {
///
/// Returns a stream of `SseEvent` values and increments/decrements the
/// connection counter on creation/drop, just like `subscribe()` does for SSE.
pub fn subscribe_raw(&self) -> impl Stream<Item = SseEvent> + Send + 'static + use<> {
///
/// Returns `None` if the maximum connection limit has been reached.
pub fn subscribe_raw(&self) -> Option<impl Stream<Item = SseEvent> + Send + 'static + use<>> {
// Atomically increment only if below the limit. This prevents
// concurrent callers from overshooting max_connections.
let counter = Arc::clone(&self.connection_count);
counter.fetch_add(1, Ordering::Relaxed);
let max = self.max_connections;
counter
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
if current < max {
Some(current + 1)
} else {
None
}
})
.ok()?;
let rx = self.tx.subscribe();
let stream = BroadcastStream::new(rx).filter_map(|result| result.ok());
CountedStream {
Some(CountedStream {
inner: stream,
counter,
}
})
}
/// Create a new SSE stream for a client connection.
///
/// Returns `None` if the maximum connection limit has been reached.
pub fn subscribe(
&self,
) -> Sse<impl Stream<Item = Result<Event, Infallible>> + Send + 'static + use<>> {
) -> Option<Sse<impl Stream<Item = Result<Event, Infallible>> + Send + 'static + use<>>> {
// Atomically increment only if below the limit.
let counter = Arc::clone(&self.connection_count);
counter.fetch_add(1, Ordering::Relaxed);
let max = self.max_connections;
counter
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
if current < max {
Some(current + 1)
} else {
None
}
})
.ok()?;
let rx = self.tx.subscribe();
let stream = BroadcastStream::new(rx)
@@ -99,8 +130,10 @@ impl SseManager {
counter,
};
Sse::new(counted_stream)
.keep_alive(KeepAlive::new().interval(Duration::from_secs(30)).text(""))
Some(
Sse::new(counted_stream)
.keep_alive(KeepAlive::new().interval(Duration::from_secs(30)).text("")),
)
}
}
@@ -175,7 +208,7 @@ mod tests {
#[tokio::test]
async fn test_subscribe_raw_receives_events() {
let manager = SseManager::new();
let mut stream = Box::pin(manager.subscribe_raw());
let mut stream = Box::pin(manager.subscribe_raw().expect("should subscribe"));
assert_eq!(manager.connection_count(), 1);
@@ -195,7 +228,7 @@ mod tests {
async fn test_subscribe_raw_decrements_on_drop() {
let manager = SseManager::new();
{
let _stream = Box::pin(manager.subscribe_raw());
let _stream = Box::pin(manager.subscribe_raw().expect("should subscribe"));
assert_eq!(manager.connection_count(), 1);
}
// Stream dropped, counter should decrement
@@ -205,8 +238,8 @@ mod tests {
#[tokio::test]
async fn test_subscribe_raw_multiple_subscribers() {
let manager = SseManager::new();
let mut s1 = Box::pin(manager.subscribe_raw());
let mut s2 = Box::pin(manager.subscribe_raw());
let mut s1 = Box::pin(manager.subscribe_raw().expect("should subscribe"));
let mut s2 = Box::pin(manager.subscribe_raw().expect("should subscribe"));
assert_eq!(manager.connection_count(), 2);
manager.broadcast(SseEvent::Heartbeat);
@@ -221,4 +254,18 @@ mod tests {
drop(s2);
assert_eq!(manager.connection_count(), 0);
}
#[tokio::test]
async fn test_subscribe_raw_rejects_over_limit() {
let mut manager = SseManager::new();
manager.max_connections = 2; // Low limit for testing
let _s1 = Box::pin(manager.subscribe_raw().expect("first should succeed"));
let _s2 = Box::pin(manager.subscribe_raw().expect("second should succeed"));
assert_eq!(manager.connection_count(), 2);
// Third should be rejected
assert!(manager.subscribe_raw().is_none());
assert!(manager.subscribe().is_none());
}
}
+24
View File
@@ -281,6 +281,8 @@ function sendApprovalAction(requestId, action) {
function renderMarkdown(text) {
if (typeof marked !== 'undefined') {
let html = marked.parse(text);
// Sanitize HTML output to prevent XSS from tool output or LLM responses.
html = sanitizeRenderedHtml(html);
// Inject copy buttons into <pre> blocks
html = html.replace(/<pre>/g, '<pre class="code-block-wrapper"><button class="copy-btn" onclick="copyCodeBlock(this)">Copy</button>');
return html;
@@ -288,6 +290,28 @@ function renderMarkdown(text) {
return escapeHtml(text);
}
// Strip dangerous HTML elements and attributes from rendered markdown.
// This prevents XSS from tool output or prompt injection in LLM responses.
function sanitizeRenderedHtml(html) {
html = html.replace(/<script\b[^<]*(?:(?!<\/script>)<[^<]*)*<\/script>/gi, '');
html = html.replace(/<iframe\b[^>]*>[\s\S]*?<\/iframe>/gi, '');
html = html.replace(/<object\b[^>]*>[\s\S]*?<\/object>/gi, '');
html = html.replace(/<embed\b[^>]*\/?>/gi, '');
html = html.replace(/<form\b[^>]*>[\s\S]*?<\/form>/gi, '');
html = html.replace(/<style\b[^>]*>[\s\S]*?<\/style>/gi, '');
html = html.replace(/<link\b[^>]*\/?>/gi, '');
html = html.replace(/<base\b[^>]*\/?>/gi, '');
html = html.replace(/<meta\b[^>]*\/?>/gi, '');
// Remove event handler attributes (onclick, onerror, onload, etc.)
html = html.replace(/\s+on\w+\s*=\s*"[^"]*"/gi, '');
html = html.replace(/\s+on\w+\s*=\s*'[^']*'/gi, '');
html = html.replace(/\s+on\w+\s*=\s*[^\s>]+/gi, '');
// Remove javascript: and data: URLs in href/src attributes
html = html.replace(/(href|src|action)\s*=\s*["']?\s*javascript\s*:/gi, '$1="');
html = html.replace(/(href|src|action)\s*=\s*["']?\s*data\s*:/gi, '$1="');
return html;
}
function copyCodeBlock(btn) {
const pre = btn.parentElement;
const code = pre.querySelector('code');
+13 -2
View File
@@ -71,8 +71,17 @@ pub async fn handle_ws_connection(socket: WebSocket, state: Arc<GatewayState>) {
}
let tracker_for_drop = state.ws_tracker.clone();
// Subscribe to broadcast events (same source as SSE)
let mut event_stream = Box::pin(state.sse.subscribe_raw());
// Subscribe to broadcast events (same source as SSE).
// Reject if we've hit the connection limit.
let Some(raw_stream) = state.sse.subscribe_raw() else {
tracing::warn!("WebSocket rejected: too many connections");
// Decrement the WS tracker we already incremented above.
if let Some(ref tracker) = tracker_for_drop {
tracker.decrement();
}
return;
};
let mut event_stream = Box::pin(raw_stream);
// Channel for the sender task to receive messages from both
// the broadcast stream and any direct sends (like Pong)
@@ -476,6 +485,8 @@ mod tests {
user_id: "test".to_string(),
shutdown_tx: tokio::sync::RwLock::new(None),
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
llm_provider: None,
chat_rate_limiter: crate::channels::web::server::RateLimiter::new(30, 60),
}
}
}
+50 -79
View File
@@ -1,7 +1,9 @@
//! Configuration management CLI commands.
//!
//! Commands for viewing and modifying settings.
//! Settings are stored in PostgreSQL (env > DB > default).
//! Settings are stored in the database (env > DB > default).
use std::sync::Arc;
use clap::Subcommand;
@@ -46,11 +48,9 @@ pub enum ConfigCommand {
/// 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),
let db: Option<Arc<dyn crate::db::Database>> = match connect_db().await {
Ok(d) => Some(d),
Err(e) => {
eprintln!(
"Warning: Could not connect to database ({}), using disk fallback",
@@ -60,41 +60,42 @@ pub async fn run_config_command(cmd: ConfigCommand) -> anyhow::Result<()> {
}
};
let db_ref = db.as_deref();
match cmd {
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()),
ConfigCommand::List { filter } => list_settings(db_ref, filter).await,
ConfigCommand::Get { path } => get_setting(db_ref, &path).await,
ConfigCommand::Set { path, value } => set_setting(db_ref, &path, &value).await,
ConfigCommand::Reset { path } => reset_setting(db_ref, &path).await,
ConfigCommand::Path => show_path(db_ref.is_some()),
}
}
/// Bootstrap a DB connection for config commands.
async fn connect_store() -> anyhow::Result<crate::history::Store> {
/// Bootstrap a DB connection for config commands (backend-agnostic).
async fn connect_db() -> anyhow::Result<Arc<dyn crate::db::Database>> {
let config = crate::config::Config::from_env()
.await
.map_err(|e| anyhow::anyhow!("{}", e))?;
let store = crate::history::Store::new(&config.database).await?;
store.run_migrations().await?;
Ok(store)
crate::db::connect_from_config(&config.database)
.await
.map_err(|e| anyhow::anyhow!("{}", e))
}
const DEFAULT_USER_ID: &str = "default";
/// Load settings: DB if available, else disk.
async fn load_settings(store: Option<&crate::history::Store>) -> Settings {
async fn load_settings(store: Option<&dyn crate::db::Database>) -> 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()
Settings::default()
}
/// List all settings.
async fn list_settings(
store: Option<&crate::history::Store>,
store: Option<&dyn crate::db::Database>,
filter: Option<String>,
) -> anyhow::Result<()> {
let settings = load_settings(store).await;
@@ -107,10 +108,10 @@ async fn list_settings(
println!();
for (key, value) in all {
if let Some(ref f) = filter {
if !key.starts_with(f) {
continue;
}
if let Some(ref f) = filter
&& !key.starts_with(f)
{
continue;
}
let display_value = if value.len() > 60 {
@@ -126,7 +127,7 @@ async fn list_settings(
}
/// Get a specific setting.
async fn get_setting(store: Option<&crate::history::Store>, path: &str) -> anyhow::Result<()> {
async fn get_setting(store: Option<&dyn crate::db::Database>, path: &str) -> anyhow::Result<()> {
let settings = load_settings(store).await;
match settings.get(path) {
@@ -142,7 +143,7 @@ async fn get_setting(store: Option<&crate::history::Store>, path: &str) -> anyho
/// Set a setting value.
async fn set_setting(
store: Option<&crate::history::Store>,
store: Option<&dyn crate::db::Database>,
path: &str,
value: &str,
) -> anyhow::Result<()> {
@@ -152,42 +153,36 @@ async fn set_setting(
.set(path, value)
.map_err(|e| anyhow::anyhow!("{}", e))?;
// 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()?;
}
let store = store.ok_or_else(|| {
anyhow::anyhow!("Database connection required to save settings. Check DATABASE_URL.")
})?;
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))?;
println!("Set {} = {}", path, value);
Ok(())
}
/// Reset a setting to default.
async fn reset_setting(store: Option<&crate::history::Store>, path: &str) -> anyhow::Result<()> {
async fn reset_setting(store: Option<&dyn crate::db::Database>, path: &str) -> anyhow::Result<()> {
let default = Settings::default();
let default_value = default
.get(path)
.ok_or_else(|| anyhow::anyhow!("Unknown setting: {}", path))?;
// 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()?;
}
let store = store.ok_or_else(|| {
anyhow::anyhow!("Database connection required to reset settings. Check DATABASE_URL.")
})?;
store
.delete_setting(DEFAULT_USER_ID, path)
.await
.map_err(|e| anyhow::anyhow!("Failed to delete setting from database: {}", e))?;
println!("Reset {} to default: {}", path, default_value);
Ok(())
@@ -196,38 +191,14 @@ async fn reset_setting(store: Option<&crate::history::Store>, path: &str) -> any
/// 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()
);
println!("Settings stored in: database (settings table)");
} else {
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)");
}
println!("Settings stored in: PostgreSQL (not connected, using defaults)");
}
println!(
"Env config: {}",
crate::bootstrap::ironclaw_env_path().display()
);
Ok(())
}
+82 -35
View File
@@ -8,8 +8,10 @@ use std::sync::Arc;
use clap::Subcommand;
use crate::config::Config;
use crate::history::Store;
use crate::secrets::{PostgresSecretsStore, SecretsCrypto, SecretsStore};
use crate::db::Database;
#[cfg(feature = "postgres")]
use crate::secrets::PostgresSecretsStore;
use crate::secrets::{SecretsCrypto, SecretsStore};
use crate::tools::mcp::{
McpClient, McpServerConfig, McpSessionManager, OAuthConfig,
auth::{authorize_mcp_server, is_authenticated},
@@ -172,10 +174,10 @@ async fn add_server(
config.validate()?;
// Save (DB if available, else disk)
let store = connect_store().await;
let mut servers = load_servers(store.as_ref()).await?;
let db = connect_db().await;
let mut servers = load_servers(db.as_deref()).await?;
servers.upsert(config);
save_servers(store.as_ref(), &servers).await?;
save_servers(db.as_deref(), &servers).await?;
println!();
println!(" ✓ Added MCP server '{}'", name);
@@ -193,12 +195,12 @@ async fn add_server(
/// Remove an MCP server.
async fn remove_server(name: String) -> anyhow::Result<()> {
let store = connect_store().await;
let mut servers = load_servers(store.as_ref()).await?;
let db = connect_db().await;
let mut servers = load_servers(db.as_deref()).await?;
if !servers.remove(&name) {
anyhow::bail!("Server '{}' not found", name);
}
save_servers(store.as_ref(), &servers).await?;
save_servers(db.as_deref(), &servers).await?;
println!();
println!(" ✓ Removed MCP server '{}'", name);
@@ -209,8 +211,8 @@ async fn remove_server(name: String) -> anyhow::Result<()> {
/// List configured MCP servers.
async fn list_servers(verbose: bool) -> anyhow::Result<()> {
let store = connect_store().await;
let servers = load_servers(store.as_ref()).await?;
let db = connect_db().await;
let servers = load_servers(db.as_deref()).await?;
if servers.servers.is_empty() {
println!();
@@ -268,8 +270,8 @@ 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 store = connect_store().await;
let servers = load_servers(store.as_ref()).await?;
let db = connect_db().await;
let servers = load_servers(db.as_deref()).await?;
let server = servers
.get(&name)
.cloned()
@@ -341,8 +343,8 @@ 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 store = connect_store().await;
let servers = load_servers(store.as_ref()).await?;
let db = connect_db().await;
let servers = load_servers(db.as_deref()).await?;
let server = servers
.get(&name)
.cloned()
@@ -437,8 +439,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 store = connect_store().await;
let mut servers = load_servers(store.as_ref()).await?;
let db = connect_db().await;
let mut servers = load_servers(db.as_deref()).await?;
let server = servers
.get_mut(&name)
@@ -453,7 +455,7 @@ async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Res
};
server.enabled = new_state;
save_servers(store.as_ref(), &servers).await?;
save_servers(db.as_deref(), &servers).await?;
let status = if new_state { "enabled" } else { "disabled" };
println!();
@@ -465,18 +467,16 @@ async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Res
const DEFAULT_USER_ID: &str = "default";
/// Try to connect to the database store for DB-backed config.
async fn connect_store() -> Option<Store> {
/// Try to connect to the database (backend-agnostic).
async fn connect_db() -> Option<Arc<dyn Database>> {
let config = Config::from_env().await.ok()?;
let store = Store::new(&config.database).await.ok()?;
store.run_migrations().await.ok()?;
Some(store)
crate::db::connect_from_config(&config.database).await.ok()
}
/// 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
async fn load_servers(db: Option<&dyn Database>) -> Result<McpServersFile, config::ConfigError> {
if let Some(db) = db {
config::load_mcp_servers_from_db(db, DEFAULT_USER_ID).await
} else {
config::load_mcp_servers().await
}
@@ -484,11 +484,11 @@ async fn load_servers(store: Option<&Store>) -> Result<McpServersFile, config::C
/// Save MCP servers (DB if available, else disk).
async fn save_servers(
store: Option<&Store>,
db: Option<&dyn Database>,
servers: &McpServersFile,
) -> Result<(), config::ConfigError> {
if let Some(store) = store {
config::save_mcp_servers_to_db(store, DEFAULT_USER_ID, servers).await
if let Some(db) = db {
config::save_mcp_servers_to_db(db, DEFAULT_USER_ID, servers).await
} else {
config::save_mcp_servers(servers).await
}
@@ -504,14 +504,61 @@ async fn get_secrets_store() -> anyhow::Result<Arc<dyn SecretsStore + Send + Syn
)
})?;
let store = Store::new(&config.database).await?;
store.run_migrations().await?;
let crypto = SecretsCrypto::new(master_key.clone())?;
Ok(Arc::new(PostgresSecretsStore::new(
store.pool(),
Arc::new(crypto),
)))
#[cfg(feature = "postgres")]
{
let store = crate::history::Store::new(&config.database).await?;
store.run_migrations().await?;
Ok(Arc::new(PostgresSecretsStore::new(
store.pool(),
Arc::new(crypto),
)))
}
#[cfg(all(feature = "libsql", not(feature = "postgres")))]
{
use crate::db::Database as _;
use crate::db::libsql_backend::LibSqlBackend;
use secrecy::ExposeSecret as _;
let default_path = crate::config::default_libsql_path();
let db_path = config
.database
.libsql_path
.as_deref()
.unwrap_or(&default_path);
let backend = if let Some(ref url) = config.database.libsql_url {
let token = config.database.libsql_auth_token.as_ref().ok_or_else(|| {
anyhow::anyhow!("LIBSQL_AUTH_TOKEN is required when LIBSQL_URL is set")
})?;
LibSqlBackend::new_remote_replica(db_path, url, token.expose_secret())
.await
.map_err(|e| anyhow::anyhow!("{}", e))?
} else {
LibSqlBackend::new_local(db_path)
.await
.map_err(|e| anyhow::anyhow!("{}", e))?
};
backend
.run_migrations()
.await
.map_err(|e| anyhow::anyhow!("{}", e))?;
return Ok(Arc::new(crate::secrets::LibSqlSecretsStore::new(
backend.shared_db(),
Arc::new(crypto),
)));
}
#[cfg(not(any(feature = "postgres", feature = "libsql")))]
{
let _ = crypto;
anyhow::bail!(
"No database backend available for secrets. Enable 'postgres' or 'libsql' feature."
);
}
}
#[cfg(test)]
+26 -1
View File
@@ -9,6 +9,30 @@ use clap::Subcommand;
use crate::workspace::{EmbeddingProvider, SearchConfig, Workspace};
/// Run a memory command using the Database trait (works with any backend).
pub async fn run_memory_command_with_db(
cmd: MemoryCommand,
db: std::sync::Arc<dyn crate::db::Database>,
embeddings: Option<Arc<dyn EmbeddingProvider>>,
) -> anyhow::Result<()> {
let mut workspace = Workspace::new_with_db("default", db);
if let Some(emb) = embeddings {
workspace = workspace.with_embeddings(emb);
}
match cmd {
MemoryCommand::Search { query, limit } => search(&workspace, &query, limit).await,
MemoryCommand::Read { path } => read(&workspace, &path).await,
MemoryCommand::Write {
path,
content,
append,
} => write(&workspace, &path, content, append).await,
MemoryCommand::Tree { path, depth } => tree(&workspace, &path, depth).await,
MemoryCommand::Status => status(&workspace).await,
}
}
#[derive(Subcommand, Debug, Clone)]
pub enum MemoryCommand {
/// Search workspace memory (hybrid full-text + semantic)
@@ -55,7 +79,8 @@ pub enum MemoryCommand {
Status,
}
/// Run a memory command.
/// Run a memory command (PostgreSQL backend).
#[cfg(feature = "postgres")]
pub async fn run_memory_command(
cmd: MemoryCommand,
pool: deadpool_postgres::Pool,
+5 -1
View File
@@ -12,13 +12,17 @@
mod config;
mod mcp;
pub mod memory;
pub mod oauth_defaults;
mod pairing;
pub mod status;
mod tool;
pub use config::{ConfigCommand, run_config_command};
pub use mcp::{McpCommand, run_mcp_command};
pub use memory::{MemoryCommand, run_memory_command};
pub use memory::MemoryCommand;
#[cfg(feature = "postgres")]
pub use memory::run_memory_command;
pub use memory::run_memory_command_with_db;
pub use pairing::{PairingCommand, run_pairing_command, run_pairing_command_with_store};
pub use status::run_status_command;
pub use tool::{ToolCommand, run_tool_command};
+342
View File
@@ -0,0 +1,342 @@
//! Shared OAuth infrastructure: built-in credentials, callback server, landing pages.
//!
//! Every OAuth flow in the codebase (WASM tool auth, MCP server auth, NEAR AI login)
//! uses the same callback port, landing page, and listener logic from this module.
//!
//! # Built-in Credentials
//!
//! Many CLI tools (gcloud, rclone, gdrive) ship with default OAuth credentials
//! so users don't need to register their own OAuth app. Google explicitly
//! documents that client_secret for "Desktop App" / "Installed App" types
//! is NOT actually secret.
//!
//! Default credentials are hardcoded below. They can be overridden at:
//!
//! - **Compile time**: Set IRONCLAW_GOOGLE_CLIENT_ID / IRONCLAW_GOOGLE_CLIENT_SECRET
//! env vars before building to replace the hardcoded defaults.
//! - **Runtime**: Users can set GOOGLE_OAUTH_CLIENT_ID / GOOGLE_OAUTH_CLIENT_SECRET
//! env vars, which take priority over built-in defaults.
use std::time::Duration;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::TcpListener;
// ── Built-in credentials ────────────────────────────────────────────────
pub struct OAuthCredentials {
pub client_id: &'static str,
pub client_secret: &'static str,
}
/// Google OAuth "Desktop App" credentials, shared across all Google tools.
/// Compile-time env vars override the hardcoded defaults below.
const GOOGLE_CLIENT_ID: &str = match option_env!("IRONCLAW_GOOGLE_CLIENT_ID") {
Some(v) => v,
None => "564604149681-efo25d43rs85v0tibdepsmdv5dsrhhr0.apps.googleusercontent.com",
};
const GOOGLE_CLIENT_SECRET: &str = match option_env!("IRONCLAW_GOOGLE_CLIENT_SECRET") {
Some(v) => v,
None => "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",
};
/// Returns built-in OAuth credentials for a provider, keyed by secret_name.
///
/// The secret_name comes from the tool's capabilities.json `auth.secret_name` field.
/// Returns `None` if no built-in credentials are configured for that provider.
pub fn builtin_credentials(secret_name: &str) -> Option<OAuthCredentials> {
match secret_name {
"google_oauth_token" => Some(OAuthCredentials {
client_id: GOOGLE_CLIENT_ID,
client_secret: GOOGLE_CLIENT_SECRET,
}),
_ => None,
}
}
// ── Shared callback server ──────────────────────────────────────────────
/// Fixed port for all OAuth callbacks.
///
/// Every redirect URI registered with providers must use this port:
/// `http://localhost:9876/callback` (or `/auth/callback` for NEAR AI).
pub const OAUTH_CALLBACK_PORT: u16 = 9876;
/// Error from the OAuth callback listener.
#[derive(Debug, thiserror::Error)]
pub enum OAuthCallbackError {
#[error("Port {0} is in use (another auth flow running?): {1}")]
PortInUse(u16, String),
#[error("Authorization denied by user")]
Denied,
#[error("Timed out waiting for authorization")]
Timeout,
#[error("IO error: {0}")]
Io(String),
}
/// Bind the OAuth callback listener on the fixed port.
///
/// Binds to IPv4 `127.0.0.1` first because callback URLs use `127.0.0.1`
/// explicitly (e.g., NEAR AI redirects to `http://127.0.0.1:9876/auth/callback`).
/// Falls back to IPv6 `[::1]` only if IPv4 binding fails for a reason other
/// than `AddrInUse`. If the port is already occupied, fails immediately.
pub async fn bind_callback_listener() -> Result<TcpListener, OAuthCallbackError> {
let ipv4_addr = format!("127.0.0.1:{}", OAUTH_CALLBACK_PORT);
match TcpListener::bind(&ipv4_addr).await {
Ok(listener) => return Ok(listener),
Err(e) if e.kind() == std::io::ErrorKind::AddrInUse => {
return Err(OAuthCallbackError::PortInUse(
OAUTH_CALLBACK_PORT,
e.to_string(),
));
}
Err(_) => {
// IPv4 not available, fall back to IPv6
}
}
TcpListener::bind(format!("[::1]:{}", OAUTH_CALLBACK_PORT))
.await
.map_err(|e| {
if e.kind() == std::io::ErrorKind::AddrInUse {
OAuthCallbackError::PortInUse(OAUTH_CALLBACK_PORT, e.to_string())
} else {
OAuthCallbackError::Io(e.to_string())
}
})
}
/// Wait for an OAuth callback and extract a query parameter value.
///
/// Listens for a GET request matching `path_prefix` (e.g., "/callback" or "/auth/callback"),
/// extracts the value of `param_name` (e.g., "code" or "token"), and shows a branded
/// landing page using `display_name` (e.g., "Google", "Notion", "NEAR AI").
///
/// Times out after 5 minutes.
pub async fn wait_for_callback(
listener: TcpListener,
path_prefix: &str,
param_name: &str,
display_name: &str,
) -> Result<String, OAuthCallbackError> {
let path_prefix = path_prefix.to_string();
let param_name = param_name.to_string();
let display_name = display_name.to_string();
tokio::time::timeout(Duration::from_secs(300), async move {
loop {
let (mut socket, _) = listener
.accept()
.await
.map_err(|e| OAuthCallbackError::Io(e.to_string()))?;
let mut reader = BufReader::new(&mut socket);
let mut request_line = String::new();
reader
.read_line(&mut request_line)
.await
.map_err(|e| OAuthCallbackError::Io(e.to_string()))?;
if let Some(path) = request_line.split_whitespace().nth(1)
&& path.starts_with(&path_prefix)
&& let Some(query) = path.split('?').nth(1)
{
// Check for error first
if query.contains("error=") {
let html = landing_html(&display_name, false);
let response = format!(
"HTTP/1.1 400 Bad Request\r\n\
Content-Type: text/html; charset=utf-8\r\n\
Connection: close\r\n\
\r\n\
{}",
html
);
let _ = socket.write_all(response.as_bytes()).await;
return Err(OAuthCallbackError::Denied);
}
// Look for the target parameter
for param in query.split('&') {
let parts: Vec<&str> = param.splitn(2, '=').collect();
if parts.len() == 2 && parts[0] == param_name {
let value = urlencoding::decode(parts[1])
.unwrap_or_else(|_| parts[1].into())
.into_owned();
let html = landing_html(&display_name, true);
let response = format!(
"HTTP/1.1 200 OK\r\n\
Content-Type: text/html; charset=utf-8\r\n\
Connection: close\r\n\
\r\n\
{}",
html
);
let _ = socket.write_all(response.as_bytes()).await;
let _ = socket.shutdown().await;
return Ok(value);
}
}
}
// Not the callback we're looking for
let response = "HTTP/1.1 404 Not Found\r\nConnection: close\r\n\r\n";
let _ = socket.write_all(response.as_bytes()).await;
}
})
.await
.map_err(|_| OAuthCallbackError::Timeout)?
}
/// Escape a string for safe interpolation into HTML content.
fn html_escape(s: &str) -> String {
let mut out = String::with_capacity(s.len());
for c in s.chars() {
match c {
'&' => out.push_str("&amp;"),
'<' => out.push_str("&lt;"),
'>' => out.push_str("&gt;"),
'"' => out.push_str("&quot;"),
'\'' => out.push_str("&#x27;"),
_ => out.push(c),
}
}
out
}
/// HTML landing page shown in the browser after an OAuth redirect.
pub fn landing_html(provider_name: &str, success: bool) -> String {
let safe_name = html_escape(provider_name);
let (icon, heading, subtitle, accent) = if success {
(
r##"<div style="width:64px;height:64px;border-radius:50%;background:#22c55e;display:flex;align-items:center;justify-content:center;margin:0 auto 24px">
<svg width="32" height="32" viewBox="0 0 24 24" fill="none" stroke="#fff" stroke-width="3" stroke-linecap="round" stroke-linejoin="round"><polyline points="20 6 9 17 4 12"/></svg>
</div>"##,
format!("{} Connected", safe_name),
"You can close this window and return to your terminal.",
"#22c55e",
)
} else {
(
r##"<div style="width:64px;height:64px;border-radius:50%;background:#ef4444;display:flex;align-items:center;justify-content:center;margin:0 auto 24px">
<svg width="32" height="32" viewBox="0 0 24 24" fill="none" stroke="#fff" stroke-width="3" stroke-linecap="round" stroke-linejoin="round"><line x1="18" y1="6" x2="6" y2="18"/><line x1="6" y1="6" x2="18" y2="18"/></svg>
</div>"##,
"Authorization Failed".to_string(),
"The request was denied. You can close this window and try again.",
"#ef4444",
)
};
format!(
r#"<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<title>IronClaw - {heading}</title>
<style>
* {{ margin:0; padding:0; box-sizing:border-box }}
body {{
font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif;
background: #0a0a0a;
color: #e5e5e5;
display: flex;
justify-content: center;
align-items: center;
min-height: 100vh;
}}
.card {{
text-align: center;
padding: 48px 40px;
max-width: 420px;
border: 1px solid #262626;
border-radius: 16px;
background: #141414;
}}
h1 {{
font-size: 22px;
font-weight: 600;
margin-bottom: 8px;
color: #fafafa;
}}
p {{
font-size: 14px;
color: #a3a3a3;
line-height: 1.5;
}}
.accent {{ color: {accent}; }}
.brand {{
margin-top: 32px;
font-size: 12px;
color: #525252;
letter-spacing: 0.5px;
text-transform: uppercase;
}}
</style>
</head>
<body>
<div class="card">
{icon}
<h1>{heading}</h1>
<p>{subtitle}</p>
<div class="brand">IronClaw</div>
</div>
</body>
</html>"#,
heading = heading,
icon = icon,
subtitle = subtitle,
accent = accent,
)
}
#[cfg(test)]
mod tests {
use crate::cli::oauth_defaults::{builtin_credentials, landing_html};
#[test]
fn test_unknown_provider_returns_none() {
assert!(builtin_credentials("unknown_token").is_none());
}
#[test]
fn test_google_returns_based_on_compile_env() {
let creds = builtin_credentials("google_oauth_token");
assert!(creds.is_some());
let creds = creds.unwrap();
assert!(!creds.client_id.is_empty());
assert!(!creds.client_secret.is_empty());
}
#[test]
fn test_landing_html_success_contains_key_elements() {
let html = landing_html("Google", true);
assert!(html.contains("Google Connected"));
assert!(html.contains("charset"));
assert!(html.contains("IronClaw"));
assert!(html.contains("#22c55e")); // green accent
assert!(!html.contains("Failed"));
}
#[test]
fn test_landing_html_escapes_provider_name() {
let html = landing_html("<script>alert(1)</script>", true);
assert!(!html.contains("<script>"));
assert!(html.contains("&lt;script&gt;"));
}
#[test]
fn test_landing_html_error_contains_key_elements() {
let html = landing_html("Notion", false);
assert!(html.contains("Authorization Failed"));
assert!(html.contains("charset"));
assert!(html.contains("IronClaw"));
assert!(html.contains("#ef4444")); // red accent
assert!(!html.contains("Connected"));
}
}
+51 -24
View File
@@ -9,7 +9,7 @@ use crate::settings::Settings;
/// Run the status command, printing system health info.
pub async fn run_status_command() -> anyhow::Result<()> {
let settings = Settings::load();
let settings = Settings::default();
println!("IronClaw Status");
println!("===============\n");
@@ -22,16 +22,36 @@ pub async fn run_status_command() -> anyhow::Result<()> {
);
// Database
let db_url_set = settings.database_url.is_some() || std::env::var("DATABASE_URL").is_ok();
print!(" Database: ");
if db_url_set {
// Try to connect
match check_database().await {
Ok(()) => println!("connected"),
Err(e) => println!("error ({})", e),
let db_backend = std::env::var("DATABASE_BACKEND")
.ok()
.unwrap_or_else(|| "postgres".to_string());
match db_backend.as_str() {
"libsql" | "turso" | "sqlite" => {
let path = std::env::var("LIBSQL_PATH")
.map(std::path::PathBuf::from)
.unwrap_or_else(|_| crate::config::default_libsql_path());
if path.exists() {
let turso = if std::env::var("LIBSQL_URL").is_ok() {
" + Turso sync"
} else {
""
};
println!("libSQL ({}{})", path.display(), turso);
} else {
println!("libSQL (file missing: {})", path.display());
}
}
_ => {
if std::env::var("DATABASE_URL").is_ok() {
match check_database().await {
Ok(()) => println!("connected (PostgreSQL)"),
Err(e) => println!("error ({})", e),
}
} else {
println!("not configured");
}
}
} else {
println!("not configured");
}
// Session / Auth
@@ -43,15 +63,17 @@ pub async fn run_status_command() -> anyhow::Result<()> {
println!("not found (run `ironclaw onboard`)");
}
// Secrets
// Secrets (auto-detect from env only; skip keychain probe to avoid
// triggering macOS system password dialogs on a simple status check)
print!(" Secrets: ");
let secrets_configured = settings.secrets_master_key_source != crate::settings::KeySource::None
|| std::env::var("SECRETS_MASTER_KEY").is_ok()
|| crate::secrets::keychain::has_master_key().await;
if secrets_configured {
println!("configured ({:?})", settings.secrets_master_key_source);
if std::env::var("SECRETS_MASTER_KEY").is_ok() {
println!("configured (env)");
} else {
println!("not configured");
// We don't probe the keychain here because get_generic_password()
// triggers macOS unlock+authorization dialogs, which is bad UX for
// a read-only status command. If onboarding completed with keychain
// storage, the key is there; we just can't cheaply verify it.
println!("env not set (keychain may be configured)");
}
// Embeddings
@@ -129,19 +151,18 @@ pub async fn run_status_command() -> anyhow::Result<()> {
Err(_) => println!("none configured"),
}
// Settings path
println!("\n Settings: {}", Settings::default_path().display());
// Config path
println!(
"\n Config: {}",
crate::bootstrap::ironclaw_env_path().display()
);
Ok(())
}
#[cfg(feature = "postgres")]
async fn check_database() -> anyhow::Result<()> {
let _ = dotenvy::dotenv();
let settings = Settings::load();
let url = std::env::var("DATABASE_URL")
.ok()
.or(settings.database_url)
.ok_or_else(|| anyhow::anyhow!("no URL"))?;
let url = std::env::var("DATABASE_URL").map_err(|_| anyhow::anyhow!("DATABASE_URL not set"))?;
let config: deadpool_postgres::Config = deadpool_postgres::Config {
url: Some(url),
@@ -167,6 +188,12 @@ async fn check_database() -> anyhow::Result<()> {
Ok(())
}
#[cfg(not(feature = "postgres"))]
async fn check_database() -> anyhow::Result<()> {
// For non-postgres backends, just report configured
Ok(())
}
fn count_wasm_files(dir: &std::path::Path) -> usize {
std::fs::read_dir(dir)
.map(|entries| {
+235 -163
View File
@@ -11,8 +11,11 @@ use clap::Subcommand;
use tokio::fs;
use crate::config::Config;
use crate::history::Store;
use crate::secrets::{CreateSecretParams, PostgresSecretsStore, SecretsCrypto, SecretsStore};
#[allow(unused_imports)]
use crate::db::Database;
#[cfg(feature = "postgres")]
use crate::secrets::PostgresSecretsStore;
use crate::secrets::{CreateSecretParams, SecretsCrypto, SecretsStore};
use crate::tools::wasm::{CapabilitiesFile, compute_binary_hash};
/// Default tools directory.
@@ -420,11 +423,11 @@ async fn extract_crate_name(cargo_toml: &Path) -> anyhow::Result<String> {
// Simple TOML parsing for [package] name
for line in content.lines() {
let line = line.trim();
if line.starts_with("name") {
if let Some((_, value)) = line.split_once('=') {
let name = value.trim().trim_matches('"').trim_matches('\'');
return Ok(name.to_string());
}
if line.starts_with("name")
&& let Some((_, value)) = line.split_once('=')
{
let name = value.trim().trim_matches('"').trim_matches('\'');
return Ok(name.to_string());
}
}
@@ -488,10 +491,10 @@ async fn list_tools(dir: Option<PathBuf>, verbose: bool) -> anyhow::Result<()> {
if has_caps {
let caps_path = path.with_extension("capabilities.json");
if let Ok(content) = fs::read_to_string(&caps_path).await {
if let Ok(caps) = CapabilitiesFile::from_json(&content) {
print_capabilities_summary(&caps);
}
if let Ok(content) = fs::read_to_string(&caps_path).await
&& let Ok(caps) = CapabilitiesFile::from_json(&content)
{
print_capabilities_summary(&caps);
}
}
println!();
@@ -604,16 +607,16 @@ fn print_capabilities_summary(caps: &CapabilitiesFile) {
}
}
if let Some(ref secrets) = caps.secrets {
if !secrets.allowed_names.is_empty() {
parts.push(format!("secrets: {}", secrets.allowed_names.len()));
}
if let Some(ref secrets) = caps.secrets
&& !secrets.allowed_names.is_empty()
{
parts.push(format!("secrets: {}", secrets.allowed_names.len()));
}
if let Some(ref ws) = caps.workspace {
if !ws.allowed_prefixes.is_empty() {
parts.push("workspace: read".to_string());
}
if let Some(ref ws) = caps.workspace
&& !ws.allowed_prefixes.is_empty()
{
parts.push("workspace: read".to_string());
}
if !parts.is_empty() {
@@ -650,30 +653,30 @@ fn print_capabilities_detail(caps: &CapabilitiesFile) {
}
}
if let Some(ref secrets) = caps.secrets {
if !secrets.allowed_names.is_empty() {
println!(" Secrets (existence check only):");
for name in &secrets.allowed_names {
println!(" {}", name);
}
if let Some(ref secrets) = caps.secrets
&& !secrets.allowed_names.is_empty()
{
println!(" Secrets (existence check only):");
for name in &secrets.allowed_names {
println!(" {}", name);
}
}
if let Some(ref tool_invoke) = caps.tool_invoke {
if !tool_invoke.aliases.is_empty() {
println!(" Tool aliases:");
for (alias, real_name) in &tool_invoke.aliases {
println!(" {} -> {}", alias, real_name);
}
if let Some(ref tool_invoke) = caps.tool_invoke
&& !tool_invoke.aliases.is_empty()
{
println!(" Tool aliases:");
for (alias, real_name) in &tool_invoke.aliases {
println!(" {} -> {}", alias, real_name);
}
}
if let Some(ref ws) = caps.workspace {
if !ws.allowed_prefixes.is_empty() {
println!(" Workspace read prefixes:");
for prefix in &ws.allowed_prefixes {
println!(" {}", prefix);
}
if let Some(ref ws) = caps.workspace
&& !ws.allowed_prefixes.is_empty()
{
println!(" Workspace read prefixes:");
for prefix in &ws.allowed_prefixes {
println!(" {}", prefix);
}
}
}
@@ -722,11 +725,58 @@ async fn auth_tool(name: String, dir: Option<PathBuf>, user_id: String) -> anyho
)
})?;
let store = Store::new(&config.database).await?;
store.run_migrations().await?;
let crypto = SecretsCrypto::new(master_key.clone())?;
let secrets_store = Arc::new(PostgresSecretsStore::new(store.pool(), Arc::new(crypto)));
let secrets_store: Arc<dyn SecretsStore + Send + Sync> = {
#[cfg(feature = "postgres")]
{
let store = crate::history::Store::new(&config.database).await?;
store.run_migrations().await?;
Arc::new(PostgresSecretsStore::new(store.pool(), Arc::new(crypto)))
}
#[cfg(all(feature = "libsql", not(feature = "postgres")))]
{
use crate::db::Database as _;
use crate::db::libsql_backend::LibSqlBackend;
use secrecy::ExposeSecret as _;
let default_path = crate::config::default_libsql_path();
let db_path = config
.database
.libsql_path
.as_deref()
.unwrap_or(&default_path);
let backend = if let Some(ref url) = config.database.libsql_url {
let token = config.database.libsql_auth_token.as_ref().ok_or_else(|| {
anyhow::anyhow!("LIBSQL_AUTH_TOKEN is required when LIBSQL_URL is set")
})?;
LibSqlBackend::new_remote_replica(db_path, url, token.expose_secret())
.await
.map_err(|e| anyhow::anyhow!("{}", e))?
} else {
LibSqlBackend::new_local(db_path)
.await
.map_err(|e| anyhow::anyhow!("{}", e))?
};
backend
.run_migrations()
.await
.map_err(|e| anyhow::anyhow!("{}", e))?;
Arc::new(crate::secrets::LibSqlSecretsStore::new(
backend.shared_db(),
Arc::new(crypto),
))
}
#[cfg(not(any(feature = "postgres", feature = "libsql")))]
{
let _ = crypto;
anyhow::bail!(
"No database backend available for secrets. Enable 'postgres' or 'libsql' feature."
);
}
};
// Check if already configured
let already_configured = secrets_store
@@ -752,51 +802,103 @@ async fn auth_tool(name: String, dir: Option<PathBuf>, user_id: String) -> anyho
}
// Check for environment variable
if let Some(ref env_var) = auth.env_var {
if let Ok(token) = std::env::var(env_var) {
if !token.is_empty() {
println!(" Found {} in environment.", env_var);
println!();
if let Some(ref env_var) = auth.env_var
&& let Ok(token) = std::env::var(env_var)
&& !token.is_empty()
{
println!(" Found {} in environment.", env_var);
println!();
// Validate if endpoint is provided
if let Some(ref validation) = auth.validation_endpoint {
print!(" Validating token...");
std::io::stdout().flush()?;
// Validate if endpoint is provided
if let Some(ref validation) = auth.validation_endpoint {
print!(" Validating token...");
std::io::stdout().flush()?;
match validate_token(&token, validation, &auth.secret_name).await {
Ok(()) => {
println!("");
}
Err(e) => {
println!("");
println!(" Validation failed: {}", e);
println!();
println!(" Falling back to manual entry...");
return auth_tool_manual(&secrets_store, &user_id, &auth).await;
}
}
match validate_token(&token, validation, &auth.secret_name).await {
Ok(()) => {
println!("");
}
Err(e) => {
println!("");
println!(" Validation failed: {}", e);
println!();
println!(" Falling back to manual entry...");
return auth_tool_manual(secrets_store.as_ref(), &user_id, &auth).await;
}
// Save the token
save_token(&secrets_store, &user_id, &auth, &token).await?;
print_success(display_name);
return Ok(());
}
}
// Save the token
save_token(secrets_store.as_ref(), &user_id, &auth, &token, None, None).await?;
print_success(display_name);
return Ok(());
}
// Check for OAuth configuration
if let Some(ref oauth) = auth.oauth {
return auth_tool_oauth(&secrets_store, &user_id, &auth, oauth).await;
// For providers with shared tokens (e.g., all Google tools share google_oauth_token),
// combine scopes from all installed tools so one auth covers everything.
let combined = combine_provider_scopes(&tools_dir, &auth.secret_name, oauth).await;
if combined.scopes.len() > oauth.scopes.len() {
let extra = combined.scopes.len() - oauth.scopes.len();
println!(
" Including scopes from {} other installed tool(s) sharing this credential.",
extra
);
println!();
}
return auth_tool_oauth(secrets_store.as_ref(), &user_id, &auth, &combined).await;
}
// Fall back to manual entry
auth_tool_manual(&secrets_store, &user_id, &auth).await
auth_tool_manual(secrets_store.as_ref(), &user_id, &auth).await
}
/// Scan the tools directory for all capabilities files sharing the same secret_name
/// and combine their OAuth scopes. This way, authing any Google tool requests scopes
/// for ALL installed Google tools, so one login covers everything.
async fn combine_provider_scopes(
tools_dir: &Path,
secret_name: &str,
base_oauth: &crate::tools::wasm::OAuthConfigSchema,
) -> crate::tools::wasm::OAuthConfigSchema {
let mut all_scopes: std::collections::HashSet<String> =
base_oauth.scopes.iter().cloned().collect();
if let Ok(mut entries) = tokio::fs::read_dir(tools_dir).await {
while let Ok(Some(entry)) = entries.next_entry().await {
let path = entry.path();
if path.extension().and_then(|e| e.to_str()) != Some("json") {
continue;
}
let name = path
.file_name()
.and_then(|n| n.to_str())
.unwrap_or_default();
if !name.ends_with(".capabilities.json") {
continue;
}
if let Ok(content) = tokio::fs::read_to_string(&path).await
&& let Ok(caps) = CapabilitiesFile::from_json(&content)
&& let Some(auth) = &caps.auth
&& auth.secret_name == secret_name
&& let Some(oauth) = &auth.oauth
{
all_scopes.extend(oauth.scopes.iter().cloned());
}
}
}
let mut combined = base_oauth.clone();
combined.scopes = all_scopes.into_iter().collect();
combined.scopes.sort(); // deterministic ordering
combined
}
/// OAuth browser-based login flow.
async fn auth_tool_oauth(
store: &PostgresSecretsStore,
store: &(dyn SecretsStore + Send + Sync),
user_id: &str,
auth: &crate::tools::wasm::AuthCapabilitySchema,
oauth: &crate::tools::wasm::OAuthConfigSchema,
@@ -804,12 +906,14 @@ async fn auth_tool_oauth(
use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
use rand::RngCore;
use sha2::{Digest, Sha256};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::TcpListener;
use crate::cli::oauth_defaults::{self, OAUTH_CALLBACK_PORT};
let display_name = auth.display_name.as_deref().unwrap_or(&auth.secret_name);
// Get client_id from config or env
// Get client_id: capabilities file > runtime env var > built-in defaults
let builtin = oauth_defaults::builtin_credentials(&auth.secret_name);
let client_id = oauth
.client_id
.clone()
@@ -819,41 +923,32 @@ async fn auth_tool_oauth(
.as_ref()
.and_then(|env| std::env::var(env).ok())
})
.or_else(|| builtin.as_ref().map(|c| c.client_id.to_string()))
.ok_or_else(|| {
anyhow::anyhow!(
"OAuth client_id not configured.\n\
Set it in the capabilities file or via environment variable."
Set {} env var, or build with IRONCLAW_GOOGLE_CLIENT_ID.",
oauth.client_id_env.as_deref().unwrap_or("the client_id")
)
})?;
// Get client_secret if provided
let client_secret = oauth.client_secret.clone().or_else(|| {
oauth
.client_secret_env
.as_ref()
.and_then(|env| std::env::var(env).ok())
});
// Get client_secret: capabilities file > runtime env var > built-in defaults
let client_secret = oauth
.client_secret
.clone()
.or_else(|| {
oauth
.client_secret_env
.as_ref()
.and_then(|env| std::env::var(env).ok())
})
.or_else(|| builtin.as_ref().map(|c| c.client_secret.to_string()));
println!(" Starting OAuth authentication...");
println!();
// Find an available port for the callback
let mut listener = None;
let mut port = 0;
for p in 9876..=9886 {
match TcpListener::bind(format!("127.0.0.1:{}", p)).await {
Ok(l) => {
listener = Some(l);
port = p;
break;
}
Err(_) => continue,
}
}
let listener = listener.ok_or_else(|| anyhow::anyhow!("Could not find available port"))?;
let redirect_uri = format!("http://localhost:{}/callback", port);
let listener = oauth_defaults::bind_callback_listener().await?;
let redirect_uri = format!("http://localhost:{}/callback", OAUTH_CALLBACK_PORT);
// Generate PKCE verifier and challenge
let (code_verifier, code_challenge) = if oauth.use_pkce {
@@ -912,65 +1007,8 @@ async fn auth_tool_oauth(
println!(" Waiting for authorization...");
// Wait for callback with timeout
let timeout = std::time::Duration::from_secs(300);
let code = tokio::time::timeout(timeout, async {
loop {
let (mut socket, _) = listener.accept().await?;
let mut reader = BufReader::new(&mut socket);
let mut request_line = String::new();
reader.read_line(&mut request_line).await?;
// Parse GET /callback?code=xxx HTTP/1.1
if let Some(path) = request_line.split_whitespace().nth(1) {
if path.starts_with("/callback") {
if let Some(query) = path.split('?').nth(1) {
for param in query.split('&') {
let parts: Vec<&str> = param.splitn(2, '=').collect();
if parts.len() == 2 && parts[0] == "code" {
let code = urlencoding::decode(parts[1])
.unwrap_or_else(|_| parts[1].into())
.into_owned();
// Send success response
let response = format!(
"HTTP/1.1 200 OK\r\n\
Content-Type: text/html\r\n\
\r\n\
<!DOCTYPE html><html><body style=\"font-family: sans-serif; \
display: flex; justify-content: center; align-items: center; \
height: 100vh; margin: 0; background: #191919; color: white;\">\
<div style=\"text-align: center;\">\
<h1> {} Connected!</h1>\
<p>You can close this window.</p>\
</div></body></html>",
display_name
);
let _ = socket.write_all(response.as_bytes()).await;
let _ = socket.shutdown().await;
return Ok::<_, anyhow::Error>(code);
}
}
// Check for error
if query.contains("error=") {
let response =
"HTTP/1.1 400 Bad Request\r\n\r\nAuthorization denied";
let _ = socket.write_all(response.as_bytes()).await;
return Err(anyhow::anyhow!("Authorization denied by user"));
}
}
}
}
let response = "HTTP/1.1 404 Not Found\r\n\r\n";
let _ = socket.write_all(response.as_bytes()).await;
}
})
.await
.map_err(|_| anyhow::anyhow!("Timed out waiting for authorization"))??;
let code =
oauth_defaults::wait_for_callback(listener, "/callback", "code", display_name).await?;
println!();
println!(" Exchanging code for token...");
@@ -1021,8 +1059,19 @@ async fn auth_tool_oauth(
)
})?;
// Save the token
save_token(store, user_id, auth, access_token).await?;
let refresh_token = token_data.get("refresh_token").and_then(|v| v.as_str());
let expires_in = token_data.get("expires_in").and_then(|v| v.as_u64());
// Save the token (with refresh token and expiry if provided)
save_token(
store,
user_id,
auth,
access_token,
refresh_token,
expires_in,
)
.await?;
// Extract any additional info for display
let workspace_name = token_data
@@ -1044,7 +1093,7 @@ async fn auth_tool_oauth(
/// Manual token entry flow.
async fn auth_tool_manual(
store: &PostgresSecretsStore,
store: &(dyn SecretsStore + Send + Sync),
user_id: &str,
auth: &crate::tools::wasm::AuthCapabilitySchema,
) -> anyhow::Result<()> {
@@ -1124,8 +1173,8 @@ async fn auth_tool_manual(
}
}
// Save the token
save_token(store, user_id, auth, &token).await?;
// Save the token (manual path: no refresh token or expiry)
save_token(store, user_id, auth, &token, None, None).await?;
print_success(display_name);
Ok(())
}
@@ -1216,11 +1265,16 @@ async fn validate_token(
}
/// Save token to secrets store.
///
/// Optionally stores a refresh token (as `{secret_name}_refresh_token`) and
/// sets `expires_at` on the access token so the runtime can auto-refresh.
async fn save_token(
store: &PostgresSecretsStore,
store: &(dyn SecretsStore + Send + Sync),
user_id: &str,
auth: &crate::tools::wasm::AuthCapabilitySchema,
token: &str,
refresh_token: Option<&str>,
expires_in: Option<u64>,
) -> anyhow::Result<()> {
let mut params = CreateSecretParams::new(&auth.secret_name, token);
@@ -1228,11 +1282,29 @@ async fn save_token(
params = params.with_provider(provider);
}
if let Some(secs) = expires_in {
let expires_at = chrono::Utc::now() + chrono::Duration::seconds(secs as i64);
params = params.with_expiry(expires_at);
}
store
.create(user_id, params)
.await
.map_err(|e| anyhow::anyhow!("Failed to save token: {}", e))?;
// Store refresh token separately (no expiry, it's long-lived)
if let Some(rt) = refresh_token {
let refresh_name = format!("{}_refresh_token", auth.secret_name);
let mut refresh_params = CreateSecretParams::new(&refresh_name, rt);
if let Some(ref provider) = auth.provider {
refresh_params = refresh_params.with_provider(provider);
}
store
.create(user_id, refresh_params)
.await
.map_err(|e| anyhow::anyhow!("Failed to save refresh token: {}", e))?;
}
Ok(())
}
+274 -76
View File
@@ -1,11 +1,13 @@
//! 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`.
//! `DATABASE_URL` lives in `~/.ironclaw/.env` (loaded via dotenvy early
//! in startup). Everything else comes from env vars, the DB settings
//! table, or auto-detection.
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::OnceLock;
use std::time::Duration;
use secrecy::{ExposeSecret, SecretString};
@@ -13,6 +15,13 @@ use secrecy::{ExposeSecret, SecretString};
use crate::error::ConfigError;
use crate::settings::Settings;
/// Thread-safe overlay for injected env vars (secrets loaded from DB).
///
/// Used by `inject_llm_keys_from_secrets()` to make API keys available to
/// `optional_env()` without unsafe `set_var` calls. `optional_env()` checks
/// real env vars first, then falls back to this overlay.
static INJECTED_VARS: OnceLock<HashMap<String, String>> = OnceLock::new();
/// Main configuration for the agent.
#[derive(Debug, Clone)]
pub struct Config {
@@ -38,11 +47,11 @@ impl Config {
/// 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,
store: &dyn crate::db::Database,
user_id: &str,
bootstrap: &crate::bootstrap::BootstrapConfig,
) -> Result<Self, ConfigError> {
let _ = dotenvy::dotenv();
crate::bootstrap::load_ironclaw_env();
// Load all settings from DB into a Settings struct
let db_settings = match store.get_all_settings(user_id).await {
@@ -53,7 +62,7 @@ impl Config {
}
};
Self::build(bootstrap, &db_settings).await
Self::build(&db_settings).await
}
/// Load configuration from environment variables only (no database).
@@ -61,20 +70,20 @@ impl Config {
/// 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.
///
/// Loads both `./.env` (standard, higher priority) and `~/.ironclaw/.env`
/// (lower priority) via dotenvy, which never overwrites existing vars.
pub async fn from_env() -> Result<Self, ConfigError> {
let _ = dotenvy::dotenv();
let bootstrap = crate::bootstrap::BootstrapConfig::load();
crate::bootstrap::load_ironclaw_env();
let settings = Settings::load();
Self::build(&bootstrap, &settings).await
Self::build(&settings).await
}
/// Build config from bootstrap + settings (shared by from_env and from_db).
async fn build(
bootstrap: &crate::bootstrap::BootstrapConfig,
settings: &Settings,
) -> Result<Self, ConfigError> {
/// Build config from settings (shared by from_env and from_db).
async fn build(settings: &Settings) -> Result<Self, ConfigError> {
Ok(Self {
database: DatabaseConfig::resolve(bootstrap)?,
database: DatabaseConfig::resolve()?,
llm: LlmConfig::resolve(settings)?,
embeddings: EmbeddingsConfig::resolve(settings)?,
tunnel: TunnelConfig::resolve(settings)?,
@@ -82,7 +91,7 @@ impl Config {
agent: AgentConfig::resolve(settings)?,
safety: SafetyConfig::resolve()?,
wasm: WasmConfig::resolve()?,
secrets: SecretsConfig::resolve(bootstrap).await?,
secrets: SecretsConfig::resolve().await?,
builder: BuilderModeConfig::resolve()?,
heartbeat: HeartbeatConfig::resolve(settings)?,
routines: RoutineConfig::resolve()?,
@@ -107,13 +116,13 @@ impl TunnelConfig {
let public_url = optional_env("TUNNEL_URL")?
.or_else(|| settings.tunnel.public_url.clone().filter(|s| !s.is_empty()));
if let Some(ref url) = public_url {
if !url.starts_with("https://") {
return Err(ConfigError::InvalidValue {
key: "TUNNEL_URL".to_string(),
message: "must start with https:// (webhooks require HTTPS)".to_string(),
});
}
if let Some(ref url) = public_url
&& !url.starts_with("https://")
{
return Err(ConfigError::InvalidValue {
key: "TUNNEL_URL".to_string(),
message: "must start with https:// (webhooks require HTTPS)".to_string(),
});
}
Ok(Self { public_url })
@@ -134,35 +143,104 @@ impl TunnelConfig {
}
}
/// Which database backend to use.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum DatabaseBackend {
/// PostgreSQL via deadpool-postgres (default).
#[default]
Postgres,
/// libSQL/Turso embedded database.
LibSql,
}
impl std::str::FromStr for DatabaseBackend {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.to_lowercase().as_str() {
"postgres" | "postgresql" | "pg" => Ok(Self::Postgres),
"libsql" | "turso" | "sqlite" => Ok(Self::LibSql),
_ => Err(format!(
"invalid database backend '{}', expected 'postgres' or 'libsql'",
s
)),
}
}
}
/// Database configuration.
#[derive(Debug, Clone)]
pub struct DatabaseConfig {
/// Which backend to use (default: Postgres).
pub backend: DatabaseBackend,
// -- PostgreSQL fields --
pub url: SecretString,
pub pool_size: usize,
// -- libSQL fields --
/// Path to local libSQL database file (default: ~/.ironclaw/ironclaw.db).
pub libsql_path: Option<PathBuf>,
/// Turso cloud URL for remote sync (optional).
pub libsql_url: Option<String>,
/// Turso auth token (required when libsql_url is set).
pub libsql_auth_token: Option<SecretString>,
}
impl DatabaseConfig {
fn resolve(bootstrap: &crate::bootstrap::BootstrapConfig) -> Result<Self, ConfigError> {
fn resolve() -> Result<Self, ConfigError> {
let backend: DatabaseBackend = if let Some(b) = optional_env("DATABASE_BACKEND")? {
b.parse().map_err(|e| ConfigError::InvalidValue {
key: "DATABASE_BACKEND".to_string(),
message: e,
})?
} else {
DatabaseBackend::default()
};
// PostgreSQL URL is required only when using the postgres backend.
// For libsql backend, default to an empty placeholder.
// DATABASE_URL is loaded from ~/.ironclaw/.env via dotenvy early in startup.
let url = optional_env("DATABASE_URL")?
.or_else(|| bootstrap.database_url.clone())
.or_else(|| {
if backend == DatabaseBackend::LibSql {
Some("unused://libsql".to_string())
} else {
None
}
})
.ok_or_else(|| ConfigError::MissingRequired {
key: "database_url".to_string(),
hint: "Run 'ironclaw onboard' or set DATABASE_URL environment variable".to_string(),
})?;
let pool_size = optional_env("DATABASE_POOL_SIZE")?
.map(|s| s.parse())
.transpose()
.map_err(|e| ConfigError::InvalidValue {
key: "DATABASE_POOL_SIZE".to_string(),
message: format!("must be a positive integer: {e}"),
})?
.or(bootstrap.database_pool_size)
.unwrap_or(10);
let pool_size = parse_optional_env("DATABASE_POOL_SIZE", 10)?;
let libsql_path = optional_env("LIBSQL_PATH")?.map(PathBuf::from).or_else(|| {
if backend == DatabaseBackend::LibSql {
Some(default_libsql_path())
} else {
None
}
});
let libsql_url = optional_env("LIBSQL_URL")?;
let libsql_auth_token = optional_env("LIBSQL_AUTH_TOKEN")?.map(SecretString::from);
if libsql_url.is_some() && libsql_auth_token.is_none() {
return Err(ConfigError::MissingRequired {
key: "LIBSQL_AUTH_TOKEN".to_string(),
hint: "LIBSQL_AUTH_TOKEN is required when LIBSQL_URL is set".to_string(),
});
}
Ok(Self {
backend,
url: SecretString::from(url),
pool_size,
libsql_path,
libsql_url,
libsql_auth_token,
})
}
@@ -172,6 +250,14 @@ impl DatabaseConfig {
}
}
/// Default libSQL database path (~/.ironclaw/ironclaw.db).
pub fn default_libsql_path() -> PathBuf {
dirs::home_dir()
.unwrap_or_else(|| PathBuf::from("."))
.join(".ironclaw")
.join("ironclaw.db")
}
/// Which LLM backend to use.
///
/// Defaults to `NearAi` to keep IronClaw close to the NEAR ecosystem.
@@ -312,16 +398,37 @@ pub struct NearAiConfig {
pub api_mode: NearAiApiMode,
/// API key for cloud-api (required for chat_completions mode)
pub api_key: Option<SecretString>,
/// Optional fallback model for failover (default: None).
/// When set, a secondary provider is created with this model and wrapped
/// in a `FailoverProvider` so transient errors on the primary model
/// automatically fall through to the fallback.
pub fallback_model: Option<String>,
/// Maximum number of retries for transient errors (default: 3).
/// With the default of 3, the provider makes up to 4 total attempts
/// (1 initial + 3 retries) before giving up.
pub max_retries: u32,
}
impl LlmConfig {
fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
// Determine backend (default: NearAi)
// Determine backend: env var > settings > default (NearAi)
let backend: LlmBackend = if let Some(b) = optional_env("LLM_BACKEND")? {
b.parse().map_err(|e| ConfigError::InvalidValue {
key: "LLM_BACKEND".to_string(),
message: e,
})?
} else if let Some(ref b) = settings.llm_backend {
match b.parse() {
Ok(backend) => backend,
Err(e) => {
tracing::warn!(
"Invalid llm_backend '{}' in settings: {}. Using default NearAi.",
b,
e
);
LlmBackend::NearAi
}
}
} else {
LlmBackend::NearAi
};
@@ -356,6 +463,8 @@ impl LlmConfig {
.unwrap_or_else(default_session_path),
api_mode,
api_key: nearai_api_key,
fallback_model: optional_env("NEARAI_FALLBACK_MODEL")?,
max_retries: parse_optional_env("NEARAI_MAX_RETRIES", 3)?,
};
// Resolve provider-specific configs based on backend
@@ -388,6 +497,7 @@ impl LlmConfig {
let ollama = if backend == LlmBackend::Ollama {
let base_url = optional_env("OLLAMA_BASE_URL")?
.or_else(|| settings.ollama_base_url.clone())
.unwrap_or_else(|| "http://localhost:11434".to_string());
let model = optional_env("OLLAMA_MODEL")?.unwrap_or_else(|| "llama3".to_string());
Some(OllamaConfig { base_url, model })
@@ -396,8 +506,9 @@ impl LlmConfig {
};
let openai_compatible = if backend == LlmBackend::OpenAiCompatible {
let base_url =
optional_env("LLM_BASE_URL")?.ok_or_else(|| ConfigError::MissingRequired {
let base_url = optional_env("LLM_BASE_URL")?
.or_else(|| settings.openai_compatible_base_url.clone())
.ok_or_else(|| ConfigError::MissingRequired {
key: "LLM_BASE_URL".to_string(),
hint: "Set LLM_BASE_URL when LLM_BACKEND=openai_compatible".to_string(),
})?;
@@ -767,52 +878,41 @@ impl std::fmt::Debug for SecretsConfig {
}
}
/// Process-wide cache for the keychain master key.
///
/// Avoids re-prompting the OS keychain on every `SecretsConfig::resolve()` call
/// (e.g. `Config::from_env()` then `Config::from_db()`). Thread-safe alternative
/// to caching in a process env var.
impl SecretsConfig {
async fn resolve(bootstrap: &crate::bootstrap::BootstrapConfig) -> Result<Self, ConfigError> {
/// Auto-detect secrets master key from env var, then OS keychain.
///
/// Sequential probe: SECRETS_MASTER_KEY env var first, then OS keychain.
/// No saved "source" needed; just try each source in order.
async fn resolve() -> Result<Self, ConfigError> {
use crate::settings::KeySource;
let (master_key, source) = if let Some(env_key) = optional_env("SECRETS_MASTER_KEY")? {
(Some(SecretString::from(env_key)), KeySource::Env)
} else {
match bootstrap.secrets_master_key_source {
KeySource::Keychain => {
// Try to load from OS keychain (async on Linux)
match crate::secrets::keychain::get_master_key().await {
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)
}
}
// Probe the OS keychain; if a key is stored, use it
match crate::secrets::keychain::get_master_key().await {
Ok(key_bytes) => {
let key_hex: String = key_bytes.iter().map(|b| format!("{:02x}", b)).collect();
(Some(SecretString::from(key_hex)), KeySource::Keychain)
}
KeySource::Env => {
tracing::warn!(
"Secrets configured for env var but SECRETS_MASTER_KEY not set."
);
(None, KeySource::None)
}
KeySource::None => (None, KeySource::None),
Err(_) => (None, KeySource::None),
}
};
let enabled = master_key.is_some();
if let Some(ref key) = master_key {
if key.expose_secret().len() < 32 {
return Err(ConfigError::InvalidValue {
key: "SECRETS_MASTER_KEY".to_string(),
message: "must be at least 32 bytes for AES-256-GCM".to_string(),
});
}
if let Some(ref key) = master_key
&& key.expose_secret().len() < 32
{
return Err(ConfigError::InvalidValue {
key: "SECRETS_MASTER_KEY".to_string(),
message: "must be at least 32 bytes for AES-256-GCM".to_string(),
});
}
Ok(Self {
@@ -1174,6 +1274,36 @@ pub struct ClaudeCodeConfig {
pub max_turns: u32,
/// Memory limit in MB for Claude Code containers (heavier than workers).
pub memory_limit_mb: u64,
/// Allowed tool patterns for Claude Code permission settings.
///
/// Written to `/workspace/.claude/settings.json` before spawning the CLI.
/// Provides defense-in-depth: only explicitly listed tools are auto-approved.
/// Any new/unknown tools would require interactive approval (which times out
/// in the non-interactive container, failing safely).
///
/// Patterns follow Claude Code syntax: `"Bash(*)"`, `"Read"`, `"Edit(*)"`, etc.
pub allowed_tools: Vec<String>,
}
/// Default allowed tools for Claude Code inside containers.
///
/// These cover all standard Claude Code tools needed for autonomous operation.
/// The Docker container provides the primary security boundary; this allowlist
/// provides defense-in-depth by preventing any future unknown tools from being
/// silently auto-approved.
fn default_claude_code_allowed_tools() -> Vec<String> {
[
"Bash(*)",
"Read",
"Edit(*)",
"Glob",
"Grep",
"WebFetch(*)",
"Task(*)",
]
.into_iter()
.map(String::from)
.collect()
}
impl Default for ClaudeCodeConfig {
@@ -1186,11 +1316,24 @@ impl Default for ClaudeCodeConfig {
model: "sonnet".to_string(),
max_turns: 50,
memory_limit_mb: 4096,
allowed_tools: default_claude_code_allowed_tools(),
}
}
}
impl ClaudeCodeConfig {
/// Load from environment variables only (used inside containers where
/// there is no database or full config).
pub fn from_env() -> Self {
match Self::resolve() {
Ok(c) => c,
Err(e) => {
tracing::warn!("Failed to resolve ClaudeCodeConfig: {e}, using defaults");
Self::default()
}
}
}
fn resolve() -> Result<Self, ConfigError> {
let defaults = Self::default();
Ok(Self {
@@ -1211,21 +1354,76 @@ impl ClaudeCodeConfig {
"CLAUDE_CODE_MEMORY_LIMIT_MB",
defaults.memory_limit_mb,
)?,
allowed_tools: optional_env("CLAUDE_CODE_ALLOWED_TOOLS")?
.map(|s| {
s.split(',')
.map(|t| t.trim().to_string())
.filter(|t| !t.is_empty())
.collect()
})
.unwrap_or(defaults.allowed_tools),
})
}
}
/// Load API keys from the encrypted secrets store into a thread-safe overlay.
///
/// This bridges the gap between secrets stored during onboarding and the
/// env-var-first resolution in `LlmConfig::resolve()`. Keys in the overlay
/// are read by `optional_env()` before falling back to `std::env::var()`,
/// so explicit env vars always win.
pub async fn inject_llm_keys_from_secrets(
secrets: &dyn crate::secrets::SecretsStore,
user_id: &str,
) {
let mappings = [
("llm_openai_api_key", "OPENAI_API_KEY"),
("llm_anthropic_api_key", "ANTHROPIC_API_KEY"),
("llm_compatible_api_key", "LLM_API_KEY"),
];
let mut injected = HashMap::new();
for (secret_name, env_var) in mappings {
match std::env::var(env_var) {
Ok(val) if !val.is_empty() => continue,
_ => {}
}
match secrets.get_decrypted(user_id, secret_name).await {
Ok(decrypted) => {
injected.insert(env_var.to_string(), decrypted.expose().to_string());
tracing::debug!("Loaded secret '{}' for env var '{}'", secret_name, env_var);
}
Err(_) => {
// Secret doesn't exist, that's fine
}
}
}
let _ = INJECTED_VARS.set(injected);
}
// Helper functions
fn optional_env(key: &str) -> Result<Option<String>, ConfigError> {
// Check real env vars first (always win over injected secrets)
match std::env::var(key) {
Ok(val) if val.is_empty() => Ok(None),
Ok(val) => Ok(Some(val)),
Err(std::env::VarError::NotPresent) => Ok(None),
Err(e) => Err(ConfigError::ParseError(format!(
"failed to read {key}: {e}"
))),
Ok(val) if val.is_empty() => {}
Ok(val) => return Ok(Some(val)),
Err(std::env::VarError::NotPresent) => {}
Err(e) => {
return Err(ConfigError::ParseError(format!(
"failed to read {key}: {e}"
)));
}
}
// Fall back to thread-safe overlay (secrets injected from DB)
if let Some(val) = INJECTED_VARS.get().and_then(|map| map.get(key)) {
return Ok(Some(val.clone()));
}
Ok(None)
}
fn parse_optional_env<T>(key: &str, default: T) -> Result<T, ConfigError>
+5 -4
View File
@@ -45,20 +45,21 @@ impl ContextManager {
title: impl Into<String>,
description: impl Into<String>,
) -> Result<Uuid, JobError> {
let contexts = self.contexts.read().await;
// Hold write lock for the entire check-insert to prevent TOCTOU races
// where two concurrent calls both pass the active_count check.
let mut contexts = self.contexts.write().await;
let active_count = contexts.values().filter(|c| c.state.is_active()).count();
if active_count >= self.max_jobs {
return Err(JobError::MaxJobsExceeded { max: self.max_jobs });
}
drop(contexts);
let context = JobContext::with_user(user_id, title, description);
let job_id = context.job_id;
contexts.insert(job_id, context);
drop(contexts);
let memory = Memory::new(job_id);
self.contexts.write().await.insert(job_id, context);
self.memories.write().await.insert(job_id, memory);
Ok(job_id)
+88
View File
@@ -119,6 +119,10 @@ pub struct JobContext {
pub estimated_duration: Option<Duration>,
/// Actual cost so far.
pub actual_cost: Decimal,
/// Total tokens consumed by LLM calls in this job.
pub total_tokens_used: u64,
/// Maximum tokens allowed per job (0 = unlimited).
pub max_tokens: u64,
/// When the job was created.
pub created_at: DateTime<Utc>,
/// When the job was started.
@@ -159,6 +163,8 @@ impl JobContext {
estimated_cost: None,
estimated_duration: None,
actual_cost: Decimal::ZERO,
total_tokens_used: 0,
max_tokens: 0,
created_at: Utc::now(),
started_at: None,
completed_at: None,
@@ -189,6 +195,14 @@ impl JobContext {
};
self.transitions.push(transition);
// Cap transition history to prevent unbounded memory growth
const MAX_TRANSITIONS: usize = 200;
if self.transitions.len() > MAX_TRANSITIONS {
let drain_count = self.transitions.len() - MAX_TRANSITIONS;
self.transitions.drain(..drain_count);
}
self.state = new_state;
// Update timestamps
@@ -210,6 +224,29 @@ impl JobContext {
self.actual_cost += cost;
}
/// Record token usage from an LLM call. Returns an error string if the
/// token budget has been exceeded after this addition.
pub fn add_tokens(&mut self, tokens: u64) -> Result<(), String> {
self.total_tokens_used += tokens;
if self.max_tokens > 0 && self.total_tokens_used > self.max_tokens {
Err(format!(
"Token budget exceeded: used {} of {} allowed tokens",
self.total_tokens_used, self.max_tokens
))
} else {
Ok(())
}
}
/// Check whether the monetary budget has been exceeded.
pub fn budget_exceeded(&self) -> bool {
if let Some(ref budget) = self.budget {
self.actual_cost > *budget
} else {
false
}
}
/// Get the duration since the job started.
pub fn elapsed(&self) -> Option<Duration> {
self.started_at.map(|start| {
@@ -274,6 +311,57 @@ mod tests {
assert_eq!(ctx.state, JobState::Completed);
}
#[test]
fn test_transition_history_capped() {
let mut ctx = JobContext::new("Test", "Transition cap test");
// Cycle through Pending -> InProgress -> Stuck -> InProgress -> Stuck ...
ctx.transition_to(JobState::InProgress, None).unwrap();
for i in 0..250 {
ctx.mark_stuck(format!("stuck {}", i)).unwrap();
ctx.attempt_recovery().unwrap();
}
// 1 initial + 250*2 = 501 transitions, should be capped at 200
assert!(
ctx.transitions.len() <= 200,
"transitions should be capped at 200, got {}",
ctx.transitions.len()
);
}
#[test]
fn test_add_tokens_enforces_budget() {
let mut ctx = JobContext::new("Test", "Budget test");
ctx.max_tokens = 1000;
assert!(ctx.add_tokens(500).is_ok());
assert_eq!(ctx.total_tokens_used, 500);
assert!(ctx.add_tokens(600).is_err());
assert_eq!(ctx.total_tokens_used, 1100); // tokens still recorded
}
#[test]
fn test_add_tokens_unlimited() {
let mut ctx = JobContext::new("Test", "No budget");
// max_tokens = 0 means unlimited
assert!(ctx.add_tokens(1_000_000).is_ok());
}
#[test]
fn test_budget_exceeded() {
let mut ctx = JobContext::new("Test", "Money test");
ctx.budget = Some(Decimal::new(100, 0)); // $100
assert!(!ctx.budget_exceeded());
ctx.add_cost(Decimal::new(50, 0));
assert!(!ctx.budget_exceeded());
ctx.add_cost(Decimal::new(60, 0));
assert!(ctx.budget_exceeded());
}
#[test]
fn test_budget_exceeded_none() {
let ctx = JobContext::new("Test", "No budget");
assert!(!ctx.budget_exceeded()); // No budget = never exceeded
}
#[test]
fn test_stuck_recovery() {
let mut ctx = JobContext::new("Test", "Test job");
File diff suppressed because it is too large Load Diff
+549
View File
@@ -0,0 +1,549 @@
//! SQLite-dialect migrations for the libSQL/Turso backend.
//!
//! Consolidates all PostgreSQL migrations (V1-V8) into a single SQLite-compatible
//! schema. Run once on database creation; idempotent via `IF NOT EXISTS`.
/// Consolidated schema for libSQL.
///
/// Translates PostgreSQL types and features:
/// - `UUID` -> `TEXT` (store as hex string)
/// - `TIMESTAMPTZ` -> `TEXT` (ISO-8601)
/// - `JSONB` -> `TEXT` (JSON encoded)
/// - `BYTEA` -> `BLOB`
/// - `NUMERIC` -> `TEXT` (preserve precision for rust_decimal)
/// - `TEXT[]` -> `TEXT` (JSON array)
/// - `VECTOR(1536)` -> `F32_BLOB(1536)` (libsql native)
/// - `TSVECTOR` -> FTS5 virtual table
/// - `BIGSERIAL` -> `INTEGER PRIMARY KEY AUTOINCREMENT`
/// - PL/pgSQL functions -> SQLite triggers
pub const SCHEMA: &str = r#"
-- ==================== Migration tracking ====================
CREATE TABLE IF NOT EXISTS _migrations (
version INTEGER PRIMARY KEY,
name TEXT NOT NULL,
applied_at TEXT NOT NULL DEFAULT (datetime('now'))
);
-- ==================== Conversations ====================
CREATE TABLE IF NOT EXISTS conversations (
id TEXT PRIMARY KEY,
channel TEXT NOT NULL,
user_id TEXT NOT NULL,
thread_id TEXT,
started_at TEXT NOT NULL DEFAULT (datetime('now')),
last_activity TEXT NOT NULL DEFAULT (datetime('now')),
metadata TEXT NOT NULL DEFAULT '{}'
);
CREATE INDEX IF NOT EXISTS idx_conversations_channel ON conversations(channel);
CREATE INDEX IF NOT EXISTS idx_conversations_user ON conversations(user_id);
CREATE INDEX IF NOT EXISTS idx_conversations_last_activity ON conversations(last_activity);
CREATE TABLE IF NOT EXISTS conversation_messages (
id TEXT PRIMARY KEY,
conversation_id TEXT NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
role TEXT NOT NULL,
content TEXT NOT NULL,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE INDEX IF NOT EXISTS idx_conversation_messages_conversation
ON conversation_messages(conversation_id);
-- ==================== Agent Jobs ====================
CREATE TABLE IF NOT EXISTS agent_jobs (
id TEXT PRIMARY KEY,
marketplace_job_id TEXT,
conversation_id TEXT REFERENCES conversations(id),
title TEXT NOT NULL,
description TEXT NOT NULL,
category TEXT,
status TEXT NOT NULL,
source TEXT NOT NULL,
user_id TEXT NOT NULL DEFAULT 'default',
project_dir TEXT,
job_mode TEXT NOT NULL DEFAULT 'worker',
budget_amount TEXT,
budget_token TEXT,
bid_amount TEXT,
estimated_cost TEXT,
estimated_time_secs INTEGER,
estimated_value TEXT,
actual_cost TEXT,
actual_time_secs INTEGER,
success INTEGER,
failure_reason TEXT,
stuck_since TEXT,
repair_attempts INTEGER NOT NULL DEFAULT 0,
created_at TEXT NOT NULL DEFAULT (datetime('now')),
started_at TEXT,
completed_at TEXT
);
CREATE INDEX IF NOT EXISTS idx_agent_jobs_status ON agent_jobs(status);
CREATE INDEX IF NOT EXISTS idx_agent_jobs_marketplace ON agent_jobs(marketplace_job_id);
CREATE INDEX IF NOT EXISTS idx_agent_jobs_conversation ON agent_jobs(conversation_id);
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);
CREATE TABLE IF NOT EXISTS job_actions (
id TEXT PRIMARY KEY,
job_id TEXT NOT NULL REFERENCES agent_jobs(id) ON DELETE CASCADE,
sequence_num INTEGER NOT NULL,
tool_name TEXT NOT NULL,
input TEXT NOT NULL,
output_raw TEXT,
output_sanitized TEXT,
sanitization_warnings TEXT,
cost TEXT,
duration_ms INTEGER,
success INTEGER NOT NULL,
error_message TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now')),
UNIQUE(job_id, sequence_num)
);
CREATE INDEX IF NOT EXISTS idx_job_actions_job_id ON job_actions(job_id);
CREATE INDEX IF NOT EXISTS idx_job_actions_tool ON job_actions(tool_name);
-- ==================== Dynamic Tools ====================
CREATE TABLE IF NOT EXISTS dynamic_tools (
id TEXT PRIMARY KEY,
name TEXT NOT NULL UNIQUE,
description TEXT NOT NULL,
parameters_schema TEXT NOT NULL,
code TEXT NOT NULL,
sandbox_config TEXT NOT NULL,
created_by_job_id TEXT REFERENCES agent_jobs(id),
success_count INTEGER NOT NULL DEFAULT 0,
failure_count INTEGER NOT NULL DEFAULT 0,
last_error TEXT,
status TEXT NOT NULL DEFAULT 'active',
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE INDEX IF NOT EXISTS idx_dynamic_tools_status ON dynamic_tools(status);
CREATE INDEX IF NOT EXISTS idx_dynamic_tools_name ON dynamic_tools(name);
-- ==================== LLM Calls ====================
CREATE TABLE IF NOT EXISTS llm_calls (
id TEXT PRIMARY KEY,
job_id TEXT REFERENCES agent_jobs(id) ON DELETE CASCADE,
conversation_id TEXT REFERENCES conversations(id),
provider TEXT NOT NULL,
model TEXT NOT NULL,
input_tokens INTEGER NOT NULL,
output_tokens INTEGER NOT NULL,
cost TEXT NOT NULL,
purpose TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE INDEX IF NOT EXISTS idx_llm_calls_job ON llm_calls(job_id);
CREATE INDEX IF NOT EXISTS idx_llm_calls_conversation ON llm_calls(conversation_id);
CREATE INDEX IF NOT EXISTS idx_llm_calls_provider ON llm_calls(provider);
-- ==================== Estimation ====================
CREATE TABLE IF NOT EXISTS estimation_snapshots (
id TEXT PRIMARY KEY,
job_id TEXT NOT NULL REFERENCES agent_jobs(id) ON DELETE CASCADE,
category TEXT NOT NULL,
tool_names TEXT NOT NULL DEFAULT '[]',
estimated_cost TEXT NOT NULL,
actual_cost TEXT,
estimated_time_secs INTEGER NOT NULL,
actual_time_secs INTEGER,
estimated_value TEXT NOT NULL,
actual_value TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE INDEX IF NOT EXISTS idx_estimation_category ON estimation_snapshots(category);
CREATE INDEX IF NOT EXISTS idx_estimation_job ON estimation_snapshots(job_id);
-- ==================== Self Repair ====================
CREATE TABLE IF NOT EXISTS repair_attempts (
id TEXT PRIMARY KEY,
target_type TEXT NOT NULL,
target_id TEXT NOT NULL,
diagnosis TEXT NOT NULL,
action_taken TEXT NOT NULL,
success INTEGER NOT NULL,
error_message TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE INDEX IF NOT EXISTS idx_repair_attempts_target ON repair_attempts(target_type, target_id);
CREATE INDEX IF NOT EXISTS idx_repair_attempts_created ON repair_attempts(created_at);
-- ==================== Workspace: Memory Documents ====================
CREATE TABLE IF NOT EXISTS memory_documents (
id TEXT PRIMARY KEY,
user_id TEXT NOT NULL,
agent_id TEXT,
path TEXT NOT NULL,
content TEXT NOT NULL,
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
metadata TEXT NOT NULL DEFAULT '{}',
UNIQUE (user_id, agent_id, path)
);
CREATE INDEX IF NOT EXISTS idx_memory_documents_user ON memory_documents(user_id);
CREATE INDEX IF NOT EXISTS idx_memory_documents_path ON memory_documents(user_id, path);
CREATE INDEX IF NOT EXISTS idx_memory_documents_updated ON memory_documents(updated_at DESC);
-- Trigger to auto-update updated_at on memory_documents
CREATE TRIGGER IF NOT EXISTS update_memory_documents_updated_at
AFTER UPDATE ON memory_documents
FOR EACH ROW
WHEN NEW.updated_at = OLD.updated_at
BEGIN
UPDATE memory_documents SET updated_at = datetime('now') WHERE id = NEW.id;
END;
-- ==================== Workspace: Memory Chunks ====================
CREATE TABLE IF NOT EXISTS memory_chunks (
_rowid INTEGER PRIMARY KEY AUTOINCREMENT,
id TEXT NOT NULL UNIQUE,
document_id TEXT NOT NULL REFERENCES memory_documents(id) ON DELETE CASCADE,
chunk_index INTEGER NOT NULL,
content TEXT NOT NULL,
embedding F32_BLOB(1536),
created_at TEXT NOT NULL DEFAULT (datetime('now')),
UNIQUE (document_id, chunk_index)
);
CREATE INDEX IF NOT EXISTS idx_memory_chunks_document ON memory_chunks(document_id);
-- Vector index for semantic search (libSQL native)
CREATE INDEX IF NOT EXISTS idx_memory_chunks_embedding
ON memory_chunks (libsql_vector_idx(embedding));
-- FTS5 virtual table for full-text search
CREATE VIRTUAL TABLE IF NOT EXISTS memory_chunks_fts USING fts5(
content,
content='memory_chunks',
content_rowid='_rowid'
);
-- Triggers to keep FTS5 in sync with memory_chunks
CREATE TRIGGER IF NOT EXISTS memory_chunks_fts_insert AFTER INSERT ON memory_chunks BEGIN
INSERT INTO memory_chunks_fts(rowid, content) VALUES (new._rowid, new.content);
END;
CREATE TRIGGER IF NOT EXISTS memory_chunks_fts_delete AFTER DELETE ON memory_chunks BEGIN
INSERT INTO memory_chunks_fts(memory_chunks_fts, rowid, content)
VALUES ('delete', old._rowid, old.content);
END;
CREATE TRIGGER IF NOT EXISTS memory_chunks_fts_update AFTER UPDATE ON memory_chunks BEGIN
INSERT INTO memory_chunks_fts(memory_chunks_fts, rowid, content)
VALUES ('delete', old._rowid, old.content);
INSERT INTO memory_chunks_fts(rowid, content) VALUES (new._rowid, new.content);
END;
-- ==================== Workspace: Heartbeat State ====================
CREATE TABLE IF NOT EXISTS heartbeat_state (
id TEXT PRIMARY KEY,
user_id TEXT NOT NULL,
agent_id TEXT,
last_run TEXT,
next_run TEXT,
interval_seconds INTEGER NOT NULL DEFAULT 1800,
enabled INTEGER NOT NULL DEFAULT 1,
consecutive_failures INTEGER NOT NULL DEFAULT 0,
last_checks TEXT NOT NULL DEFAULT '{}',
UNIQUE (user_id, agent_id)
);
CREATE INDEX IF NOT EXISTS idx_heartbeat_user ON heartbeat_state(user_id);
-- ==================== Secrets ====================
CREATE TABLE IF NOT EXISTS secrets (
id TEXT PRIMARY KEY,
user_id TEXT NOT NULL,
name TEXT NOT NULL,
encrypted_value BLOB NOT NULL,
key_salt BLOB NOT NULL,
provider TEXT,
expires_at TEXT,
last_used_at TEXT,
usage_count INTEGER NOT NULL DEFAULT 0,
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
UNIQUE (user_id, name)
);
CREATE INDEX IF NOT EXISTS idx_secrets_user ON secrets(user_id);
-- ==================== WASM Tools ====================
CREATE TABLE IF NOT EXISTS wasm_tools (
id TEXT PRIMARY KEY,
user_id TEXT NOT NULL,
name TEXT NOT NULL,
version TEXT NOT NULL DEFAULT '1.0.0',
description TEXT NOT NULL,
wasm_binary BLOB NOT NULL,
binary_hash BLOB NOT NULL,
parameters_schema TEXT NOT NULL,
source_url TEXT,
trust_level TEXT NOT NULL DEFAULT 'user',
status TEXT NOT NULL DEFAULT 'active',
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
UNIQUE (user_id, name, version)
);
CREATE INDEX IF NOT EXISTS idx_wasm_tools_user ON wasm_tools(user_id);
CREATE INDEX IF NOT EXISTS idx_wasm_tools_name ON wasm_tools(user_id, name);
CREATE INDEX IF NOT EXISTS idx_wasm_tools_status ON wasm_tools(status);
-- ==================== Tool Capabilities ====================
CREATE TABLE IF NOT EXISTS tool_capabilities (
id TEXT PRIMARY KEY,
wasm_tool_id TEXT NOT NULL REFERENCES wasm_tools(id) ON DELETE CASCADE,
http_allowlist TEXT NOT NULL DEFAULT '[]',
allowed_secrets TEXT NOT NULL DEFAULT '[]',
tool_aliases TEXT NOT NULL DEFAULT '{}',
requests_per_minute INTEGER NOT NULL DEFAULT 60,
requests_per_hour INTEGER NOT NULL DEFAULT 1000,
max_request_body_bytes INTEGER NOT NULL DEFAULT 1048576,
max_response_body_bytes INTEGER NOT NULL DEFAULT 10485760,
workspace_read_prefixes TEXT NOT NULL DEFAULT '[]',
http_timeout_secs INTEGER NOT NULL DEFAULT 30,
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
UNIQUE (wasm_tool_id)
);
-- ==================== Leak Detection Patterns ====================
CREATE TABLE IF NOT EXISTS leak_detection_patterns (
id TEXT PRIMARY KEY,
name TEXT NOT NULL UNIQUE,
pattern TEXT NOT NULL,
severity TEXT NOT NULL DEFAULT 'high',
action TEXT NOT NULL DEFAULT 'block',
enabled INTEGER NOT NULL DEFAULT 1,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
-- ==================== Rate Limit State ====================
CREATE TABLE IF NOT EXISTS tool_rate_limit_state (
id TEXT PRIMARY KEY,
wasm_tool_id TEXT NOT NULL REFERENCES wasm_tools(id) ON DELETE CASCADE,
user_id TEXT NOT NULL,
minute_window_start TEXT NOT NULL DEFAULT (datetime('now')),
minute_count INTEGER NOT NULL DEFAULT 0,
hour_window_start TEXT NOT NULL DEFAULT (datetime('now')),
hour_count INTEGER NOT NULL DEFAULT 0,
UNIQUE (wasm_tool_id, user_id)
);
-- ==================== Secret Usage Audit Log ====================
CREATE TABLE IF NOT EXISTS secret_usage_log (
id TEXT PRIMARY KEY,
secret_id TEXT NOT NULL REFERENCES secrets(id) ON DELETE CASCADE,
wasm_tool_id TEXT REFERENCES wasm_tools(id) ON DELETE SET NULL,
user_id TEXT NOT NULL,
target_host TEXT NOT NULL,
target_path TEXT,
success INTEGER NOT NULL,
error_message TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE INDEX IF NOT EXISTS idx_secret_usage_user ON secret_usage_log(user_id);
-- ==================== Leak Detection Events ====================
CREATE TABLE IF NOT EXISTS leak_detection_events (
id TEXT PRIMARY KEY,
pattern_id TEXT REFERENCES leak_detection_patterns(id) ON DELETE SET NULL,
wasm_tool_id TEXT REFERENCES wasm_tools(id) ON DELETE SET NULL,
user_id TEXT NOT NULL,
source TEXT NOT NULL,
action_taken TEXT NOT NULL,
context_preview TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
-- ==================== Tool Failures ====================
CREATE TABLE IF NOT EXISTS tool_failures (
id TEXT PRIMARY KEY,
tool_name TEXT NOT NULL UNIQUE,
error_message TEXT,
error_count INTEGER DEFAULT 1,
first_failure TEXT DEFAULT (datetime('now')),
last_failure TEXT DEFAULT (datetime('now')),
last_build_result TEXT,
repaired_at TEXT,
repair_attempts INTEGER DEFAULT 0
);
CREATE INDEX IF NOT EXISTS idx_tool_failures_name ON tool_failures(tool_name);
-- ==================== Job Events ====================
CREATE TABLE IF NOT EXISTS job_events (
id INTEGER PRIMARY KEY AUTOINCREMENT,
job_id TEXT NOT NULL REFERENCES agent_jobs(id),
event_type TEXT NOT NULL,
data TEXT NOT NULL,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE INDEX IF NOT EXISTS idx_job_events_job ON job_events(job_id, id);
-- ==================== Routines ====================
CREATE TABLE IF NOT EXISTS routines (
id TEXT PRIMARY KEY,
name TEXT NOT NULL,
description TEXT NOT NULL DEFAULT '',
user_id TEXT NOT NULL,
enabled INTEGER NOT NULL DEFAULT 1,
trigger_type TEXT NOT NULL,
trigger_config TEXT NOT NULL,
action_type TEXT NOT NULL,
action_config TEXT NOT NULL,
cooldown_secs INTEGER NOT NULL DEFAULT 300,
max_concurrent INTEGER NOT NULL DEFAULT 1,
dedup_window_secs INTEGER,
notify_channel TEXT,
notify_user TEXT NOT NULL DEFAULT 'default',
notify_on_success INTEGER NOT NULL DEFAULT 0,
notify_on_failure INTEGER NOT NULL DEFAULT 1,
notify_on_attention INTEGER NOT NULL DEFAULT 1,
state TEXT NOT NULL DEFAULT '{}',
last_run_at TEXT,
next_fire_at TEXT,
run_count INTEGER NOT NULL DEFAULT 0,
consecutive_failures INTEGER NOT NULL DEFAULT 0,
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
UNIQUE (user_id, name)
);
CREATE INDEX IF NOT EXISTS idx_routines_user ON routines(user_id);
-- ==================== Routine Runs ====================
CREATE TABLE IF NOT EXISTS routine_runs (
id TEXT PRIMARY KEY,
routine_id TEXT NOT NULL REFERENCES routines(id) ON DELETE CASCADE,
trigger_type TEXT NOT NULL,
trigger_detail TEXT,
started_at TEXT NOT NULL DEFAULT (datetime('now')),
completed_at TEXT,
status TEXT NOT NULL DEFAULT 'running',
result_summary TEXT,
tokens_used INTEGER,
job_id TEXT REFERENCES agent_jobs(id),
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE INDEX IF NOT EXISTS idx_routine_runs_routine ON routine_runs(routine_id);
-- ==================== Settings ====================
CREATE TABLE IF NOT EXISTS settings (
user_id TEXT NOT NULL,
key TEXT NOT NULL,
value TEXT NOT NULL,
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
PRIMARY KEY (user_id, key)
);
CREATE INDEX IF NOT EXISTS idx_settings_user ON settings(user_id);
-- ==================== Missing indexes (parity with PostgreSQL) ====================
-- agent_jobs
CREATE INDEX IF NOT EXISTS idx_agent_jobs_stuck ON agent_jobs(stuck_since);
-- secrets
CREATE INDEX IF NOT EXISTS idx_secrets_provider ON secrets(provider);
CREATE INDEX IF NOT EXISTS idx_secrets_expires ON secrets(expires_at);
-- wasm_tools
CREATE INDEX IF NOT EXISTS idx_wasm_tools_trust ON wasm_tools(trust_level);
-- tool_capabilities
CREATE INDEX IF NOT EXISTS idx_tool_capabilities_tool ON tool_capabilities(wasm_tool_id);
-- leak_detection_patterns
CREATE INDEX IF NOT EXISTS idx_leak_patterns_enabled ON leak_detection_patterns(enabled);
-- tool_rate_limit_state
CREATE INDEX IF NOT EXISTS idx_rate_limit_tool ON tool_rate_limit_state(wasm_tool_id);
-- secret_usage_log
CREATE INDEX IF NOT EXISTS idx_secret_usage_secret ON secret_usage_log(secret_id);
CREATE INDEX IF NOT EXISTS idx_secret_usage_tool ON secret_usage_log(wasm_tool_id);
CREATE INDEX IF NOT EXISTS idx_secret_usage_created ON secret_usage_log(created_at DESC);
-- leak_detection_events
CREATE INDEX IF NOT EXISTS idx_leak_events_pattern ON leak_detection_events(pattern_id);
CREATE INDEX IF NOT EXISTS idx_leak_events_tool ON leak_detection_events(wasm_tool_id);
CREATE INDEX IF NOT EXISTS idx_leak_events_user ON leak_detection_events(user_id);
CREATE INDEX IF NOT EXISTS idx_leak_events_created ON leak_detection_events(created_at DESC);
-- tool_failures
CREATE INDEX IF NOT EXISTS idx_tool_failures_count ON tool_failures(error_count DESC);
CREATE INDEX IF NOT EXISTS idx_tool_failures_unrepaired ON tool_failures(tool_name);
-- routines
CREATE INDEX IF NOT EXISTS idx_routines_next_fire ON routines(next_fire_at);
CREATE INDEX IF NOT EXISTS idx_routines_event_triggers ON routines(user_id);
-- routine_runs
CREATE INDEX IF NOT EXISTS idx_routine_runs_status ON routine_runs(status);
-- heartbeat_state
CREATE INDEX IF NOT EXISTS idx_heartbeat_next_run ON heartbeat_state(next_run);
-- ==================== Seed data ====================
-- Pre-populate leak detection patterns (matches PostgreSQL V2 migration).
INSERT OR IGNORE INTO leak_detection_patterns (id, name, pattern, severity, action, enabled, created_at) VALUES
('550e8400-e29b-41d4-a716-446655440001', 'openai_api_key', 'sk-(?:proj-)?[a-zA-Z0-9]{20,}(?:T3BlbkFJ[a-zA-Z0-9_-]*)?', 'critical', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440002', 'anthropic_api_key', 'sk-ant-api[a-zA-Z0-9_-]{90,}', 'critical', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440003', 'aws_access_key', 'AKIA[0-9A-Z]{16}', 'critical', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440004', 'aws_secret_key', '(?<![A-Za-z0-9/+=])[A-Za-z0-9/+=]{40}(?![A-Za-z0-9/+=])', 'high', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440005', 'github_token', 'gh[pousr]_[A-Za-z0-9_]{36,}', 'critical', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440006', 'github_fine_grained_pat', 'github_pat_[a-zA-Z0-9]{22}_[a-zA-Z0-9]{59}', 'critical', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440007', 'stripe_api_key', 'sk_(?:live|test)_[a-zA-Z0-9]{24,}', 'critical', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440008', 'nearai_session', 'sess_[a-zA-Z0-9]{32,}', 'critical', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440009', 'bearer_token', 'Bearer\s+[a-zA-Z0-9_-]{20,}', 'high', 'redact', 1, datetime('now')),
('550e8400-e29b-41d4-a716-44665544000a', 'pem_private_key', '-----BEGIN\s+(?:RSA\s+)?PRIVATE\s+KEY-----', 'critical', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-44665544000b', 'ssh_private_key', '-----BEGIN\s+(?:OPENSSH|EC|DSA)\s+PRIVATE\s+KEY-----', 'critical', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-44665544000c', 'google_api_key', 'AIza[0-9A-Za-z_-]{35}', 'high', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-44665544000d', 'slack_token', 'xox[baprs]-[0-9a-zA-Z-]{10,}', 'high', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-44665544000e', 'discord_token', '[MN][A-Za-z\d]{23,}\.[\w-]{6}\.[\w-]{27}', 'high', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-44665544000f', 'twilio_api_key', 'SK[a-fA-F0-9]{32}', 'high', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440010', 'sendgrid_api_key', 'SG\.[a-zA-Z0-9_-]{22}\.[a-zA-Z0-9_-]{43}', 'high', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440011', 'mailchimp_api_key', '[a-f0-9]{32}-us[0-9]{1,2}', 'medium', 'block', 1, datetime('now')),
('550e8400-e29b-41d4-a716-446655440012', 'high_entropy_hex', '(?<![a-fA-F0-9])[a-fA-F0-9]{64}(?![a-fA-F0-9])', 'medium', 'warn', 1, datetime('now'));
"#;
+538
View File
@@ -0,0 +1,538 @@
//! Database abstraction layer.
//!
//! Provides a backend-agnostic `Database` trait that unifies all persistence
//! operations. Two implementations exist behind feature flags:
//!
//! - `postgres` (default): Uses `deadpool-postgres` + `tokio-postgres`
//! - `libsql`: Uses libSQL (Turso's SQLite fork) for embedded/edge deployment
//!
//! The existing `Store`, `Repository`, `SecretsStore`, and `WasmToolStore`
//! types become thin wrappers that delegate to `Arc<dyn Database>`.
#[cfg(feature = "postgres")]
pub mod postgres;
#[cfg(feature = "libsql")]
pub mod libsql_backend;
#[cfg(feature = "libsql")]
pub mod libsql_migrations;
use std::collections::HashMap;
use std::sync::Arc;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use rust_decimal::Decimal;
use uuid::Uuid;
use crate::agent::BrokenTool;
use crate::agent::routine::{Routine, RoutineRun, RunStatus};
use crate::context::{ActionRecord, JobContext, JobState};
use crate::error::DatabaseError;
use crate::error::WorkspaceError;
use crate::history::{
ConversationMessage, ConversationSummary, JobEventRecord, LlmCallRecord, SandboxJobRecord,
SandboxJobSummary, SettingRow,
};
use crate::workspace::{MemoryChunk, MemoryDocument, WorkspaceEntry};
use crate::workspace::{SearchConfig, SearchResult};
/// Create a database backend from configuration, run migrations, and return it.
///
/// This is the shared helper for CLI commands and other call sites that need
/// a simple `Arc<dyn Database>` without retaining backend-specific handles
/// (e.g., `pg_pool` or `libsql_conn` for the secrets store). The main agent
/// startup in `main.rs` uses its own initialization block because it also
/// captures those backend-specific handles.
pub async fn connect_from_config(
config: &crate::config::DatabaseConfig,
) -> Result<Arc<dyn Database>, DatabaseError> {
match config.backend {
#[cfg(feature = "libsql")]
crate::config::DatabaseBackend::LibSql => {
use secrecy::ExposeSecret as _;
let default_path = crate::config::default_libsql_path();
let db_path = config.libsql_path.as_deref().unwrap_or(&default_path);
let backend = if let Some(ref url) = config.libsql_url {
let token = config.libsql_auth_token.as_ref().ok_or_else(|| {
DatabaseError::Pool(
"LIBSQL_AUTH_TOKEN required when LIBSQL_URL is set".to_string(),
)
})?;
libsql_backend::LibSqlBackend::new_remote_replica(
db_path,
url,
token.expose_secret(),
)
.await
.map_err(|e| DatabaseError::Pool(e.to_string()))?
} else {
libsql_backend::LibSqlBackend::new_local(db_path)
.await
.map_err(|e| DatabaseError::Pool(e.to_string()))?
};
backend.run_migrations().await?;
Ok(Arc::new(backend))
}
#[cfg(feature = "postgres")]
_ => {
let pg = postgres::PgBackend::new(config)
.await
.map_err(|e| DatabaseError::Pool(e.to_string()))?;
pg.run_migrations().await?;
Ok(Arc::new(pg))
}
#[cfg(not(feature = "postgres"))]
_ => Err(DatabaseError::Pool(
"No database backend available. Enable 'postgres' or 'libsql' feature.".to_string(),
)),
}
}
/// Backend-agnostic database trait.
///
/// Combines all persistence operations from Store, Repository, and related
/// stores into a single trait that can be implemented for different backends.
#[async_trait]
pub trait Database: Send + Sync {
/// Run schema migrations for this backend.
async fn run_migrations(&self) -> Result<(), DatabaseError>;
// ==================== Conversations ====================
/// Create a new conversation.
async fn create_conversation(
&self,
channel: &str,
user_id: &str,
thread_id: Option<&str>,
) -> Result<Uuid, DatabaseError>;
/// Update conversation last activity.
async fn touch_conversation(&self, id: Uuid) -> Result<(), DatabaseError>;
/// Add a message to a conversation.
async fn add_conversation_message(
&self,
conversation_id: Uuid,
role: &str,
content: &str,
) -> Result<Uuid, DatabaseError>;
/// Ensure a conversation row exists (upsert).
async fn ensure_conversation(
&self,
id: Uuid,
channel: &str,
user_id: &str,
thread_id: Option<&str>,
) -> Result<(), DatabaseError>;
/// List conversations with a title preview.
async fn list_conversations_with_preview(
&self,
user_id: &str,
channel: &str,
limit: i64,
) -> Result<Vec<ConversationSummary>, DatabaseError>;
/// Get or create the singleton assistant conversation.
async fn get_or_create_assistant_conversation(
&self,
user_id: &str,
channel: &str,
) -> Result<Uuid, DatabaseError>;
/// Create a conversation with specific metadata.
async fn create_conversation_with_metadata(
&self,
channel: &str,
user_id: &str,
metadata: &serde_json::Value,
) -> Result<Uuid, DatabaseError>;
/// Load messages with cursor-based pagination.
async fn list_conversation_messages_paginated(
&self,
conversation_id: Uuid,
before: Option<DateTime<Utc>>,
limit: i64,
) -> Result<(Vec<ConversationMessage>, bool), DatabaseError>;
/// Merge a single key into conversation metadata.
async fn update_conversation_metadata_field(
&self,
id: Uuid,
key: &str,
value: &serde_json::Value,
) -> Result<(), DatabaseError>;
/// Read conversation metadata.
async fn get_conversation_metadata(
&self,
id: Uuid,
) -> Result<Option<serde_json::Value>, DatabaseError>;
/// Load all messages for a conversation.
async fn list_conversation_messages(
&self,
conversation_id: Uuid,
) -> Result<Vec<ConversationMessage>, DatabaseError>;
/// Check if a conversation belongs to a specific user.
async fn conversation_belongs_to_user(
&self,
conversation_id: Uuid,
user_id: &str,
) -> Result<bool, DatabaseError>;
// ==================== Jobs ====================
/// Save a job context.
async fn save_job(&self, ctx: &JobContext) -> Result<(), DatabaseError>;
/// Get a job by ID.
async fn get_job(&self, id: Uuid) -> Result<Option<JobContext>, DatabaseError>;
/// Update job status.
async fn update_job_status(
&self,
id: Uuid,
status: JobState,
failure_reason: Option<&str>,
) -> Result<(), DatabaseError>;
/// Mark job as stuck.
async fn mark_job_stuck(&self, id: Uuid) -> Result<(), DatabaseError>;
/// Get stuck jobs.
async fn get_stuck_jobs(&self) -> Result<Vec<Uuid>, DatabaseError>;
// ==================== Actions ====================
/// Save a job action.
async fn save_action(&self, job_id: Uuid, action: &ActionRecord) -> Result<(), DatabaseError>;
/// Get actions for a job.
async fn get_job_actions(&self, job_id: Uuid) -> Result<Vec<ActionRecord>, DatabaseError>;
// ==================== LLM Calls ====================
/// Record an LLM call.
async fn record_llm_call(&self, record: &LlmCallRecord<'_>) -> Result<Uuid, DatabaseError>;
// ==================== Estimation Snapshots ====================
/// Save an estimation snapshot.
async fn save_estimation_snapshot(
&self,
job_id: Uuid,
category: &str,
tool_names: &[String],
estimated_cost: Decimal,
estimated_time_secs: i32,
estimated_value: Decimal,
) -> Result<Uuid, DatabaseError>;
/// Update estimation snapshot with actual values.
async fn update_estimation_actuals(
&self,
id: Uuid,
actual_cost: Decimal,
actual_time_secs: i32,
actual_value: Option<Decimal>,
) -> Result<(), DatabaseError>;
// ==================== Sandbox Jobs ====================
/// Insert a new sandbox job.
async fn save_sandbox_job(&self, job: &SandboxJobRecord) -> Result<(), DatabaseError>;
/// Get a sandbox job by ID.
async fn get_sandbox_job(&self, id: Uuid) -> Result<Option<SandboxJobRecord>, DatabaseError>;
/// List all sandbox jobs, most recent first.
async fn list_sandbox_jobs(&self) -> Result<Vec<SandboxJobRecord>, DatabaseError>;
/// Update sandbox job status.
async fn update_sandbox_job_status(
&self,
id: Uuid,
status: &str,
success: Option<bool>,
message: Option<&str>,
started_at: Option<DateTime<Utc>>,
completed_at: Option<DateTime<Utc>>,
) -> Result<(), DatabaseError>;
/// Mark stale sandbox jobs as interrupted.
async fn cleanup_stale_sandbox_jobs(&self) -> Result<u64, DatabaseError>;
/// Get sandbox job summary.
async fn sandbox_job_summary(&self) -> Result<SandboxJobSummary, DatabaseError>;
/// List sandbox jobs for a specific user, most recent first.
async fn list_sandbox_jobs_for_user(
&self,
user_id: &str,
) -> Result<Vec<SandboxJobRecord>, DatabaseError>;
/// Get sandbox job summary for a specific user.
async fn sandbox_job_summary_for_user(
&self,
user_id: &str,
) -> Result<SandboxJobSummary, DatabaseError>;
/// Check if a sandbox job belongs to a specific user.
async fn sandbox_job_belongs_to_user(
&self,
job_id: Uuid,
user_id: &str,
) -> Result<bool, DatabaseError>;
/// Update sandbox job mode.
async fn update_sandbox_job_mode(&self, id: Uuid, mode: &str) -> Result<(), DatabaseError>;
/// Get sandbox job mode.
async fn get_sandbox_job_mode(&self, id: Uuid) -> Result<Option<String>, DatabaseError>;
// ==================== Job Events ====================
/// Persist a job event.
async fn save_job_event(
&self,
job_id: Uuid,
event_type: &str,
data: &serde_json::Value,
) -> Result<(), DatabaseError>;
/// Load all job events.
async fn list_job_events(&self, job_id: Uuid) -> Result<Vec<JobEventRecord>, DatabaseError>;
// ==================== Routines ====================
/// Create a new routine.
async fn create_routine(&self, routine: &Routine) -> Result<(), DatabaseError>;
/// Get a routine by ID.
async fn get_routine(&self, id: Uuid) -> Result<Option<Routine>, DatabaseError>;
/// Get a routine by user_id and name.
async fn get_routine_by_name(
&self,
user_id: &str,
name: &str,
) -> Result<Option<Routine>, DatabaseError>;
/// List routines for a user.
async fn list_routines(&self, user_id: &str) -> Result<Vec<Routine>, DatabaseError>;
/// List all enabled event routines.
async fn list_event_routines(&self) -> Result<Vec<Routine>, DatabaseError>;
/// List due cron routines.
async fn list_due_cron_routines(&self) -> Result<Vec<Routine>, DatabaseError>;
/// Update a routine.
async fn update_routine(&self, routine: &Routine) -> Result<(), DatabaseError>;
/// Update runtime state after a routine fires.
async fn update_routine_runtime(
&self,
id: Uuid,
last_run_at: DateTime<Utc>,
next_fire_at: Option<DateTime<Utc>>,
run_count: u64,
consecutive_failures: u32,
state: &serde_json::Value,
) -> Result<(), DatabaseError>;
/// Delete a routine.
async fn delete_routine(&self, id: Uuid) -> Result<bool, DatabaseError>;
// ==================== Routine Runs ====================
/// Record a routine run starting.
async fn create_routine_run(&self, run: &RoutineRun) -> Result<(), DatabaseError>;
/// Complete a routine run.
async fn complete_routine_run(
&self,
id: Uuid,
status: RunStatus,
result_summary: Option<&str>,
tokens_used: Option<i32>,
) -> Result<(), DatabaseError>;
/// List recent runs for a routine.
async fn list_routine_runs(
&self,
routine_id: Uuid,
limit: i64,
) -> Result<Vec<RoutineRun>, DatabaseError>;
/// Count currently running runs for a routine.
async fn count_running_routine_runs(&self, routine_id: Uuid) -> Result<i64, DatabaseError>;
// ==================== Tool Failures ====================
/// Record a tool failure (upsert).
async fn record_tool_failure(
&self,
tool_name: &str,
error_message: &str,
) -> Result<(), DatabaseError>;
/// Get broken tools exceeding threshold.
async fn get_broken_tools(&self, threshold: i32) -> Result<Vec<BrokenTool>, DatabaseError>;
/// Mark a tool as repaired.
async fn mark_tool_repaired(&self, tool_name: &str) -> Result<(), DatabaseError>;
/// Increment repair attempts.
async fn increment_repair_attempts(&self, tool_name: &str) -> Result<(), DatabaseError>;
// ==================== Settings ====================
/// Get a single setting.
async fn get_setting(
&self,
user_id: &str,
key: &str,
) -> Result<Option<serde_json::Value>, DatabaseError>;
/// Get a single setting with metadata.
async fn get_setting_full(
&self,
user_id: &str,
key: &str,
) -> Result<Option<SettingRow>, DatabaseError>;
/// Set a single setting (upsert).
async fn set_setting(
&self,
user_id: &str,
key: &str,
value: &serde_json::Value,
) -> Result<(), DatabaseError>;
/// Delete a single setting.
async fn delete_setting(&self, user_id: &str, key: &str) -> Result<bool, DatabaseError>;
/// List all settings for a user.
async fn list_settings(&self, user_id: &str) -> Result<Vec<SettingRow>, DatabaseError>;
/// Get all settings as a flat map.
async fn get_all_settings(
&self,
user_id: &str,
) -> Result<HashMap<String, serde_json::Value>, DatabaseError>;
/// Bulk-write settings atomically.
async fn set_all_settings(
&self,
user_id: &str,
settings: &HashMap<String, serde_json::Value>,
) -> Result<(), DatabaseError>;
/// Check if settings exist for a user.
async fn has_settings(&self, user_id: &str) -> Result<bool, DatabaseError>;
// ==================== Workspace: Documents ====================
/// Get a document by path.
async fn get_document_by_path(
&self,
user_id: &str,
agent_id: Option<Uuid>,
path: &str,
) -> Result<MemoryDocument, WorkspaceError>;
/// Get a document by ID.
async fn get_document_by_id(&self, id: Uuid) -> Result<MemoryDocument, WorkspaceError>;
/// Get or create a document by path.
async fn get_or_create_document_by_path(
&self,
user_id: &str,
agent_id: Option<Uuid>,
path: &str,
) -> Result<MemoryDocument, WorkspaceError>;
/// Update a document's content.
async fn update_document(&self, id: Uuid, content: &str) -> Result<(), WorkspaceError>;
/// Delete a document by path.
async fn delete_document_by_path(
&self,
user_id: &str,
agent_id: Option<Uuid>,
path: &str,
) -> Result<(), WorkspaceError>;
/// List files and directories in a directory path.
async fn list_directory(
&self,
user_id: &str,
agent_id: Option<Uuid>,
directory: &str,
) -> Result<Vec<WorkspaceEntry>, WorkspaceError>;
/// List all file paths in the workspace.
async fn list_all_paths(
&self,
user_id: &str,
agent_id: Option<Uuid>,
) -> Result<Vec<String>, WorkspaceError>;
/// List all documents for a user.
async fn list_documents(
&self,
user_id: &str,
agent_id: Option<Uuid>,
) -> Result<Vec<MemoryDocument>, WorkspaceError>;
// ==================== Workspace: Chunks ====================
/// Delete all chunks for a document.
async fn delete_chunks(&self, document_id: Uuid) -> Result<(), WorkspaceError>;
/// Insert a chunk.
async fn insert_chunk(
&self,
document_id: Uuid,
chunk_index: i32,
content: &str,
embedding: Option<&[f32]>,
) -> Result<Uuid, WorkspaceError>;
/// Update a chunk's embedding.
async fn update_chunk_embedding(
&self,
chunk_id: Uuid,
embedding: &[f32],
) -> Result<(), WorkspaceError>;
/// Get chunks without embeddings for backfilling.
async fn get_chunks_without_embeddings(
&self,
user_id: &str,
agent_id: Option<Uuid>,
limit: usize,
) -> Result<Vec<MemoryChunk>, WorkspaceError>;
// ==================== Workspace: Search ====================
/// Perform hybrid search combining FTS and vector similarity.
async fn hybrid_search(
&self,
user_id: &str,
agent_id: Option<Uuid>,
query: &str,
embedding: Option<&[f32]>,
config: &SearchConfig,
) -> Result<Vec<SearchResult>, WorkspaceError>;
}
+627
View File
@@ -0,0 +1,627 @@
//! PostgreSQL backend for the Database trait.
//!
//! Delegates to the existing `Store` (history) and `Repository` (workspace)
//! implementations, avoiding SQL duplication.
use std::collections::HashMap;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use deadpool_postgres::Pool;
use rust_decimal::Decimal;
use uuid::Uuid;
use crate::agent::BrokenTool;
use crate::agent::routine::{Routine, RoutineRun, RunStatus};
use crate::config::DatabaseConfig;
use crate::context::{ActionRecord, JobContext, JobState};
use crate::db::Database;
use crate::error::{DatabaseError, WorkspaceError};
use crate::history::{
ConversationMessage, ConversationSummary, JobEventRecord, LlmCallRecord, SandboxJobRecord,
SandboxJobSummary, SettingRow, Store,
};
use crate::workspace::{
MemoryChunk, MemoryDocument, Repository, SearchConfig, SearchResult, WorkspaceEntry,
};
/// PostgreSQL database backend.
///
/// Wraps the existing `Store` (for history/conversations/jobs/routines/settings)
/// and `Repository` (for workspace documents/chunks/search) to implement the
/// unified `Database` trait.
pub struct PgBackend {
store: Store,
repo: Repository,
}
impl PgBackend {
/// Create a new PostgreSQL backend from configuration.
pub async fn new(config: &DatabaseConfig) -> Result<Self, DatabaseError> {
let store = Store::new(config).await?;
let repo = Repository::new(store.pool());
Ok(Self { store, repo })
}
/// Get a clone of the connection pool.
///
/// Useful for sharing with components that still need raw pool access.
pub fn pool(&self) -> Pool {
self.store.pool()
}
}
#[async_trait]
impl Database for PgBackend {
async fn run_migrations(&self) -> Result<(), DatabaseError> {
self.store.run_migrations().await
}
// ==================== Conversations ====================
async fn create_conversation(
&self,
channel: &str,
user_id: &str,
thread_id: Option<&str>,
) -> Result<Uuid, DatabaseError> {
self.store
.create_conversation(channel, user_id, thread_id)
.await
}
async fn touch_conversation(&self, id: Uuid) -> Result<(), DatabaseError> {
self.store.touch_conversation(id).await
}
async fn add_conversation_message(
&self,
conversation_id: Uuid,
role: &str,
content: &str,
) -> Result<Uuid, DatabaseError> {
self.store
.add_conversation_message(conversation_id, role, content)
.await
}
async fn ensure_conversation(
&self,
id: Uuid,
channel: &str,
user_id: &str,
thread_id: Option<&str>,
) -> Result<(), DatabaseError> {
self.store
.ensure_conversation(id, channel, user_id, thread_id)
.await
}
async fn list_conversations_with_preview(
&self,
user_id: &str,
channel: &str,
limit: i64,
) -> Result<Vec<ConversationSummary>, DatabaseError> {
self.store
.list_conversations_with_preview(user_id, channel, limit)
.await
}
async fn get_or_create_assistant_conversation(
&self,
user_id: &str,
channel: &str,
) -> Result<Uuid, DatabaseError> {
self.store
.get_or_create_assistant_conversation(user_id, channel)
.await
}
async fn create_conversation_with_metadata(
&self,
channel: &str,
user_id: &str,
metadata: &serde_json::Value,
) -> Result<Uuid, DatabaseError> {
self.store
.create_conversation_with_metadata(channel, user_id, metadata)
.await
}
async fn list_conversation_messages_paginated(
&self,
conversation_id: Uuid,
before: Option<DateTime<Utc>>,
limit: i64,
) -> Result<(Vec<ConversationMessage>, bool), DatabaseError> {
self.store
.list_conversation_messages_paginated(conversation_id, before, limit)
.await
}
async fn update_conversation_metadata_field(
&self,
id: Uuid,
key: &str,
value: &serde_json::Value,
) -> Result<(), DatabaseError> {
self.store
.update_conversation_metadata_field(id, key, value)
.await
}
async fn get_conversation_metadata(
&self,
id: Uuid,
) -> Result<Option<serde_json::Value>, DatabaseError> {
self.store.get_conversation_metadata(id).await
}
async fn list_conversation_messages(
&self,
conversation_id: Uuid,
) -> Result<Vec<ConversationMessage>, DatabaseError> {
self.store.list_conversation_messages(conversation_id).await
}
async fn conversation_belongs_to_user(
&self,
conversation_id: Uuid,
user_id: &str,
) -> Result<bool, DatabaseError> {
self.store
.conversation_belongs_to_user(conversation_id, user_id)
.await
}
// ==================== Jobs ====================
async fn save_job(&self, ctx: &JobContext) -> Result<(), DatabaseError> {
self.store.save_job(ctx).await
}
async fn get_job(&self, id: Uuid) -> Result<Option<JobContext>, DatabaseError> {
self.store.get_job(id).await
}
async fn update_job_status(
&self,
id: Uuid,
status: JobState,
failure_reason: Option<&str>,
) -> Result<(), DatabaseError> {
self.store
.update_job_status(id, status, failure_reason)
.await
}
async fn mark_job_stuck(&self, id: Uuid) -> Result<(), DatabaseError> {
self.store.mark_job_stuck(id).await
}
async fn get_stuck_jobs(&self) -> Result<Vec<Uuid>, DatabaseError> {
self.store.get_stuck_jobs().await
}
// ==================== Actions ====================
async fn save_action(&self, job_id: Uuid, action: &ActionRecord) -> Result<(), DatabaseError> {
self.store.save_action(job_id, action).await
}
async fn get_job_actions(&self, job_id: Uuid) -> Result<Vec<ActionRecord>, DatabaseError> {
self.store.get_job_actions(job_id).await
}
// ==================== LLM Calls ====================
async fn record_llm_call(&self, record: &LlmCallRecord<'_>) -> Result<Uuid, DatabaseError> {
self.store.record_llm_call(record).await
}
// ==================== Estimation Snapshots ====================
async fn save_estimation_snapshot(
&self,
job_id: Uuid,
category: &str,
tool_names: &[String],
estimated_cost: Decimal,
estimated_time_secs: i32,
estimated_value: Decimal,
) -> Result<Uuid, DatabaseError> {
self.store
.save_estimation_snapshot(
job_id,
category,
tool_names,
estimated_cost,
estimated_time_secs,
estimated_value,
)
.await
}
async fn update_estimation_actuals(
&self,
id: Uuid,
actual_cost: Decimal,
actual_time_secs: i32,
actual_value: Option<Decimal>,
) -> Result<(), DatabaseError> {
self.store
.update_estimation_actuals(id, actual_cost, actual_time_secs, actual_value)
.await
}
// ==================== Sandbox Jobs ====================
async fn save_sandbox_job(&self, job: &SandboxJobRecord) -> Result<(), DatabaseError> {
self.store.save_sandbox_job(job).await
}
async fn get_sandbox_job(&self, id: Uuid) -> Result<Option<SandboxJobRecord>, DatabaseError> {
self.store.get_sandbox_job(id).await
}
async fn list_sandbox_jobs(&self) -> Result<Vec<SandboxJobRecord>, DatabaseError> {
self.store.list_sandbox_jobs().await
}
async fn update_sandbox_job_status(
&self,
id: Uuid,
status: &str,
success: Option<bool>,
message: Option<&str>,
started_at: Option<DateTime<Utc>>,
completed_at: Option<DateTime<Utc>>,
) -> Result<(), DatabaseError> {
self.store
.update_sandbox_job_status(id, status, success, message, started_at, completed_at)
.await
}
async fn cleanup_stale_sandbox_jobs(&self) -> Result<u64, DatabaseError> {
self.store.cleanup_stale_sandbox_jobs().await
}
async fn sandbox_job_summary(&self) -> Result<SandboxJobSummary, DatabaseError> {
self.store.sandbox_job_summary().await
}
async fn list_sandbox_jobs_for_user(
&self,
user_id: &str,
) -> Result<Vec<SandboxJobRecord>, DatabaseError> {
self.store.list_sandbox_jobs_for_user(user_id).await
}
async fn sandbox_job_summary_for_user(
&self,
user_id: &str,
) -> Result<SandboxJobSummary, DatabaseError> {
self.store.sandbox_job_summary_for_user(user_id).await
}
async fn sandbox_job_belongs_to_user(
&self,
job_id: Uuid,
user_id: &str,
) -> Result<bool, DatabaseError> {
self.store
.sandbox_job_belongs_to_user(job_id, user_id)
.await
}
async fn update_sandbox_job_mode(&self, id: Uuid, mode: &str) -> Result<(), DatabaseError> {
self.store.update_sandbox_job_mode(id, mode).await
}
async fn get_sandbox_job_mode(&self, id: Uuid) -> Result<Option<String>, DatabaseError> {
self.store.get_sandbox_job_mode(id).await
}
// ==================== Job Events ====================
async fn save_job_event(
&self,
job_id: Uuid,
event_type: &str,
data: &serde_json::Value,
) -> Result<(), DatabaseError> {
self.store.save_job_event(job_id, event_type, data).await
}
async fn list_job_events(&self, job_id: Uuid) -> Result<Vec<JobEventRecord>, DatabaseError> {
self.store.list_job_events(job_id).await
}
// ==================== Routines ====================
async fn create_routine(&self, routine: &Routine) -> Result<(), DatabaseError> {
self.store.create_routine(routine).await
}
async fn get_routine(&self, id: Uuid) -> Result<Option<Routine>, DatabaseError> {
self.store.get_routine(id).await
}
async fn get_routine_by_name(
&self,
user_id: &str,
name: &str,
) -> Result<Option<Routine>, DatabaseError> {
self.store.get_routine_by_name(user_id, name).await
}
async fn list_routines(&self, user_id: &str) -> Result<Vec<Routine>, DatabaseError> {
self.store.list_routines(user_id).await
}
async fn list_event_routines(&self) -> Result<Vec<Routine>, DatabaseError> {
self.store.list_event_routines().await
}
async fn list_due_cron_routines(&self) -> Result<Vec<Routine>, DatabaseError> {
self.store.list_due_cron_routines().await
}
async fn update_routine(&self, routine: &Routine) -> Result<(), DatabaseError> {
self.store.update_routine(routine).await
}
async fn update_routine_runtime(
&self,
id: Uuid,
last_run_at: DateTime<Utc>,
next_fire_at: Option<DateTime<Utc>>,
run_count: u64,
consecutive_failures: u32,
state: &serde_json::Value,
) -> Result<(), DatabaseError> {
self.store
.update_routine_runtime(
id,
last_run_at,
next_fire_at,
run_count,
consecutive_failures,
state,
)
.await
}
async fn delete_routine(&self, id: Uuid) -> Result<bool, DatabaseError> {
self.store.delete_routine(id).await
}
// ==================== Routine Runs ====================
async fn create_routine_run(&self, run: &RoutineRun) -> Result<(), DatabaseError> {
self.store.create_routine_run(run).await
}
async fn complete_routine_run(
&self,
id: Uuid,
status: RunStatus,
result_summary: Option<&str>,
tokens_used: Option<i32>,
) -> Result<(), DatabaseError> {
self.store
.complete_routine_run(id, status, result_summary, tokens_used)
.await
}
async fn list_routine_runs(
&self,
routine_id: Uuid,
limit: i64,
) -> Result<Vec<RoutineRun>, DatabaseError> {
self.store.list_routine_runs(routine_id, limit).await
}
async fn count_running_routine_runs(&self, routine_id: Uuid) -> Result<i64, DatabaseError> {
self.store.count_running_routine_runs(routine_id).await
}
// ==================== Tool Failures ====================
async fn record_tool_failure(
&self,
tool_name: &str,
error_message: &str,
) -> Result<(), DatabaseError> {
self.store
.record_tool_failure(tool_name, error_message)
.await
}
async fn get_broken_tools(&self, threshold: i32) -> Result<Vec<BrokenTool>, DatabaseError> {
self.store.get_broken_tools(threshold).await
}
async fn mark_tool_repaired(&self, tool_name: &str) -> Result<(), DatabaseError> {
self.store.mark_tool_repaired(tool_name).await
}
async fn increment_repair_attempts(&self, tool_name: &str) -> Result<(), DatabaseError> {
self.store.increment_repair_attempts(tool_name).await
}
// ==================== Settings ====================
async fn get_setting(
&self,
user_id: &str,
key: &str,
) -> Result<Option<serde_json::Value>, DatabaseError> {
self.store.get_setting(user_id, key).await
}
async fn get_setting_full(
&self,
user_id: &str,
key: &str,
) -> Result<Option<SettingRow>, DatabaseError> {
self.store.get_setting_full(user_id, key).await
}
async fn set_setting(
&self,
user_id: &str,
key: &str,
value: &serde_json::Value,
) -> Result<(), DatabaseError> {
self.store.set_setting(user_id, key, value).await
}
async fn delete_setting(&self, user_id: &str, key: &str) -> Result<bool, DatabaseError> {
self.store.delete_setting(user_id, key).await
}
async fn list_settings(&self, user_id: &str) -> Result<Vec<SettingRow>, DatabaseError> {
self.store.list_settings(user_id).await
}
async fn get_all_settings(
&self,
user_id: &str,
) -> Result<HashMap<String, serde_json::Value>, DatabaseError> {
self.store.get_all_settings(user_id).await
}
async fn set_all_settings(
&self,
user_id: &str,
settings: &HashMap<String, serde_json::Value>,
) -> Result<(), DatabaseError> {
self.store.set_all_settings(user_id, settings).await
}
async fn has_settings(&self, user_id: &str) -> Result<bool, DatabaseError> {
self.store.has_settings(user_id).await
}
// ==================== Workspace: Documents ====================
async fn get_document_by_path(
&self,
user_id: &str,
agent_id: Option<Uuid>,
path: &str,
) -> Result<MemoryDocument, WorkspaceError> {
self.repo
.get_document_by_path(user_id, agent_id, path)
.await
}
async fn get_document_by_id(&self, id: Uuid) -> Result<MemoryDocument, WorkspaceError> {
self.repo.get_document_by_id(id).await
}
async fn get_or_create_document_by_path(
&self,
user_id: &str,
agent_id: Option<Uuid>,
path: &str,
) -> Result<MemoryDocument, WorkspaceError> {
self.repo
.get_or_create_document_by_path(user_id, agent_id, path)
.await
}
async fn update_document(&self, id: Uuid, content: &str) -> Result<(), WorkspaceError> {
self.repo.update_document(id, content).await
}
async fn delete_document_by_path(
&self,
user_id: &str,
agent_id: Option<Uuid>,
path: &str,
) -> Result<(), WorkspaceError> {
self.repo
.delete_document_by_path(user_id, agent_id, path)
.await
}
async fn list_directory(
&self,
user_id: &str,
agent_id: Option<Uuid>,
directory: &str,
) -> Result<Vec<WorkspaceEntry>, WorkspaceError> {
self.repo.list_directory(user_id, agent_id, directory).await
}
async fn list_all_paths(
&self,
user_id: &str,
agent_id: Option<Uuid>,
) -> Result<Vec<String>, WorkspaceError> {
self.repo.list_all_paths(user_id, agent_id).await
}
async fn list_documents(
&self,
user_id: &str,
agent_id: Option<Uuid>,
) -> Result<Vec<MemoryDocument>, WorkspaceError> {
self.repo.list_documents(user_id, agent_id).await
}
// ==================== Workspace: Chunks ====================
async fn delete_chunks(&self, document_id: Uuid) -> Result<(), WorkspaceError> {
self.repo.delete_chunks(document_id).await
}
async fn insert_chunk(
&self,
document_id: Uuid,
chunk_index: i32,
content: &str,
embedding: Option<&[f32]>,
) -> Result<Uuid, WorkspaceError> {
self.repo
.insert_chunk(document_id, chunk_index, content, embedding)
.await
}
async fn update_chunk_embedding(
&self,
chunk_id: Uuid,
embedding: &[f32],
) -> Result<(), WorkspaceError> {
self.repo.update_chunk_embedding(chunk_id, embedding).await
}
async fn get_chunks_without_embeddings(
&self,
user_id: &str,
agent_id: Option<Uuid>,
limit: usize,
) -> Result<Vec<MemoryChunk>, WorkspaceError> {
self.repo
.get_chunks_without_embeddings(user_id, agent_id, limit)
.await
}
// ==================== Workspace: Search ====================
async fn hybrid_search(
&self,
user_id: &str,
agent_id: Option<Uuid>,
query: &str,
embedding: Option<&[f32]>,
config: &SearchConfig,
) -> Result<Vec<SearchResult>, WorkspaceError> {
self.repo
.hybrid_search(user_id, agent_id, query, embedding, config)
.await
}
}
+7
View File
@@ -87,14 +87,21 @@ pub enum DatabaseError {
#[error("Serialization error: {0}")]
Serialization(String),
#[cfg(feature = "postgres")]
#[error("PostgreSQL error: {0}")]
Postgres(#[from] tokio_postgres::Error),
#[cfg(feature = "postgres")]
#[error("Pool build error: {0}")]
PoolBuild(#[from] deadpool_postgres::BuildError),
#[cfg(feature = "postgres")]
#[error("Pool runtime error: {0}")]
PoolRuntime(#[from] deadpool_postgres::PoolError),
#[cfg(feature = "libsql")]
#[error("LibSQL error: {0}")]
LibSql(#[from] libsql::Error),
}
/// Channel-related errors.
+5 -6
View File
@@ -144,12 +144,11 @@ impl SuccessEvaluator for RuleBasedEvaluator {
// Check for critical errors
for action in actions.iter().filter(|a| !a.success) {
if let Some(ref error) = action.error {
if error.to_lowercase().contains("critical")
|| error.to_lowercase().contains("fatal")
{
issues.push(format!("Critical error in {}: {}", action.tool_name, error));
}
if let Some(ref error) = action.error
&& (error.to_lowercase().contains("critical")
|| error.to_lowercase().contains("fatal"))
{
issues.push(format!("Critical error in {}: {}", action.tool_name, error));
}
}
+64 -27
View File
@@ -57,7 +57,7 @@ pub struct ExtensionManager {
_tunnel_url: Option<String>,
user_id: String,
/// Optional database store for DB-backed MCP config.
store: Option<Arc<crate::history::Store>>,
store: Option<Arc<dyn crate::db::Database>>,
}
impl ExtensionManager {
@@ -71,7 +71,7 @@ impl ExtensionManager {
wasm_channels_dir: PathBuf,
tunnel_url: Option<String>,
user_id: String,
store: Option<Arc<crate::history::Store>>,
store: Option<Arc<dyn crate::db::Database>>,
) -> Self {
Self {
registry: ExtensionRegistry::new(),
@@ -351,7 +351,7 @@ impl ExtensionManager {
) -> 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
crate::tools::mcp::config::load_mcp_servers_from_db(store.as_ref(), &self.user_id).await
} else {
crate::tools::mcp::config::load_mcp_servers().await
}
@@ -375,7 +375,8 @@ impl ExtensionManager {
) -> 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
crate::tools::mcp::config::add_mcp_server_db(store.as_ref(), &self.user_id, config)
.await
} else {
crate::tools::mcp::config::add_mcp_server(config).await
}
@@ -386,7 +387,8 @@ impl ExtensionManager {
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
crate::tools::mcp::config::remove_mcp_server_db(store.as_ref(), &self.user_id, name)
.await
} else {
crate::tools::mcp::config::remove_mcp_server(name).await
}
@@ -461,7 +463,16 @@ impl ExtensionManager {
name: &str,
url: &str,
) -> Result<InstallResult, ExtensionError> {
// Download the WASM binary
// Require HTTPS to prevent downgrade attacks
if !url.starts_with("https://") {
return Err(ExtensionError::InstallFailed(
"Only HTTPS URLs are allowed for extension downloads".to_string(),
));
}
// 50 MB cap to prevent disk-fill DoS
const MAX_WASM_SIZE: usize = 50 * 1024 * 1024;
let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(60))
.build()
@@ -480,11 +491,36 @@ impl ExtensionManager {
)));
}
// Check Content-Length header before downloading the full body
if let Some(len) = response.content_length()
&& len as usize > MAX_WASM_SIZE
{
return Err(ExtensionError::InstallFailed(format!(
"WASM binary too large ({} bytes, max {} bytes)",
len, MAX_WASM_SIZE
)));
}
let bytes = response
.bytes()
.await
.map_err(|e| ExtensionError::DownloadFailed(e.to_string()))?;
if bytes.len() > MAX_WASM_SIZE {
return Err(ExtensionError::InstallFailed(format!(
"WASM binary too large ({} bytes, max {} bytes)",
bytes.len(),
MAX_WASM_SIZE
)));
}
// Basic WASM magic number check (\0asm)
if bytes.len() < 4 || &bytes[..4] != b"\0asm" {
return Err(ExtensionError::InstallFailed(
"Downloaded file is not a valid WASM binary (bad magic number)".to_string(),
));
}
// Ensure tools directory exists
tokio::fs::create_dir_all(&self.wasm_tools_dir)
.await
@@ -497,9 +533,10 @@ impl ExtensionManager {
.map_err(|e| ExtensionError::InstallFailed(e.to_string()))?;
tracing::info!(
"Installed WASM tool '{}' ({} bytes) to {}",
"Installed WASM tool '{}' ({} bytes) from {} to {}",
name,
bytes.len(),
url,
wasm_path.display()
);
@@ -731,27 +768,27 @@ impl ExtensionManager {
};
// Check env var first
if let Some(ref env_var) = auth.env_var {
if let Ok(value) = std::env::var(env_var) {
// Store the env var value as a secret
let params = CreateSecretParams::new(&auth.secret_name, &value)
.with_provider(name.to_string());
self.secrets
.create(&self.user_id, params)
.await
.map_err(|e| ExtensionError::AuthFailed(e.to_string()))?;
if let Some(ref env_var) = auth.env_var
&& let Ok(value) = std::env::var(env_var)
{
// Store the env var value as a secret
let params =
CreateSecretParams::new(&auth.secret_name, &value).with_provider(name.to_string());
self.secrets
.create(&self.user_id, params)
.await
.map_err(|e| ExtensionError::AuthFailed(e.to_string()))?;
return Ok(AuthResult {
name: name.to_string(),
kind: ExtensionKind::WasmTool,
auth_url: None,
callback_type: None,
instructions: None,
setup_url: None,
awaiting_token: false,
status: "authenticated".to_string(),
});
}
return Ok(AuthResult {
name: name.to_string(),
kind: ExtensionKind::WasmTool,
auth_url: None,
callback_type: None,
instructions: None,
setup_url: None,
awaiting_token: false,
status: "authenticated".to_string(),
});
}
// Check if already authenticated
+5 -1
View File
@@ -5,11 +5,15 @@
//! - Learning from past executions
//! - Analytics and metrics
#[cfg(feature = "postgres")]
mod analytics;
mod store;
#[cfg(feature = "postgres")]
pub use analytics::{JobStats, ToolStats};
#[cfg(feature = "postgres")]
pub use store::Store;
pub use store::{
ConversationMessage, ConversationSummary, JobEventRecord, LlmCallRecord, SandboxJobRecord,
SandboxJobSummary, Store,
SandboxJobSummary, SettingRow,
};
+130
View File
@@ -1,13 +1,18 @@
//! PostgreSQL store for persisting agent data.
use chrono::{DateTime, Utc};
#[cfg(feature = "postgres")]
use deadpool_postgres::{Config, Pool, Runtime};
use rust_decimal::Decimal;
#[cfg(feature = "postgres")]
use tokio_postgres::NoTls;
use uuid::Uuid;
#[cfg(feature = "postgres")]
use crate::config::DatabaseConfig;
#[cfg(feature = "postgres")]
use crate::context::{ActionRecord, JobContext, JobState};
#[cfg(feature = "postgres")]
use crate::error::DatabaseError;
/// Record for an LLM call to be persisted.
@@ -24,11 +29,18 @@ pub struct LlmCallRecord<'a> {
}
/// Database store for the agent.
#[cfg(feature = "postgres")]
pub struct Store {
pool: Pool,
}
#[cfg(feature = "postgres")]
impl Store {
/// Wrap an existing pool (useful when the caller already has a connection).
pub fn from_pool(pool: Pool) -> Self {
Self { pool }
}
/// Create a new store and connect to the database.
pub async fn new(config: &DatabaseConfig) -> Result<Self, DatabaseError> {
let mut cfg = Config::new();
@@ -144,7 +156,12 @@ impl Store {
actual_cost, repair_attempts, created_at, started_at, completed_at
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17)
ON CONFLICT (id) DO UPDATE SET
title = EXCLUDED.title,
description = EXCLUDED.description,
category = EXCLUDED.category,
status = EXCLUDED.status,
estimated_cost = EXCLUDED.estimated_cost,
estimated_time_secs = EXCLUDED.estimated_time_secs,
actual_cost = EXCLUDED.actual_cost,
repair_attempts = EXCLUDED.repair_attempts,
started_at = EXCLUDED.started_at,
@@ -220,6 +237,8 @@ impl Store {
completed_at: row.get("completed_at"),
transitions: Vec::new(), // Not loaded from DB for now
metadata: serde_json::Value::Null,
total_tokens_used: 0,
max_tokens: 0,
}))
}
None => Ok(None),
@@ -464,6 +483,7 @@ pub struct SandboxJobSummary {
pub interrupted: usize,
}
#[cfg(feature = "postgres")]
impl Store {
/// Insert a new sandbox job into `agent_jobs`.
pub async fn save_sandbox_job(&self, job: &SandboxJobRecord) -> Result<(), DatabaseError> {
@@ -565,6 +585,90 @@ impl Store {
.collect())
}
/// List sandbox jobs for a specific user, most recent first.
pub async fn list_sandbox_jobs_for_user(
&self,
user_id: &str,
) -> Result<Vec<SandboxJobRecord>, DatabaseError> {
let conn = self.conn().await?;
let rows = conn
.query(
r#"
SELECT id, title, status, user_id, project_dir,
success, failure_reason, created_at, started_at, completed_at
FROM agent_jobs WHERE source = 'sandbox' AND user_id = $1
ORDER BY created_at DESC
"#,
&[&user_id],
)
.await?;
Ok(rows
.iter()
.map(|r| SandboxJobRecord {
id: r.get("id"),
task: r.get("title"),
status: r.get("status"),
user_id: r.get("user_id"),
project_dir: r
.get::<_, Option<String>>("project_dir")
.unwrap_or_default(),
success: r.get("success"),
failure_reason: r.get("failure_reason"),
created_at: r.get("created_at"),
started_at: r.get("started_at"),
completed_at: r.get("completed_at"),
})
.collect())
}
/// Get a summary of sandbox job counts by status for a specific user.
pub async fn sandbox_job_summary_for_user(
&self,
user_id: &str,
) -> Result<SandboxJobSummary, DatabaseError> {
let conn = self.conn().await?;
let rows = conn
.query(
"SELECT status, COUNT(*) as cnt FROM agent_jobs WHERE source = 'sandbox' AND user_id = $1 GROUP BY status",
&[&user_id],
)
.await?;
let mut summary = SandboxJobSummary::default();
for row in &rows {
let status: String = row.get("status");
let count: i64 = row.get("cnt");
let c = count as usize;
summary.total += c;
match status.as_str() {
"creating" => summary.creating += c,
"running" => summary.running += c,
"completed" => summary.completed += c,
"failed" => summary.failed += c,
"interrupted" => summary.interrupted += c,
_ => {}
}
}
Ok(summary)
}
/// Check if a sandbox job belongs to a specific user.
pub async fn sandbox_job_belongs_to_user(
&self,
job_id: Uuid,
user_id: &str,
) -> Result<bool, DatabaseError> {
let conn = self.conn().await?;
let row = conn
.query_opt(
"SELECT 1 FROM agent_jobs WHERE id = $1 AND user_id = $2 AND source = 'sandbox'",
&[&job_id, &user_id],
)
.await?;
Ok(row.is_some())
}
/// Update sandbox job status and optional timestamps/result.
pub async fn update_sandbox_job_status(
&self,
@@ -656,6 +760,7 @@ pub struct JobEventRecord {
pub created_at: DateTime<Utc>,
}
#[cfg(feature = "postgres")]
impl Store {
/// Persist a job event (fire-and-forget from orchestrator handler).
pub async fn save_job_event(
@@ -728,10 +833,12 @@ impl Store {
// ==================== Routines ====================
#[cfg(feature = "postgres")]
use crate::agent::routine::{
NotifyConfig, Routine, RoutineAction, RoutineGuardrails, RoutineRun, RunStatus, Trigger,
};
#[cfg(feature = "postgres")]
impl Store {
/// Create a new routine.
pub async fn create_routine(&self, routine: &Routine) -> Result<(), DatabaseError> {
@@ -1032,6 +1139,7 @@ impl Store {
}
}
#[cfg(feature = "postgres")]
fn row_to_routine(row: &tokio_postgres::Row) -> Result<Routine, DatabaseError> {
let trigger_type: String = row.get("trigger_type");
let trigger_config: serde_json::Value = row.get("trigger_config");
@@ -1076,6 +1184,7 @@ fn row_to_routine(row: &tokio_postgres::Row) -> Result<Routine, DatabaseError> {
})
}
#[cfg(feature = "postgres")]
fn row_to_routine_run(row: &tokio_postgres::Row) -> Result<RoutineRun, DatabaseError> {
let status_str: String = row.get("status");
let status: RunStatus = status_str
@@ -1121,6 +1230,7 @@ pub struct ConversationMessage {
pub created_at: DateTime<Utc>,
}
#[cfg(feature = "postgres")]
impl Store {
/// Ensure a conversation row exists for a given UUID.
///
@@ -1258,6 +1368,22 @@ impl Store {
Ok(id)
}
/// Check whether a conversation belongs to the given user.
pub async fn conversation_belongs_to_user(
&self,
conversation_id: Uuid,
user_id: &str,
) -> Result<bool, DatabaseError> {
let conn = self.conn().await?;
let row = conn
.query_opt(
"SELECT 1 FROM conversations WHERE id = $1 AND user_id = $2",
&[&conversation_id, &user_id],
)
.await?;
Ok(row.is_some())
}
/// Load messages for a conversation with cursor-based pagination.
///
/// Returns `(messages_oldest_first, has_more)`.
@@ -1375,6 +1501,7 @@ impl Store {
}
}
#[cfg(feature = "postgres")]
fn parse_job_state(s: &str) -> JobState {
match s {
"pending" => JobState::Pending,
@@ -1391,8 +1518,10 @@ fn parse_job_state(s: &str) -> JobState {
// ==================== Tool Failures ====================
#[cfg(feature = "postgres")]
use crate::agent::BrokenTool;
#[cfg(feature = "postgres")]
impl Store {
/// Record a tool failure (upsert: increment count if exists).
pub async fn record_tool_failure(
@@ -1486,6 +1615,7 @@ pub struct SettingRow {
pub updated_at: DateTime<Utc>,
}
#[cfg(feature = "postgres")]
impl Store {
/// Get a single setting by key.
pub async fn get_setting(
+3
View File
@@ -44,6 +44,7 @@ pub mod channels;
pub mod cli;
pub mod config;
pub mod context;
pub mod db;
pub mod error;
pub mod estimation;
pub mod evaluation;
@@ -58,6 +59,8 @@ pub mod secrets;
pub mod settings;
pub mod setup;
pub mod tools;
pub mod tracing_fmt;
pub mod util;
pub mod worker;
pub mod workspace;
+483
View File
@@ -0,0 +1,483 @@
//! Multi-provider LLM failover.
//!
//! Wraps multiple LlmProvider instances and tries each in sequence
//! until one succeeds. Transparent to callers --- same LlmProvider trait.
use std::future::Future;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use async_trait::async_trait;
use rust_decimal::Decimal;
use crate::error::LlmError;
use crate::llm::provider::{
CompletionRequest, CompletionResponse, LlmProvider, ToolCompletionRequest,
ToolCompletionResponse,
};
/// Returns `true` if the error is transient and the request should be retried
/// on the next provider in the failover chain.
///
/// Retryable: `RequestFailed`, `RateLimited`, `InvalidResponse`,
/// `SessionRenewalFailed`, `ModelNotAvailable`, `Http`, `Io`.
///
/// `ModelNotAvailable` is retryable because the next provider in the chain may
/// offer a different model, so it's worth trying.
///
/// Non-retryable errors (`AuthFailed`, `SessionExpired`, `ContextLengthExceeded`)
/// propagate immediately because a different provider won't fix them.
fn is_retryable(err: &LlmError) -> bool {
matches!(
err,
LlmError::RequestFailed { .. }
| LlmError::RateLimited { .. }
| LlmError::InvalidResponse { .. }
| LlmError::SessionRenewalFailed { .. }
// ModelNotAvailable is retryable: the next provider may offer a different model.
| LlmError::ModelNotAvailable { .. }
| LlmError::Http(_)
| LlmError::Io(_)
)
}
/// An LLM provider that wraps multiple providers and tries each in sequence
/// on transient failures.
///
/// The first provider in the list is the primary. If it fails with a retryable
/// error, the next provider is tried, and so on. Non-retryable errors
/// (e.g. `AuthFailed`, `ContextLengthExceeded`) propagate immediately.
pub struct FailoverProvider {
providers: Vec<Arc<dyn LlmProvider>>,
/// Index of the provider that last handled a request successfully.
/// Used by `model_name()` and `cost_per_token()` so downstream cost
/// tracking reflects the provider that actually served the request.
last_used: AtomicUsize,
}
impl FailoverProvider {
/// Create a new failover provider.
///
/// Returns an error if `providers` is empty.
pub fn new(providers: Vec<Arc<dyn LlmProvider>>) -> Result<Self, LlmError> {
if providers.is_empty() {
return Err(LlmError::RequestFailed {
provider: "failover".to_string(),
reason: "FailoverProvider requires at least one provider".to_string(),
});
}
Ok(Self {
providers,
last_used: AtomicUsize::new(0),
})
}
/// Try each provider in sequence until one succeeds or all fail.
async fn try_providers<T, F, Fut>(&self, mut call: F) -> Result<T, LlmError>
where
F: FnMut(Arc<dyn LlmProvider>) -> Fut,
Fut: Future<Output = Result<T, LlmError>>,
{
let mut last_error: Option<LlmError> = None;
for (i, provider) in self.providers.iter().enumerate() {
let result = call(Arc::clone(provider)).await;
match result {
Ok(response) => {
self.last_used.store(i, Ordering::Relaxed);
return Ok(response);
}
Err(err) => {
if !is_retryable(&err) {
return Err(err);
}
if i + 1 < self.providers.len() {
tracing::warn!(
provider = %provider.model_name(),
error = %err,
next_provider = %self.providers[i + 1].model_name(),
"Provider failed with retryable error, trying next provider"
);
}
last_error = Some(err);
}
}
}
// SAFETY: providers is non-empty (checked in `new`), so at least one
// iteration ran and `last_error` is `Some`.
Err(last_error.expect("providers list is non-empty"))
}
}
#[async_trait]
impl LlmProvider for FailoverProvider {
fn model_name(&self) -> &str {
self.providers[self.last_used.load(Ordering::Relaxed)].model_name()
}
fn cost_per_token(&self) -> (Decimal, Decimal) {
self.providers[self.last_used.load(Ordering::Relaxed)].cost_per_token()
}
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
self.try_providers(|provider| {
let req = request.clone();
async move { provider.complete(req).await }
})
.await
}
async fn complete_with_tools(
&self,
request: ToolCompletionRequest,
) -> Result<ToolCompletionResponse, LlmError> {
self.try_providers(|provider| {
let req = request.clone();
async move { provider.complete_with_tools(req).await }
})
.await
}
async fn list_models(&self) -> Result<Vec<String>, LlmError> {
let mut all_models = Vec::new();
for provider in &self.providers {
match provider.list_models().await {
Ok(models) => all_models.extend(models),
Err(err) => {
tracing::warn!(
provider = %provider.model_name(),
error = %err,
"Failed to list models from provider, skipping"
);
}
}
}
all_models.sort();
all_models.dedup();
Ok(all_models)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Mutex;
use std::time::Duration;
use crate::llm::provider::{CompletionResponse, FinishReason, ToolCompletionResponse};
/// A mock LLM provider that returns a predetermined result.
struct MockProvider {
name: String,
input_cost: Decimal,
output_cost: Decimal,
complete_result: Mutex<Option<Result<CompletionResponse, LlmError>>>,
tool_complete_result: Mutex<Option<Result<ToolCompletionResponse, LlmError>>>,
}
impl MockProvider {
fn succeeding(name: &str, content: &str) -> Self {
Self {
name: name.to_string(),
input_cost: Decimal::ZERO,
output_cost: Decimal::ZERO,
complete_result: Mutex::new(Some(Ok(CompletionResponse {
content: content.to_string(),
input_tokens: 10,
output_tokens: 5,
finish_reason: FinishReason::Stop,
response_id: None,
}))),
tool_complete_result: Mutex::new(Some(Ok(ToolCompletionResponse {
content: Some(content.to_string()),
tool_calls: vec![],
input_tokens: 10,
output_tokens: 5,
finish_reason: FinishReason::Stop,
response_id: None,
}))),
}
}
fn succeeding_with_cost(
name: &str,
content: &str,
input_cost: Decimal,
output_cost: Decimal,
) -> Self {
Self {
input_cost,
output_cost,
..Self::succeeding(name, content)
}
}
fn failing_retryable(name: &str) -> Self {
Self {
name: name.to_string(),
input_cost: Decimal::ZERO,
output_cost: Decimal::ZERO,
complete_result: Mutex::new(Some(Err(LlmError::RequestFailed {
provider: name.to_string(),
reason: "server error".to_string(),
}))),
tool_complete_result: Mutex::new(Some(Err(LlmError::RequestFailed {
provider: name.to_string(),
reason: "server error".to_string(),
}))),
}
}
fn failing_non_retryable(name: &str) -> Self {
Self {
name: name.to_string(),
input_cost: Decimal::ZERO,
output_cost: Decimal::ZERO,
complete_result: Mutex::new(Some(Err(LlmError::AuthFailed {
provider: name.to_string(),
}))),
tool_complete_result: Mutex::new(Some(Err(LlmError::AuthFailed {
provider: name.to_string(),
}))),
}
}
fn failing_rate_limited(name: &str) -> Self {
Self {
name: name.to_string(),
input_cost: Decimal::ZERO,
output_cost: Decimal::ZERO,
complete_result: Mutex::new(Some(Err(LlmError::RateLimited {
provider: name.to_string(),
retry_after: Some(Duration::from_secs(30)),
}))),
tool_complete_result: Mutex::new(Some(Err(LlmError::RateLimited {
provider: name.to_string(),
retry_after: Some(Duration::from_secs(30)),
}))),
}
}
}
#[async_trait]
impl LlmProvider for MockProvider {
fn model_name(&self) -> &str {
&self.name
}
fn cost_per_token(&self) -> (Decimal, Decimal) {
(self.input_cost, self.output_cost)
}
async fn complete(
&self,
_request: CompletionRequest,
) -> Result<CompletionResponse, LlmError> {
self.complete_result
.lock()
.unwrap()
.take()
.expect("MockProvider::complete called more than once")
}
async fn complete_with_tools(
&self,
_request: ToolCompletionRequest,
) -> Result<ToolCompletionResponse, LlmError> {
self.tool_complete_result
.lock()
.unwrap()
.take()
.expect("MockProvider::complete_with_tools called more than once")
}
async fn list_models(&self) -> Result<Vec<String>, LlmError> {
Ok(vec![self.name.clone()])
}
}
fn make_request() -> CompletionRequest {
CompletionRequest::new(vec![crate::llm::ChatMessage::user("hello")])
}
fn make_tool_request() -> ToolCompletionRequest {
ToolCompletionRequest::new(vec![crate::llm::ChatMessage::user("hello")], vec![])
}
// Test 1: Primary succeeds, no failover occurs.
#[tokio::test]
async fn primary_succeeds_no_failover() {
let primary = Arc::new(MockProvider::succeeding("primary", "primary response"));
let fallback = Arc::new(MockProvider::succeeding("fallback", "fallback response"));
let failover = FailoverProvider::new(vec![primary, fallback]).unwrap();
let response = failover.complete(make_request()).await.unwrap();
assert_eq!(response.content, "primary response");
}
// Test 2: Primary fails with retryable error, fallback succeeds.
#[tokio::test]
async fn primary_fails_retryable_fallback_succeeds() {
let primary = Arc::new(MockProvider::failing_retryable("primary"));
let fallback = Arc::new(MockProvider::succeeding("fallback", "fallback response"));
let failover = FailoverProvider::new(vec![primary, fallback]).unwrap();
let response = failover.complete(make_request()).await.unwrap();
assert_eq!(response.content, "fallback response");
}
// Test 3: All providers fail, returns last error.
#[tokio::test]
async fn all_providers_fail_returns_last_error() {
let primary = Arc::new(MockProvider::failing_retryable("primary"));
let fallback = Arc::new(MockProvider::failing_retryable("fallback"));
let failover = FailoverProvider::new(vec![primary, fallback]).unwrap();
let err = failover.complete(make_request()).await.unwrap_err();
match err {
LlmError::RequestFailed { provider, .. } => {
assert_eq!(provider, "fallback");
}
other => panic!("expected RequestFailed, got: {other:?}"),
}
}
// Test 4: Non-retryable error fails immediately, no failover.
#[tokio::test]
async fn non_retryable_error_fails_immediately() {
let primary = Arc::new(MockProvider::failing_non_retryable("primary"));
let fallback = Arc::new(MockProvider::succeeding("fallback", "fallback response"));
let failover = FailoverProvider::new(vec![primary, fallback]).unwrap();
let err = failover.complete(make_request()).await.unwrap_err();
match err {
LlmError::AuthFailed { provider } => {
assert_eq!(provider, "primary");
}
other => panic!("expected AuthFailed, got: {other:?}"),
}
}
// Test 5: Three providers, first two fail (retryable), third succeeds.
#[tokio::test]
async fn three_providers_first_two_fail_third_succeeds() {
let p1 = Arc::new(MockProvider::failing_retryable("provider-1"));
let p2 = Arc::new(MockProvider::failing_rate_limited("provider-2"));
let p3 = Arc::new(MockProvider::succeeding("provider-3", "third time lucky"));
let failover = FailoverProvider::new(vec![p1, p2, p3]).unwrap();
let response = failover.complete(make_request()).await.unwrap();
assert_eq!(response.content, "third time lucky");
}
// Test: complete_with_tools follows same failover logic.
#[tokio::test]
async fn complete_with_tools_failover() {
let primary = Arc::new(MockProvider::failing_retryable("primary"));
let fallback = Arc::new(MockProvider::succeeding("fallback", "tools fallback"));
let failover = FailoverProvider::new(vec![primary, fallback]).unwrap();
let response = failover
.complete_with_tools(make_tool_request())
.await
.unwrap();
assert_eq!(response.content.as_deref(), Some("tools fallback"));
}
// Test: model_name and cost_per_token reflect the last-used provider.
#[tokio::test]
async fn model_name_and_cost_track_last_used_provider() {
let fallback_cost = Decimal::new(15, 6); // 0.000015
let primary = Arc::new(MockProvider::failing_retryable("primary-model"));
let fallback = Arc::new(MockProvider::succeeding_with_cost(
"fallback-model",
"ok",
fallback_cost,
fallback_cost,
));
let failover = FailoverProvider::new(vec![primary, fallback]).unwrap();
// Before any call, defaults to primary (index 0).
assert_eq!(failover.model_name(), "primary-model");
assert_eq!(failover.cost_per_token(), (Decimal::ZERO, Decimal::ZERO));
// After failover, should reflect the fallback provider.
let _ = failover.complete(make_request()).await.unwrap();
assert_eq!(failover.model_name(), "fallback-model");
assert_eq!(failover.cost_per_token(), (fallback_cost, fallback_cost));
}
// Test: list_models aggregates from all providers.
#[tokio::test]
async fn list_models_aggregates_all() {
let p1 = Arc::new(MockProvider::succeeding("model-a", "ok"));
let p2 = Arc::new(MockProvider::succeeding("model-b", "ok"));
let failover = FailoverProvider::new(vec![p1, p2]).unwrap();
let models = failover.list_models().await.unwrap();
assert!(models.contains(&"model-a".to_string()));
assert!(models.contains(&"model-b".to_string()));
}
// Test: is_retryable correctly classifies errors.
#[test]
fn retryable_classification() {
// Retryable
assert!(is_retryable(&LlmError::RequestFailed {
provider: "p".into(),
reason: "err".into(),
}));
assert!(is_retryable(&LlmError::RateLimited {
provider: "p".into(),
retry_after: None,
}));
assert!(is_retryable(&LlmError::InvalidResponse {
provider: "p".into(),
reason: "bad json".into(),
}));
assert!(is_retryable(&LlmError::SessionRenewalFailed {
provider: "p".into(),
reason: "timeout".into(),
}));
assert!(is_retryable(&LlmError::Io(std::io::Error::new(
std::io::ErrorKind::ConnectionReset,
"reset"
))));
assert!(is_retryable(&LlmError::ModelNotAvailable {
provider: "p".into(),
model: "m".into(),
}));
// Non-retryable
assert!(!is_retryable(&LlmError::AuthFailed {
provider: "p".into(),
}));
assert!(!is_retryable(&LlmError::SessionExpired {
provider: "p".into(),
}));
assert!(!is_retryable(&LlmError::ContextLengthExceeded {
used: 100_000,
limit: 50_000,
}));
}
// Test: empty providers list returns error (not panic).
#[test]
fn empty_providers_returns_error() {
let result = FailoverProvider::new(vec![]);
assert!(result.is_err());
}
}
+26 -13
View File
@@ -8,20 +8,26 @@
//! - **OpenAI-compatible**: Any endpoint that speaks the OpenAI API
mod costs;
pub mod failover;
mod nearai;
mod nearai_chat;
mod provider;
mod reasoning;
mod retry;
mod rig_adapter;
pub mod session;
pub use failover::FailoverProvider;
pub use nearai::{ModelInfo, NearAiProvider};
pub use nearai_chat::NearAiChatProvider;
pub use provider::{
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, ModelMetadata,
Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse, ToolDefinition, ToolResult,
};
pub use reasoning::{ActionPlan, Reasoning, ReasoningContext, RespondResult, ToolSelection};
pub use reasoning::{
ActionPlan, Reasoning, ReasoningContext, RespondOutput, RespondResult, TokenUsage,
ToolSelection,
};
pub use rig_adapter::RigAdapter;
pub use session::{SessionConfig, SessionManager, create_session_manager};
@@ -30,7 +36,7 @@ use std::sync::Arc;
use rig::client::CompletionClient;
use secrecy::ExposeSecret;
use crate::config::{LlmBackend, LlmConfig, NearAiApiMode};
use crate::config::{LlmBackend, LlmConfig, NearAiApiMode, NearAiConfig};
use crate::error::LlmError;
/// Create an LLM provider based on configuration.
@@ -43,7 +49,7 @@ pub fn create_llm_provider(
session: Arc<SessionManager>,
) -> Result<Arc<dyn LlmProvider>, LlmError> {
match config.backend {
LlmBackend::NearAi => create_nearai_provider(config, session),
LlmBackend::NearAi => create_llm_provider_with_config(&config.nearai, session),
LlmBackend::OpenAi => create_openai_provider(config),
LlmBackend::Anthropic => create_anthropic_provider(config),
LlmBackend::Ollama => create_ollama_provider(config),
@@ -51,21 +57,28 @@ pub fn create_llm_provider(
}
}
fn create_nearai_provider(
config: &LlmConfig,
/// Create an LLM provider from a `NearAiConfig` directly.
///
/// This is useful when constructing additional providers for failover,
/// where only the model name differs from the primary config.
pub fn create_llm_provider_with_config(
config: &NearAiConfig,
session: Arc<SessionManager>,
) -> Result<Arc<dyn LlmProvider>, LlmError> {
match config.nearai.api_mode {
match config.api_mode {
NearAiApiMode::Responses => {
tracing::info!("Using NEAR AI Responses API (chat-api) with session auth");
Ok(Arc::new(NearAiProvider::new(
config.nearai.clone(),
session,
)))
tracing::info!(
model = %config.model,
"Using Responses API (chat-api) with session auth"
);
Ok(Arc::new(NearAiProvider::new(config.clone(), session)))
}
NearAiApiMode::ChatCompletions => {
tracing::info!("Using NEAR AI Chat Completions API (cloud-api) with API key auth");
Ok(Arc::new(NearAiChatProvider::new(config.nearai.clone())?))
tracing::info!(
model = %config.model,
"Using Chat Completions API (cloud-api) with API key auth"
);
Ok(Arc::new(NearAiChatProvider::new(config.clone())?))
}
}
}
+135 -93
View File
@@ -19,6 +19,7 @@ use crate::llm::provider::{
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, Role, ToolCall,
ToolCompletionRequest, ToolCompletionResponse,
};
use crate::llm::retry::{is_retryable_status, retry_backoff_delay};
use crate::llm::session::SessionManager;
/// Information about an available model from NEAR AI API.
@@ -209,20 +210,20 @@ impl NearAiProvider {
data: Option<Vec<ModelEntry>>,
}
if let Ok(resp) = serde_json::from_str::<ModelsResponse>(&response_text) {
if let Some(entries) = resp.models.or(resp.data) {
let models: Vec<ModelInfo> = entries
.into_iter()
.filter_map(|e| {
e.get_name().map(|name| ModelInfo {
name,
provider: None,
})
if let Ok(resp) = serde_json::from_str::<ModelsResponse>(&response_text)
&& let Some(entries) = resp.models.or(resp.data)
{
let models: Vec<ModelInfo> = entries
.into_iter()
.filter_map(|e| {
e.get_name().map(|name| ModelInfo {
name,
provider: None,
})
.collect();
if !models.is_empty() {
return Ok(models);
}
})
.collect();
if !models.is_empty() {
return Ok(models);
}
}
@@ -270,88 +271,139 @@ impl NearAiProvider {
}
}
/// Inner request implementation without retry logic.
/// Inner request implementation with retry logic for transient errors.
///
/// Retries on HTTP 429, 500, 502, 503, 504 with exponential backoff.
/// Does not retry on client errors (400, 401, 403, 404) or parse errors.
async fn send_request_inner<T: Serialize + std::fmt::Debug, R: for<'de> Deserialize<'de>>(
&self,
path: &str,
body: &T,
) -> Result<R, LlmError> {
let url = self.api_url(path);
let token = self.session.get_token().await?;
let max_retries = self.config.max_retries;
tracing::debug!("Sending request to NEAR AI: {}", url);
tracing::debug!("Request body: {:?}", body);
for attempt in 0..=max_retries {
let token = self.session.get_token().await?;
let response = self
.client
.post(&url)
.header("Authorization", format!("Bearer {}", token.expose_secret()))
.header("Content-Type", "application/json")
.json(body)
.send()
.await
.map_err(|e| {
tracing::error!("NEAR AI request failed: {}", e);
e
})?;
tracing::debug!(
"Sending request to NEAR AI: {} (attempt {})",
url,
attempt + 1
);
tracing::debug!("Request body: {:?}", body);
let status = response.status();
let response_text = response.text().await.unwrap_or_default();
let response = self
.client
.post(&url)
.header("Authorization", format!("Bearer {}", token.expose_secret()))
.header("Content-Type", "application/json")
.json(body)
.send()
.await;
tracing::debug!("NEAR AI response status: {}", status);
tracing::debug!("NEAR AI response body: {}", response_text);
let response = match response {
Ok(r) => r,
Err(e) => {
tracing::error!("NEAR AI request failed: {}", e);
// Network errors (timeout, connection refused) are transient
if attempt < max_retries {
let delay = retry_backoff_delay(attempt);
tracing::warn!(
"NEAR AI request error (attempt {}/{}), retrying in {:?}: {}",
attempt + 1,
max_retries + 1,
delay,
e,
);
tokio::time::sleep(delay).await;
continue;
}
return Err(e.into());
}
};
if !status.is_success() {
// Check for session expiration (401 with specific message patterns)
if status.as_u16() == 401 {
let is_session_expired = response_text.to_lowercase().contains("session")
&& (response_text.to_lowercase().contains("expired")
|| response_text.to_lowercase().contains("invalid"));
let status = response.status();
let response_text = response.text().await.unwrap_or_default();
if is_session_expired {
return Err(LlmError::SessionExpired {
tracing::debug!("NEAR AI response status: {}", status);
tracing::debug!("NEAR AI response body: {}", response_text);
if !status.is_success() {
let status_code = status.as_u16();
// Check for session expiration (401 with specific message patterns)
if status_code == 401 {
let lower = response_text.to_lowercase();
let is_session_expired = lower.contains("session")
&& (lower.contains("expired") || lower.contains("invalid"));
if is_session_expired {
return Err(LlmError::SessionExpired {
provider: "nearai".to_string(),
});
}
// Generic 401 -- not retryable
return Err(LlmError::AuthFailed {
provider: "nearai".to_string(),
});
}
// Generic 401 without session expiration indication
return Err(LlmError::AuthFailed {
provider: "nearai".to_string(),
});
}
// Check if this is a transient error worth retrying
if is_retryable_status(status_code) && attempt < max_retries {
let delay = retry_backoff_delay(attempt);
tracing::warn!(
"NEAR AI returned HTTP {} (attempt {}/{}), retrying in {:?}",
status_code,
attempt + 1,
max_retries + 1,
delay,
);
tokio::time::sleep(delay).await;
continue;
}
// Try to parse as JSON error
if let Ok(error) = serde_json::from_str::<NearAiErrorResponse>(&response_text) {
if status.as_u16() == 429 {
return Err(LlmError::RateLimited {
// Non-retryable error or exhausted retries
if let Ok(error) = serde_json::from_str::<NearAiErrorResponse>(&response_text) {
if status_code == 429 {
return Err(LlmError::RateLimited {
provider: "nearai".to_string(),
retry_after: None,
});
}
return Err(LlmError::RequestFailed {
provider: "nearai".to_string(),
retry_after: None,
reason: error.error,
});
}
return Err(LlmError::RequestFailed {
provider: "nearai".to_string(),
reason: error.error,
reason: format!("HTTP {}: {}", status, response_text),
});
}
return Err(LlmError::RequestFailed {
provider: "nearai".to_string(),
reason: format!("HTTP {}: {}", status, response_text),
});
// Success -- parse the response
return match serde_json::from_str::<R>(&response_text) {
Ok(parsed) => Ok(parsed),
Err(e) => {
tracing::debug!("Response is not expected JSON format: {}", e);
tracing::debug!("Will try alternative parsing in caller");
Err(LlmError::InvalidResponse {
provider: "nearai".to_string(),
reason: format!("Parse error: {}. Raw: {}", e, response_text),
})
}
};
}
// Try to parse as our expected type
match serde_json::from_str::<R>(&response_text) {
Ok(parsed) => Ok(parsed),
Err(e) => {
tracing::debug!("Response is not expected JSON format: {}", e);
tracing::debug!("Will try alternative parsing in caller");
Err(LlmError::InvalidResponse {
provider: "nearai".to_string(),
reason: format!("Parse error: {}. Raw: {}", e, response_text),
})
}
}
// This is unreachable because the loop always returns, but the compiler
// cannot prove that. Return a generic error as a safety net.
Err(LlmError::RequestFailed {
provider: "nearai".to_string(),
reason: "retry loop exited unexpectedly".to_string(),
})
}
}
@@ -456,7 +508,7 @@ impl LlmProvider for NearAiProvider {
Err(e) => return Err(e),
};
tracing::debug!("NEAR AI response: {:?}", response);
tracing::debug!("NEAR AI response: output_items={}", response.output.len());
// Extract text from response output
// Try multiple formats since API response shape may vary
@@ -464,11 +516,6 @@ impl LlmProvider for NearAiProvider {
.output
.iter()
.filter_map(|item| {
tracing::debug!(
"Processing output item: type={}, text={:?}",
item.item_type,
item.text
);
if item.item_type == "message" {
// First check for direct text field on item
if let Some(ref text) = item.text {
@@ -479,11 +526,6 @@ impl LlmProvider for NearAiProvider {
contents
.iter()
.filter_map(|c| {
tracing::debug!(
"Content item: type={}, text={:?}",
c.content_type,
c.text
);
// Accept various content types that might contain text
match c.content_type.as_str() {
"output_text" | "text" => c.text.clone(),
@@ -694,21 +736,21 @@ impl LlmProvider for NearAiProvider {
}
}
}
} else if item.item_type == "function_call" {
if let (Some(name), Some(call_id)) = (&item.name, &item.call_id) {
// Parse arguments JSON string into Value
let arguments = item
.arguments
.as_ref()
.and_then(|s| serde_json::from_str(s).ok())
.unwrap_or(serde_json::Value::Object(Default::default()));
} else if item.item_type == "function_call"
&& let (Some(name), Some(call_id)) = (&item.name, &item.call_id)
{
// Parse arguments JSON string into Value
let arguments = item
.arguments
.as_ref()
.and_then(|s| serde_json::from_str(s).ok())
.unwrap_or(serde_json::Value::Object(Default::default()));
tool_calls.push(ToolCall {
id: call_id.clone(),
name: name.clone(),
arguments,
});
}
tool_calls.push(ToolCall {
id: call_id.clone(),
name: name.clone(),
arguments,
});
}
}
+270 -37
View File
@@ -16,6 +16,7 @@ use crate::llm::provider::{
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, ModelMetadata,
Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse,
};
use crate::llm::retry::{is_retryable_status, retry_backoff_delay};
/// NEAR AI Chat Completions API provider.
pub struct NearAiChatProvider {
@@ -62,63 +63,116 @@ impl NearAiChatProvider {
.unwrap_or_default()
}
/// Send a request to the chat completions API.
/// Send a request to the chat completions API with retry on transient errors.
///
/// Retries on HTTP 429, 500, 502, 503, 504 with exponential backoff.
/// Does not retry on client errors (400, 401, 403, 404) or parse errors.
async fn send_request<T: Serialize, R: for<'de> Deserialize<'de>>(
&self,
body: &T,
) -> Result<R, LlmError> {
let url = self.api_url("chat/completions");
let max_retries = self.config.max_retries;
tracing::debug!("Sending request to NEAR AI Chat: {}", url);
for attempt in 0..=max_retries {
tracing::debug!(
"Sending request to NEAR AI Chat: {} (attempt {})",
url,
attempt + 1,
);
// 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);
}
if tracing::enabled!(tracing::Level::DEBUG)
&& let Ok(json) = serde_json::to_string(body)
{
tracing::debug!("NEAR AI Chat request body: {}", json);
}
let response = self
.client
.post(&url)
.header("Authorization", format!("Bearer {}", self.api_key()))
.header("Content-Type", "application/json")
.json(body)
.send()
.await
.map_err(|e| {
tracing::error!("NEAR AI Chat request failed: {}", e);
LlmError::RequestFailed {
provider: "nearai_chat".to_string(),
reason: e.to_string(),
let response = self
.client
.post(&url)
.header("Authorization", format!("Bearer {}", self.api_key()))
.header("Content-Type", "application/json")
.json(body)
.send()
.await;
let response = match response {
Ok(r) => r,
Err(e) => {
tracing::error!("NEAR AI Chat request failed: {}", e);
if attempt < max_retries {
let delay = retry_backoff_delay(attempt);
tracing::warn!(
"NEAR AI Chat request error (attempt {}/{}), retrying in {:?}: {}",
attempt + 1,
max_retries + 1,
delay,
e,
);
tokio::time::sleep(delay).await;
continue;
}
return Err(LlmError::RequestFailed {
provider: "nearai_chat".to_string(),
reason: e.to_string(),
});
}
})?;
};
let status = response.status();
let response_text = response.text().await.unwrap_or_default();
let status = response.status();
let response_text = response.text().await.unwrap_or_default();
tracing::debug!("NEAR AI Chat response status: {}", status);
tracing::debug!("NEAR AI Chat response body: {}", response_text);
tracing::debug!("NEAR AI Chat response status: {}", status);
tracing::debug!("NEAR AI Chat response body: {}", response_text);
if !status.is_success() {
if status.as_u16() == 401 {
return Err(LlmError::AuthFailed {
if !status.is_success() {
let status_code = status.as_u16();
// Auth errors are not retryable
if status_code == 401 {
return Err(LlmError::AuthFailed {
provider: "nearai_chat".to_string(),
});
}
// Transient errors: retry with backoff
if is_retryable_status(status_code) && attempt < max_retries {
let delay = retry_backoff_delay(attempt);
tracing::warn!(
"NEAR AI Chat returned HTTP {} (attempt {}/{}), retrying in {:?}",
status_code,
attempt + 1,
max_retries + 1,
delay,
);
tokio::time::sleep(delay).await;
continue;
}
// Non-retryable or exhausted retries
if status_code == 429 {
return Err(LlmError::RateLimited {
provider: "nearai_chat".to_string(),
retry_after: None,
});
}
return Err(LlmError::RequestFailed {
provider: "nearai_chat".to_string(),
reason: format!("HTTP {}: {}", status, response_text),
});
}
if status.as_u16() == 429 {
return Err(LlmError::RateLimited {
provider: "nearai_chat".to_string(),
retry_after: None,
});
}
return Err(LlmError::RequestFailed {
// Success — parse the response
return serde_json::from_str(&response_text).map_err(|e| LlmError::InvalidResponse {
provider: "nearai_chat".to_string(),
reason: format!("HTTP {}: {}", status, response_text),
reason: format!("JSON parse error: {}. Raw: {}", e, response_text),
});
}
serde_json::from_str(&response_text).map_err(|e| LlmError::InvalidResponse {
// Safety net: unreachable because the loop always returns
Err(LlmError::RequestFailed {
provider: "nearai_chat".to_string(),
reason: format!("JSON parse error: {}. Raw: {}", e, response_text),
reason: "retry loop exited unexpectedly".to_string(),
})
}
@@ -222,6 +276,12 @@ impl LlmProvider for NearAiChatProvider {
let messages: Vec<ChatCompletionMessage> =
req.messages.into_iter().map(|m| m.into()).collect();
// NEAR AI cloud-api does not support multi-turn tool calling (rejects
// any request containing role:"tool" messages with HTTP 400). Rewrite
// tool-call / tool-result pairs into plain text so the conversation
// history is preserved without using unsupported message roles.
let messages = flatten_tool_messages(messages);
let tools: Vec<ChatCompletionTool> = req
.tools
.into_iter()
@@ -367,6 +427,64 @@ struct ChatCompletionMessage {
tool_calls: Option<Vec<ChatCompletionToolCall>>,
}
/// Rewrite tool-call / tool-result messages into plain assistant/user text.
///
/// NEAR AI cloud-api does not support the OpenAI multi-turn tool-calling
/// protocol (`role: "tool"` messages). This function converts:
/// - Assistant messages with `tool_calls` → assistant text describing the calls
/// - Tool result messages (`role: "tool"`) → user messages with the result
///
/// Non-tool messages pass through unchanged.
fn flatten_tool_messages(messages: Vec<ChatCompletionMessage>) -> Vec<ChatCompletionMessage> {
let has_tool_msgs = messages.iter().any(|m| m.role == "tool");
if !has_tool_msgs {
return messages;
}
tracing::debug!("Flattening tool messages for NEAR AI compatibility");
messages
.into_iter()
.map(|msg| {
if let (true, Some(calls)) = (msg.role == "assistant", &msg.tool_calls) {
// Convert assistant tool_calls into descriptive text
let mut parts: Vec<String> = Vec::new();
if let Some(ref text) = msg.content
&& !text.is_empty()
{
parts.push(text.clone());
}
for tc in calls {
parts.push(format!(
"[Called tool `{}` with arguments: {}]",
tc.function.name, tc.function.arguments
));
}
ChatCompletionMessage {
role: "assistant".to_string(),
content: Some(parts.join("\n")),
tool_call_id: None,
name: None,
tool_calls: None,
}
} else if msg.role == "tool" {
// Convert tool result into a user message
let tool_name = msg.name.as_deref().unwrap_or("unknown");
let result = msg.content.as_deref().unwrap_or("");
ChatCompletionMessage {
role: "user".to_string(),
content: Some(format!("[Tool `{}` returned: {}]", tool_name, result)),
tool_call_id: None,
name: None,
tool_calls: None,
}
} else {
msg
}
})
.collect()
}
impl From<ChatMessage> for ChatCompletionMessage {
fn from(msg: ChatMessage) -> Self {
let role = match msg.role {
@@ -544,4 +662,119 @@ mod tests {
serde_json::from_str(&calls[0].function.arguments).expect("valid JSON string");
assert_eq!(parsed["key"], "value");
}
#[test]
fn test_flatten_no_tool_messages_passthrough() {
let messages = vec![
ChatCompletionMessage {
role: "system".to_string(),
content: Some("You are helpful.".to_string()),
tool_call_id: None,
name: None,
tool_calls: None,
},
ChatCompletionMessage {
role: "user".to_string(),
content: Some("Hello".to_string()),
tool_call_id: None,
name: None,
tool_calls: None,
},
];
let result = flatten_tool_messages(messages);
assert_eq!(result.len(), 2);
assert_eq!(result[0].role, "system");
assert_eq!(result[1].role, "user");
}
#[test]
fn test_flatten_tool_call_and_result() {
let messages = vec![
ChatCompletionMessage {
role: "user".to_string(),
content: Some("test".to_string()),
tool_call_id: None,
name: None,
tool_calls: None,
},
ChatCompletionMessage {
role: "assistant".to_string(),
content: None,
tool_call_id: None,
name: None,
tool_calls: Some(vec![ChatCompletionToolCall {
id: "call_1".to_string(),
call_type: "function".to_string(),
function: ChatCompletionToolCallFunction {
name: "echo".to_string(),
arguments: r#"{"message":"hi"}"#.to_string(),
},
}]),
},
ChatCompletionMessage {
role: "tool".to_string(),
content: Some("hi".to_string()),
tool_call_id: Some("call_1".to_string()),
name: Some("echo".to_string()),
tool_calls: None,
},
];
let result = flatten_tool_messages(messages);
assert_eq!(result.len(), 3);
// Assistant tool_calls → plain assistant text
assert_eq!(result[1].role, "assistant");
assert!(result[1].tool_calls.is_none());
assert!(
result[1]
.content
.as_ref()
.unwrap()
.contains("[Called tool `echo`")
);
// Tool result → user message
assert_eq!(result[2].role, "user");
assert!(result[2].tool_call_id.is_none());
assert!(
result[2]
.content
.as_ref()
.unwrap()
.contains("[Tool `echo` returned: hi]")
);
}
#[test]
fn test_flatten_preserves_assistant_text_with_tool_calls() {
let messages = vec![
ChatCompletionMessage {
role: "assistant".to_string(),
content: Some("Let me check that.".to_string()),
tool_call_id: None,
name: None,
tool_calls: Some(vec![ChatCompletionToolCall {
id: "call_1".to_string(),
call_type: "function".to_string(),
function: ChatCompletionToolCallFunction {
name: "search".to_string(),
arguments: r#"{"q":"test"}"#.to_string(),
},
}]),
},
ChatCompletionMessage {
role: "tool".to_string(),
content: Some("found it".to_string()),
tool_call_id: Some("call_1".to_string()),
name: Some("search".to_string()),
tool_calls: None,
},
];
let result = flatten_tool_messages(messages);
let text = result[0].content.as_ref().unwrap();
assert!(text.starts_with("Let me check that."));
assert!(text.contains("[Called tool `search`"));
}
}
+77 -32
View File
@@ -113,6 +113,25 @@ pub struct ToolSelection {
pub reasoning: String,
/// Alternative tools considered.
pub alternatives: Vec<String>,
/// The tool call ID from the LLM response.
///
/// OpenAI-compatible providers assign each tool call a unique ID that must
/// be echoed back in the corresponding tool result message. Without this,
/// the provider cannot match results to their originating calls.
pub tool_call_id: String,
}
/// Token usage from a single LLM call.
#[derive(Debug, Clone, Copy, Default)]
pub struct TokenUsage {
pub input_tokens: u32,
pub output_tokens: u32,
}
impl TokenUsage {
pub fn total(&self) -> u32 {
self.input_tokens + self.output_tokens
}
}
/// Result of a response with potential tool calls.
@@ -131,6 +150,13 @@ pub enum RespondResult {
},
}
/// A `RespondResult` bundled with the token usage from the LLM call that produced it.
#[derive(Debug, Clone)]
pub struct RespondOutput {
pub result: RespondResult,
pub usage: TokenUsage,
}
/// Reasoning engine for the agent.
pub struct Reasoning {
llm: Arc<dyn LlmProvider>,
@@ -224,6 +250,7 @@ impl Reasoning {
parameters: tool_call.arguments,
reasoning: reasoning.clone(),
alternatives: vec![],
tool_call_id: tool_call.id,
})
.collect();
@@ -284,7 +311,8 @@ Respond in JSON format:
/// tool calls as text for simple cases. Use `respond_with_tools()` when you
/// need to actually execute tool calls in an agentic loop.
pub async fn respond(&self, context: &ReasoningContext) -> Result<String, LlmError> {
match self.respond_with_tools(context).await? {
let output = self.respond_with_tools(context).await?;
match output.result {
RespondResult::Text(text) => Ok(text),
RespondResult::ToolCalls {
tool_calls: calls, ..
@@ -299,15 +327,14 @@ Respond in JSON format:
}
}
/// Generate a response that may include tool calls.
/// Generate a response that may include tool calls, with token usage tracking.
///
/// Returns `RespondResult::ToolCalls` if the model wants to call tools,
/// allowing the caller to execute them and continue the conversation.
/// Returns `RespondResult::Text` when the model has a final text response.
/// Returns `RespondOutput` containing the result and token usage from the LLM call.
/// The caller should use `usage` to track cost/budget against the job.
pub async fn respond_with_tools(
&self,
context: &ReasoningContext,
) -> Result<RespondResult, LlmError> {
) -> Result<RespondOutput, LlmError> {
let system_prompt = self.build_conversation_prompt(context);
let mut messages = vec![ChatMessage::system(system_prompt)];
@@ -322,12 +349,19 @@ Respond in JSON format:
request.metadata = context.metadata.clone();
let response = self.llm.complete_with_tools(request).await?;
let usage = TokenUsage {
input_tokens: response.input_tokens,
output_tokens: response.output_tokens,
};
// If there were tool calls, return them for execution
if !response.tool_calls.is_empty() {
return Ok(RespondResult::ToolCalls {
tool_calls: response.tool_calls,
content: response.content,
return Ok(RespondOutput {
result: RespondResult::ToolCalls {
tool_calls: response.tool_calls,
content: response.content,
},
usage,
});
}
@@ -341,17 +375,23 @@ Respond in JSON format:
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)
return Ok(RespondOutput {
result: RespondResult::ToolCalls {
tool_calls: recovered,
content: if cleaned.is_empty() {
None
} else {
Some(cleaned)
},
},
usage,
});
}
Ok(RespondResult::Text(clean_response(&content)))
Ok(RespondOutput {
result: RespondResult::Text(clean_response(&content)),
usage,
})
} else {
// No tools, use simple completion
let mut request = CompletionRequest::new(messages)
@@ -360,7 +400,13 @@ Respond in JSON format:
request.metadata = context.metadata.clone();
let response = self.llm.complete(request).await?;
Ok(RespondResult::Text(clean_response(&response.content)))
Ok(RespondOutput {
result: RespondResult::Text(clean_response(&response.content)),
usage: TokenUsage {
input_tokens: response.input_tokens,
output_tokens: response.output_tokens,
},
})
}
}
@@ -542,21 +588,20 @@ fn recover_tool_calls_from_content(
}
// 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;
}
}
if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(inner)
&& let Some(name) = parsed.get("name").and_then(|v| v.as_str())
&& 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>")
+96
View File
@@ -0,0 +1,96 @@
//! Shared retry helpers for LLM providers.
//!
//! Provides exponential backoff with jitter and retryable status classification
//! used by both `NearAiProvider` and `NearAiChatProvider`.
use std::time::Duration;
use rand::Rng;
/// Returns `true` if the HTTP status code is transient and worth retrying.
pub(crate) fn is_retryable_status(status: u16) -> bool {
matches!(status, 429 | 500 | 502 | 503 | 504)
}
/// Calculate exponential backoff delay with random jitter.
///
/// Base delay is 1 second, doubled each attempt, with +/-25% jitter.
/// - attempt 0: ~1s (0.75s - 1.25s)
/// - attempt 1: ~2s (1.5s - 2.5s)
/// - attempt 2: ~4s (3.0s - 5.0s)
pub(crate) fn retry_backoff_delay(attempt: u32) -> Duration {
let base_ms: u64 = 1000u64.saturating_mul(2u64.saturating_pow(attempt));
let jitter_range = base_ms / 4; // 25%
let jitter = if jitter_range > 0 {
let offset = rand::thread_rng().gen_range(0..=jitter_range * 2);
offset as i64 - jitter_range as i64
} else {
0
};
let delay_ms = (base_ms as i64 + jitter).max(100) as u64;
Duration::from_millis(delay_ms)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_is_retryable_status() {
// Transient errors should be retryable
assert!(is_retryable_status(429));
assert!(is_retryable_status(500));
assert!(is_retryable_status(502));
assert!(is_retryable_status(503));
assert!(is_retryable_status(504));
// Client errors should not be retryable
assert!(!is_retryable_status(400));
assert!(!is_retryable_status(401));
assert!(!is_retryable_status(403));
assert!(!is_retryable_status(404));
assert!(!is_retryable_status(422));
// Success codes should not be retryable
assert!(!is_retryable_status(200));
assert!(!is_retryable_status(201));
}
#[test]
fn test_retry_backoff_delay_exponential_growth() {
// Run multiple samples to verify the range, accounting for jitter
for _ in 0..20 {
let d0 = retry_backoff_delay(0);
let d1 = retry_backoff_delay(1);
let d2 = retry_backoff_delay(2);
// Attempt 0: base 1000ms, jitter +/-250ms -> [750, 1250]
assert!(d0.as_millis() >= 750, "attempt 0 too low: {:?}", d0);
assert!(d0.as_millis() <= 1250, "attempt 0 too high: {:?}", d0);
// Attempt 1: base 2000ms, jitter +/-500ms -> [1500, 2500]
assert!(d1.as_millis() >= 1500, "attempt 1 too low: {:?}", d1);
assert!(d1.as_millis() <= 2500, "attempt 1 too high: {:?}", d1);
// Attempt 2: base 4000ms, jitter +/-1000ms -> [3000, 5000]
assert!(d2.as_millis() >= 3000, "attempt 2 too low: {:?}", d2);
assert!(d2.as_millis() <= 5000, "attempt 2 too high: {:?}", d2);
}
}
#[test]
fn test_retry_backoff_delay_minimum() {
// Even at attempt 0, delay should be at least 100ms (the minimum floor)
for _ in 0..20 {
let delay = retry_backoff_delay(0);
assert!(delay.as_millis() >= 100);
}
}
#[test]
fn test_retry_backoff_delay_no_overflow() {
// Very high attempt numbers should not panic from overflow
let delay = retry_backoff_delay(30);
assert!(delay.as_millis() >= 100);
}
}
+56 -182
View File
@@ -31,8 +31,6 @@ pub struct SessionConfig {
pub auth_base_url: String,
/// Path to session file (e.g., ~/.ironclaw/session.json).
pub session_path: PathBuf,
/// Port range for OAuth callback server.
pub callback_port_range: (u16, u16),
}
impl Default for SessionConfig {
@@ -40,7 +38,6 @@ impl Default for SessionConfig {
Self {
auth_base_url: "https://private.near.ai".to_string(),
session_path: default_session_path(),
callback_port_range: (9876, 9886),
}
}
}
@@ -62,7 +59,7 @@ pub struct SessionManager {
/// 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>>>,
store: RwLock<Option<Arc<dyn crate::db::Database>>>,
/// User ID for DB settings (default: "default").
user_id: RwLock<String>,
}
@@ -83,16 +80,16 @@ impl SessionManager {
};
// Try to load existing session synchronously during construction
if let Ok(data) = std::fs::read_to_string(&manager.config.session_path) {
if let Ok(session) = serde_json::from_str::<SessionData>(&data) {
// We can't await here, so we use try_write
if let Ok(mut guard) = manager.token.try_write() {
*guard = Some(SecretString::from(session.session_token));
tracing::info!(
"Loaded session token from {}",
manager.config.session_path.display()
);
}
if let Ok(data) = std::fs::read_to_string(&manager.config.session_path)
&& let Ok(session) = serde_json::from_str::<SessionData>(&data)
{
// We can't await here, so we use try_write
if let Ok(mut guard) = manager.token.try_write() {
*guard = Some(SecretString::from(session.session_token));
tracing::info!(
"Loaded session token from {}",
manager.config.session_path.display()
);
}
}
@@ -125,7 +122,7 @@ impl SessionManager {
/// 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) {
pub async fn attach_store(&self, store: Arc<dyn crate::db::Database>, user_id: &str) {
*self.store.write().await = Some(store);
*self.user_id.write().await = user_id.to_string();
@@ -222,38 +219,21 @@ impl SessionManager {
/// Start the OAuth login flow.
///
/// 1. Find an available port for the callback server
/// 1. Bind the fixed callback port
/// 2. Print the auth URL and attempt to open browser
/// 3. Wait for OAuth callback with session token
/// 4. Save and return the token
async fn initiate_login(&self) -> Result<(), LlmError> {
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::TcpListener;
use crate::cli::oauth_defaults::{self, OAUTH_CALLBACK_PORT};
// Find an available port
let mut listener = None;
let mut port = 0;
let listener = oauth_defaults::bind_callback_listener()
.await
.map_err(|e| LlmError::SessionRenewalFailed {
provider: "nearai".to_string(),
reason: e.to_string(),
})?;
for p in self.config.callback_port_range.0..=self.config.callback_port_range.1 {
match TcpListener::bind(format!("127.0.0.1:{}", p)).await {
Ok(l) => {
listener = Some(l);
port = p;
break;
}
Err(_) => continue,
}
}
let listener = listener.ok_or_else(|| LlmError::SessionRenewalFailed {
provider: "nearai".to_string(),
reason: format!(
"Could not find available port in range {}-{}",
self.config.callback_port_range.0, self.config.callback_port_range.1
),
})?;
let callback_url = format!("http://127.0.0.1:{}", port);
let callback_url = format!("http://127.0.0.1:{}", OAUTH_CALLBACK_PORT);
// Show auth provider menu
println!();
@@ -333,138 +313,16 @@ impl SessionManager {
println!();
println!("Waiting for authentication...");
// Wait for callback with timeout
// The API redirects to: {frontend_callback}/auth/callback?token=X&session_id=X&expires_at=X&is_new_user=X
let timeout = std::time::Duration::from_secs(300); // 5 minutes
let selected_provider = auth_provider.to_string();
let (session_token, auth_provider) = tokio::time::timeout(timeout, async move {
loop {
let (mut socket, _) = listener.accept().await.map_err(|e| {
LlmError::SessionRenewalFailed {
provider: "nearai".to_string(),
reason: format!("Failed to accept connection: {}", e),
}
// The NEAR AI API redirects to: {frontend_callback}/auth/callback?token=X&...
let session_token =
oauth_defaults::wait_for_callback(listener, "/auth/callback", "token", "NEAR AI")
.await
.map_err(|e| LlmError::SessionRenewalFailed {
provider: "nearai".to_string(),
reason: e.to_string(),
})?;
let mut reader = BufReader::new(&mut socket);
let mut request_line = String::new();
reader.read_line(&mut request_line).await.map_err(|e| {
LlmError::SessionRenewalFailed {
provider: "nearai".to_string(),
reason: format!("Failed to read request: {}", e),
}
})?;
// Parse GET /auth/callback?token=xxx&session_id=xxx&expires_at=xxx&is_new_user=xxx HTTP/1.1
if let Some(path) = request_line.split_whitespace().nth(1) {
if path.starts_with("/auth/callback") {
// Parse query parameters
if let Some(query) = path.split('?').nth(1) {
let mut token = None;
for param in query.split('&') {
let parts: Vec<&str> = param.splitn(2, '=').collect();
if parts.len() == 2 && parts[0] == "token" {
token = Some(
urlencoding::decode(parts[1])
.unwrap_or_else(|_| parts[1].into())
.into_owned(),
);
}
}
if let Some(token) = token {
// Send success response with nice styling
let response = concat!(
"HTTP/1.1 200 OK\r\n",
"Content-Type: text/html; charset=utf-8\r\n",
"Connection: close\r\n",
"\r\n",
"<!DOCTYPE html>\n",
"<html>\n",
"<head>\n",
" <meta charset=\"utf-8\">\n",
" <title>NEAR AI - Authentication Successful</title>\n",
" <style>\n",
" * { margin: 0; padding: 0; box-sizing: border-box; }\n",
" body {\n",
" font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;\n",
" background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%);\n",
" min-height: 100vh;\n",
" display: flex;\n",
" align-items: center;\n",
" justify-content: center;\n",
" color: #fff;\n",
" }\n",
" .container {\n",
" text-align: center;\n",
" padding: 3rem;\n",
" background: rgba(255,255,255,0.05);\n",
" border-radius: 16px;\n",
" backdrop-filter: blur(10px);\n",
" border: 1px solid rgba(255,255,255,0.1);\n",
" max-width: 400px;\n",
" }\n",
" .checkmark {\n",
" width: 80px;\n",
" height: 80px;\n",
" background: linear-gradient(135deg, #00d9a5 0%, #00b386 100%);\n",
" border-radius: 50%;\n",
" display: flex;\n",
" align-items: center;\n",
" justify-content: center;\n",
" margin: 0 auto 1.5rem;\n",
" font-size: 40px;\n",
" }\n",
" h1 {\n",
" font-size: 1.5rem;\n",
" font-weight: 600;\n",
" margin-bottom: 0.75rem;\n",
" }\n",
" p {\n",
" color: rgba(255,255,255,0.7);\n",
" font-size: 0.95rem;\n",
" line-height: 1.5;\n",
" }\n",
" .brand {\n",
" margin-top: 2rem;\n",
" padding-top: 1.5rem;\n",
" border-top: 1px solid rgba(255,255,255,0.1);\n",
" font-size: 0.8rem;\n",
" color: rgba(255,255,255,0.4);\n",
" }\n",
" </style>\n",
"</head>\n",
"<body>\n",
" <div class=\"container\">\n",
" <div class=\"checkmark\">&#10003;</div>\n",
" <h1>Authentication Successful</h1>\n",
" <p>You can close this window and return to the terminal.</p>\n",
" <div class=\"brand\">NEAR AI Agent</div>\n",
" </div>\n",
"</body>\n",
"</html>"
);
let _ = socket.write_all(response.as_bytes()).await;
let _ = socket.shutdown().await;
return Ok::<_, LlmError>((token, Some(selected_provider.clone())));
}
}
}
}
// Not the callback we're looking for, send 404
let response = "HTTP/1.1 404 Not Found\r\nConnection: close\r\n\r\n";
let _ = socket.write_all(response.as_bytes()).await;
}
})
.await
.map_err(|_| LlmError::SessionRenewalFailed {
provider: "nearai".to_string(),
reason: "Authentication timed out after 5 minutes".to_string(),
})??;
let auth_provider = Some(auth_provider.to_string());
// Save the token
self.save_session(&session_token, auth_provider.as_deref())
@@ -520,6 +378,25 @@ impl SessionManager {
))
})?;
// Restrictive permissions: session file contains a secret token
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let perms = std::fs::Permissions::from_mode(0o600);
tokio::fs::set_permissions(&self.config.session_path, perms)
.await
.map_err(|e| {
LlmError::Io(std::io::Error::new(
e.kind(),
format!(
"Failed to set permissions on {}: {}",
self.config.session_path.display(),
e
),
))
})?;
}
tracing::debug!("Session saved to {}", self.config.session_path.display());
// Also save to DB if a store is attached
@@ -623,15 +500,14 @@ pub async fn create_session_manager(config: SessionConfig) -> Arc<SessionManager
let manager = SessionManager::new_async(config).await;
// Check for legacy env var and migrate if present and no file token
if !manager.has_token().await {
if let Ok(token) = std::env::var("NEARAI_SESSION_TOKEN") {
if !token.is_empty() {
tracing::info!("Migrating session token from NEARAI_SESSION_TOKEN env var to file");
manager.set_token(SecretString::from(token.clone())).await;
if let Err(e) = manager.save_session(&token, None).await {
tracing::warn!("Failed to save migrated session: {}", e);
}
}
if !manager.has_token().await
&& let Ok(token) = std::env::var("NEARAI_SESSION_TOKEN")
&& !token.is_empty()
{
tracing::info!("Migrating session token from NEARAI_SESSION_TOKEN env var to file");
manager.set_token(SecretString::from(token.clone())).await;
if let Err(e) = manager.save_session(&token, None).await {
tracing::warn!("Failed to save migrated session: {}", e);
}
}
@@ -652,7 +528,6 @@ mod tests {
let config = SessionConfig {
auth_base_url: "https://example.com".to_string(),
session_path: session_path.clone(),
callback_port_range: (9900, 9910),
};
let manager = SessionManager::new_async(config.clone()).await;
@@ -693,7 +568,6 @@ mod tests {
let config = SessionConfig {
auth_base_url: "https://example.com".to_string(),
session_path: dir.path().join("nonexistent.json"),
callback_port_range: (9900, 9910),
};
let manager = SessionManager::new_async(config).await;
+270 -123
View File
@@ -17,22 +17,22 @@ use ironclaw::{
web::log_layer::{LogBroadcaster, WebLogLayer},
},
cli::{
Cli, Command, run_mcp_command, run_memory_command, run_pairing_command, run_status_command,
run_tool_command,
Cli, Command, run_mcp_command, run_pairing_command, run_status_command, run_tool_command,
},
config::Config,
context::ContextManager,
extensions::ExtensionManager,
history::Store,
llm::{SessionConfig, create_llm_provider, create_session_manager},
llm::{
FailoverProvider, LlmProvider, SessionConfig, create_llm_provider,
create_llm_provider_with_config, create_session_manager,
},
orchestrator::{
ContainerJobConfig, ContainerJobManager, OrchestratorApi, TokenStore,
api::OrchestratorState,
},
pairing::PairingStore,
safety::SafetyLayer,
secrets::{PostgresSecretsStore, SecretsCrypto, SecretsStore},
setup::{SetupConfig, SetupWizard},
secrets::SecretsStore,
tools::{
ToolRegistry,
mcp::{McpClient, McpSessionManager, config::load_mcp_servers_from_db, is_authenticated},
@@ -41,6 +41,13 @@ use ironclaw::{
workspace::{EmbeddingProvider, NearAiEmbeddings, OpenAiEmbeddings, Workspace},
};
#[cfg(feature = "libsql")]
use ironclaw::secrets::LibSqlSecretsStore;
#[cfg(feature = "postgres")]
use ironclaw::secrets::PostgresSecretsStore;
use ironclaw::secrets::SecretsCrypto;
#[cfg(any(feature = "postgres", feature = "libsql"))]
use ironclaw::setup::{SetupConfig, SetupWizard};
#[tokio::main]
async fn main() -> anyhow::Result<()> {
let cli = Cli::parse();
@@ -85,18 +92,14 @@ async fn main() -> anyhow::Result<()> {
.init();
// Memory commands need database (and optionally embeddings)
let _ = dotenvy::dotenv();
let config = Config::from_env()
.await
.map_err(|e| anyhow::anyhow!("{}", e))?;
let store = ironclaw::history::Store::new(&config.database).await?;
store.run_migrations().await?;
// Set up embeddings if available
let session = ironclaw::llm::create_session_manager(ironclaw::llm::SessionConfig {
auth_base_url: config.llm.nearai.auth_base_url.clone(),
session_path: config.llm.nearai.session_path.clone(),
..Default::default()
})
.await;
@@ -130,7 +133,14 @@ async fn main() -> anyhow::Result<()> {
None
};
return run_memory_command(mem_cmd.clone(), store.pool(), embeddings).await;
// Create a Database-trait-backed workspace for the memory command
let db: Arc<dyn ironclaw::db::Database> =
ironclaw::db::connect_from_config(&config.database)
.await
.map_err(|e| anyhow::anyhow!("{}", e))?;
return ironclaw::cli::run_memory_command_with_db(mem_cmd.clone(), db, embeddings)
.await;
}
Some(Command::Pairing(pairing_cmd)) => {
tracing_subscriber::fmt()
@@ -142,7 +152,6 @@ async fn main() -> anyhow::Result<()> {
return run_pairing_command(pairing_cmd.clone()).map_err(|e| anyhow::anyhow!("{}", e));
}
Some(Command::Status) => {
let _ = dotenvy::dotenv();
tracing_subscriber::fmt()
.with_env_filter(
EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")),
@@ -216,6 +225,7 @@ async fn main() -> anyhow::Result<()> {
max_turns: *max_turns,
model: model.clone(),
timeout: std::time::Duration::from_secs(1800),
allowed_tools: Vec::new(),
};
let runtime = ironclaw::worker::ClaudeBridgeRuntime::new(config)
@@ -232,15 +242,25 @@ async fn main() -> anyhow::Result<()> {
skip_auth,
channels_only,
}) => {
// Load .env before running onboarding wizard
// Load .env files before running onboarding wizard.
// Standard ./.env first (higher priority), then ~/.ironclaw/.env.
let _ = dotenvy::dotenv();
ironclaw::bootstrap::load_ironclaw_env();
let config = SetupConfig {
skip_auth: *skip_auth,
channels_only: *channels_only,
};
let mut wizard = SetupWizard::with_config(config);
wizard.run().await?;
#[cfg(any(feature = "postgres", feature = "libsql"))]
{
let config = SetupConfig {
skip_auth: *skip_auth,
channels_only: *channels_only,
};
let mut wizard = SetupWizard::with_config(config);
wizard.run().await?;
}
#[cfg(not(any(feature = "postgres", feature = "libsql")))]
{
let _ = (skip_auth, channels_only);
eprintln!("Onboarding wizard requires the 'postgres' or 'libsql' feature.");
}
return Ok(());
}
None | Some(Command::Run) => {
@@ -248,22 +268,23 @@ async fn main() -> anyhow::Result<()> {
}
}
// Load .env if present
// Load .env files early so DATABASE_URL (and any other vars) are
// available to all subsequent env-based config resolution.
// Standard ./.env first (higher priority), then ~/.ironclaw/.env.
let _ = dotenvy::dotenv();
ironclaw::bootstrap::load_ironclaw_env();
// Enhanced first-run detection
if !cli.no_onboard {
if let Some(reason) = check_onboard_needed().await {
println!("Onboarding needed: {}", reason);
println!();
let mut wizard = SetupWizard::new();
wizard.run().await?;
}
#[cfg(any(feature = "postgres", feature = "libsql"))]
if !cli.no_onboard
&& let Some(reason) = check_onboard_needed()
{
println!("Onboarding needed: {}", reason);
println!();
let mut wizard = SetupWizard::new();
wizard.run().await?;
}
// 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().await {
Ok(c) => c,
@@ -283,7 +304,6 @@ async fn main() -> anyhow::Result<()> {
let session_config = SessionConfig {
auth_base_url: config.llm.nearai.auth_base_url.clone(),
session_path: config.llm.nearai.session_path.clone(),
..Default::default()
};
let session = create_session_manager(session_config).await;
@@ -294,7 +314,7 @@ async fn main() -> anyhow::Result<()> {
// Initialize tracing
let env_filter = EnvFilter::try_from_default_env()
.unwrap_or_else(|_| EnvFilter::new("ironclaw=info,tower_http=debug"));
.unwrap_or_else(|_| EnvFilter::new("ironclaw=info,tower_http=warn"));
// Create log broadcaster before tracing init so the WebLogLayer can capture all events.
// This gets wired to the gateway's /api/logs/events SSE endpoint later.
@@ -302,7 +322,11 @@ async fn main() -> anyhow::Result<()> {
tracing_subscriber::registry()
.with(env_filter)
.with(tracing_subscriber::fmt::layer().with_target(false))
.with(
tracing_subscriber::fmt::layer()
.with_target(false)
.with_writer(ironclaw::tracing_fmt::TruncatingStderr::default()),
)
.with(WebLogLayer::new(Arc::clone(&log_broadcaster)))
.init();
@@ -319,23 +343,86 @@ async fn main() -> anyhow::Result<()> {
tracing::info!("Loaded configuration for agent: {}", config.agent.name);
tracing::info!("LLM backend: {}", config.llm.backend);
// Initialize database store (optional for testing)
let store = if cli.no_db {
// Initialize database backend.
//
// Creates an `Arc<dyn Database>` that all consumers share.
// Backend is selected by the `DATABASE_BACKEND` env var / config.
//
// NOTE: For simpler call sites (CLI commands, Memory handler) use the shared
// helper `ironclaw::db::connect_from_config()`. This block is kept inline
// because it also captures backend-specific handles (`pg_pool`, `libsql_db`)
// needed by the secrets store.
#[cfg(feature = "postgres")]
let mut pg_pool: Option<deadpool_postgres::Pool> = None;
#[cfg(feature = "libsql")]
let mut libsql_db: Option<std::sync::Arc<libsql::Database>> = None;
let db: Option<Arc<dyn ironclaw::db::Database>> = if cli.no_db {
tracing::warn!("Running without database connection");
None
} else {
let store = Store::new(&config.database).await?;
store.run_migrations().await?;
tracing::info!("Database connected and migrations applied");
match config.database.backend {
#[cfg(feature = "libsql")]
ironclaw::config::DatabaseBackend::LibSql => {
use ironclaw::db::Database as _;
use ironclaw::db::libsql_backend::LibSqlBackend;
use secrecy::ExposeSecret as _;
let default_path = ironclaw::config::default_libsql_path();
let db_path = config
.database
.libsql_path
.as_deref()
.unwrap_or(&default_path);
let backend = if let Some(ref url) = config.database.libsql_url {
let token = config.database.libsql_auth_token.as_ref().ok_or_else(|| {
anyhow::anyhow!("LIBSQL_AUTH_TOKEN is required when LIBSQL_URL is set")
})?;
LibSqlBackend::new_remote_replica(db_path, url, token.expose_secret()).await?
} else {
LibSqlBackend::new_local(db_path).await?
};
backend.run_migrations().await?;
tracing::info!("libSQL database connected and migrations applied");
// Capture the Database handle for SecretsStore (connection-per-op)
libsql_db = Some(backend.shared_db());
Some(Arc::new(backend) as Arc<dyn ironclaw::db::Database>)
}
#[cfg(feature = "postgres")]
_ => {
use ironclaw::db::Database as _;
let pg = ironclaw::db::postgres::PgBackend::new(&config.database)
.await
.map_err(|e| anyhow::anyhow!("{}", e))?;
pg.run_migrations()
.await
.map_err(|e| anyhow::anyhow!("{}", e))?;
tracing::info!("PostgreSQL database connected and migrations applied");
pg_pool = Some(pg.pool());
Some(Arc::new(pg) as Arc<dyn ironclaw::db::Database>)
}
#[cfg(not(feature = "postgres"))]
_ => {
anyhow::bail!(
"No database backend available. Enable 'postgres' or 'libsql' feature."
);
}
}
};
// Post-init operations using the database
if let Some(ref db) = db {
// 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 {
if let Err(e) = ironclaw::bootstrap::migrate_disk_to_db(db.as_ref(), "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 {
match Config::from_db(db.as_ref(), "default").await {
Ok(db_config) => {
config = db_config;
tracing::info!("Configuration reloaded from database");
@@ -348,23 +435,105 @@ async fn main() -> anyhow::Result<()> {
}
}
let store = Arc::new(store);
// Attach store to session manager so tokens save to DB too
session.attach_store(Arc::clone(&store), "default").await;
// Attach DB to session manager so tokens save to DB too
session.attach_store(Arc::clone(db), "default").await;
// Mark any jobs left in "running" or "creating" state as "interrupted".
if let Err(e) = store.cleanup_stale_sandbox_jobs().await {
if let Err(e) = db.cleanup_stale_sandbox_jobs().await {
tracing::warn!("Failed to cleanup stale sandbox jobs: {}", e);
}
}
Some(store)
};
// Create secrets store early: needed for injecting LLM API keys from encrypted
// storage before creating the LLM provider, and later for MCP auth + WASM channels.
//
// When both `postgres` and `libsql` features are compiled, the runtime-selected
// backend determines which store is created: whichever DB init branch ran will
// have set its handle (pg_pool or libsql_db), and the or_else chain picks it up.
let secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>> =
if let Some(master_key) = config.secrets.master_key() {
match SecretsCrypto::new(master_key.clone()) {
Ok(crypto) => {
let crypto = Arc::new(crypto);
let store: Option<Arc<dyn SecretsStore + Send + Sync>> = None;
#[cfg(feature = "libsql")]
let store = store.or_else(|| {
libsql_db.take().map(|db| {
Arc::new(LibSqlSecretsStore::new(db, Arc::clone(&crypto)))
as Arc<dyn SecretsStore + Send + Sync>
})
});
#[cfg(feature = "postgres")]
let store = store.or_else(|| {
pg_pool.as_ref().map(|pool| {
Arc::new(PostgresSecretsStore::new(pool.clone(), Arc::clone(&crypto)))
as Arc<dyn SecretsStore + Send + Sync>
})
});
store
}
Err(e) => {
tracing::warn!("Failed to initialize secrets crypto: {}", e);
#[cfg(feature = "libsql")]
let _ = libsql_db.take();
None
}
}
} else {
#[cfg(feature = "libsql")]
let _ = libsql_db.take();
None
};
// Inject LLM API keys from the encrypted secrets store into a thread-safe
// overlay so that optional_env() (used by LlmConfig::resolve()) picks them
// up. Then re-resolve LlmConfig with the newly available keys (backend may
// have been set during onboarding but the API key is in the secrets store).
if let Some(ref secrets) = secrets_store {
ironclaw::config::inject_llm_keys_from_secrets(secrets.as_ref(), "default").await;
// Re-resolve LlmConfig now that secrets overlay has been populated
if let Some(ref db_ref) = db {
match Config::from_db(db_ref.as_ref(), "default").await {
Ok(refreshed) => {
config = refreshed;
tracing::debug!("LlmConfig re-resolved after secret injection");
}
Err(e) => {
tracing::warn!("Failed to re-resolve config after secret injection: {}", e);
}
}
}
}
// Initialize LLM provider (clone session so we can reuse it for embeddings)
let llm = create_llm_provider(&config.llm, session.clone())?;
tracing::info!("LLM provider initialized: {}", llm.model_name());
// Wrap in failover if a fallback model is configured
let llm: Arc<dyn LlmProvider> =
if let Some(fallback_model) = config.llm.nearai.fallback_model.as_ref() {
if fallback_model == &config.llm.nearai.model {
tracing::warn!(
"fallback_model is the same as primary model, failover may not be effective"
);
}
let mut fallback_config = config.llm.nearai.clone();
fallback_config.model = fallback_model.clone();
let fallback = create_llm_provider_with_config(&fallback_config, session.clone())?;
tracing::info!(
primary = %llm.model_name(),
fallback = %fallback.model_name(),
"LLM failover enabled"
);
Arc::new(FailoverProvider::new(vec![llm, fallback])?)
} else {
llm
};
// Initialize safety layer
let safety = Arc::new(SafetyLayer::new(&config.safety));
tracing::info!("Safety layer initialized");
@@ -414,8 +583,8 @@ async fn main() -> anyhow::Result<()> {
};
// Register memory tools if database is available
if let Some(ref store) = store {
let mut workspace = Workspace::new("default", store.pool());
if let Some(ref db) = db {
let mut workspace = Workspace::new_with_db("default", Arc::clone(db));
if let Some(ref emb) = embeddings {
workspace = workspace.with_embeddings(emb.clone());
}
@@ -438,23 +607,6 @@ async fn main() -> anyhow::Result<()> {
tracing::info!("Builder mode enabled");
}
// Create secrets store if master key is configured (needed for MCP auth and WASM channels)
let secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>> =
if let (Some(store), Some(master_key)) = (&store, config.secrets.master_key()) {
match SecretsCrypto::new(master_key.clone()) {
Ok(crypto) => Some(Arc::new(PostgresSecretsStore::new(
store.pool(),
Arc::new(crypto),
))),
Err(e) => {
tracing::warn!("Failed to initialize secrets crypto: {}", e);
None
}
}
} else {
None
};
let mcp_session_manager = Arc::new(McpSessionManager::new());
// Create WASM tool runtime (sync, just builds the wasmtime engine)
@@ -475,7 +627,10 @@ async fn main() -> anyhow::Result<()> {
// Both register into the shared ToolRegistry (RwLock-based) so concurrent writes are safe.
let wasm_tools_future = async {
if let Some(ref runtime) = wasm_tool_runtime {
let loader = WasmToolLoader::new(Arc::clone(runtime), Arc::clone(&tools));
let mut loader = WasmToolLoader::new(Arc::clone(runtime), Arc::clone(&tools));
if let Some(ref secrets) = secrets_store {
loader = loader.with_secrets_store(Arc::clone(secrets));
}
// Load installed tools from ~/.ironclaw/tools/
match loader.load_from_dir(&config.wasm.tools_dir).await {
@@ -515,8 +670,8 @@ async fn main() -> anyhow::Result<()> {
let mcp_servers_future = async {
if let Some(ref secrets) = secrets_store {
let servers_result = if let Some(ref s) = store {
load_mcp_servers_from_db(s, "default").await
let servers_result = if let Some(ref d) = db {
load_mcp_servers_from_db(d.as_ref(), "default").await
} else {
ironclaw::tools::mcp::config::load_mcp_servers().await
};
@@ -629,7 +784,7 @@ async fn main() -> anyhow::Result<()> {
config.channels.wasm_channels_dir.clone(),
config.tunnel.public_url.clone(),
"default".to_string(),
store.clone(),
db.clone(),
));
tools.register_extension_tools(Arc::clone(&manager));
tracing::info!("Extension manager initialized with in-chat discovery tools");
@@ -681,6 +836,7 @@ async fn main() -> anyhow::Result<()> {
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,
claude_code_allowed_tools: config.claude_code.allowed_tools.clone(),
};
let jm = Arc::new(ContainerJobManager::new(job_config, token_store.clone()));
@@ -691,7 +847,7 @@ async fn main() -> anyhow::Result<()> {
token_store,
job_event_tx: job_event_tx.clone(),
prompt_queue: Arc::clone(&prompt_queue),
store: store.clone(),
store: db.clone(),
};
tokio::spawn(async move {
@@ -797,13 +953,13 @@ 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 channel_name == "telegram"
&& 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() {
@@ -894,23 +1050,23 @@ async fn main() -> anyhow::Result<()> {
// Extract its routes for the unified server; the channel itself just
// provides the mpsc stream.
let mut webhook_server_addr: Option<std::net::SocketAddr> = None;
if !cli.cli_only {
if let Some(ref http_config) = config.channels.http {
let http_channel = HttpChannel::new(http_config.clone());
webhook_routes.push(http_channel.routes());
let (host, port) = http_channel.addr();
webhook_server_addr = Some(
format!("{}:{}", host, port)
.parse()
.expect("HttpConfig host:port must be a valid SocketAddr"),
);
channels.add(Box::new(http_channel));
tracing::info!(
"HTTP channel enabled on {}:{}",
http_config.host,
http_config.port
);
}
if !cli.cli_only
&& let Some(ref http_config) = config.channels.http
{
let http_channel = HttpChannel::new(http_config.clone());
webhook_routes.push(http_channel.routes());
let (host, port) = http_channel.addr();
webhook_server_addr = Some(
format!("{}:{}", host, port)
.parse()
.expect("HttpConfig host:port must be a valid SocketAddr"),
);
channels.add(Box::new(http_channel));
tracing::info!(
"HTTP channel enabled on {}:{}",
http_config.host,
http_config.port
);
}
// Start the unified webhook server if any routes were registered.
@@ -928,13 +1084,15 @@ async fn main() -> anyhow::Result<()> {
};
// Create workspace for agent (shared with memory tools)
let workspace = store.as_ref().map(|s| {
let mut ws = Workspace::new("default", s.pool());
let workspace = if let Some(ref db_ref) = db {
let mut ws = Workspace::new_with_db("default", Arc::clone(db_ref));
if let Some(ref emb) = embeddings {
ws = ws.with_embeddings(emb.clone());
}
Arc::new(ws)
});
Some(Arc::new(ws))
} else {
None
};
// Seed workspace with core identity files on first boot
if let Some(ref ws) = workspace {
@@ -972,7 +1130,7 @@ async fn main() -> anyhow::Result<()> {
tools.register_job_tools(
Arc::clone(&context_manager),
container_job_manager.clone(),
store.clone(),
db.clone(),
);
// Add web gateway channel if configured
@@ -987,8 +1145,8 @@ async fn main() -> anyhow::Result<()> {
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 d) = db {
gw = gw.with_store(Arc::clone(d));
}
if let Some(ref jm) = container_job_manager {
gw = gw.with_job_manager(Arc::clone(jm));
@@ -1025,7 +1183,7 @@ async fn main() -> anyhow::Result<()> {
// Create and run the agent
let deps = AgentDeps {
store,
store: db,
llm,
safety,
tools,
@@ -1059,29 +1217,18 @@ async fn main() -> anyhow::Result<()> {
/// Check if onboarding is needed and return the reason.
///
/// Returns `Some(reason)` if onboarding should be triggered, `None` otherwise.
async fn check_onboard_needed() -> Option<&'static str> {
let bootstrap = ironclaw::bootstrap::BootstrapConfig::load();
/// Called after `load_ironclaw_env()`, so DATABASE_URL from `~/.ironclaw/.env`
/// is already in the environment.
#[cfg(any(feature = "postgres", feature = "libsql"))]
fn check_onboard_needed() -> Option<&'static str> {
let has_db = std::env::var("DATABASE_URL").is_ok()
|| std::env::var("LIBSQL_PATH").is_ok()
|| ironclaw::config::default_libsql_path().exists();
// Database not configured (and not in env)
if bootstrap.database_url.is_none() && std::env::var("DATABASE_URL").is_err() {
if !has_db {
return Some("Database not configured");
}
// Secrets not configured (and not in env)
if bootstrap.secrets_master_key_source == ironclaw::settings::KeySource::None
&& std::env::var("SECRETS_MASTER_KEY").is_err()
&& !ironclaw::secrets::keychain::has_master_key().await
{
// Only require secrets setup if user hasn't explicitly disabled it
// For now, we don't require it for first run
}
// First run (onboarding never completed and no session)
let session_path = ironclaw::llm::session::default_session_path();
if !bootstrap.onboard_completed && !session_path.exists() {
return Some("First run");
}
None
}
+12 -12
View File
@@ -15,7 +15,7 @@ use tokio::sync::{Mutex, broadcast};
use uuid::Uuid;
use crate::channels::web::types::SseEvent;
use crate::history::Store;
use crate::db::Database;
use crate::llm::{CompletionRequest, LlmProvider, ToolCompletionRequest};
use crate::orchestrator::auth::{TokenStore, worker_auth_middleware};
use crate::orchestrator::job_manager::ContainerJobManager;
@@ -43,7 +43,7 @@ pub struct OrchestratorState {
/// 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>>,
pub store: Option<Arc<dyn Database>>,
}
/// The orchestrator's internal API server.
@@ -339,16 +339,16 @@ async fn get_prompt_handler(
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,
})),
));
}
if let Some(prompts) = queue.get_mut(&job_id)
&& 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
+3 -2
View File
@@ -14,6 +14,7 @@ use axum::http::StatusCode;
use axum::middleware::Next;
use axum::response::Response;
use rand::Rng;
use subtle::ConstantTimeEq;
use tokio::sync::RwLock;
use uuid::Uuid;
@@ -38,13 +39,13 @@ impl TokenStore {
token
}
/// Validate a token for a specific job.
/// Validate a token for a specific job (constant-time comparison).
pub async fn validate(&self, job_id: Uuid, token: &str) -> bool {
self.tokens
.read()
.await
.get(&job_id)
.map(|stored| stored == token)
.map(|stored| stored.as_bytes().ct_eq(token.as_bytes()).into())
.unwrap_or(false)
}
+75 -29
View File
@@ -58,6 +58,8 @@ pub struct ContainerJobConfig {
pub claude_code_max_turns: u32,
/// Memory limit in MB for Claude Code containers (heavier than workers).
pub claude_code_memory_limit_mb: u64,
/// Allowed tool patterns for Claude Code (passed as CLAUDE_CODE_ALLOWED_TOOLS env var).
pub claude_code_allowed_tools: Vec<String>,
}
impl Default for ContainerJobConfig {
@@ -71,6 +73,7 @@ impl Default for ContainerJobConfig {
claude_code_model: "sonnet".to_string(),
claude_code_max_turns: 50,
claude_code_memory_limit_mb: 4096,
claude_code_allowed_tools: crate::config::ClaudeCodeConfig::default().allowed_tools,
}
}
}
@@ -161,6 +164,29 @@ impl ContainerJobManager {
};
self.containers.write().await.insert(job_id, handle);
// Run the actual container creation. On any failure, revoke the token
// and remove the handle so we don't leak resources.
match self
.create_job_inner(job_id, &token, project_dir, mode)
.await
{
Ok(()) => Ok(token),
Err(e) => {
self.token_store.revoke(job_id).await;
self.containers.write().await.remove(&job_id);
Err(e)
}
}
}
/// Inner implementation of container creation (separated for cleanup).
async fn create_job_inner(
&self,
job_id: Uuid,
token: &str,
project_dir: Option<PathBuf>,
mode: JobMode,
) -> Result<(), OrchestratorError> {
// Connect to Docker
let docker = connect_docker()
.await
@@ -203,27 +229,34 @@ impl ContainerJobManager {
.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()
),
});
}
if let Ok(canonical_base) = projects_base.canonicalize()
&& !canonical.starts_with(&canonical_base)
{
return Err(OrchestratorError::ContainerCreationFailed {
job_id,
reason: format!(
"project directory {} is outside allowed base {}",
canonical.display(),
canonical_base.display()
),
});
}
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
// Claude Code mode: mount host ~/.claude read-only for auth,
// and pass the tool allowlist so the bridge can write settings.json.
if mode == JobMode::ClaudeCode {
if let Some(ref claude_dir) = self.config.claude_config_dir {
binds.push(format!("{}:/home/sandbox/.claude:ro", claude_dir.display()));
}
if !self.config.claude_code_allowed_tools.is_empty() {
env_vec.push(format!(
"CLAUDE_CODE_ALLOWED_TOOLS={}",
self.config.claude_code_allowed_tools.join(",")
));
}
}
// Memory limit: Claude Code gets more memory
@@ -243,11 +276,7 @@ impl ContainerJobManager {
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(),
]),
cap_add: Some(vec!["CHOWN".to_string()]),
security_opt: Some(vec!["no-new-privileges:true".to_string()]),
tmpfs: Some(
[("/tmp".to_string(), "size=512M".to_string())]
@@ -328,7 +357,7 @@ impl ContainerJobManager {
"Created and started worker container"
);
Ok(token)
Ok(())
}
/// Stop a running container job.
@@ -355,15 +384,18 @@ impl ContainerJobManager {
})?;
// Stop the container (10 second grace period)
let _ = docker
if let Err(e) = docker
.stop_container(
&container_id,
Some(bollard::container::StopContainerOptions { t: 10 }),
)
.await;
.await
{
tracing::warn!(job_id = %job_id, error = %e, "Failed to stop container (may already be stopped)");
}
// Remove the container
let _ = docker
if let Err(e) = docker
.remove_container(
&container_id,
Some(bollard::container::RemoveContainerOptions {
@@ -371,7 +403,10 @@ impl ContainerJobManager {
..Default::default()
}),
)
.await;
.await
{
tracing::warn!(job_id = %job_id, error = %e, "Failed to remove container (may require manual cleanup)");
}
// Update state
if let Some(handle) = self.containers.write().await.get_mut(&job_id) {
@@ -407,16 +442,21 @@ impl ContainerJobManager {
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
if let Some(cid) = container_id
&& !cid.is_empty()
{
match connect_docker().await {
Ok(docker) => {
if let Err(e) = docker
.stop_container(
&cid,
Some(bollard::container::StopContainerOptions { t: 5 }),
)
.await;
let _ = docker
.await
{
tracing::warn!(job_id = %job_id, error = %e, "Failed to stop completed container");
}
if let Err(e) = docker
.remove_container(
&cid,
Some(bollard::container::RemoveContainerOptions {
@@ -424,7 +464,13 @@ impl ContainerJobManager {
..Default::default()
}),
)
.await;
.await
{
tracing::warn!(job_id = %job_id, error = %e, "Failed to remove completed container");
}
}
Err(e) => {
tracing::warn!(job_id = %job_id, error = %e, "Failed to connect to Docker for container cleanup");
}
}
}
+11 -3
View File
@@ -320,19 +320,27 @@ impl PairingStore {
fn record_failed_approve(&self, channel: &str) -> Result<(), PairingStoreError> {
let path = approve_attempts_path(&self.base_dir, channel)?;
fs::create_dir_all(path.parent().unwrap())?;
// Open (or create) and lock before reading so concurrent callers
// don't clobber each other's writes.
let file = fs::OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(true)
.truncate(false)
.open(&path)?;
file.lock_exclusive()?;
let content = fs::read_to_string(&path).unwrap_or_default();
let mut data: ApproveAttemptsFile = serde_json::from_str(&content).unwrap_or_default();
let mut data: ApproveAttemptsFile = fs::read_to_string(&path)
.ok()
.and_then(|c| serde_json::from_str(&c).ok())
.unwrap_or_default();
let now = now_secs();
data.failed_at.push(now);
let cutoff = now.saturating_sub(PAIRING_APPROVE_RATE_WINDOW_SECS);
data.failed_at.retain(|&t| t >= cutoff);
let json = serde_json::to_string_pretty(&data)?;
fs::write(&path, json)?;
fs4::FileExt::unlock(&file)?;
+21 -9
View File
@@ -147,10 +147,10 @@ impl LeakDetector {
// Build prefix matcher for patterns that start with a known prefix
let mut prefixes = Vec::new();
for (idx, pattern) in patterns.iter().enumerate() {
if let Some(prefix) = extract_literal_prefix(pattern.regex.as_str()) {
if prefix.len() >= 3 {
prefixes.push((prefix, idx));
}
if let Some(prefix) = extract_literal_prefix(pattern.regex.as_str())
&& prefix.len() >= 3
{
prefixes.push((prefix, idx));
}
}
@@ -306,12 +306,11 @@ impl LeakDetector {
})?;
}
// Scan body if present and valid UTF-8
// Scan body if present. Use lossy UTF-8 conversion so a leading
// non-UTF8 byte can't be used to skip scanning entirely.
if let Some(body_bytes) = body {
if let Ok(body_str) = std::str::from_utf8(body_bytes) {
self.scan_and_clean(body_str)?;
}
// Binary bodies are not scanned (could add hex pattern detection later)
let body_str = String::from_utf8_lossy(body_bytes);
self.scan_and_clean(&body_str)?;
}
Ok(())
@@ -705,4 +704,17 @@ mod tests {
let result = detector.scan_http_request("https://api.example.com/webhook", &[], Some(body));
assert!(result.is_err());
}
#[test]
fn test_scan_http_request_blocks_secret_in_binary_body() {
let detector = LeakDetector::new();
// Attacker prepends a non-UTF8 byte to bypass strict from_utf8 check.
// The lossy conversion should still detect the secret.
let mut body = vec![0xFF]; // invalid UTF-8 leading byte
body.extend_from_slice(b"sk-proj-test1234567890abcdefghij");
let result = detector.scan_http_request("https://api.example.com/exfil", &[], Some(&body));
assert!(result.is_err(), "binary body should still be scanned");
}
}
+21 -5
View File
@@ -98,15 +98,15 @@ impl SafetyLayer {
was_modified: true,
};
}
if violations
let force_sanitize = violations
.iter()
.any(|rule| rule.action == crate::safety::PolicyAction::Sanitize)
{
.any(|rule| rule.action == crate::safety::PolicyAction::Sanitize);
if force_sanitize {
was_modified = true;
}
// Run sanitization if enabled
if self.config.injection_check_enabled {
// Run sanitization once: if injection_check is enabled OR policy requires it
if self.config.injection_check_enabled || force_sanitize {
let mut sanitized = self.sanitizer.sanitize(&content);
sanitized.was_modified = sanitized.was_modified || was_modified;
sanitized
@@ -190,4 +190,20 @@ mod tests {
assert!(wrapped.contains("sanitized=\"true\""));
assert!(wrapped.contains("Hello &lt;world&gt;"));
}
#[test]
fn test_sanitize_action_forces_sanitization_when_injection_check_disabled() {
let config = SafetyConfig {
max_output_length: 100_000,
injection_check_enabled: false,
};
let safety = SafetyLayer::new(&config);
// Content with an injection-like pattern that a policy might flag
let output = safety.sanitize_tool_output("test", "normal text");
// With injection_check disabled and no policy violations, content
// should pass through unmodified
assert_eq!(output.content, "normal text");
assert!(!output.was_modified);
}
}
+7 -12
View File
@@ -279,11 +279,7 @@ impl ContainerRunner {
network_mode: Some("bridge".to_string()),
// Security: drop all capabilities and add back only what's needed
cap_drop: Some(vec!["ALL".to_string()]),
cap_add: Some(vec![
"CHOWN".to_string(),
"SETUID".to_string(),
"SETGID".to_string(),
]),
cap_add: Some(vec!["CHOWN".to_string()]),
// Prevent privilege escalation
security_opt: Some(vec!["no-new-privileges:true".to_string()]),
// Read-only root filesystem (workspace is still writable if policy allows)
@@ -498,10 +494,10 @@ impl ContainerRunner {
/// 3. `~/.docker/run/docker.sock` (Docker Desktop on macOS)
pub async fn connect_docker() -> Result<Docker> {
// 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);
}
if let Ok(docker) = Docker::connect_with_local_defaults()
&& docker.ping().await.is_ok()
{
return Ok(docker);
}
// Try Docker Desktop socket (macOS)
@@ -511,10 +507,9 @@ pub async fn connect_docker() -> Result<Docker> {
let sock_str = desktop_sock.to_string_lossy();
if let Ok(docker) =
Docker::connect_with_socket(&sock_str, 120, bollard::API_DEFAULT_VERSION)
&& docker.ping().await.is_ok()
{
if docker.ping().await.is_ok() {
return Ok(docker);
}
return Ok(docker);
}
}
}
+9 -9
View File
@@ -259,11 +259,11 @@ async fn handle_connect(
let decision = state.decider.decide(&network_req).await;
if !decision.is_allowed() {
if let NetworkDecision::Deny { reason } = decision {
tracing::info!("Proxy: blocked CONNECT {} - {}", host, reason);
return error_response(StatusCode::FORBIDDEN, reason);
}
if !decision.is_allowed()
&& let NetworkDecision::Deny { reason } = decision
{
tracing::info!("Proxy: blocked CONNECT {} - {}", host, reason);
return error_response(StatusCode::FORBIDDEN, reason);
}
tracing::debug!("Proxy: allowing CONNECT to {}", host);
@@ -294,10 +294,10 @@ async fn forward_request(
// Copy headers (except hop-by-hop headers)
for (name, value) in req.headers() {
if !is_hop_by_hop_header(name.as_str()) {
if let Ok(v) = value.to_str() {
builder = builder.header(name.as_str(), v);
}
if !is_hop_by_hop_header(name.as_str())
&& let Ok(v) = value.to_str()
{
builder = builder.header(name.as_str(), v);
}
}
+4 -5
View File
@@ -109,12 +109,11 @@ impl NetworkPolicyDecider for DefaultPolicyDecider {
async fn decide(&self, request: &NetworkRequest) -> NetworkDecision {
// First check if the domain is allowed
let validation = self.allowlist.is_allowed(&request.host);
if !validation.is_allowed() {
if let crate::sandbox::proxy::allowlist::DomainValidationResult::Denied(reason) =
if !validation.is_allowed()
&& let crate::sandbox::proxy::allowlist::DomainValidationResult::Denied(reason) =
validation
{
return NetworkDecision::Deny { reason };
}
{
return NetworkDecision::Deny { reason };
}
// Check if we need to inject credentials
+1 -1
View File
@@ -261,7 +261,7 @@ pub use platform::{delete_master_key, get_master_key, has_master_key, store_mast
/// Parse a hex string to bytes.
fn hex_to_bytes(hex: &str) -> Result<Vec<u8>, SecretError> {
if hex.len() % 2 != 0 {
if !hex.len().is_multiple_of(2) {
return Err(SecretError::KeychainError(
"Invalid hex string length".to_string(),
));
+5 -1
View File
@@ -64,7 +64,11 @@ mod store;
mod types;
pub use crypto::SecretsCrypto;
pub use store::{PostgresSecretsStore, SecretsStore};
#[cfg(feature = "libsql")]
pub use store::LibSqlSecretsStore;
#[cfg(feature = "postgres")]
pub use store::PostgresSecretsStore;
pub use store::SecretsStore;
pub use types::{
CreateSecretParams, CredentialLocation, CredentialMapping, DecryptedSecret, Secret,
SecretError, SecretRef,
+382 -14
View File
@@ -10,6 +10,7 @@ use std::sync::Arc;
use async_trait::async_trait;
use chrono::Utc;
#[cfg(feature = "postgres")]
use deadpool_postgres::Pool;
use secrecy::ExposeSecret;
use uuid::Uuid;
@@ -61,11 +62,13 @@ pub trait SecretsStore: Send + Sync {
}
/// PostgreSQL implementation of SecretsStore.
#[cfg(feature = "postgres")]
pub struct PostgresSecretsStore {
pool: Pool,
crypto: Arc<SecretsCrypto>,
}
#[cfg(feature = "postgres")]
impl PostgresSecretsStore {
/// Create a new store with the given database pool and crypto instance.
pub fn new(pool: Pool, crypto: Arc<SecretsCrypto>) -> Self {
@@ -73,6 +76,7 @@ impl PostgresSecretsStore {
}
}
#[cfg(feature = "postgres")]
#[async_trait]
impl SecretsStore for PostgresSecretsStore {
async fn create(
@@ -149,10 +153,10 @@ impl SecretsStore for PostgresSecretsStore {
let secret = row_to_secret(&r);
// Check expiration
if let Some(expires_at) = secret.expires_at {
if expires_at < Utc::now() {
return Err(SecretError::Expired);
}
if let Some(expires_at) = secret.expires_at
&& expires_at < Utc::now()
{
return Err(SecretError::Expired);
}
Ok(secret)
@@ -272,10 +276,10 @@ impl SecretsStore for PostgresSecretsStore {
}
// Simple glob: * matches any suffix
if let Some(prefix) = pattern.strip_suffix('*') {
if secret_name.starts_with(prefix) {
return Ok(true);
}
if let Some(prefix) = pattern.strip_suffix('*')
&& secret_name.starts_with(prefix)
{
return Ok(true);
}
}
@@ -283,6 +287,7 @@ impl SecretsStore for PostgresSecretsStore {
}
}
#[cfg(feature = "postgres")]
fn row_to_secret(row: &tokio_postgres::Row) -> Secret {
Secret {
id: row.get("id"),
@@ -299,6 +304,332 @@ fn row_to_secret(row: &tokio_postgres::Row) -> Secret {
}
}
// ==================== libSQL implementation ====================
/// libSQL/Turso implementation of SecretsStore.
///
/// Holds an `Arc<Database>` handle and creates a fresh connection per operation,
/// matching the connection-per-request pattern used by the main `LibSqlBackend`.
#[cfg(feature = "libsql")]
pub struct LibSqlSecretsStore {
db: Arc<libsql::Database>,
crypto: Arc<SecretsCrypto>,
}
#[cfg(feature = "libsql")]
impl LibSqlSecretsStore {
/// Create a new store with the given shared libsql database handle and crypto instance.
pub fn new(db: Arc<libsql::Database>, crypto: Arc<SecretsCrypto>) -> Self {
Self { db, crypto }
}
fn connect(&self) -> Result<libsql::Connection, SecretError> {
self.db
.connect()
.map_err(|e| SecretError::Database(format!("Connection failed: {}", e)))
}
}
#[cfg(feature = "libsql")]
#[async_trait]
impl SecretsStore for LibSqlSecretsStore {
async fn create(
&self,
user_id: &str,
params: CreateSecretParams,
) -> Result<Secret, SecretError> {
let plaintext = params.value.expose_secret().as_bytes();
let (encrypted_value, key_salt) = self.crypto.encrypt(plaintext)?;
let id = Uuid::new_v4();
let now = Utc::now();
let now_str = now.to_rfc3339_opts(chrono::SecondsFormat::Millis, true);
let expires_at_str = params
.expires_at
.map(|dt| dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true));
// Start transaction for atomic upsert + read-back
let conn = self.connect()?;
let tx = conn
.transaction()
.await
.map_err(|e| SecretError::Database(e.to_string()))?;
tx.execute(
r#"
INSERT INTO secrets (id, user_id, name, encrypted_value, key_salt, provider, expires_at, created_at, updated_at)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?8)
ON CONFLICT (user_id, name) DO UPDATE SET
encrypted_value = excluded.encrypted_value,
key_salt = excluded.key_salt,
provider = excluded.provider,
expires_at = excluded.expires_at,
updated_at = ?8
"#,
libsql::params![
id.to_string(),
user_id,
params.name.as_str(),
libsql::Value::Blob(encrypted_value.clone()),
libsql::Value::Blob(key_salt.clone()),
libsql_opt_text(params.provider.as_deref()),
libsql_opt_text(expires_at_str.as_deref()),
now_str.as_str(),
],
)
.await
.map_err(|e| SecretError::Database(e.to_string()))?;
// Read back the row (may have been upserted)
let mut rows = tx
.query(
r#"
SELECT id, user_id, name, encrypted_value, key_salt, provider, expires_at,
last_used_at, usage_count, created_at, updated_at
FROM secrets
WHERE user_id = ?1 AND name = ?2
"#,
libsql::params![user_id, params.name.as_str()],
)
.await
.map_err(|e| SecretError::Database(e.to_string()))?;
let row = rows
.next()
.await
.map_err(|e| SecretError::Database(e.to_string()))?
.ok_or_else(|| SecretError::Database("Insert succeeded but row not found".into()))?;
let secret = libsql_row_to_secret(&row)?;
tx.commit()
.await
.map_err(|e| SecretError::Database(e.to_string()))?;
Ok(secret)
}
async fn get(&self, user_id: &str, name: &str) -> Result<Secret, SecretError> {
let conn = self.connect()?;
let mut rows = conn
.query(
r#"
SELECT id, user_id, name, encrypted_value, key_salt, provider, expires_at,
last_used_at, usage_count, created_at, updated_at
FROM secrets
WHERE user_id = ?1 AND name = ?2
"#,
libsql::params![user_id, name],
)
.await
.map_err(|e| SecretError::Database(e.to_string()))?;
match rows
.next()
.await
.map_err(|e| SecretError::Database(e.to_string()))?
{
Some(row) => {
let secret = libsql_row_to_secret(&row)?;
if let Some(expires_at) = secret.expires_at
&& expires_at < Utc::now()
{
return Err(SecretError::Expired);
}
Ok(secret)
}
None => Err(SecretError::NotFound(name.to_string())),
}
}
async fn get_decrypted(
&self,
user_id: &str,
name: &str,
) -> Result<DecryptedSecret, SecretError> {
let secret = self.get(user_id, name).await?;
self.crypto
.decrypt(&secret.encrypted_value, &secret.key_salt)
}
async fn exists(&self, user_id: &str, name: &str) -> Result<bool, SecretError> {
let conn = self.connect()?;
let mut rows = conn
.query(
"SELECT 1 FROM secrets WHERE user_id = ?1 AND name = ?2",
libsql::params![user_id, name],
)
.await
.map_err(|e| SecretError::Database(e.to_string()))?;
Ok(rows
.next()
.await
.map_err(|e| SecretError::Database(e.to_string()))?
.is_some())
}
async fn list(&self, user_id: &str) -> Result<Vec<SecretRef>, SecretError> {
let conn = self.connect()?;
let mut rows = conn
.query(
"SELECT name, provider FROM secrets WHERE user_id = ?1 ORDER BY name",
libsql::params![user_id],
)
.await
.map_err(|e| SecretError::Database(e.to_string()))?;
let mut refs = Vec::new();
while let Some(row) = rows
.next()
.await
.map_err(|e| SecretError::Database(e.to_string()))?
{
refs.push(SecretRef {
name: row.get::<String>(0).unwrap_or_default(),
provider: row.get::<String>(1).ok(),
});
}
Ok(refs)
}
async fn delete(&self, user_id: &str, name: &str) -> Result<bool, SecretError> {
let conn = self.connect()?;
let affected = conn
.execute(
"DELETE FROM secrets WHERE user_id = ?1 AND name = ?2",
libsql::params![user_id, name],
)
.await
.map_err(|e| SecretError::Database(e.to_string()))?;
Ok(affected > 0)
}
async fn record_usage(&self, secret_id: Uuid) -> Result<(), SecretError> {
let now = Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true);
let conn = self.connect()?;
conn.execute(
r#"
UPDATE secrets
SET last_used_at = ?1, usage_count = usage_count + 1
WHERE id = ?2
"#,
libsql::params![now.as_str(), secret_id.to_string()],
)
.await
.map_err(|e| SecretError::Database(e.to_string()))?;
Ok(())
}
async fn is_accessible(
&self,
user_id: &str,
secret_name: &str,
allowed_secrets: &[String],
) -> Result<bool, SecretError> {
if !self.exists(user_id, secret_name).await? {
return Ok(false);
}
for pattern in allowed_secrets {
if pattern == secret_name {
return Ok(true);
}
if let Some(prefix) = pattern.strip_suffix('*')
&& secret_name.starts_with(prefix)
{
return Ok(true);
}
}
Ok(false)
}
}
#[cfg(feature = "libsql")]
fn libsql_opt_text(s: Option<&str>) -> libsql::Value {
match s {
Some(s) => libsql::Value::Text(s.to_string()),
None => libsql::Value::Null,
}
}
#[cfg(feature = "libsql")]
fn libsql_parse_timestamp(s: &str) -> Result<chrono::DateTime<Utc>, SecretError> {
if let Ok(dt) = chrono::DateTime::parse_from_rfc3339(s) {
return Ok(dt.with_timezone(&Utc));
}
if let Ok(ndt) = chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f") {
return Ok(ndt.and_utc());
}
if let Ok(ndt) = chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S") {
return Ok(ndt.and_utc());
}
Err(SecretError::Database(format!(
"unparseable timestamp: {:?}",
s
)))
}
#[cfg(feature = "libsql")]
fn libsql_row_to_secret(row: &libsql::Row) -> Result<Secret, SecretError> {
let id_str: String = row
.get(0)
.map_err(|e| SecretError::Database(e.to_string()))?;
let user_id: String = row
.get(1)
.map_err(|e| SecretError::Database(e.to_string()))?;
let name: String = row
.get(2)
.map_err(|e| SecretError::Database(e.to_string()))?;
let encrypted_value: Vec<u8> = row
.get(3)
.map_err(|e| SecretError::Database(e.to_string()))?;
let key_salt: Vec<u8> = row
.get(4)
.map_err(|e| SecretError::Database(e.to_string()))?;
let provider: Option<String> = row.get::<String>(5).ok().filter(|s| !s.is_empty());
let expires_at = row
.get::<String>(6)
.ok()
.filter(|s| !s.is_empty())
.and_then(|s| libsql_parse_timestamp(&s).ok());
let last_used_at = row
.get::<String>(7)
.ok()
.filter(|s| !s.is_empty())
.and_then(|s| libsql_parse_timestamp(&s).ok());
let usage_count: i64 = row.get::<i64>(8).unwrap_or(0);
let created_at_str: String = row
.get(9)
.map_err(|e| SecretError::Database(e.to_string()))?;
let updated_at_str: String = row
.get(10)
.map_err(|e| SecretError::Database(e.to_string()))?;
Ok(Secret {
id: id_str
.parse()
.map_err(|e: uuid::Error| SecretError::Database(e.to_string()))?,
user_id,
name,
encrypted_value,
key_salt,
provider,
expires_at,
last_used_at,
usage_count,
created_at: libsql_parse_timestamp(&created_at_str)?,
updated_at: libsql_parse_timestamp(&updated_at_str)?,
})
}
/// In-memory implementation for testing.
#[cfg(test)]
pub mod testing {
@@ -364,12 +695,21 @@ pub mod testing {
}
async fn get(&self, user_id: &str, name: &str) -> Result<Secret, SecretError> {
self.secrets
let secret = self
.secrets
.read()
.await
.get(&(user_id.to_string(), name.to_string()))
.cloned()
.ok_or_else(|| SecretError::NotFound(name.to_string()))
.ok_or_else(|| SecretError::NotFound(name.to_string()))?;
if let Some(expires_at) = secret.expires_at
&& expires_at < Utc::now()
{
return Err(SecretError::Expired);
}
Ok(secret)
}
async fn get_decrypted(
@@ -430,10 +770,10 @@ pub mod testing {
if pattern == secret_name {
return Ok(true);
}
if let Some(prefix) = pattern.strip_suffix('*') {
if secret_name.starts_with(prefix) {
return Ok(true);
}
if let Some(prefix) = pattern.strip_suffix('*')
&& secret_name.starts_with(prefix)
{
return Ok(true);
}
}
Ok(false)
@@ -558,6 +898,34 @@ mod tests {
);
}
#[tokio::test]
async fn test_expired_secret_returns_error() {
let store = test_store();
let expires_at = chrono::Utc::now() - chrono::Duration::hours(1);
let params = CreateSecretParams::new("expired_key", "value").with_expiry(expires_at);
store.create("user1", params).await.unwrap();
let result = store.get("user1", "expired_key").await;
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
crate::secrets::SecretError::Expired
));
}
#[tokio::test]
async fn test_non_expired_secret_succeeds() {
let store = test_store();
let expires_at = chrono::Utc::now() + chrono::Duration::hours(1);
let params = CreateSecretParams::new("fresh_key", "value").with_expiry(expires_at);
store.create("user1", params).await.unwrap();
let result = store.get("user1", "fresh_key").await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_user_isolation() {
let store = test_store();
+87 -83
View File
@@ -15,6 +15,10 @@ pub struct Settings {
pub onboard_completed: bool,
// === Step 1: Database ===
/// Database backend: "postgres" or "libsql".
#[serde(default)]
pub database_backend: Option<String>,
/// Database connection URL (postgres://...).
#[serde(default)]
pub database_url: Option<String>,
@@ -23,13 +27,31 @@ pub struct Settings {
#[serde(default)]
pub database_pool_size: Option<usize>,
/// Path to local libSQL database file.
#[serde(default)]
pub libsql_path: Option<String>,
/// Turso cloud URL for remote replica sync.
#[serde(default)]
pub libsql_url: Option<String>,
// === Step 2: Security ===
/// Source for the secrets master key.
#[serde(default)]
pub secrets_master_key_source: KeySource,
// === Step 3: NEAR AI Auth ===
// Session stored separately in session.json
// === Step 3: Inference Provider ===
/// LLM backend: "nearai", "anthropic", "openai", "ollama", "openai_compatible".
#[serde(default)]
pub llm_backend: Option<String>,
/// Ollama base URL (when llm_backend = "ollama").
#[serde(default)]
pub ollama_base_url: Option<String>,
/// OpenAI-compatible endpoint base URL (when llm_backend = "openai_compatible").
#[serde(default)]
pub openai_compatible_base_url: Option<String>,
// === Step 4: Model Selection ===
/// Currently selected model.
@@ -487,20 +509,16 @@ impl Default for BuilderSettings {
}
impl Settings {
/// Get the default settings file path (~/.ironclaw/settings.json).
pub fn default_path() -> PathBuf {
dirs::home_dir()
.unwrap_or_else(|| PathBuf::from("."))
.join(".ironclaw")
.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
// Start with defaults, then overlay each DB setting.
//
// The settings table stores both Settings struct fields and app-specific
// data (e.g. nearai.session_token). Skip keys that don't correspond to
// a known Settings path.
let mut settings = Self::default();
for (key, value) in map {
@@ -509,17 +527,23 @@ impl Settings {
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(),
serde_json::Value::Null => continue, // null means default, skip
other => other.to_string(),
};
if let Err(e) = settings.set(key, &value_str) {
tracing::warn!(
"Failed to apply DB setting '{}' = '{}': {}",
key,
value_str,
e
);
match settings.set(key, &value_str) {
Ok(()) => {}
// The settings table stores both Settings fields and app-specific
// data (e.g. nearai.session_token). Silently skip unknown paths.
Err(e) if e.starts_with("Path not found") => {}
Err(e) => {
tracing::warn!(
"Failed to apply DB setting '{}' = '{}': {}",
key,
value_str,
e
);
}
}
}
@@ -540,50 +564,27 @@ impl Settings {
map
}
/// Get the default settings file path (~/.ironclaw/settings.json).
pub fn default_path() -> std::path::PathBuf {
dirs::home_dir()
.unwrap_or_else(|| std::path::PathBuf::from("."))
.join(".ironclaw")
.join("settings.json")
}
/// Load settings from disk, returning default if not found.
pub fn load() -> Self {
Self::load_from(&Self::default_path())
}
/// Load settings from a specific path.
pub fn load_from(path: &PathBuf) -> Self {
/// Load settings from a specific path (used by bootstrap legacy migration).
pub fn load_from(path: &std::path::Path) -> Self {
match std::fs::read_to_string(path) {
Ok(data) => serde_json::from_str(&data).unwrap_or_default(),
Err(_) => Self::default(),
}
}
/// Save settings to disk.
pub fn save(&self) -> std::io::Result<()> {
self.save_to(&Self::default_path())
}
/// Save settings to a specific path.
pub fn save_to(&self, path: &PathBuf) -> std::io::Result<()> {
// Ensure parent directory exists
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)
}
/// Get the selected model, falling back to the provided default.
pub fn model_or(&self, default: &str) -> String {
self.selected_model
.clone()
.unwrap_or_else(|| default.to_string())
}
/// Set the selected model and save.
pub fn set_model(&mut self, model: &str) -> std::io::Result<()> {
self.selected_model = Some(model.to_string());
self.save()
}
/// Get a setting value by dotted path (e.g., "agent.max_parallel_jobs").
pub fn get(&self, path: &str) -> Option<String> {
let json = serde_json::to_value(self).ok()?;
@@ -768,42 +769,22 @@ fn collect_settings(
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
#[test]
fn test_settings_save_load() {
let dir = tempdir().unwrap();
let path = dir.path().join("settings.json");
fn test_db_map_round_trip() {
let settings = Settings {
selected_model: Some("claude-3-5-sonnet-20241022".to_string()),
..Default::default()
};
settings.save_to(&path).unwrap();
let loaded = Settings::load_from(&path);
let map = settings.to_db_map();
let restored = Settings::from_db_map(&map);
assert_eq!(
loaded.selected_model,
restored.selected_model,
Some("claude-3-5-sonnet-20241022".to_string())
);
}
#[test]
fn test_model_or_default() {
let settings = Settings::default();
assert_eq!(
settings.model_or("default-model"),
"default-model".to_string()
);
let settings = Settings {
selected_model: Some("my-model".to_string()),
..Default::default()
};
assert_eq!(settings.model_or("default-model"), "my-model".to_string());
}
#[test]
fn test_get_setting() {
let settings = Settings::default();
@@ -874,16 +855,13 @@ mod tests {
}
#[test]
fn test_telegram_owner_id_round_trip() {
let dir = tempdir().unwrap();
let path = dir.path().join("settings.json");
fn test_telegram_owner_id_db_round_trip() {
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));
let map = settings.to_db_map();
let restored = Settings::from_db_map(&map);
assert_eq!(restored.channels.telegram_owner_id, Some(123456789));
}
#[test]
@@ -900,4 +878,30 @@ mod tests {
.unwrap();
assert_eq!(settings.channels.telegram_owner_id, Some(987654321));
}
#[test]
fn test_llm_backend_round_trip() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("settings.json");
let settings = Settings {
llm_backend: Some("anthropic".to_string()),
ollama_base_url: Some("http://localhost:11434".to_string()),
openai_compatible_base_url: Some("http://my-vllm:8000/v1".to_string()),
..Default::default()
};
let json = serde_json::to_string_pretty(&settings).unwrap();
std::fs::write(&path, json).unwrap();
let loaded = Settings::load_from(&path);
assert_eq!(loaded.llm_backend, Some("anthropic".to_string()));
assert_eq!(
loaded.ollama_base_url,
Some("http://localhost:11434".to_string())
);
assert_eq!(
loaded.openai_compatible_base_url,
Some("http://my-vllm:8000/v1".to_string())
);
}
}
+199 -149
View File
@@ -12,54 +12,88 @@ use reqwest::Client;
use secrecy::{ExposeSecret, SecretString};
use serde::Deserialize;
use crate::secrets::{CreateSecretParams, PostgresSecretsStore, SecretsCrypto, SecretsStore};
use crate::settings::Settings;
#[cfg(feature = "postgres")]
use crate::secrets::SecretsCrypto;
use crate::secrets::{CreateSecretParams, SecretsStore};
use crate::settings::{Settings, TunnelSettings};
use crate::setup::prompts::{
confirm, input, optional_input, print_error, print_info, print_success, secret_input,
};
/// Typed errors for channel setup flows.
#[derive(Debug, thiserror::Error)]
pub enum ChannelSetupError {
#[error("I/O error: {0}")]
Io(#[from] std::io::Error),
#[error("{0}")]
Network(String),
#[error("{0}")]
Secrets(String),
#[error("{0}")]
Validation(String),
}
/// Context for saving secrets during setup.
pub struct SecretsContext {
store: PostgresSecretsStore,
store: Arc<dyn SecretsStore>,
user_id: String,
}
impl SecretsContext {
/// Create a new secrets context.
/// Create a new secrets context from a trait-object store.
pub fn from_store(store: Arc<dyn SecretsStore>, user_id: &str) -> Self {
Self {
store,
user_id: user_id.to_string(),
}
}
/// Create a new secrets context from a PostgreSQL pool and crypto.
#[cfg(feature = "postgres")]
pub fn new(pool: deadpool_postgres::Pool, crypto: Arc<SecretsCrypto>, user_id: &str) -> Self {
Self {
store: PostgresSecretsStore::new(pool, crypto),
store: Arc::new(crate::secrets::PostgresSecretsStore::new(pool, crypto)),
user_id: user_id.to_string(),
}
}
/// Save a secret to the database.
pub async fn save_secret(&self, name: &str, value: &SecretString) -> Result<(), String> {
pub async fn save_secret(
&self,
name: &str,
value: &SecretString,
) -> Result<(), ChannelSetupError> {
let params = CreateSecretParams::new(name, value.expose_secret());
self.store
.create(&self.user_id, params)
.await
.map_err(|e| format!("Failed to save secret: {}", e))?;
.map_err(|e| ChannelSetupError::Secrets(format!("Failed to save secret: {}", e)))?;
Ok(())
}
/// Check if a secret exists.
pub async fn secret_exists(&self, name: &str) -> bool {
self.store
.exists(&self.user_id, name)
.await
.unwrap_or(false)
match self.store.exists(&self.user_id, name).await {
Ok(exists) => exists,
Err(e) => {
tracing::warn!(secret = name, error = %e, "Failed to check if secret exists, assuming absent");
false
}
}
}
/// Read a secret from the database (decrypted).
pub async fn get_secret(&self, name: &str) -> Result<SecretString, String> {
pub async fn get_secret(&self, name: &str) -> Result<SecretString, ChannelSetupError> {
let decrypted = self
.store
.get_decrypted(&self.user_id, name)
.await
.map_err(|e| format!("Failed to read secret: {}", e))?;
.map_err(|e| ChannelSetupError::Secrets(format!("Failed to read secret: {}", e)))?;
Ok(SecretString::from(decrypted.expose().to_string()))
}
}
@@ -96,7 +130,6 @@ struct TelegramGetUpdatesResponse {
#[derive(Debug, Deserialize)]
struct TelegramUpdate {
#[allow(dead_code)]
update_id: i64,
message: Option<TelegramUpdateMessage>,
}
@@ -120,7 +153,10 @@ struct TelegramUpdateUser {
/// 2. Entering the bot token
/// 3. Validating the token
/// 4. Saving the token to the database
pub async fn setup_telegram(secrets: &SecretsContext) -> Result<TelegramSetupResult, String> {
pub async fn setup_telegram(
secrets: &SecretsContext,
settings: &Settings,
) -> Result<TelegramSetupResult, ChannelSetupError> {
println!("Telegram Setup:");
println!();
print_info("To create a Telegram bot:");
@@ -132,10 +168,10 @@ pub async fn setup_telegram(secrets: &SecretsContext) -> Result<TelegramSetupRes
// Check if token already exists
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())? {
if !confirm("Replace existing token?", false)? {
// 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?;
let webhook_secret = setup_telegram_webhook_secret(secrets, &settings.tunnel).await?;
let owner_id = bind_telegram_owner_flow(secrets, settings).await?;
return Ok(TelegramSetupResult {
enabled: true,
bot_username: None,
@@ -145,47 +181,48 @@ pub async fn setup_telegram(secrets: &SecretsContext) -> Result<TelegramSetupRes
}
}
let token = secret_input("Bot token (from @BotFather)").map_err(|e| e.to_string())?;
loop {
let token = secret_input("Bot token (from @BotFather)")?;
// Validate the token
print_info("Validating bot token...");
// Validate the token
print_info("Validating bot token...");
match validate_telegram_token(&token).await {
Ok(username) => {
print_success(&format!(
"Bot validated: @{}",
username.as_deref().unwrap_or("unknown")
));
match validate_telegram_token(&token).await {
Ok(username) => {
print_success(&format!(
"Bot validated: @{}",
username.as_deref().unwrap_or("unknown")
));
// Save to database
secrets.save_secret("telegram_bot_token", &token).await?;
print_success("Token saved to database");
// Save to database
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?;
// 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?;
// Offer webhook secret configuration
let webhook_secret =
setup_telegram_webhook_secret(secrets, &settings.tunnel).await?;
Ok(TelegramSetupResult {
enabled: true,
bot_username: username,
webhook_secret,
owner_id,
})
}
Err(e) => {
print_error(&format!("Token validation failed: {}", e));
return Ok(TelegramSetupResult {
enabled: true,
bot_username: username,
webhook_secret,
owner_id,
});
}
Err(e) => {
print_error(&format!("Token validation failed: {}", e));
if confirm("Try again?", true).map_err(|e| e.to_string())? {
Box::pin(setup_telegram(secrets)).await
} else {
Ok(TelegramSetupResult {
enabled: false,
bot_username: None,
webhook_secret: None,
owner_id: None,
})
if !confirm("Try again?", true)? {
return Ok(TelegramSetupResult {
enabled: false,
bot_username: None,
webhook_secret: None,
owner_id: None,
});
}
}
}
}
@@ -195,14 +232,14 @@ pub async fn setup_telegram(secrets: &SecretsContext) -> Result<TelegramSetupRes
///
/// 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> {
async fn bind_telegram_owner(token: &SecretString) -> Result<Option<i64>, ChannelSetupError> {
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())? {
if !confirm("Bind bot to your Telegram account?", true)? {
print_info("Skipping account binding. Bot will accept messages from all users.");
return Ok(None);
}
@@ -213,14 +250,16 @@ async fn bind_telegram_owner(token: &SecretString) -> Result<Option<i64>, String
let client = Client::builder()
.timeout(std::time::Duration::from_secs(35))
.build()
.map_err(|e| format!("Failed to create HTTP client: {}", e))?;
.map_err(|e| ChannelSetupError::Network(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;
if let Err(e) = client.post(&delete_url).send().await {
tracing::warn!("Failed to delete webhook (getUpdates may not work): {e}");
}
let updates_url = format!(
"https://api.telegram.org/bot{}/getUpdates",
@@ -235,49 +274,56 @@ async fn bind_telegram_owner(token: &SecretString) -> Result<Option<i64>, String
.query(&[("timeout", "30"), ("allowed_updates", "[\"message\"]")])
.send()
.await
.map_err(|e| format!("getUpdates request failed: {}", e))?;
.map_err(|e| ChannelSetupError::Network(format!("getUpdates request failed: {}", e)))?;
if !response.status().is_success() {
return Err(format!("getUpdates returned status {}", response.status()));
return Err(ChannelSetupError::Network(format!(
"getUpdates returned status {}",
response.status()
)));
}
let body: TelegramGetUpdatesResponse = response
.json()
.await
.map_err(|e| format!("Failed to parse getUpdates response: {}", e))?;
let body: TelegramGetUpdatesResponse = response.json().await.map_err(|e| {
ChannelSetupError::Network(format!("Failed to parse getUpdates response: {}", e))
})?;
if !body.ok {
return Err("Telegram API returned error for getUpdates".to_string());
return Err(ChannelSetupError::Network(
"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());
if let Some(ref msg) = update.message
&& 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
));
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));
// Acknowledge the update so it doesn't pile up
let ack_url = format!(
"https://api.telegram.org/bot{}/getUpdates",
token.expose_secret()
);
if let Err(e) = client
.get(&ack_url)
.query(&[("offset", &(update.update_id + 1).to_string())])
.send()
.await
{
tracing::warn!("Failed to acknowledge Telegram update: {e}");
}
return Ok(Some(from.id));
}
}
}
@@ -290,12 +336,13 @@ async fn bind_telegram_owner(token: &SecretString) -> Result<Option<i64>, String
/// 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();
async fn bind_telegram_owner_flow(
secrets: &SecretsContext,
settings: &Settings,
) -> Result<Option<i64>, ChannelSetupError> {
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())? {
if !confirm("Re-bind to a different account?", false)? {
return Ok(settings.channels.telegram_owner_id);
}
}
@@ -310,12 +357,10 @@ async fn bind_telegram_owner_flow(secrets: &SecretsContext) -> Result<Option<i64
///
/// This is shared across all channels that need webhook endpoints.
/// Returns the tunnel URL if configured.
pub fn setup_tunnel() -> Result<Option<String>, String> {
// Check if already configured
let settings = Settings::load();
pub fn setup_tunnel(settings: &Settings) -> Result<Option<String>, ChannelSetupError> {
if let Some(ref url) = settings.tunnel.public_url {
print_info(&format!("Existing tunnel configured: {}", url));
if !confirm("Change tunnel configuration?", false).map_err(|e| e.to_string())? {
if !confirm("Change tunnel configuration?", false)? {
return Ok(Some(url.clone()));
}
}
@@ -335,30 +380,24 @@ pub fn setup_tunnel() -> Result<Option<String>, String> {
print_info("Security comes from provider-specific secrets (e.g., Telegram webhook secret).");
println!();
if !confirm("Configure a tunnel?", false).map_err(|e| e.to_string())? {
if !confirm("Configure a tunnel?", false)? {
return Ok(None);
}
let tunnel_url =
input("Tunnel URL (e.g., https://abc123.ngrok.io)").map_err(|e| e.to_string())?;
let tunnel_url = input("Tunnel URL (e.g., https://abc123.ngrok.io)")?;
// Validate URL format
if !tunnel_url.starts_with("https://") {
print_error("URL must start with https:// (webhooks require HTTPS)");
return Err("Invalid tunnel URL: must use HTTPS".to_string());
return Err(ChannelSetupError::Validation(
"Invalid tunnel URL: must use HTTPS".to_string(),
));
}
// Remove trailing slash if present
let tunnel_url = tunnel_url.trim_end_matches('/').to_string();
// Save to settings
let mut settings = Settings::load();
settings.tunnel.public_url = Some(tunnel_url.clone());
settings
.save()
.map_err(|e| format!("Failed to save settings: {}", e))?;
print_success(&format!("Tunnel URL saved: {}", tunnel_url));
print_success(&format!("Tunnel URL configured: {}", tunnel_url));
print_info("");
print_info("Make sure your tunnel is running before starting the agent.");
print_info("You can also set TUNNEL_URL environment variable to override.");
@@ -369,10 +408,11 @@ pub fn setup_tunnel() -> Result<Option<String>, String> {
/// Set up Telegram webhook secret for signature validation.
///
/// Returns the webhook secret if configured.
async fn setup_telegram_webhook_secret(secrets: &SecretsContext) -> Result<Option<String>, String> {
// Check if tunnel is configured
let settings = Settings::load();
if settings.tunnel.public_url.is_none() {
async fn setup_telegram_webhook_secret(
secrets: &SecretsContext,
tunnel: &TunnelSettings,
) -> Result<Option<String>, ChannelSetupError> {
if tunnel.public_url.is_none() {
print_info("");
print_info("No tunnel configured. Telegram will use polling mode (30s+ delay).");
print_info("Run setup again to configure a tunnel for instant delivery.");
@@ -384,7 +424,7 @@ async fn setup_telegram_webhook_secret(secrets: &SecretsContext) -> Result<Optio
print_info("A webhook secret adds an extra layer of security by validating");
print_info("that requests actually come from Telegram's servers.");
if !confirm("Generate a webhook secret?", true).map_err(|e| e.to_string())? {
if !confirm("Generate a webhook secret?", true)? {
return Ok(None);
}
@@ -403,11 +443,13 @@ async fn setup_telegram_webhook_secret(secrets: &SecretsContext) -> Result<Optio
/// Validate a Telegram bot token by calling the getMe API.
///
/// Returns the bot's username if valid.
pub async fn validate_telegram_token(token: &SecretString) -> Result<Option<String>, String> {
pub async fn validate_telegram_token(
token: &SecretString,
) -> Result<Option<String>, ChannelSetupError> {
let client = Client::builder()
.timeout(std::time::Duration::from_secs(10))
.build()
.map_err(|e| format!("Failed to create HTTP client: {}", e))?;
.map_err(|e| ChannelSetupError::Network(format!("Failed to create HTTP client: {}", e)))?;
let url = format!(
"https://api.telegram.org/bot{}/getMe",
@@ -418,21 +460,26 @@ pub async fn validate_telegram_token(token: &SecretString) -> Result<Option<Stri
.get(&url)
.send()
.await
.map_err(|e| format!("Request failed: {}", e))?;
.map_err(|e| ChannelSetupError::Network(format!("Request failed: {}", e)))?;
if !response.status().is_success() {
return Err(format!("API returned status {}", response.status()));
return Err(ChannelSetupError::Network(format!(
"API returned status {}",
response.status()
)));
}
let body: TelegramGetMeResponse = response
.json()
.await
.map_err(|e| format!("Failed to parse response: {}", e))?;
.map_err(|e| ChannelSetupError::Network(format!("Failed to parse response: {}", e)))?;
if body.ok {
Ok(body.result.and_then(|u| u.username))
} else {
Err("Telegram API returned error".to_string())
Err(ChannelSetupError::Network(
"Telegram API returned error".to_string(),
))
}
}
@@ -445,38 +492,34 @@ pub struct HttpSetupResult {
}
/// Set up HTTP webhook channel.
pub async fn setup_http(secrets: &SecretsContext) -> Result<HttpSetupResult, String> {
pub async fn setup_http(secrets: &SecretsContext) -> Result<HttpSetupResult, ChannelSetupError> {
println!("HTTP Webhook Setup:");
println!();
print_info("The HTTP webhook allows external services to send messages to the agent.");
println!();
let port_str = optional_input("Port", Some("default: 8080")).map_err(|e| e.to_string())?;
let port_str = optional_input("Port", Some("default: 8080"))?;
let port: u16 = port_str
.as_deref()
.unwrap_or("8080")
.parse()
.map_err(|e| format!("Invalid port: {}", e))?;
.map_err(|e| ChannelSetupError::Validation(format!("Invalid port: {}", e)))?;
if port < 1024 {
print_info("Note: Ports below 1024 may require root privileges");
}
let host = optional_input("Host", Some("default: 0.0.0.0"))
.map_err(|e| e.to_string())?
.unwrap_or_else(|| "0.0.0.0".to_string());
let host =
optional_input("Host", Some("default: 0.0.0.0"))?.unwrap_or_else(|| "0.0.0.0".to_string());
// Generate a webhook secret
if confirm("Generate a webhook secret for authentication?", true).map_err(|e| e.to_string())? {
if confirm("Generate a webhook secret for authentication?", true)? {
let secret = generate_webhook_secret();
secrets
.save_secret("http_webhook_secret", &SecretString::from(secret.clone()))
.save_secret("http_webhook_secret", &SecretString::from(secret))
.await?;
print_success("Webhook secret generated and saved to database");
print_info(&format!(
"Secret: {} (store this for your webhook clients)",
secret
));
print_info("Retrieve it later with: ironclaw secret get http_webhook_secret");
}
print_success(&format!("HTTP webhook will listen on {}:{}", host, port));
@@ -490,11 +533,7 @@ pub async fn setup_http(secrets: &SecretsContext) -> Result<HttpSetupResult, Str
/// Generate a random webhook secret.
pub fn generate_webhook_secret() -> String {
use rand::RngCore;
let mut rng = rand::thread_rng();
let mut bytes = [0u8; 32];
rng.fill_bytes(&mut bytes);
bytes.iter().map(|b| format!("{:02x}", b)).collect()
generate_secret_with_length(32)
}
/// Result of WASM channel setup.
@@ -512,7 +551,7 @@ pub async fn setup_wasm_channel(
secrets: &SecretsContext,
channel_name: &str,
setup: &crate::channels::wasm::SetupSchema,
) -> Result<WasmChannelSetupResult, String> {
) -> Result<WasmChannelSetupResult, ChannelSetupError> {
println!("{} Setup:", channel_name);
println!();
@@ -523,7 +562,7 @@ pub async fn setup_wasm_channel(
"Existing {} found in database.",
secret_config.name
));
if !confirm("Replace existing value?", false).map_err(|e| e.to_string())? {
if !confirm("Replace existing value?", false)? {
continue;
}
}
@@ -531,8 +570,7 @@ pub async fn setup_wasm_channel(
// Get the value from user or auto-generate
let value = if secret_config.optional {
let input_value =
optional_input(&secret_config.prompt, Some("leave empty to auto-generate"))
.map_err(|e| e.to_string())?;
optional_input(&secret_config.prompt, Some("leave empty to auto-generate"))?;
if let Some(v) = input_value {
if !v.is_empty() {
@@ -559,18 +597,21 @@ pub async fn setup_wasm_channel(
}
} else {
// Required secret
let input_value = secret_input(&secret_config.prompt).map_err(|e| e.to_string())?;
let input_value = secret_input(&secret_config.prompt)?;
// Validate if pattern is provided
if let Some(ref pattern) = secret_config.validation {
let re = regex::Regex::new(pattern)
.map_err(|e| format!("Invalid validation pattern: {}", e))?;
let re = regex::Regex::new(pattern).map_err(|e| {
ChannelSetupError::Validation(format!("Invalid validation pattern: {}", e))
})?;
if !re.is_match(input_value.expose_secret()) {
print_error(&format!(
"Value does not match expected format: {}",
pattern
));
return Err("Validation failed".to_string());
return Err(ChannelSetupError::Validation(
"Validation failed".to_string(),
));
}
}
@@ -582,14 +623,11 @@ pub async fn setup_wasm_channel(
print_success(&format!("{} saved to database", secret_config.name));
}
// Optionally validate the configuration
// TODO: Substitute secrets into the validation URL and make a
// GET request to verify the configured credentials actually work.
if let Some(ref validation_endpoint) = setup.validation_endpoint {
print_info("Validating configuration...");
// The validation endpoint may contain placeholders like {telegram_bot_token}
// For now, we skip validation since we'd need to substitute secrets
// A full implementation would fetch secrets and substitute them
print_info(&format!(
"Validation endpoint configured: {} (validation skipped)",
"Validation endpoint configured: {} (validation not yet implemented)",
validation_endpoint
));
}
@@ -613,11 +651,23 @@ fn generate_secret_with_length(length: usize) -> String {
#[cfg(test)]
mod tests {
use super::*;
use crate::setup::channels::generate_webhook_secret;
#[test]
fn test_generate_webhook_secret() {
let secret = generate_webhook_secret();
assert_eq!(secret.len(), 64); // 32 bytes = 64 hex chars
}
#[test]
fn test_generate_secret_with_length() {
use super::generate_secret_with_length;
let s = generate_secret_with_length(16);
assert_eq!(s.len(), 32); // 16 bytes = 32 hex chars
assert!(s.chars().all(|c| c.is_ascii_hexdigit()));
let s2 = generate_secret_with_length(1);
assert_eq!(s2.len(), 2);
}
}
+5 -2
View File
@@ -3,7 +3,7 @@
//! Provides a guided setup experience for:
//! 1. Database connection
//! 2. Security (secrets master key)
//! 3. NEAR AI authentication
//! 3. Inference provider selection
//! 4. Model selection
//! 5. Embeddings
//! 6. Channel configuration (HTTP, Telegram, etc.)
@@ -20,13 +20,16 @@
mod channels;
mod prompts;
#[cfg(any(feature = "postgres", feature = "libsql"))]
mod wizard;
pub use channels::{
SecretsContext, setup_http, setup_telegram, setup_tunnel, validate_telegram_token,
ChannelSetupError, SecretsContext, setup_http, setup_telegram, setup_tunnel,
validate_telegram_token,
};
pub use prompts::{
confirm, input, optional_input, print_error, print_header, print_info, print_step,
print_success, secret_input, select_many, select_one,
};
#[cfg(any(feature = "postgres", feature = "libsql"))]
pub use wizard::{SetupConfig, SetupWizard};
+10 -4
View File
@@ -21,6 +21,7 @@ use secrecy::SecretString;
/// Display a numbered menu and get user selection.
///
/// Returns the index (0-based) of the selected option.
/// Pressing Enter without input selects the first option (index 0).
///
/// # Example
///
@@ -54,10 +55,11 @@ pub fn select_one(prompt: &str, options: &[&str]) -> io::Result<usize> {
}
// Parse number
if let Ok(num) = input.parse::<usize>() {
if num >= 1 && num <= options.len() {
return Ok(num - 1);
}
if let Ok(num) = input.parse::<usize>()
&& num >= 1
&& num <= options.len()
{
return Ok(num - 1);
}
writeln!(
@@ -83,6 +85,10 @@ pub fn select_one(prompt: &str, options: &[&str]) -> io::Result<usize> {
/// ])?;
/// ```
pub fn select_many(prompt: &str, options: &[(&str, bool)]) -> io::Result<Vec<usize>> {
if options.is_empty() {
return Ok(vec![]);
}
let mut stdout = io::stdout();
let mut selected: Vec<bool> = options.iter().map(|(_, s)| *s).collect();
let mut cursor_pos = 0;
+1230 -147
View File
File diff suppressed because it is too large Load Diff
+1 -1
View File
@@ -595,7 +595,7 @@ Create alongside the .wasm file to grant capabilities:
AgentToolError::BuilderFailed(format!("LLM response failed: {}", e))
})?;
match result {
match result.result {
RespondResult::Text(response) => {
reason_ctx.messages.push(ChatMessage::assistant(&response));
+27 -27
View File
@@ -326,20 +326,20 @@ impl TestHarness {
}
// Verify expected output
if let Some(ref expected) = test.expected_output {
if &actual != expected {
return TestResult {
name: test.name.clone(),
passed: false,
duration,
error: Some(format!(
"Output mismatch:\nExpected: {}\nActual: {}",
serde_json::to_string_pretty(expected).unwrap_or_default(),
serde_json::to_string_pretty(&actual).unwrap_or_default()
)),
actual_output: Some(actual),
};
}
if let Some(ref expected) = test.expected_output
&& &actual != expected
{
return TestResult {
name: test.name.clone(),
passed: false,
duration,
error: Some(format!(
"Output mismatch:\nExpected: {}\nActual: {}",
serde_json::to_string_pretty(expected).unwrap_or_default(),
serde_json::to_string_pretty(&actual).unwrap_or_default()
)),
actual_output: Some(actual),
};
}
// Verify expected fields
@@ -357,19 +357,19 @@ impl TestHarness {
};
}
if let Some(ref expected_value) = field.value {
if field_value != Some(expected_value) {
return TestResult {
name: test.name.clone(),
passed: false,
duration,
error: Some(format!(
"Field '{}' mismatch: expected {:?}, got {:?}",
field.path, expected_value, field_value
)),
actual_output: Some(actual),
};
}
if let Some(ref expected_value) = field.value
&& field_value != Some(expected_value)
{
return TestResult {
name: test.name.clone(),
passed: false,
duration,
error: Some(format!(
"Field '{}' mismatch: expected {:?}, got {:?}",
field.path, expected_value, field_value
)),
actual_output: Some(actual),
};
}
}
}
+137 -35
View File
@@ -50,65 +50,90 @@ const MAX_WRITE_SIZE: usize = 5 * 1024 * 1024;
/// Maximum directory listing entries.
const MAX_DIR_ENTRIES: usize = 500;
/// Validate that a path is safe (no traversal attacks).
fn validate_path(path_str: &str, base_dir: Option<&Path>) -> Result<PathBuf, ToolError> {
let path = PathBuf::from(path_str);
// Reject paths with suspicious components (validation only, no action needed)
/// Normalize a path by resolving `.` and `..` components lexically (no filesystem access).
///
/// This is critical for security: `std::fs::canonicalize` only works on paths that exist,
/// so for new files we must normalize without touching the filesystem.
fn normalize_lexical(path: &Path) -> PathBuf {
let mut components = Vec::new();
for component in path.components() {
match component {
std::path::Component::ParentDir => {
// Allow .. but validate final path is within sandbox
}
std::path::Component::Normal(s) => {
let s = s.to_string_lossy();
if s.starts_with('.') && s != "." && s != ".." && !s.starts_with(".git") {
// Hidden files are OK for .git, .gitignore, etc.
// Only pop if there's a normal component to pop (don't escape root/prefix)
if components
.last()
.is_some_and(|c| matches!(c, std::path::Component::Normal(_)))
{
components.pop();
}
}
_ => {}
std::path::Component::CurDir => {}
other => components.push(other),
}
}
components.iter().collect()
}
/// Validate that a path is safe (no traversal attacks).
///
/// For sandboxed paths (base_dir is set), we normalize the joined path lexically
/// and then verify it lives under the canonical base. This prevents escapes through
/// non-existent parent directories where `canonicalize()` would fall back to the
/// raw (un-normalized) path.
fn validate_path(path_str: &str, base_dir: Option<&Path>) -> Result<PathBuf, ToolError> {
let path = PathBuf::from(path_str);
// Resolve to absolute path
let resolved = if path.is_absolute() {
path.canonicalize().unwrap_or_else(|_| path.clone())
path.canonicalize()
.unwrap_or_else(|_| normalize_lexical(&path))
} else if let Some(base) = base_dir {
base.join(&path)
let joined = base.join(&path);
joined
.canonicalize()
.unwrap_or_else(|_| base.join(&path))
.unwrap_or_else(|_| normalize_lexical(&joined))
} else {
std::env::current_dir()
let joined = std::env::current_dir()
.unwrap_or_else(|_| PathBuf::from("."))
.join(&path)
.join(&path);
normalize_lexical(&joined)
};
// If base_dir is set, ensure path is within it
// If base_dir is set, ensure the resolved path is within it
if let Some(base) = base_dir {
// Canonicalize the base to handle symlinks (e.g., /var -> /private/var on macOS)
let base_canonical = base.canonicalize().unwrap_or_else(|_| base.to_path_buf());
let base_canonical = base
.canonicalize()
.unwrap_or_else(|_| normalize_lexical(base));
// For files that don't exist yet, we need to check the parent directory
// and ensure the resolved path would be within the base
// For existing paths, canonicalize to resolve symlinks.
// For non-existent paths, the lexical normalization above already removed
// all `..` components, so starts_with is reliable.
let check_path = if resolved.exists() {
resolved.canonicalize().unwrap_or_else(|_| resolved.clone())
} else {
// For non-existent files, canonicalize the parent and append the filename
if let Some(parent) = resolved.parent() {
if parent.exists() {
let canonical_parent = parent
// Walk up to the nearest existing ancestor directory, canonicalize it,
// then re-append the remaining tail. This handles the case where a
// symlink sits above the new file.
let mut ancestor = resolved.as_path();
let mut tail_parts: Vec<&std::ffi::OsStr> = Vec::new();
loop {
if ancestor.exists() {
let canonical_ancestor = ancestor
.canonicalize()
.unwrap_or_else(|_| parent.to_path_buf());
if let Some(filename) = resolved.file_name() {
canonical_parent.join(filename)
} else {
resolved.clone()
.unwrap_or_else(|_| ancestor.to_path_buf());
let mut result = canonical_ancestor;
for part in tail_parts.into_iter().rev() {
result = result.join(part);
}
} else {
resolved.clone()
break result;
}
if let Some(name) = ancestor.file_name() {
tail_parts.push(name);
}
match ancestor.parent() {
Some(parent) if parent != ancestor => ancestor = parent,
_ => break resolved.clone(),
}
} else {
resolved.clone()
}
};
@@ -871,4 +896,81 @@ mod tests {
let entries = result.result.get("entries").unwrap().as_array().unwrap();
assert!(entries.len() >= 2);
}
#[test]
fn test_normalize_lexical() {
// Basic .. resolution
assert_eq!(
normalize_lexical(Path::new("/a/b/../c")),
PathBuf::from("/a/c")
);
// Multiple .. components
assert_eq!(
normalize_lexical(Path::new("/a/b/c/../../d")),
PathBuf::from("/a/d")
);
// . components stripped
assert_eq!(
normalize_lexical(Path::new("/a/./b/./c")),
PathBuf::from("/a/b/c")
);
// Cannot escape root
assert_eq!(
normalize_lexical(Path::new("/a/../../..")),
PathBuf::from("/")
);
}
#[test]
fn test_validate_path_rejects_traversal_nonexistent_parent() {
// The critical test: writing to ../../outside/newdir/file with base_dir
// set should be rejected even when the parent directory does not exist
// (i.e. canonicalize() cannot resolve it).
let dir = TempDir::new().unwrap();
let evil_path = format!(
"{}/../../outside/newdir/file.txt",
dir.path().to_str().unwrap()
);
let result = validate_path(&evil_path, Some(dir.path()));
assert!(
result.is_err(),
"Should reject traversal via non-existent parent, got: {:?}",
result
);
}
#[test]
fn test_validate_path_rejects_relative_traversal() {
let dir = TempDir::new().unwrap();
let result = validate_path("../../etc/passwd", Some(dir.path()));
assert!(
result.is_err(),
"Should reject relative traversal, got: {:?}",
result
);
}
#[test]
fn test_validate_path_allows_valid_nested_write() {
let dir = TempDir::new().unwrap();
let result = validate_path("subdir/newfile.txt", Some(dir.path()));
assert!(
result.is_ok(),
"Should allow nested writes within sandbox: {:?}",
result
);
}
#[test]
fn test_validate_path_allows_dot_dot_within_sandbox() {
// a/b/../c resolves to a/c which is still inside the sandbox
let dir = TempDir::new().unwrap();
std::fs::create_dir_all(dir.path().join("a/b")).unwrap();
let result = validate_path("a/b/../c.txt", Some(dir.path()));
assert!(
result.is_ok(),
"Should allow .. that stays within sandbox: {:?}",
result
);
}
}
+85 -9
View File
@@ -1,7 +1,7 @@
//! HTTP request tool.
use std::collections::HashMap;
use std::net::IpAddr;
use std::net::{IpAddr, ToSocketAddrs};
use std::time::Duration;
use async_trait::async_trait;
@@ -11,6 +11,9 @@ use crate::context::JobContext;
use crate::safety::LeakDetector;
use crate::tools::tool::{Tool, ToolError, ToolOutput};
/// Maximum response body size (5 MB). Prevents OOM from unbounded responses.
const MAX_RESPONSE_SIZE: usize = 5 * 1024 * 1024;
/// Tool for making HTTP requests.
pub struct HttpTool {
client: Client,
@@ -21,6 +24,7 @@ impl HttpTool {
pub fn new() -> Self {
let client = Client::builder()
.timeout(Duration::from_secs(30))
.redirect(reqwest::redirect::Policy::none())
.build()
.expect("Failed to create HTTP client");
@@ -49,11 +53,28 @@ fn validate_url(url: &str) -> Result<reqwest::Url, ToolError> {
));
}
if let Ok(ip) = host.parse::<IpAddr>() {
if is_disallowed_ip(&ip) {
return Err(ToolError::NotAuthorized(
"private or local IPs are not allowed".to_string(),
));
// Check literal IP addresses
if let Ok(ip) = host.parse::<IpAddr>()
&& is_disallowed_ip(&ip)
{
return Err(ToolError::NotAuthorized(
"private or local IPs are not allowed".to_string(),
));
}
// Resolve hostname and check all resolved IPs against the blocklist.
// This prevents DNS rebinding where a hostname resolves to a private IP.
let port = parsed.port_or_known_default().unwrap_or(443);
let socket_addr = format!("{}:{}", host, port);
if let Ok(addrs) = socket_addr.to_socket_addrs() {
for addr in addrs {
if is_disallowed_ip(&addr.ip()) {
return Err(ToolError::NotAuthorized(format!(
"hostname '{}' resolves to disallowed IP {}",
host,
addr.ip()
)));
}
}
}
@@ -202,17 +223,36 @@ impl Tool for HttpTool {
})?;
let status = response.status().as_u16();
// Block redirects: the server tried to send us elsewhere (potential SSRF)
if (300..400).contains(&status) {
return Err(ToolError::NotAuthorized(format!(
"request returned redirect (HTTP {}), which is blocked to prevent SSRF",
status
)));
}
let headers: HashMap<String, String> = response
.headers()
.iter()
.filter_map(|(k, v)| v.to_str().ok().map(|v| (k.to_string(), v.to_string())))
.collect();
// Get response body
let body_text = response.text().await.map_err(|e| {
// Get response body with size cap to prevent OOM
let body_bytes = response.bytes().await.map_err(|e| {
ToolError::ExternalService(format!("failed to read response body: {}", e))
})?;
if body_bytes.len() > MAX_RESPONSE_SIZE {
return Err(ToolError::ExecutionFailed(format!(
"Response body too large ({} bytes, max {})",
body_bytes.len(),
MAX_RESPONSE_SIZE
)));
}
let body_text = String::from_utf8_lossy(&body_bytes).into_owned();
// Try to parse as JSON, fall back to string
let body: serde_json::Value = serde_json::from_str(&body_text)
.unwrap_or_else(|_| serde_json::Value::String(body_text.clone()));
@@ -241,7 +281,7 @@ impl Tool for HttpTool {
#[cfg(test)]
mod tests {
use super::validate_url;
use super::*;
#[test]
fn test_validate_url_rejects_http() {
@@ -260,4 +300,40 @@ mod tests {
let url = validate_url("https://example.com").unwrap();
assert_eq!(url.host_str(), Some("example.com"));
}
#[test]
fn test_validate_url_rejects_private_ip_literal() {
let err = validate_url("https://192.168.1.1/api").unwrap_err();
assert!(err.to_string().contains("private"));
}
#[test]
fn test_validate_url_rejects_loopback_ip() {
let err = validate_url("https://127.0.0.1/api").unwrap_err();
assert!(err.to_string().contains("private"));
}
#[test]
fn test_validate_url_rejects_link_local() {
let err = validate_url("https://169.254.169.254/latest/meta-data/").unwrap_err();
assert!(err.to_string().contains("private"));
}
#[test]
fn test_is_disallowed_ip_covers_ranges() {
use std::net::Ipv4Addr;
// Private ranges
assert!(is_disallowed_ip(&IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1))));
assert!(is_disallowed_ip(&IpAddr::V4(Ipv4Addr::new(172, 16, 0, 1))));
assert!(is_disallowed_ip(&IpAddr::V4(Ipv4Addr::new(192, 168, 0, 1))));
// Loopback
assert!(is_disallowed_ip(&IpAddr::V4(Ipv4Addr::LOCALHOST)));
// Cloud metadata
assert!(is_disallowed_ip(&IpAddr::V4(Ipv4Addr::new(
169, 254, 169, 254
))));
// Public
assert!(!is_disallowed_ip(&IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8))));
}
}
+16 -15
View File
@@ -15,7 +15,8 @@ use chrono::Utc;
use uuid::Uuid;
use crate::context::{ContextManager, JobContext, JobState};
use crate::history::{SandboxJobRecord, Store};
use crate::db::Database;
use crate::history::SandboxJobRecord;
use crate::orchestrator::job_manager::{ContainerJobManager, JobMode};
use crate::tools::tool::{Tool, ToolError, ToolOutput};
@@ -27,7 +28,7 @@ use crate::tools::tool::{Tool, ToolError, ToolOutput};
pub struct CreateJobTool {
context_manager: Arc<ContextManager>,
job_manager: Option<Arc<ContainerJobManager>>,
store: Option<Arc<Store>>,
store: Option<Arc<dyn Database>>,
}
impl CreateJobTool {
@@ -43,7 +44,7 @@ impl CreateJobTool {
pub fn with_sandbox(
mut self,
job_manager: Arc<ContainerJobManager>,
store: Option<Arc<Store>>,
store: Option<Arc<dyn Database>>,
) -> Self {
self.job_manager = Some(job_manager);
self.store = store;
@@ -157,18 +158,18 @@ impl CreateJobTool {
});
// 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);
}
});
}
if mode == JobMode::ClaudeCode
&& 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.
+31 -1
View File
@@ -20,6 +20,12 @@ use crate::context::JobContext;
use crate::tools::tool::{Tool, ToolError, ToolOutput};
use crate::workspace::{Workspace, paths};
/// Identity files that the LLM must not overwrite via tool calls.
/// These are loaded into the system prompt and could be used for prompt
/// injection if an attacker tricks the agent into overwriting them.
const PROTECTED_IDENTITY_FILES: &[&str] =
&[paths::IDENTITY, paths::SOUL, paths::AGENTS, paths::USER];
/// Tool for searching workspace memory.
///
/// Performs hybrid search (FTS + semantic) across all memory documents.
@@ -188,6 +194,16 @@ impl Tool for MemoryWriteTool {
.and_then(|v| v.as_str())
.unwrap_or("daily_log");
// Reject writes to identity files that are loaded into the system prompt.
// An attacker could use prompt injection to trick the agent into overwriting
// these, poisoning future conversations.
if PROTECTED_IDENTITY_FILES.contains(&target) {
return Err(ToolError::NotAuthorized(format!(
"writing to '{}' is not allowed (identity file protected from tool writes)",
target,
)));
}
let append = params
.get("append")
.and_then(|v| v.as_bool())
@@ -230,6 +246,20 @@ impl Tool for MemoryWriteTool {
paths::HEARTBEAT.to_string()
}
path => {
// Protect identity files from LLM overwrites (prompt injection defense).
// These files are injected into the system prompt, so poisoning them
// would let an attacker rewrite the agent's core instructions.
let normalized = path.trim_start_matches('/');
if PROTECTED_IDENTITY_FILES
.iter()
.any(|p| normalized.eq_ignore_ascii_case(p))
{
return Err(ToolError::NotAuthorized(format!(
"writing to '{}' is not allowed (identity file protected from tool access)",
path
)));
}
if append {
self.workspace
.append(path, content)
@@ -452,7 +482,7 @@ impl Tool for MemoryTreeTool {
}
}
#[cfg(test)]
#[cfg(all(test, feature = "postgres"))]
mod tests {
use super::*;
+1 -1
View File
@@ -11,7 +11,7 @@ mod marketplace;
mod memory;
mod restaurant;
pub mod routine;
mod shell;
pub(crate) mod shell;
mod taskrabbit;
mod time;
+11 -11
View File
@@ -19,18 +19,18 @@ use crate::agent::routine::{
};
use crate::agent::routine_engine::RoutineEngine;
use crate::context::JobContext;
use crate::history::Store;
use crate::db::Database;
use crate::tools::tool::{Tool, ToolError, ToolOutput};
// ==================== routine_create ====================
pub struct RoutineCreateTool {
store: Arc<Store>,
store: Arc<dyn Database>,
engine: Arc<RoutineEngine>,
}
impl RoutineCreateTool {
pub fn new(store: Arc<Store>, engine: Arc<RoutineEngine>) -> Self {
pub fn new(store: Arc<dyn Database>, engine: Arc<RoutineEngine>) -> Self {
Self { store, engine }
}
}
@@ -277,11 +277,11 @@ impl Tool for RoutineCreateTool {
// ==================== routine_list ====================
pub struct RoutineListTool {
store: Arc<Store>,
store: Arc<dyn Database>,
}
impl RoutineListTool {
pub fn new(store: Arc<Store>) -> Self {
pub fn new(store: Arc<dyn Database>) -> Self {
Self { store }
}
}
@@ -351,12 +351,12 @@ impl Tool for RoutineListTool {
// ==================== routine_update ====================
pub struct RoutineUpdateTool {
store: Arc<Store>,
store: Arc<dyn Database>,
engine: Arc<RoutineEngine>,
}
impl RoutineUpdateTool {
pub fn new(store: Arc<Store>, engine: Arc<RoutineEngine>) -> Self {
pub fn new(store: Arc<dyn Database>, engine: Arc<RoutineEngine>) -> Self {
Self { store, engine }
}
}
@@ -474,12 +474,12 @@ impl Tool for RoutineUpdateTool {
// ==================== routine_delete ====================
pub struct RoutineDeleteTool {
store: Arc<Store>,
store: Arc<dyn Database>,
engine: Arc<RoutineEngine>,
}
impl RoutineDeleteTool {
pub fn new(store: Arc<Store>, engine: Arc<RoutineEngine>) -> Self {
pub fn new(store: Arc<dyn Database>, engine: Arc<RoutineEngine>) -> Self {
Self { store, engine }
}
}
@@ -551,11 +551,11 @@ impl Tool for RoutineDeleteTool {
// ==================== routine_history ====================
pub struct RoutineHistoryTool {
store: Arc<Store>,
store: Arc<dyn Database>,
}
impl RoutineHistoryTool {
pub fn new(store: Arc<Store>) -> Self {
pub fn new(store: Arc<dyn Database>) -> Self {
Self { store }
}
}
+129 -18
View File
@@ -74,6 +74,58 @@ static DANGEROUS_PATTERNS: LazyLock<Vec<&'static str>> = LazyLock::new(|| {
]
});
/// Patterns that should NEVER be auto-approved, even if the user chose "always approve"
/// for the shell tool. These require explicit per-invocation approval because they are
/// destructive or security-sensitive.
static NEVER_AUTO_APPROVE_PATTERNS: LazyLock<Vec<&'static str>> = LazyLock::new(|| {
vec![
"rm -rf",
"rm -fr",
"chmod -r 777",
"chmod 777",
"chown -r",
"shutdown",
"reboot",
"poweroff",
"init 0",
"init 6",
"iptables",
"nft ",
"useradd",
"userdel",
"passwd",
"visudo",
"crontab",
"systemctl disable",
"launchctl unload",
"kill -9",
"killall",
"pkill",
"docker rm",
"docker rmi",
"docker system prune",
"git push --force",
"git push -f",
"git reset --hard",
"git clean -f",
"DROP TABLE",
"DROP DATABASE",
"TRUNCATE",
"DELETE FROM",
]
});
/// Check whether a shell command contains patterns that must never be auto-approved.
///
/// Even when the user has chosen "always approve" for the shell tool, these commands
/// require explicit per-invocation approval because they are destructive.
pub fn requires_explicit_approval(command: &str) -> bool {
let lower = command.to_lowercase();
NEVER_AUTO_APPROVE_PATTERNS
.iter()
.any(|p| lower.contains(&p.to_lowercase()))
}
/// Shell command execution tool.
pub struct ShellTool {
/// Working directory for commands (if None, uses job's working dir or cwd).
@@ -289,23 +341,17 @@ impl ShellTool {
// Determine timeout
let timeout_duration = timeout.map(Duration::from_secs).unwrap_or(self.timeout);
// Try sandbox execution if available
if let Some(ref sandbox) = self.sandbox {
if sandbox.is_initialized() || sandbox.config().enabled {
match self
.execute_sandboxed(sandbox, cmd, &cwd, timeout_duration)
.await
{
Ok((output, code)) => return Ok((output, code)),
Err(e) => {
// Log sandbox failure and fall through to direct execution
tracing::warn!("Sandbox execution failed, falling back to direct: {}", e);
}
}
}
// Use sandbox if configured; fail-closed (never silently fall through
// to unsandboxed execution when sandbox was intended).
if let Some(ref sandbox) = self.sandbox
&& (sandbox.is_initialized() || sandbox.config().enabled)
{
return self
.execute_sandboxed(sandbox, cmd, &cwd, timeout_duration)
.await;
}
// Fallback to direct execution
// Only execute directly when no sandbox was configured at all.
let (output, code) = self.execute_direct(cmd, &cwd, timeout_duration).await?;
Ok((output, code as i64))
}
@@ -392,17 +438,19 @@ impl Tool for ShellTool {
}
}
/// Truncate output to fit within limits.
/// Truncate output to fit within limits (UTF-8 safe).
fn truncate_output(s: &str) -> String {
if s.len() <= MAX_OUTPUT_SIZE {
s.to_string()
} else {
let half = MAX_OUTPUT_SIZE / 2;
let head_end = crate::util::floor_char_boundary(s, half);
let tail_start = crate::util::floor_char_boundary(s, s.len() - half);
format!(
"{}\n\n... [truncated {} bytes] ...\n\n{}",
&s[..half],
&s[..head_end],
s.len() - MAX_OUTPUT_SIZE,
&s[s.len() - half..]
&s[tail_start..]
)
}
}
@@ -458,6 +506,69 @@ mod tests {
assert!(matches!(result, Err(ToolError::Timeout(_))));
}
#[test]
fn test_requires_explicit_approval() {
// Destructive commands should require explicit approval
assert!(requires_explicit_approval("rm -rf /tmp/stuff"));
assert!(requires_explicit_approval("git push --force origin main"));
assert!(requires_explicit_approval("git reset --hard HEAD~5"));
assert!(requires_explicit_approval("docker rm container_name"));
assert!(requires_explicit_approval("kill -9 12345"));
assert!(requires_explicit_approval("DROP TABLE users;"));
// Safe commands should not
assert!(!requires_explicit_approval("cargo build"));
assert!(!requires_explicit_approval("git status"));
assert!(!requires_explicit_approval("ls -la"));
assert!(!requires_explicit_approval("echo hello"));
assert!(!requires_explicit_approval("cat file.txt"));
assert!(!requires_explicit_approval(
"git push origin feature-branch"
));
}
/// Replicate the extraction logic from agent_loop.rs to prove it works
/// when `arguments` is a `serde_json::Value::Object` (the common case
/// that was previously broken because `Value::Object.as_str()` returns None).
#[test]
fn test_destructive_command_extraction_from_object_args() {
let arguments = serde_json::json!({"command": "rm -rf /tmp/stuff"});
let cmd = arguments
.get("command")
.and_then(|c| c.as_str().map(String::from))
.or_else(|| {
arguments
.as_str()
.and_then(|s| serde_json::from_str::<serde_json::Value>(s).ok())
.and_then(|v| v.get("command").and_then(|c| c.as_str().map(String::from)))
});
assert_eq!(cmd.as_deref(), Some("rm -rf /tmp/stuff"));
assert!(requires_explicit_approval(cmd.as_deref().unwrap()));
}
/// Verify extraction still works when `arguments` is a JSON string
/// (rare, but possible if the LLM provider returns string-encoded JSON).
#[test]
fn test_destructive_command_extraction_from_string_args() {
let arguments =
serde_json::Value::String(r#"{"command": "git push --force origin main"}"#.to_string());
let cmd = arguments
.get("command")
.and_then(|c| c.as_str().map(String::from))
.or_else(|| {
arguments
.as_str()
.and_then(|s| serde_json::from_str::<serde_json::Value>(s).ok())
.and_then(|v| v.get("command").and_then(|c| c.as_str().map(String::from)))
});
assert_eq!(cmd.as_deref(), Some("git push --force origin main"));
assert!(requires_explicit_approval(cmd.as_deref().unwrap()));
}
#[test]
fn test_sandbox_policy_builder() {
let tool = ShellTool::new()
+15 -72
View File
@@ -11,9 +11,9 @@ use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
use rand::RngCore;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::TcpListener;
use crate::cli::oauth_defaults::{self, OAUTH_CALLBACK_PORT};
use crate::secrets::{CreateSecretParams, SecretsStore};
use crate::tools::mcp::config::McpServerConfig;
@@ -466,14 +466,12 @@ pub async fn authorize_mcp_server(
Ok(token)
}
/// Find an available port for the OAuth callback.
/// Bind the OAuth callback listener on the shared fixed port.
pub async fn find_available_port() -> Result<(TcpListener, u16), AuthError> {
for port in 9876..=9886 {
if let Ok(listener) = TcpListener::bind(format!("127.0.0.1:{}", port)).await {
return Ok((listener, port));
}
}
Err(AuthError::PortUnavailable)
let listener = oauth_defaults::bind_callback_listener()
.await
.map_err(|_| AuthError::PortUnavailable)?;
Ok((listener, OAUTH_CALLBACK_PORT))
}
/// Build the authorization URL with all required parameters.
@@ -522,71 +520,16 @@ pub async fn wait_for_authorization_callback(
listener: TcpListener,
server_name: &str,
) -> Result<String, AuthError> {
let timeout = Duration::from_secs(300);
tokio::time::timeout(timeout, async {
loop {
let (mut socket, _) = listener
.accept()
.await
.map_err(|e| AuthError::Http(e.to_string()))?;
let mut reader = BufReader::new(&mut socket);
let mut request_line = String::new();
reader
.read_line(&mut request_line)
.await
.map_err(|e| AuthError::Http(e.to_string()))?;
// Parse GET /callback?code=xxx HTTP/1.1
if let Some(path) = request_line.split_whitespace().nth(1) {
if path.starts_with("/callback") {
if let Some(query) = path.split('?').nth(1) {
// Check for error first
if query.contains("error=") {
let response = "HTTP/1.1 400 Bad Request\r\n\r\nAuthorization denied";
let _ = socket.write_all(response.as_bytes()).await;
return Err(AuthError::AuthorizationDenied);
}
// Look for code
for param in query.split('&') {
let parts: Vec<&str> = param.splitn(2, '=').collect();
if parts.len() == 2 && parts[0] == "code" {
let code = urlencoding::decode(parts[1])
.unwrap_or_else(|_| parts[1].into())
.into_owned();
// Send success response
let response = format!(
"HTTP/1.1 200 OK\r\n\
Content-Type: text/html\r\n\
\r\n\
<!DOCTYPE html><html><body style=\"font-family: sans-serif; \
display: flex; justify-content: center; align-items: center; \
height: 100vh; margin: 0; background: #191919; color: white;\">\
<div style=\"text-align: center;\">\
<h1> {} Connected!</h1>\
<p>You can close this window.</p>\
</div></body></html>",
server_name
);
let _ = socket.write_all(response.as_bytes()).await;
let _ = socket.shutdown().await;
return Ok(code);
}
}
}
}
oauth_defaults::wait_for_callback(listener, "/callback", "code", server_name)
.await
.map_err(|e| match e {
oauth_defaults::OAuthCallbackError::Denied => AuthError::AuthorizationDenied,
oauth_defaults::OAuthCallbackError::Timeout => AuthError::Timeout,
oauth_defaults::OAuthCallbackError::PortInUse(_, msg) => {
AuthError::Http(format!("Port error: {}", msg))
}
let response = "HTTP/1.1 404 Not Found\r\n\r\n";
let _ = socket.write_all(response.as_bytes()).await;
}
})
.await
.map_err(|_| AuthError::Timeout)?
oauth_defaults::OAuthCallbackError::Io(msg) => AuthError::Http(msg),
})
}
/// Exchange the authorization code for an access token.
+44 -43
View File
@@ -184,44 +184,46 @@ impl McpClient {
}
// Add Mcp-Session-Id header if we have a session
if let Some(ref session_manager) = self.session_manager {
if let Some(session_id) = session_manager.get_session_id(&self.server_name).await {
req_builder = req_builder.header("Mcp-Session-Id", session_id);
}
if let Some(ref session_manager) = self.session_manager
&& let Some(session_id) = session_manager.get_session_id(&self.server_name).await
{
req_builder = req_builder.header("Mcp-Session-Id", session_id);
}
let response = req_builder
.send()
.await
.map_err(|e| ToolError::ExternalService(format!("MCP request failed: {}", e)))?;
let response = req_builder.send().await.map_err(|e| {
let mut chain = format!("MCP request failed: {}", e);
let mut source = std::error::Error::source(&e);
while let Some(cause) = source {
chain.push_str(&format!(" -> {}", cause));
source = cause.source();
}
ToolError::ExternalService(chain)
})?;
// Check for 401 Unauthorized - try to refresh token on first attempt
if response.status() == reqwest::StatusCode::UNAUTHORIZED {
if attempt == 0 {
// Try to refresh the token
if let Some(ref secrets) = self.secrets {
if let Some(ref config) = self.server_config {
tracing::debug!(
"MCP token expired, attempting refresh for '{}'",
self.server_name
);
match refresh_access_token(config, secrets, &self.user_id).await {
Ok(_) => {
tracing::info!(
"MCP token refreshed for '{}'",
self.server_name
);
// Continue to next iteration to retry with new token
continue;
}
Err(e) => {
tracing::debug!(
"Token refresh failed for '{}': {}",
self.server_name,
e
);
// Fall through to return auth error
}
if let Some(ref secrets) = self.secrets
&& let Some(ref config) = self.server_config
{
tracing::debug!(
"MCP token expired, attempting refresh for '{}'",
self.server_name
);
match refresh_access_token(config, secrets, &self.user_id).await {
Ok(_) => {
tracing::info!("MCP token refreshed for '{}'", self.server_name);
// Continue to next iteration to retry with new token
continue;
}
Err(e) => {
tracing::debug!(
"Token refresh failed for '{}': {}",
self.server_name,
e
);
// Fall through to return auth error
}
}
}
@@ -245,16 +247,15 @@ impl McpClient {
/// Parse the HTTP response into an MCP response.
async fn parse_response(&self, response: reqwest::Response) -> Result<McpResponse, ToolError> {
// Extract session ID from response header
if let Some(ref session_manager) = self.session_manager {
if let Some(session_id) = response
if let Some(ref session_manager) = self.session_manager
&& let Some(session_id) = response
.headers()
.get("Mcp-Session-Id")
.and_then(|v| v.to_str().ok())
{
session_manager
.update_session_id(&self.server_name, Some(session_id.to_string()))
.await;
}
{
session_manager
.update_session_id(&self.server_name, Some(session_id.to_string()))
.await;
}
if !response.status().is_success() {
@@ -316,11 +317,11 @@ impl McpClient {
/// This should be called once per session to establish capabilities.
pub async fn initialize(&self) -> Result<InitializeResult, ToolError> {
// Check if already initialized
if let Some(ref session_manager) = self.session_manager {
if session_manager.is_initialized(&self.server_name).await {
// Return cached/default capabilities
return Ok(InitializeResult::default());
}
if let Some(ref session_manager) = self.session_manager
&& session_manager.is_initialized(&self.server_name).await
{
// Return cached/default capabilities
return Ok(InitializeResult::default());
}
// Ensure we have a session
+86 -5
View File
@@ -88,8 +88,18 @@ impl McpServerConfig {
}
/// Check if this server requires authentication.
///
/// Returns true if OAuth is pre-configured OR if this is a remote HTTPS server
/// (which likely supports Dynamic Client Registration even without pre-configured OAuth).
pub fn requires_auth(&self) -> bool {
self.oauth.is_some()
if self.oauth.is_some() {
return true;
}
// Remote HTTPS servers need auth handling (DCR, token refresh, 401 detection).
// Localhost/127.0.0.1 servers are assumed to be dev servers without auth.
let url_lower = self.url.to_lowercase();
let is_localhost = is_localhost_url(&url_lower);
url_lower.starts_with("https://") && !is_localhost
}
/// Get the secret name used to store the access token.
@@ -333,7 +343,7 @@ pub async fn get_mcp_server(name: &str) -> Result<McpServerConfig, ConfigError>
///
/// Falls back to the disk file if DB has no entry.
pub async fn load_mcp_servers_from_db(
store: &crate::history::Store,
store: &dyn crate::db::Database,
user_id: &str,
) -> Result<McpServersFile, ConfigError> {
match store.get_setting(user_id, "mcp_servers").await {
@@ -357,7 +367,7 @@ pub async fn load_mcp_servers_from_db(
/// Save MCP server configurations to the database settings table.
pub async fn save_mcp_servers_to_db(
store: &crate::history::Store,
store: &dyn crate::db::Database,
user_id: &str,
config: &McpServersFile,
) -> Result<(), ConfigError> {
@@ -371,7 +381,7 @@ pub async fn save_mcp_servers_to_db(
/// Add a new MCP server configuration (DB-backed).
pub async fn add_mcp_server_db(
store: &crate::history::Store,
store: &dyn crate::db::Database,
user_id: &str,
config: McpServerConfig,
) -> Result<(), ConfigError> {
@@ -386,7 +396,7 @@ pub async fn add_mcp_server_db(
/// Remove an MCP server by name (DB-backed).
pub async fn remove_mcp_server_db(
store: &crate::history::Store,
store: &dyn crate::db::Database,
user_id: &str,
name: &str,
) -> Result<(), ConfigError> {
@@ -402,11 +412,43 @@ pub async fn remove_mcp_server_db(
Ok(())
}
/// Check if a URL points to a loopback address (localhost, 127.0.0.1, [::1]).
///
/// Uses `url::Url` for proper parsing so edge cases (IPv6, userinfo, ports)
/// are handled correctly without manual string splitting.
fn is_localhost_url(url: &str) -> bool {
let Ok(parsed) = url::Url::parse(url) else {
return false;
};
match parsed.host() {
Some(url::Host::Domain(d)) => d.eq_ignore_ascii_case("localhost"),
Some(url::Host::Ipv4(ip)) => ip.is_loopback(),
Some(url::Host::Ipv6(ip)) => ip.is_loopback(),
None => false,
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
#[test]
fn test_is_localhost_url() {
assert!(is_localhost_url("http://localhost:3000/path"));
assert!(is_localhost_url("https://localhost/path"));
assert!(is_localhost_url("http://127.0.0.1:8080"));
assert!(is_localhost_url("http://127.0.0.1"));
assert!(!is_localhost_url("https://notlocalhost.com/path"));
assert!(!is_localhost_url("https://example-localhost.io"));
assert!(!is_localhost_url("https://mcp.notion.com"));
assert!(is_localhost_url("http://user:pass@localhost:3000/path"));
// IPv6 loopback
assert!(is_localhost_url("http://[::1]:8080/path"));
assert!(is_localhost_url("http://[::1]/path"));
assert!(!is_localhost_url("http://[::2]:8080/path"));
}
#[test]
fn test_server_config_validation() {
// Valid HTTPS server
@@ -514,4 +556,43 @@ mod tests {
"mcp_notion_refresh_token"
);
}
#[test]
fn test_requires_auth_with_oauth() {
let config = McpServerConfig::new("notion", "https://mcp.notion.com")
.with_oauth(OAuthConfig::new("client-123"));
assert!(config.requires_auth());
}
#[test]
fn test_requires_auth_remote_https_without_oauth() {
// Remote HTTPS servers need auth even without pre-configured OAuth (DCR)
let config = McpServerConfig::new("github-copilot", "https://api.githubcopilot.com/mcp/");
assert!(config.requires_auth());
let config = McpServerConfig::new("notion", "https://mcp.notion.com");
assert!(config.requires_auth());
}
#[test]
fn test_requires_auth_localhost_no_auth() {
// Localhost servers are dev servers, no auth needed
let config = McpServerConfig::new("local", "http://localhost:8080");
assert!(!config.requires_auth());
let config = McpServerConfig::new("local", "http://127.0.0.1:3000/mcp");
assert!(!config.requires_auth());
// Even HTTPS localhost doesn't require auth
let config = McpServerConfig::new("local", "https://localhost:8443");
assert!(!config.requires_auth());
}
#[test]
fn test_requires_auth_http_remote_no_auth() {
// HTTP remote servers won't pass validation, but if they existed
// they wouldn't trigger HTTPS auth detection
let config = McpServerConfig::new("bad", "http://mcp.example.com");
assert!(!config.requires_auth());
}
}

Some files were not shown because too many files have changed in this diff Show More