Compare commits

..
Author SHA1 Message Date
Henry ParkandClaude Sonnet 4.6 dcd8bc3697 fix(ci): secrets can't be used in step if conditions [skip-regression-check]
GitHub Actions step-level `if:` doesn't have access to `secrets` context.
Replace `if: secrets.X != ''` with `continue-on-error: true` and let
the Set token step handle the fallback.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-03-09 16:41:31 -07:00
Henry ParkandClaude Sonnet 4.6 b125832787 fix(ci): remove Sonnet reference from prompt — model is Haiku [skip-regression-check]
The orchestrator model is Haiku but the prompt said "Launch 4 parallel
Sonnet agents". Haiku can't spawn Sonnet subagents. Changed to just
"4 parallel agents" so it matches the configured model.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-03-09 16:15:28 -07:00
Henry ParkandClaude Sonnet 4.6 294449065b fix(ci): fetch tags for staging-tested + fix issue counter [skip-regression-check]
1. Add fetch-tags: true to checkout — staging-tested is a tag and
   without it every run behaves like a first run (no incremental batching).

2. Only increment ISSUES_CREATED when gh issue create succeeds.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-03-09 16:00:17 -07:00
Henry ParkandClaude Sonnet 4.6 8a3f0c0870 fix(ci): 20-min polling + timeout-minutes on gate job [skip-regression-check]
- Set polling timeout to 20 min (down from 30)
- Add timeout-minutes: 25 on gate job for hard kill on cancel
- Revert 1s sleep loop to sleep $INTERVAL (1s loop doesn't help
  with cancellation — GitHub only cancels between steps)

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-03-09 15:55:32 -07:00
Henry ParkandClaude Sonnet 4.6 6685f2ca9d fix(ci): gate long-running tests by target branch, not label [skip-regression-check]
Replace label-based gating with branch-based: run full suite unless
PR targets staging. Fixes timing bug where gh pr create --label adds
the label after the opened event, causing gated jobs to never run.

- Developer PRs (-> staging): fast tests only
- Promotion PRs (-> main): full suite
- Chained PRs (-> staging-promote/*): full suite
- Push to main: full suite
- workflow_call: fast tests only

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-03-09 15:44:49 -07:00
Henry ParkandClaude Sonnet 4.6 fb8ac7455e fix(ci): handle "No issues found" comment + restore version-check [skip-regression-check]
1. Handle Claude's "No issues found" comment as a clean pass instead
   of treating empty BODY as blocking. Checks for this before looking
   for "Found N issues" comments.

2. Restore version-check to run on all PRs (not just staging-promotion).
   It's 4-8 seconds and catches missing version bumps early.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-03-09 15:33:15 -07:00
Henry ParkandClaude Sonnet 4.6 b5b9934df0 fix(ci): Haiku model, 30-min timeout, HAS_BLOCKING injection fix [skip-regression-check]
- Switch orchestrator to claude-haiku-4-5-20251001 (cheaper per review)
- Increase polling timeout from 20 to 30 minutes for larger diffs
- Move HAS_BLOCKING from ${{ }} interpolation to env var (same
  injection fix applied to FORCE_RUN and SKIP_GATE earlier)

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-03-09 15:23:52 -07:00
Henry ParkandClaude Sonnet 4.6 0e4d50454b fix(ci): remove --delete-branch — it closes chained PRs [skip-regression-check]
Verified in ironclaw-ci-test: when a promotion branch is deleted on
merge, GitHub CLOSES (not retargets) any chained PRs using it as
their base. This silently breaks the promotion chain.

Remove --delete-branch from gh pr merge. Stale promotion branches
will be cleaned up separately.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-03-09 15:04:49 -07:00
Henry ParkandClaude Sonnet 4.6 ab37518008 fix(ci): address review round 2 — injection, cancellable polling, concurrency [skip-regression-check]
1. Script injection: use env vars for workflow_dispatch inputs instead
   of ${{ inputs.* }} interpolation in run blocks.

2. Cancellable polling: sleep in 1s increments so GitHub can cancel
   between iterations instead of blocking on a single long sleep.

3. Remove workflow_dispatch from claude-review — fails without PR
   context. Re-trigger by removing/re-adding staging-promotion label.

4. Add concurrency group to claude-review to prevent duplicate
   reviews consuming API credits.

5. Fix dead code in test.yml roll-up (redundant variable init).

6. Add comment explaining --delete-branch safety (GitHub auto-retargets
   chained PRs, verified in ironclaw-ci-test).

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-03-09 14:58:15 -07:00
Henry ParkandClaude Sonnet 4.6 e55b160a2e fix(ci): add labeled trigger — gh pr create adds label after opened event [skip-regression-check]
gh pr create --label adds the label AFTER the PR is created, so the
opened event fires without the label. The if condition checking for
staging-promotion label fails. Adding labeled trigger type fixes this.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-03-09 14:50:51 -07:00
Henry ParkandClaude Sonnet 4.6 fdfc5fbccc fix(ci): add workflow_dispatch to claude-review for manual re-trigger [skip-regression-check]
Allows manually re-triggering a Claude review without closing and
reopening the PR. Addresses review feedback from @zmanian.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-03-09 14:44:27 -07:00
Henry ParkandClaude Sonnet 4.6 402c71d642 fix(ci): address Copilot review — workflow_call, branch collision, merge failure [skip-regression-check]
1. Add workflow_call trigger to test.yml and e2e.yml so staging-ci
   can invoke them as reusable workflows.

2. Add github.run_id suffix to promotion branch name to prevent
   collision on workflow reruns.

3. Merge step now fails hard instead of swallowing errors. If merge
   fails, the step fails and the PR stays open.

4. update-tag intentionally advances regardless of gate outcome —
   this is by design so the next batch covers only new commits.
   Blocked batches stay as open PRs with review comments.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-03-09 14:41:35 -07:00
Henry ParkandClaude Sonnet 4.6 831484abab fix(ci): poll check-run status + split merge step [skip-regression-check]
1. Replace 15-min sleep with polling: check Claude Code Review
   check-run status every 30s with 20-min timeout. Exits as soon
   as job completes. Process-findings then reads the last comment.

2. Split merge into its own step conditioned on
   steps.evaluate.outputs.passed == 'true'. Prevents accidental
   merge if gate exit logic is refactored.

3. Add checks:read permission for check-runs API.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-03-09 14:38:32 -07:00
Henry ParkandClaude Sonnet 4.6 ecc026843c refactor(ci): gate long-running checks to promotion PRs + push to main [skip-regression-check]
Skip Docker Build, WASM WIT, Windows Build, Telegram tests, and
version-check on developer PRs (branch→staging). These only run on
promotion PRs (staging→main, labeled staging-promotion) and push
to main.

Developer PRs only run unit tests for fast feedback. Full integration
suite runs at promotion time as the quality gate.

Roll-up job updated to accept skipped as passing for gated jobs.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-03-09 14:27:32 -07:00
Henry ParkandClaude Sonnet 4.6 87674420c5 fix(ci): address Copilot review — fail-closed gate + subshell fix [skip-regression-check]
1. Gate fails closed: when no Claude review comment is found (API key
   missing, workflow failure, etc.), gate now blocks instead of
   passing. Use skip_claude_gate to override.

2. Process substitution: use `< <(...)` instead of pipe so the while
   loop runs in the current shell. ISSUES_CREATED counter and
   HAS_BLOCKING flag now propagate correctly. Removes /tmp file hack.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-03-09 14:06:37 -07:00
Henry ParkandClaude Sonnet 4.6 c474972a49 feat(ci): chained promotion PRs with multi-agent Claude review [skip-regression-check]
Staging CI workflow with batched promotion PRs:
- Creates staging-promote/<sha> branches per batch
- Chains PRs onto previous promotion branch (incremental diffs)
- Claude Code reviews only the incremental changes per batch
- Blocked PRs stay open as records of findings
- staging-tested tag advances regardless of gate outcome
- Runs every 60 min on cron + manual dispatch

Multi-agent Claude review (Sonnet orchestrator + Haiku agents):
- 4 parallel Sonnet review agents (security, architecture, bugs, performance)
- Haiku agents for severity/confidence scoring
- [SEVERITY:CONFIDENCE] output format
- Severity/confidence matrix for issue creation and gate blocking:
  CRITICAL: always create issue, block if confidence >=80
  HIGH: create issue if confidence >=50
  MEDIUM/LOW: create issue if confidence >=80

Gate waits 15 min for Claude review then processes comments.
Validated end-to-end in nearai/ironclaw-ci-test.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-03-09 13:55:46 -07:00
fcb152e408 feat(wasm): lazy schema injection on WASM tool errors (#638)
* feat(wasm): lazy schema injection on WASM tool errors

When a WASM tool returns an error (ToolReturnedError), call the module's
description() and schema() WIT exports and append them as a hint in the
error message. This lets the LLM retry with correct parameters without
us including large schemas in every request's tools array.

- Change ToolReturnedError from tuple to struct variant with hint field
- Add build_tool_hint() that calls WASM description()/schema() exports
- Cap description at 500 chars, schema at 3000 chars to limit context
- Hint flows automatically through Display → ToolError → ChatMessage

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: use floor_char_boundary for UTF-8 safe truncation in tool hints

Use existing crate::util::floor_char_boundary() to avoid panicking
when truncation lands mid-multibyte character. Addresses review
feedback on PR #638.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-09 11:27:46 -07:00
e86b372fa6 fix: prevent irreversible context loss when compaction archive write fails (#754)
* fix(compaction): preserve turns when archival write fails

* style: cargo fmt

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

---------

Co-authored-by: Zaki <[email protected]>
Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-03-09 10:13:13 -07:00
Nick PismenkovandGitHub 63f140d391 fix: button styles (#637) 2026-03-09 09:58:33 -07:00
ab0a2e05de fix(mcp): JSON-RPC spec compliance — flexible id, correct notification format (#685)
* fix(mcp): JSON-RPC spec compliance — flexible id, correct notification format

- McpRequest.id is now Option<u64> with skip_serializing_if, so
  notifications omit the id field as required by JSON-RPC 2.0 spec.
  Previously sent id: 0 which violates the spec.

- McpResponse.id uses flexible deserialization that accepts number,
  string, or null — fixes interop with non-standard MCP servers that
  return string ids or missing id fields on error responses.

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

* Fix review feedback: remove serde(default) from McpResponse.id, fix test assertions

- Remove #[serde(default)] from McpResponse.id so notifications (no id field)
  don't incorrectly parse as responses — prevents DoS/spoofing via SSE
- Update test assertions to use Some(value) after id became Option<u64>

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

* fix: update new transport files for Option<u64> id after rebase

Upstream #721 added stdio/unix/transport modules that use McpRequest.id
and McpResponse.id as u64. After our rebase (which changes id to
Option<u64>), these need .unwrap_or(0) for HashMap keys and Some()
wrapping in tests.

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

* test: add regression tests for JSON-RPC spec compliance

Tests for notification serialization without id field,
flexible id deserialization (string, null, non-numeric).

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

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-03-09 08:37:42 -07:00
ReidandGitHub 290d925c7f fix: preserve tool-call history across thread hydration (#568) (#670)
Prevent model re-attempts and data inconsistencies when rebuilding
  conversation context from persisted tool-call records.

  - Remove raw tool parameters from persisted tool_calls JSON to prevent
    unredacted sensitive data from being stored in the database. The LLM
    context rebuild only needs call_id + name + result.
  - Make record_tool_error/record_tool_result mutually exclusive in all
    three execution paths (dispatcher, approval, deferred). Previously
    error cases called both methods, violating the TurnToolCall invariant
    and sending contradictory outcomes to the LLM.
  - Unify call_id format to turn{N}_{i} between live sessions and
    persisted hydration to eliminate ID mismatch in the LLM context.
  - Auto-close </tool_output> XML tags after truncate_preview truncation
    to prevent malformed tool output reaching the LLM.

  [skip-regression-check]
2026-03-09 08:36:07 -07:00
d73e35cfb0 feat: add AWS Bedrock LLM provider via native Converse API (#713)
* feat: add AWS Bedrock LLM provider via native Converse API

* fix: use JSON parsing for tool result error detection instead of brittle substring matching

* refactor: extract duplicated inference config builder into helper function

* fix: address review feedback — safe casts, input validation, and tests

- Safe u32→i32 cast for max_tokens using try_from with clamp
- Remove brittle string-based error detection fallback for tool results
- Validate BEDROCK_CROSS_REGION against allowed values (us/eu/apac/global)
- Validate message list is non-empty before Converse API call
- Log when using default us-east-1 region
- Update llm_backend doc comment to list all backends
- Add tests for build_inference_config and empty message handling

* fix: persist AWS_PROFILE for Bedrock named profile auth

The wizard collected the profile name but only printed a hint to set
it manually. Now it saves to settings and writes AWS_PROFILE to the
bootstrap .env, consistent with how BEDROCK_REGION and other Bedrock
settings are persisted.

* feat: gate AWS Bedrock behind optional `bedrock` feature flag

The AWS SDK dependencies (aws-config, aws-sdk-bedrockruntime,
aws-smithy-types) require cmake and a C compiler to build aws-lc-sys.
Gate them behind an opt-in `bedrock` feature flag so default builds
are unaffected.

Build with: cargo build --features bedrock
All config, settings, and wizard code stays unconditional (no AWS deps)
so users can configure Bedrock even without the feature compiled — they
get a clear error at startup directing them to rebuild.

* fix: address review feedback and adapt Bedrock provider to registry architecture (takeover #345)

- Resolve merge conflicts with main's registry-based provider system
- Add missing cache_creation_input_tokens/cache_read_input_tokens fields
- Add missing content_parts field in test ChatMessage
- Fix string literal type mismatches in wizard env_vars (.to_string())
- Remove non-functional bearer token auth (AWS_BEARER_TOKEN_BEDROCK) from
  wizard and documentation per reviewer feedback from @zmanian and @serrrfirat
- Remove stale BEDROCK_ACCESS_KEY proxy entry from provider table
- Update Bedrock provider to use is_bedrock string check (LlmBackend enum removed)
- Add bedrock_profile fallback from settings in config resolution

[skip-regression-check]

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

* fix: use main's Cargo.lock as base to preserve dependency versions

Regenerating Cargo.lock from scratch caused transitive dependency version
drift that broke the html_to_markdown fixture test in CI.

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

* fix: bedrock config bugs — spurious warning, alias normalization, profile fallback

- Move is_bedrock check before unknown-backend warning to prevent
  spurious "unknown backend" log for bedrock users
- Normalize backend aliases ("aws", "aws_bedrock") to "bedrock" so
  the provider factory matches correctly
- Add settings.bedrock_profile fallback for AWS_PROFILE, consistent
  with region and cross_region resolution

[skip-regression-check]

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

* fix: address Copilot review feedback — bearer token cleanup, stop_sequences, model dedup

- Remove stale bearer token refs from setup README and CHANGELOG
- Remove dead bedrock_api_key secret injection mapping
- Pass stop_sequences through to Bedrock InferenceConfiguration
- Remove "API key" from wizard menu description (bearer token removed)
- Skip duplicate LLM_MODEL write for bedrock backend in wizard
- Fix cargo fmt formatting

[skip-regression-check]

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

* fix: address review feedback — async new(), remove LiteLLM entry, wizard fixes

- Remove dead LiteLLM-based bedrock entry from providers.json (native
  Converse API intercepts before registry lookup)
- Make BedrockProvider::new() async to avoid block_in_place panic in
  current_thread runtimes; propagate async to create_llm_provider,
  build_provider_chain, and init_llm
- Document CMake build prerequisite in docs/LLM_PROVIDERS.md
- Clear bedrock_profile when user selects "default credentials" in wizard
- Fix selected_model clearing to match established pattern (conditional
  on provider switch, not unconditional)
- Add regression tests for bedrock model preservation and profile clearing

Addresses review feedback from @zmanian on PR #713.
Streaming support tracked in #741.

[skip-regression-check]

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

* fix: address remaining review comments — CLAUDE.md backends, wizard UX

- Add `bedrock` to CLAUDE.md inline backend list (#10)
- Skip full setup re-run when keeping existing Bedrock config (#11)
- Clear stale bedrock_profile on empty named-profile input (#12)
- Add regression test for empty profile clearing

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

---------

Co-authored-by: Chris Gorski <[email protected]>
Co-authored-by: cgorski <[email protected]>
Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-03-09 07:10:25 +00:00
30d81fcdee docs: add simplified Chinese (zh-CN) README translation (#488)
Add README.zh-CN.md with full simplified Chinese translation of the
README, and add language switcher links to the original README.

Co-authored-by: smartchoice <[email protected]>
Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-03-09 07:06:14 +00:00
d8dcc34319 fix: CLI commands ignore runtime DATABASE_BACKEND when both features compiled (#740)
* fix: CLI commands ignore runtime DATABASE_BACKEND when both features compiled

`tool auth` and `mcp` CLI subcommands used compile-time `#[cfg]` gates
to select the database backend for secrets storage. When the binary is
compiled with both `postgres` and `libsql` features, the `#[cfg(feature
= "postgres")]` block always wins regardless of the runtime
`DATABASE_BACKEND` setting. This causes `tool auth` and `mcp auth` to
fail with a connection error for users running the libsql backend.

Switch both functions to `match config.database.backend { ... }` with
inner `#[cfg]` guards on each arm, matching the pattern already used in
`main.rs` and `app.rs`.

Also adds a top-level `auth` section to the Telegram channel
capabilities file so `ironclaw tool auth telegram` works for channels
(previously only tools had this section).

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

* fix: extract create_secrets_store factory into src/db, bump telegram version

- Move duplicated DB backend selection logic from cli/tool.rs and
  cli/mcp.rs into a shared db::create_secrets_store() factory, following
  the existing db::connect_from_config() pattern.
- Bump telegram channel version 0.2.0 → 0.2.1 to fix CI Version Bump Check.
- Add regression test for create_secrets_store with libsql backend.

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

* fix: address review feedback — wizard.rs pattern, formatting, version bump

- Convert setup/wizard.rs secrets store creation from tri-branch #[cfg]
  to runtime match on selected_backend (same pattern as CLI fix).
- Fix formatting (assert! line wrapping caught by CI).
- Bump telegram version to 0.2.2 (main already has 0.2.1).
- Merge latest main.

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

* chore: fix regression test doc comment formatting

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

[skip-regression-check]

* fix: address Copilot review — wizard default backend, error chain preservation

- Fix wizard.rs: default selected_backend to "libsql" in libsql-only
  builds so create_libsql_secrets_store is not skipped.
- Preserve error chain: replace .map_err(|e| anyhow!("{}", e)) with ?
  in cli/tool.rs and cli/mcp.rs since DatabaseError implements
  std::error::Error.

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

---------

Co-authored-by: Tiny Tim <[email protected]>
Co-authored-by: Claude Opus 4.6 <[email protected]>
Co-authored-by: firat.sertgoz <[email protected]>
2026-03-09 07:01:22 +00:00
652f30a826 fix(web): prevent fetch error when hostname is an IP address in TEE check (#672)
Currently, teeApiBase() splits the hostname by '.' and incorrectly parses IP addresses like 127.0.0.1 or localhost into invalid URLs (e.g., http://api.0.0.1/), which causes the fetch API to throw a 'Failed to construct Request' TypeError and crashes the web UI.

This fix:
- Skips TEE checks if the hostname is an IP address or localhost.
- Wraps checkTeeStatus() and fetchTeeReport() with try...catch to gracefully handle any unforeseen fetch errors without bubbling up to the global scope.

Co-authored-by: lighterEB <[email protected]>
2026-03-09 03:50:07 +00:00
Protocol ZeroandGitHub 98e9a40762 test(job): cover job tool validation and state transitions (#681)
Add focused coverage for create/list/status/cancel job tools so validation errors, summary formatting, and cancellation behavior stay stable. This locks in the current user-facing responses for running and completed jobs without changing production code.

Made-with: Cursor
2026-03-09 03:49:58 +00:00
553c306c52 feat: full image support across all channels (#725)
* feat: full image support across all channels

End-to-end image handling: upload, generation, analysis, editing, and
rendering across web gateway, HTTP webhook, WASM (Telegram/Slack), and
REPL channels. Builds on the attachment infrastructure from #596 and
draws inspiration from PR #641's image pipeline approach — credit to
that PR's author for the sentinel JSON pattern and base64-in-JSON
upload design.

Key changes:
- Image upload in web UI (file picker, paste, preview strip)
- Image generation tool (FLUX/DALL-E via /v1/images/generations)
- Image edit tool (multipart /v1/images/edits with fallback)
- Image analysis tool (vision model for workspace images)
- Model detection utilities (image_models.rs, vision_models.rs)
- Sentinel JSON detection in dispatcher for generated image rendering
- StatusUpdate::ImageGenerated → SSE/WS/REPL/WASM broadcast
- HTTP webhook attachment support (base64, 5MB/file, 10MB total)
- WASM channel image download (Telegram via file API, Slack via host HTTP)
- Tool registration wiring in app.rs

[skip-regression-check]

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

* fix: address PR #725 review comments (16 issues)

- SecretString for API keys in all image tools (image_gen, image_edit, image_analyze)
- Binary image read via tokio::fs::read instead of DB-backed workspace.read()
- Replace Arc<Workspace> with Option<PathBuf> base_dir (workspace has no filesystem API)
- ApprovalRequirement::UnlessAutoApproved for cost-sensitive image tools
- Scope sentinel detection to image_generate/image_edit tool names only
- Skip ToolResult preview broadcast for image sentinels (avoids multi-MB base64 in SSE)
- Extract shared media_type_from_path() to builtin/mod.rs
- Rename fallback_chat_edit → fallback_generate with tracing::warn
- Increase gateway body limit from 1MB to 10MB for image uploads
- Increase webhook body limit to 15MB (base64 overhead)
- Log warning on invalid base64 in images_to_attachments
- Client-side image size limits (5MB/file, 5 images max) in app.js
- aria-label on attach button for accessibility
- Update body_too_large test for new 10MB limit

[skip-regression-check]

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

* fix: add Slack file size check before download (PR review item #15)

Skip downloading files larger than 20 MB in the Slack WASM channel to
avoid excessive memory use and slow downloads in the WASM runtime.
Logs a warning when a file is skipped. Also bumps channel versions
for Slack and Telegram (prior branch changes).

[skip-regression-check]

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

* style: cargo fmt

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

* fix(security): add path validation and approval requirement to image tools

Add sandbox path validation via validate_path() to both ImageAnalyzeTool
and ImageEditTool to prevent path traversal attacks that could exfiltrate
arbitrary files through external vision/edit APIs. Also fix
ImageAnalyzeTool::requires_approval to return UnlessAutoApproved,
consistent with ImageEditTool and ImageGenerateTool.

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

* fix: post-download size guards and empty data_url sentinel check

- Slack: add post-download size check on actual bytes when metadata
  size_bytes is absent, preventing bypass of the 20MB limit
- Telegram: add 20MB download size limit (matching Slack) enforced
  in download_telegram_file() after receiving response bytes
- Dispatcher: skip broadcasting ImageGenerated SSE event when
  data_url is empty from unwrap_or_default(), log warning instead

Closes correctness issues #3, #4, #5 from PR #725 review.

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

* fix: use mime_guess for media type detection, add alt attrs and media_type validation

- Replace hardcoded media type mapping with mime_guess crate (already in deps)
- Add alt attributes to img elements in web UI for accessibility
- Validate media_type starts with "image/" in images_to_attachments()
- Update bmp test assertion to match mime_guess behavior

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

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
Co-authored-by: Zaki <[email protected]>
2026-03-09 03:41:27 +00:00
7fb2f47999 feat(skills): exclude_keywords veto in skill activation scoring (#688)
* feat(skills): exclude_keywords veto in skill activation scoring

Add exclude_keywords field to ActivationCriteria. If any exclude
keyword is present in the user message, the skill scores 0 regardless
of keyword or pattern matches — prevents cross-skill interference.

Behaviour: exclude_keywords is a hard veto. Even an exact skill name
match gets vetoed if an exclude keyword is also present. This is
intentional; partial exclusion (score reduction) would create
unpredictable interference behaviour.

Example use case: a writing skill with keywords ["write", "draft"]
and exclude_keywords ["route", "redirect"] will not activate on
messages like "don't route this to the writing agent".

Changes:
- ActivationCriteria: new exclude_keywords field (serde default)
- LoadedSkill: new lowercased_exclude_keywords (preprocessed at load)
- selector.rs: early-return 0 in score_skill() on veto match
- registry.rs: populate lowercased_exclude_keywords during loading
- Test helpers updated across mod.rs, selector.rs, attenuation.rs

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

* Fix review feedback: enforce limits on exclude_keywords, extract helper, use any()

- Add exclude_keywords to enforce_limits() with same min-length and cap
  rules as keywords — prevents empty string always-match and unbounded lists
- Extract to_lowercase_vec() helper to deduplicate three identical blocks
- Use idiomatic any() iterator instead of for loop in score_skill veto check

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

* test(skills): add exclude_keywords veto tests

Adds 4 tests for the exclude_keywords veto behavior as requested in review:

1. test_exclude_keyword_vetos_match — skill scores 0 when exclude keyword present
2. test_exclude_keyword_absent_does_not_block — skill activates normally without it
3. test_exclude_keyword_veto_wins_over_positive_match — veto wins even with multiple keyword hits
4. test_exclude_keyword_case_insensitive — veto fires regardless of message case

Also adds make_skill_with_excludes() test helper to avoid repeating the
LoadedSkill construction boilerplate in each test.

Note on substring matching: exclude_keywords uses message_lower.contains(excl)
(substring match), consistent with the existing positive keyword scoring path.
This means "red" would veto "redirect". This is documented behaviour — if
word-boundary semantics are needed, that's a follow-up change.

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

* style: run cargo fmt on selector.rs

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

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-03-08 20:23:41 -07:00
02f85a8ad5 feat(mcp): transport abstraction, stdio/UDS transports, and OAuth fixes (#721)
* feat(mcp): transport abstraction, stdio/UDS transports, and OAuth fixes

Extract McpTransport trait from HTTP-coupled McpClient, enabling pluggable
transport backends. Implements stdio and Unix domain socket transports for
local MCP server integration, fixes OAuth discovery per RFC 9728, and adds
SSRF protection.

Transport abstraction (Step 2):
- McpTransport trait with send(), shutdown(), supports_http_features()
- HttpMcpTransport extracted from McpClient with SSE parsing, session tracking
- Shared JSON-RPC framing helpers (write_jsonrpc_line, spawn_jsonrpc_reader)
- McpClient refactored to hold Arc<dyn McpTransport>

Stdio transport (#652, Step 4):
- StdioMcpTransport spawns child process, communicates via stdin/stdout
- McpProcessManager for lifecycle management with exponential backoff restart
- Background stderr drain task for debug logging

Unix domain socket transport (#134, Step 5):
- UnixMcpTransport connects to existing Unix sockets
- Reuses shared JSON-RPC framing from transport.rs

HTML error body sanitization (#263, Step 1):
- sanitize_error_body() detects HTML, strips control chars, truncates to 500

Custom headers (#639, Step 3):
- headers field on McpServerConfig, merged into every HTTP request
- --header CLI arg for `mcp add`

Config and CLI updates (Step 6):
- McpTransportConfig tagged enum (Http/Stdio/Unix) with serde support
- EffectiveTransport for zero-copy config dispatch
- CLI: --transport, --command, --arg, --env, --socket flags for `mcp add`
- `mcp list` shows transport type

OAuth fixes (#299, Step 8):
- Multi-strategy discovery (401-based, RFC 9728, direct)
- RFC 8707 resource parameter in auth and refresh flows
- SSRF protection with IPv4-mapped IPv6 bypass detection
- Well-known URI construction per RFC 8414

Closes #652, #134, #639, #263, #299

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

* fix(mcp): address audit findings from crate review

- Fix SSRF bypass: make validate_url_safe async with DNS resolution to
  block hostnames that resolve to private/link-local IPs
- Fix UTF-8 truncation: use char-based truncation in sanitize_error_body
  to avoid panicking on multi-byte characters
- Fix SSE parser: process only complete lines to handle chunks split
  across boundaries, add 10MB buffer size limit
- Add debug_assert for transport type mismatch in new_with_config
- Propagate custom headers in new_with_transport constructor
- Deduplicate effective_transport() calls in CLI list command
- Gate test-only accessors with #[cfg(test)] to eliminate dead_code warnings
- Document JSON-RPC notification id:0 limitation in protocol.rs
- Document total backoff wait time (31s) in process.rs
- Add regression test for multi-byte UTF-8 truncation

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

* fix(mcp): address PR review findings from Copilot, Gemini, and zmanian

Moderate/High fixes:
- Plumb custom headers through new_authenticated constructor
- Restrict HTTP to localhost only in validate_url_safe (prevent
  plaintext credential leaks over non-localhost HTTP)
- Add mcp_process_manager.shutdown_all() to app shutdown path to
  prevent orphaning stdio child processes
- Validate discovered authorization_url before opening browser
  (prevent malicious MCP server redirecting to phishing page)

Medium fixes:
- Upgrade debug_assert to assert in new_with_config (fires in release)
- Remove pending map entry on Ok(Err(_)) in stdio/unix send() to avoid
  stale entries and unnecessary 30s waits
- Shut down old transport in try_restart() before spawning replacement
- Redact env var values in mcp list --verbose (may contain secrets)
- Drain pending requests on shutdown to wake waiters immediately
- Add IPv6 link-local, site-local, unique-local, and documentation
  ranges to is_dangerous_ip SSRF protection

Low fixes:
- Truncate logged JSON parse error lines to 200 chars (prevent
  sensitive data in logs)
- Remove misleading shutdown comment in unix_transport
- Use tempfile::tempdir() instead of hardcoded /tmp/ path in test
- Adopt main's improved sanitize_error_body (HTML tag stripping,
  200-char truncation with char_indices)

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

* fix(mcp): gate unix_transport with #[cfg(unix)] for Windows compat

- Add #[cfg(unix)] to unix_transport module declaration
- Add #[cfg(unix)]/#[cfg(not(unix))] branches in app.rs for Unix
  socket MCP server setup
- Remove unused sanitize_error_body import in client.rs tests

[skip-regression-check]

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-03-09 02:47:42 +00:00
FrankandGitHub 9401ab0d58 fix: add timezone conversion support to time tool (#687) 2026-03-08 21:17:07 +00:00
7d1461fc74 fix: standardize libSQL timestamps as RFC 3339 UTC (#683)
* fix: standardize libsql timestamps

* style: fix formatting in libsql/mod.rs

[skip-regression-check]

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

---------

Co-authored-by: Zaki <[email protected]>
Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-03-08 21:16:46 +00:00
605a4ba46e fix(docker): bind postgres to localhost only (#686)
5432:5432 → 127.0.0.1:5432:5432 — the default docker-compose.yml
exposed postgres on all interfaces, making it reachable from the
local network in any docker compose deployment.

Co-authored-by: Claude Opus 4.6 <[email protected]>
Co-authored-by: Zaki Manian <[email protected]>
2026-03-08 20:55:57 +00:00
fe91ba2ab4 fix(repl): skip /quit on EOF when stdin is not a TTY (#724)
When running as a launchd/systemd daemon, stdin is /dev/null.
rustyline reads EOF immediately and the REPL thread was sending
a /quit message, causing the agent to shut down right after
startup — making service mode non-functional on both macOS and Linux.

Fix: check std::io::stdin().is_terminal() before sending /quit on
EOF. In daemon mode (no TTY) the REPL thread exits silently, leaving
other channels (gateway, telegram, …) running as expected.

Fixes #723

Co-authored-by: Claude Sonnet 4.6 <[email protected]>
Co-authored-by: Zaki Manian <[email protected]>
2026-03-08 20:40:56 +00:00
da2569bb77 fix(web): prevent Enter key from sending message during IME composition (#715)
Co-authored-by: Zaki Manian <[email protected]>
2026-03-08 20:40:31 +00:00
732b3ecfeb test(agent): wire TestRig job tools through the scheduler (#716)
Align TestRig with the production agent wiring so create_job exercises the real scheduler path instead of silently falling back to an unscheduled context-only job. Tighten the e2e assertion to lock in the in-progress scheduler behavior for future refactors.

Made-with: Cursor

Co-authored-by: Zaki Manian <[email protected]>
2026-03-08 20:40:03 +00:00
461d7712e8 fix(config): init_secrets no longer overwrites entire config (#726)
* fix(config): init_secrets no longer overwrites entire config

init_secrets() was calling Config::from_db_with_toml() to re-resolve
config after injecting credentials. This rebuilt the entire config from
env/DB/defaults, nuking all other config fields (agent, safety, tools,
etc.) even though only LlmConfig depends on injected credentials.

This caused 5 CI test failures: the test rig's carefully chosen config
values (max_tool_iterations, allow_local_tools, etc.) were silently
overwritten with production defaults after secret injection.

Fix: add Config::re_resolve_llm() that re-resolves only the LLM config
after credential injection, leaving all other config fields untouched.
Also fix TraceLlm::complete() to skip ToolCalls steps when called in
force_text mode (iteration limit).

[skip-regression-check]

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

* fix(test): update test to match TraceLlm::complete() skip-tool-calls behavior [skip-regression-check]

TraceLlm::complete() now skips ToolCalls steps (force_text mode) instead
of erroring. Update the test to verify it skips past a ToolCalls step and
returns the subsequent Text step.

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

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
Co-authored-by: Zaki <[email protected]>
2026-03-08 13:32:42 -07:00
ReidandGitHub 1c5117eded feat: add PID-based gateway lock to prevent multiple instances (#717) 2026-03-08 13:17:46 -07:00
ReidandGitHub 33b02eabb7 fix(cli): status command ignores config.toml and settings.json (#354) (#734) 2026-03-08 13:17:43 -07:00
ReidandGitHub 068ad2d4b7 Fix single-message mode to exit after one turn when background channels are enabled (#719) 2026-03-08 12:54:18 -07:00
56b7218897 fix(setup): preserve model name when re-running onboarding with same provider (#600) (#694)
Each provider setup function unconditionally cleared selected_model,
so re-running the wizard with "Keep current provider? Yes" would lose
the model name, forcing the user to re-select it every time.

Now only clears selected_model when the backend actually changes
(old model may be invalid for the new provider). When keeping the
same provider, the model is preserved and Step 4 shows the
"Keep current model" prompt.

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-03-08 08:32:02 +00:00
200aed16cd feat: configurable LLM request timeout via LLM_REQUEST_TIMEOUT_SECS (#615) (#630)
Add LLM_REQUEST_TIMEOUT_SECS env var (default: 120) to configure the
HTTP request timeout for LLM API calls. Primarily useful for local
models (Ollama, vLLM, LM Studio) that need more time for prompt
evaluation on consumer hardware.

The timeout is applied to the NearAI provider's HTTP client. Other
providers (Anthropic, OpenAI) use rig-core's default client.

- Add request_timeout_secs field to LlmConfig
- Thread timeout through create_llm_provider -> NearAiChatProvider
- Add NearAiChatProvider::new_with_timeout constructor
- Add .env.example documentation
- 2 regression tests for default and custom timeout values

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-03-08 08:30:52 +00:00
4c0275bcdc fix(setup): initialize secrets crypto for env-var security option (#666) (#706)
The "Environment variable" option in the setup wizard's security step
generated a master key but never initialized `secrets_crypto`, causing
subsequent API key saves to fail silently. Fix by:

1. Creating SecretsCrypto from the generated key (matching keychain path)
2. Storing the key hex in settings for write_bootstrap_env to persist
3. Auto-writing SECRETS_MASTER_KEY to ~/.ironclaw/.env
4. Using inject_single_var for thread-safe env overlay
5. Fixing misleading message (shell profiles don't work, only .env)

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-03-08 08:30:02 +00:00
272d31797e chore: remove dead code (#648) (#703)
* chore: remove dead code (LlmEvaluator, chunk_by_paragraphs, bundled channel installer, Reasoning::safety)

Delete unused code flagged in #648:
- evaluation/success.rs: delete LlmEvaluator struct/impl, remove #[allow(dead_code)] from RuleBasedEvaluator methods
- workspace/chunker.rs: delete chunk_by_paragraphs() and its tests (zero production callers)
- extensions/manager.rs: delete install_bundled_channel_from_artifacts() (hot-activation never shipped)
- llm/reasoning.rs: remove unused safety field from Reasoning struct; cascade removal through ContextCompactor, HeartbeatRunner, LlmSoftwareBuilder, and all callers

Closes #648

[skip-regression-check]

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

* fix: move RuleBasedEvaluator into test module to fix dead_code warning

RuleBasedEvaluator has no production callers -- it was only used in
tests of itself. Moving it into #[cfg(test)] eliminates the clippy
dead_code error that broke CI.

[skip-regression-check]

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

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-03-08 08:26:04 +00:00
edff54b0b1 fix: persist /model selection across restarts (#707)
* fix: persist /model selection across restarts

The /model command called set_model() on the LLM provider but never
saved the choice to settings, so the model reverted on restart. Now
persists to both the DB settings store and config.toml.

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

* fix: address CI clippy lint and use spawn_blocking for TOML I/O

- Use struct init syntax instead of field reassignment in test (clippy)
- Wrap sync filesystem operations in spawn_blocking to avoid blocking
  the tokio executor

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

* style: fix rustfmt formatting

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

* fix: address review feedback — handle JoinError, remove exists() guard

- Log warning if spawn_blocking task panics/is cancelled (JoinError)
- Remove toml_path.exists() guard; load_toml already returns Ok(None)
  for missing files, so permission errors are no longer silently skipped

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

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-03-08 08:10:46 +00:00
4d61d3eedf fix(routines): resolve message tool channel/target from per-job metadata (#708)
* fix(routines): resolve message tool channel/target from per-job metadata

When a routine's notify.channel is None, the message tool had no way to
resolve channel/target for full-job workers, causing "No target specified"
errors. The previous approach mutated shared global state via
set_message_tool_context(), which also raced with concurrent jobs.

Now the routine's notify config (channel + user) is carried in the job's
metadata JSON, and MessageTool::execute falls back to ctx.metadata when
neither explicit params nor conversation defaults are available. This
eliminates both the None-channel bug and the concurrent-job race.

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

* style: apply cargo fmt

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

* fix(message): broadcast to all channels when notify.channel is None

Address review feedback:
- Fix stale "see above" comment → "populated below"
- When notify.channel is None, use broadcast_all instead of erroring
  with "No channel specified". This matches NotifyConfig semantics
  where channel=None means "broadcast to all channels"
- Channel resolution is now Option<String>: param → default → metadata → None
- When None, MessageTool uses ChannelManager::broadcast_all(target, response)
  and reports which channels succeeded/failed
- Add regression test for broadcast-all behavior

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

* fix: use failed channels in error message, remove redundant comment

Address review feedback:
- Use `failed` vec in error message instead of re-querying channel_names
- Remove redundant orphaned comment block in routine_engine.rs

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

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-03-08 08:04:19 +00:00
df3635d6be feat(timezone): add timezone-aware session context (#671)
* feat(timezone): add timezone-aware session context (#661)

All timestamps were UTC-only, causing daily logs to split at UTC midnight,
cron schedules to fire in UTC, and no quiet hours for heartbeat. This adds
timezone as a per-session property flowing from the client.

Key changes:
- New `src/timezone.rs` module with resolution chain, parsing, and detection
- `IncomingMessage` carries optional timezone from client
- `JobContext.user_timezone` flows timezone to tools
- `next_cron_fire()` accepts timezone for schedule evaluation
- `Trigger::Cron` stores optional timezone (backward-compatible)
- Workspace gains `_tz` variants for daily logs and system prompt
- Heartbeat supports quiet hours (`HEARTBEAT_QUIET_START/END`)
- Web frontend sends `Intl.DateTimeFormat().resolvedOptions().timeZone`
- REPL auto-detects system timezone
- `DEFAULT_TIMEZONE` env var / settings for server-wide default

Storage stays UTC. Conversion happens at display boundaries.

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

* fix(timezone): address review feedback on timezone-aware sessions

- Validate quiet hours values (0-23) in HeartbeatConfig::resolve()
- Fall back to settings values when env vars are unset for quiet hours
- Validate IANA timezone strings in routine_create/update with parse_timezone
- Add timezone field to routine_create tool schema
- Allow standalone timezone update on cron routines without changing schedule
- Return path from append_daily_log_tz to avoid TOCTOU race at midnight
- Delegate append_daily_log to append_daily_log_tz(entry, UTC) to avoid drift
- Preserve timezone through approval flow via PendingApproval.user_timezone
- Improve test_today_in_tz to not depend on hardcoded year
- Add 3 regression tests for quiet hours config validation

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

* style: fix formatting in routine.rs

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

* fix(timezone): address second round of review feedback

- Remove .claude/scheduled_tasks.lock from repo and add to .gitignore
- Store resolved timezone (not raw message.timezone) in PendingApproval
- Carry forward user_timezone through chained approvals in thread_ops
- Wire quiet_hours_start/end from config to HeartbeatRunner
- Support X-Timezone header as fallback in chat_send_handler

[skip-regression-check]

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

* fix(timezone): include user's local time in time tool response

The time tool's "now" operation now returns local_iso and timezone
fields based on ctx.user_timezone, so the LLM can report time in
the user's timezone instead of always UTC.

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

* style: fix formatting in time.rs

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

* fix(timezone): address Copilot review round 3 — validation, deterministic tests, schema fixes

- Validate DEFAULT_TIMEZONE and HEARTBEAT_TIMEZONE at config load time
- Add timezone field to HeartbeatSettings and config::HeartbeatConfig
- Wire heartbeat timezone from config through agent_loop to HeartbeatRunner
- Add timezone to routine_update tool schema (was accepted but not advertised)
- Error on schedule/timezone update for non-cron routines
- Validate timezone in Trigger::from_db (coerce invalid to None with warning)
- Validate timezone in approval path (thread_ops.rs) before overwriting
- Time tool always includes timezone/local_iso fields (fallback to UTC)
- Make quiet hours tests deterministic using current UTC hour
- Add regression tests for config validation

[skip-regression-check]

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

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-03-08 08:01:56 +00:00
ReidandGitHub a20e19ab16 fix: sanitize HTML error bodies from MCP servers to prevent web UI white screen (#263) (#656)
* fix: sanitize HTML error bodies from MCP servers to prevent web UI white screen (#263)

* style: fix cargo fmt formatting in sanitize_error_body tests
2026-03-08 02:53:02 +00:00
117 changed files with 11788 additions and 1602 deletions
+1
View File
@@ -5,6 +5,7 @@ DATABASE_POOL_SIZE=10
# LLM Provider
# LLM_BACKEND=nearai # default
# Possible values: nearai, ollama, openai_compatible, openai, anthropic, tinfoil
# LLM_REQUEST_TIMEOUT_SECS=120 # Increase for local LLMs (Ollama, vLLM, LM Studio)
# === Anthropic Direct ===
# Two auth modes:
+99
View File
@@ -0,0 +1,99 @@
name: Claude Code Review
on:
pull_request:
types: [opened, labeled]
permissions:
contents: read
pull-requests: write
issues: write
id-token: write
concurrency:
group: claude-review-${{ github.event.pull_request.number || github.run_id }}
cancel-in-progress: true
jobs:
review:
name: Claude Code Review
if: contains(github.event.pull_request.labels.*.name, 'staging-promotion')
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
with:
fetch-depth: 0
- name: Run Claude Code review
uses: anthropics/claude-code-action@v1
with:
anthropic_api_key: ${{ secrets.ANTHROPIC_API_KEY }}
claude_args: "--max-turns 50 --model claude-haiku-4-5-20251001 --allowedTools 'Bash(gh pr comment:*),Bash(gh pr diff:*),Bash(gh pr view:*),Bash(gh pr list:*),Bash(gh issue view:*),Bash(gh issue list:*),Bash(gh search:*),Bash(git blame:*),Bash(git log:*),Bash(git diff:*)'"
prompt: |
Code review this pull request. Follow these steps precisely:
1. Use a Haiku agent to find relevant CLAUDE.md files: the root CLAUDE.md
and any CLAUDE.md files in directories whose files this PR modifies.
2. Use a Haiku agent to summarize the PR change (use `gh pr diff`).
3. Launch 4 parallel agents to review the change independently. Each agent should
read the PR diff with `gh pr diff` and the full source files for changed
code, then return a list of issues found:
Agent 1 — Security & Safety
Check for: command injection, path traversal, SSRF, XSS, auth bypass,
secrets in logs, .unwrap()/.expect() in production code (not tests),
race conditions, TOCTOU, unsafe blocks, panics in async, unbounded allocations.
Agent 2 — Architecture & Patterns
Check for: extensible design (traits/enums over nested conditionals),
clean abstractions, proper error types (thiserror), CLAUDE.md compliance,
type-driven design over stringly-typed code, DRY violations.
Agent 3 — Bug Scan
Shallow diff-only scan for obvious bugs: logic errors, off-by-one,
missing error handling, division by zero, incorrect return values.
Ignore nitpicks and likely false positives. Do NOT read extra context
beyond the diff — focus only on the changes.
Agent 4 — Performance & Production
Check for: blocking in async, N+1 queries, unbounded loops, missing
timeouts, resource leaks (file handles, connections), large allocations
in hot paths.
4. For each issue found, launch a parallel Haiku agent to:
a. Assign a severity:
- CRITICAL: security vulns, panics in prod (.unwrap/.expect), data exfiltration, race conditions
- HIGH: logic bugs, missing error handling, breaking API/schema changes
- MEDIUM: missing tests, unnecessary complexity, performance issues
- LOW: documentation gaps, naming suggestions
b. Score confidence 0-100 (give this rubric verbatim):
0: False positive, doesn't stand up to scrutiny, or pre-existing issue.
25: Might be real, but may be false positive. Stylistic issues not in CLAUDE.md.
50: Real issue but nitpick or rare in practice. Not very important.
75: Verified real issue, will be hit in practice. Directly impacts functionality
or explicitly mentioned in CLAUDE.md.
100: Certain, confirmed, will happen frequently. Evidence directly confirms.
5. Post a single comment on the PR using `gh pr comment` with this format.
If no issues were found, post "No issues found." instead:
### Code review
Found N issues:
1. [SEVERITY:CONFIDENCE] <brief description>
<permalink to file:line using full SHA, eg https://github.com/owner/repo/blob/abc123def/src/file.rs#L10-L15>
Example: [CRITICAL:92] `.unwrap()` can panic in production when config is missing
You MUST use the full git SHA in links (not HEAD or branch name).
Provide 1 line of context before and after each linked range.
Notes:
- Use `gh` for all GitHub interactions, not web fetch
- Do NOT check build signal or attempt to build/test the code
- Ignore pre-existing issues not introduced by this PR
- Ignore issues a linter/compiler would catch (formatting, imports, types)
+1
View File
@@ -1,5 +1,6 @@
name: E2E Tests
on:
workflow_call:
schedule:
- cron: "0 6 * * 1" # Weekly Monday 6 AM UTC
workflow_dispatch:
+473
View File
@@ -0,0 +1,473 @@
name: Staging CI (Batched)
on:
schedule:
- cron: "0 * * * *" # Every 60 minutes
workflow_dispatch:
inputs:
force:
description: "Force run even if no new commits"
type: boolean
default: false
skip_claude_gate:
description: "Skip Claude review gate (bypass blocking findings)"
type: boolean
default: false
permissions:
contents: write
issues: write
pull-requests: write
checks: read
concurrency:
group: staging-ci
cancel-in-progress: false # Let running suites finish
jobs:
# ── Check for new commits ──────────────────────────────────────
check-changes:
name: Check for new commits
runs-on: ubuntu-latest
outputs:
has_changes: ${{ steps.check.outputs.has_changes }}
current_head: ${{ steps.check.outputs.current_head }}
diff_range: ${{ steps.check.outputs.diff_range }}
steps:
- uses: actions/checkout@v6
with:
ref: staging
fetch-depth: 0
fetch-tags: true
- name: Check for changes since last tested
id: check
env:
FORCE_RUN: ${{ inputs.force }}
run: |
CURRENT_HEAD=$(git rev-parse HEAD)
echo "current_head=${CURRENT_HEAD}" >> "$GITHUB_OUTPUT"
if git rev-parse staging-tested >/dev/null 2>&1; then
LAST_TESTED=$(git rev-parse staging-tested)
else
LAST_TESTED=""
fi
DIFF_RANGE=""
if [ -n "$LAST_TESTED" ] && [ "$LAST_TESTED" = "$CURRENT_HEAD" ]; then
echo "No new commits since last tested (${CURRENT_HEAD})"
HAS_CHANGES=false
else
HAS_CHANGES=true
if [ -n "$LAST_TESTED" ]; then
COMMIT_COUNT=$(git rev-list --count "${LAST_TESTED}..HEAD")
echo "Found ${COMMIT_COUNT} new commit(s) since last tested"
DIFF_RANGE="${LAST_TESTED}..${CURRENT_HEAD}"
else
git fetch origin main
MERGE_BASE=$(git merge-base origin/main HEAD)
echo "First run -- reviewing from merge-base ${MERGE_BASE}"
DIFF_RANGE="${MERGE_BASE}..${CURRENT_HEAD}"
fi
fi
# Force override from workflow_dispatch
if [ "$FORCE_RUN" = "true" ]; then
echo "Force run requested"
HAS_CHANGES=true
if [ -z "$DIFF_RANGE" ]; then
DIFF_RANGE="${CURRENT_HEAD}..${CURRENT_HEAD}"
fi
fi
echo "has_changes=${HAS_CHANGES}" >> "$GITHUB_OUTPUT"
echo "diff_range=${DIFF_RANGE}" >> "$GITHUB_OUTPUT"
# ── Run full test suite ──────────────────────────────────────────
tests:
name: Test Suite
needs: check-changes
if: needs.check-changes.outputs.has_changes == 'true'
uses: ./.github/workflows/test.yml
# ── Run E2E browser tests ────────────────────────────────────────
e2e:
name: E2E Browser Tests
needs: check-changes
if: needs.check-changes.outputs.has_changes == 'true'
uses: ./.github/workflows/e2e.yml
# ── Create promotion PR (triggers claude-review.yml on the PR) ──
create-promotion-pr:
name: Create Promotion PR
needs: check-changes
if: needs.check-changes.outputs.has_changes == 'true'
runs-on: ubuntu-latest
outputs:
pr_number: ${{ steps.create-pr.outputs.pr_number }}
promotion_branch: ${{ steps.branch.outputs.branch }}
steps:
- uses: actions/checkout@v6
with:
ref: staging
fetch-depth: 0
- name: Generate GitHub App token
id: app-token
continue-on-error: true
uses: actions/create-github-app-token@v2
with:
app-id: ${{ secrets.GH_RELEASES_MANAGER_APP_ID }}
private-key: ${{ secrets.GH_RELEASES_MANAGER_APP_PRIVATE_KEY }}
- name: Set token
id: token
run: |
if [ -n "${{ steps.app-token.outputs.token }}" ]; then
echo "token=${{ steps.app-token.outputs.token }}" >> "$GITHUB_OUTPUT"
else
echo "token=${{ github.token }}" >> "$GITHUB_OUTPUT"
fi
- name: Check if staging is ahead of main
id: ahead-check
env:
GH_TOKEN: ${{ steps.token.outputs.token }}
run: |
git fetch origin main
AHEAD=$(git rev-list --count origin/main..origin/staging)
echo "commits_ahead=${AHEAD}" >> "$GITHUB_OUTPUT"
if [ "$AHEAD" -eq 0 ]; then
echo "Staging is not ahead of main. Nothing to promote."
else
echo "Staging is ${AHEAD} commits ahead of main."
fi
- name: Create promotion branch
id: branch
if: steps.ahead-check.outputs.commits_ahead != '0'
run: |
SHORT_SHA=$(echo "${{ needs.check-changes.outputs.current_head }}" | cut -c1-8)
BRANCH="staging-promote/${SHORT_SHA}-${{ github.run_id }}"
git checkout -b "$BRANCH"
git push origin "$BRANCH"
echo "branch=${BRANCH}" >> "$GITHUB_OUTPUT"
echo "Created promotion branch: ${BRANCH}"
- name: Find base branch
id: find-base
if: steps.ahead-check.outputs.commits_ahead != '0'
env:
GH_TOKEN: ${{ steps.token.outputs.token }}
run: |
# Find the newest open promotion PR with a staging-promote/* head branch
LATEST=$(gh pr list --label staging-promotion --state open \
--json headRefName,createdAt \
--jq '[.[] | select(.headRefName | startswith("staging-promote/"))] | sort_by(.createdAt) | last | .headRefName // empty')
if [ -n "$LATEST" ]; then
echo "base=${LATEST}" >> "$GITHUB_OUTPUT"
echo "Chaining onto existing promotion branch: ${LATEST}"
else
echo "base=main" >> "$GITHUB_OUTPUT"
echo "No existing promotion PR — targeting main"
fi
- name: Create promotion PR
id: create-pr
if: steps.ahead-check.outputs.commits_ahead != '0'
env:
GH_TOKEN: ${{ steps.token.outputs.token }}
run: |
RANGE="${{ needs.check-changes.outputs.diff_range }}"
TIMESTAMP=$(date -u +"%Y-%m-%d %H:%M UTC")
BRANCH="${{ steps.branch.outputs.branch }}"
BASE="${{ steps.find-base.outputs.base }}"
PR_URL=$(gh pr create \
--base "$BASE" \
--head "$BRANCH" \
--title "chore: promote staging to main (${TIMESTAMP})" \
--body "## Auto-promotion from staging CI
**Batch range:** \`${RANGE}\`
**Promotion branch:** \`${BRANCH}\`
**Base:** \`${BASE}\`
**Triggered by:** Staging CI batch at ${TIMESTAMP}
Waiting for gates:
- Tests: pending
- E2E: pending
- Claude Code review: pending (will post comments on this PR)
---
*Auto-created by staging-ci workflow*" \
--label "staging-promotion")
PR_NUM=$(echo "$PR_URL" | grep -oE '[0-9]+$')
echo "pr_number=${PR_NUM}" >> "$GITHUB_OUTPUT"
echo "Created promotion PR #${PR_NUM}"
# ── Gate: wait for review, process findings, merge or block ─────
gate:
name: Staging Gate
needs: [check-changes, tests, e2e, create-promotion-pr]
if: >
always() &&
needs.check-changes.outputs.has_changes == 'true' &&
needs.tests.result == 'success' &&
needs.e2e.result == 'success' &&
needs.create-promotion-pr.result == 'success'
runs-on: ubuntu-latest
timeout-minutes: 25
outputs:
gate_passed: ${{ steps.evaluate.outputs.passed }}
steps:
- uses: actions/checkout@v6
with:
ref: staging
fetch-depth: 1
- name: Generate GitHub App token
id: app-token
continue-on-error: true
uses: actions/create-github-app-token@v2
with:
app-id: ${{ secrets.GH_RELEASES_MANAGER_APP_ID }}
private-key: ${{ secrets.GH_RELEASES_MANAGER_APP_PRIVATE_KEY }}
- name: Set token
id: token
run: |
if [ -n "${{ steps.app-token.outputs.token }}" ]; then
echo "token=${{ steps.app-token.outputs.token }}" >> "$GITHUB_OUTPUT"
else
echo "token=${{ github.token }}" >> "$GITHUB_OUTPUT"
fi
- name: Wait for Claude review job
env:
GH_TOKEN: ${{ steps.token.outputs.token }}
PR_NUMBER: ${{ needs.create-promotion-pr.outputs.pr_number }}
REPO: ${{ github.repository }}
run: |
if [ -z "$PR_NUMBER" ]; then
echo "No PR number — skipping wait"
exit 0
fi
PR_SHA=$(gh pr view "$PR_NUMBER" --json headRefOid --jq '.headRefOid' || echo "")
if [ -z "$PR_SHA" ]; then
echo "::warning::Could not get PR head SHA"
exit 0
fi
echo "Polling for Claude Code Review job on PR #${PR_NUMBER} (SHA: ${PR_SHA})..."
TIMEOUT=1200 # 20 minutes
ELAPSED=0
INTERVAL=30
while [ "$ELAPSED" -lt "$TIMEOUT" ]; do
STATUS=$(gh api "repos/${REPO}/commits/${PR_SHA}/check-runs" \
--jq '[.check_runs[] | select(.name == "Claude Code Review") | .conclusion // .status] | first // "pending"' 2>/dev/null || echo "pending")
if [ "$STATUS" = "success" ] || [ "$STATUS" = "failure" ] || [ "$STATUS" = "cancelled" ]; then
echo "Claude review job completed with status: ${STATUS} (${ELAPSED}s)"
exit 0
fi
echo "Claude review status: ${STATUS} (${ELAPSED}s elapsed)"
sleep "$INTERVAL"
ELAPSED=$((ELAPSED + INTERVAL))
done
echo "::warning::Claude review job not completed after ${TIMEOUT}s"
- name: Process Claude review comments and create issues
id: process-findings
env:
GH_TOKEN: ${{ steps.token.outputs.token }}
PR_NUMBER: ${{ needs.create-promotion-pr.outputs.pr_number }}
REPO: ${{ github.repository }}
run: |
HAS_BLOCKING=false
ISSUES_CREATED=0
if [ -z "$PR_NUMBER" ]; then
echo "No PR — skipping finding processing"
echo "has_blocking=false" >> "$GITHUB_OUTPUT"
exit 0
fi
# Check for "No issues found" first (clean pass)
NO_ISSUES=$(gh api "repos/${REPO}/issues/${PR_NUMBER}/comments" \
--jq '[.[] | select(.user.login == "claude[bot]") | select(.body | test("No issues found"))] | length' 2>/dev/null || echo "0")
if [ "$NO_ISSUES" -gt 0 ]; then
echo "Claude review found no issues — gate passes"
echo "has_blocking=false" >> "$GITHUB_OUTPUT"
exit 0
fi
# Get the last Claude comment that contains findings
JQ_FILTER='[.[] | select(.user.login == "claude[bot]") | select(.body | test("Found [0-9]+ issue"))] | last'
BODY=$(gh api "repos/${REPO}/issues/${PR_NUMBER}/comments" \
--jq "${JQ_FILTER} | .body // empty" 2>/dev/null || echo "")
COMMENT_URL=$(gh api "repos/${REPO}/issues/${PR_NUMBER}/comments" \
--jq "${JQ_FILTER} | .html_url // empty" 2>/dev/null || echo "")
if [ -z "$BODY" ]; then
echo "::warning::No Claude review comment found for PR #${PR_NUMBER} — treating as blocking"
echo "has_blocking=true" >> "$GITHUB_OUTPUT"
exit 0
fi
# Parse [SEVERITY:CONFIDENCE] tags from each numbered finding
# Matrix: CRITICAL always→issue, ≥80→block. HIGH ≥50→issue. MEDIUM ≥80→issue. LOW ≥80→issue.
# Use process substitution so variables propagate to parent shell
while read -r line; do
TAG=$(echo "$line" | grep -oE '^\[(CRITICAL|HIGH|MEDIUM|LOW):[0-9]+\]')
SEVERITY=$(echo "$TAG" | sed 's/\[\(.*\):\(.*\)\]/\1/')
CONFIDENCE=$(echo "$TAG" | sed 's/\[\(.*\):\(.*\)\]/\2/')
DESC=$(echo "$line" | sed "s/\[${SEVERITY}:${CONFIDENCE}\] *//" | head -1)
echo "Found: [${SEVERITY}:${CONFIDENCE}] ${DESC}"
# Check if blocking (CRITICAL ≥80)
if [ "$SEVERITY" = "CRITICAL" ] && [ "$CONFIDENCE" -ge 80 ]; then
HAS_BLOCKING=true
fi
# Determine if this should create an issue
CREATE_ISSUE=false
case "$SEVERITY" in
CRITICAL) CREATE_ISSUE=true ;;
HIGH) [ "$CONFIDENCE" -ge 50 ] && CREATE_ISSUE=true ;;
MEDIUM) [ "$CONFIDENCE" -ge 80 ] && CREATE_ISSUE=true ;;
LOW) [ "$CONFIDENCE" -ge 80 ] && CREATE_ISSUE=true ;;
esac
if [ "$CREATE_ISSUE" = "true" ]; then
case "$SEVERITY" in
CRITICAL) LABELS="bug,risk: high,staging-ci-review" ;;
HIGH) LABELS="bug,risk: medium,staging-ci-review" ;;
MEDIUM) LABELS="risk: medium,staging-ci-review" ;;
LOW) LABELS="risk: low,staging-ci-review" ;;
esac
TITLE=$(echo "$DESC" | cut -c1-80)
{
echo "## [${SEVERITY}:${CONFIDENCE}] Issue Found by Staging CI Review"
echo ""
echo "**Severity:** ${SEVERITY}"
echo "**Confidence:** ${CONFIDENCE}/100"
echo "**PR comment:** ${COMMENT_URL}"
echo ""
echo "### Description"
echo "$DESC"
echo ""
echo "---"
echo "*Auto-created by staging-ci Claude Code review*"
} > /tmp/issue-body.md
if gh issue create \
--title "[${SEVERITY}] ${TITLE}" \
--body-file /tmp/issue-body.md \
--label "${LABELS}"; then
ISSUES_CREATED=$((ISSUES_CREATED + 1))
else
echo "::warning::Failed to create issue for ${SEVERITY} finding"
fi
fi
done < <(echo "$BODY" | grep -oE '\[(CRITICAL|HIGH|MEDIUM|LOW):[0-9]+\].*')
echo "Created ${ISSUES_CREATED} issues"
echo "has_blocking=${HAS_BLOCKING}" >> "$GITHUB_OUTPUT"
- name: Evaluate gate
id: evaluate
env:
PR_NUMBER: ${{ needs.create-promotion-pr.outputs.pr_number }}
SKIP_GATE: ${{ inputs.skip_claude_gate }}
HAS_BLOCKING: ${{ steps.process-findings.outputs.has_blocking }}
run: |
SKIP_INPUT="$SKIP_GATE"
if [ "$HAS_BLOCKING" = "true" ]; then
echo "::warning::Claude review found blocking issues (CRITICAL ≥80 confidence)"
if [ "$SKIP_INPUT" = "true" ]; then
echo "::warning::Gate overridden by skip_claude_gate workflow input"
echo "passed=true" >> "$GITHUB_OUTPUT"
else
echo "::error::Blocking promotion due to CRITICAL findings (≥80 confidence)"
echo "::error::PR #${PR_NUMBER} left open with review comments"
echo "passed=false" >> "$GITHUB_OUTPUT"
exit 1
fi
else
echo "No blocking findings. Gate passed."
echo "passed=true" >> "$GITHUB_OUTPUT"
fi
- name: Merge promotion PR
id: merge
if: steps.evaluate.outputs.passed == 'true'
env:
GH_TOKEN: ${{ steps.token.outputs.token }}
PR_NUMBER: ${{ needs.create-promotion-pr.outputs.pr_number }}
run: |
if [ -n "$PR_NUMBER" ]; then
echo "Merging promotion PR #${PR_NUMBER}"
# Do NOT use --delete-branch: deleting a promotion branch closes
# any chained PRs that use it as their base (verified in ironclaw-ci-test).
# Stale promotion branches are cleaned up separately.
gh pr merge "$PR_NUMBER" --merge
echo "merged=true" >> "$GITHUB_OUTPUT"
fi
# ── Update tested tag (always, so next batch covers only new commits) ──
update-tag:
name: Update staging-tested tag
needs: [check-changes, tests, e2e, create-promotion-pr, gate]
if: >
always() &&
needs.check-changes.outputs.has_changes == 'true' &&
needs.tests.result == 'success' &&
needs.e2e.result == 'success' &&
needs.create-promotion-pr.result == 'success'
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
with:
ref: staging
fetch-depth: 1
- name: Update staging-tested tag
run: |
git tag -f staging-tested "${{ needs.check-changes.outputs.current_head }}"
git push origin staging-tested --force
echo "Updated staging-tested tag to ${{ needs.check-changes.outputs.current_head }}"
# ── Report ───────────────────────────────────────────────────────
report:
name: Staging CI Summary
needs: [check-changes, tests, e2e, create-promotion-pr, gate, update-tag]
if: always() && needs.check-changes.outputs.has_changes == 'true'
runs-on: ubuntu-latest
steps:
- name: Summary
run: |
echo "## Staging CI Batch Results" >> "$GITHUB_STEP_SUMMARY"
echo "" >> "$GITHUB_STEP_SUMMARY"
echo "| Check | Result |" >> "$GITHUB_STEP_SUMMARY"
echo "|-------|--------|" >> "$GITHUB_STEP_SUMMARY"
echo "| Tests | ${{ needs.tests.result }} |" >> "$GITHUB_STEP_SUMMARY"
echo "| E2E | ${{ needs.e2e.result }} |" >> "$GITHUB_STEP_SUMMARY"
echo "| Promotion PR | ${{ needs.create-promotion-pr.result }} |" >> "$GITHUB_STEP_SUMMARY"
echo "| Gate | ${{ needs.gate.result }} |" >> "$GITHUB_STEP_SUMMARY"
echo "| Tag Updated | ${{ needs.update-tag.result }} |" >> "$GITHUB_STEP_SUMMARY"
echo "" >> "$GITHUB_STEP_SUMMARY"
echo "Range: ${{ needs.check-changes.outputs.diff_range }}" >> "$GITHUB_STEP_SUMMARY"
PR_NUM="${{ needs.create-promotion-pr.outputs.pr_number }}"
if [ -n "$PR_NUM" ]; then
echo "Promotion PR: #${PR_NUM}" >> "$GITHUB_STEP_SUMMARY"
fi
+30 -7
View File
@@ -1,5 +1,6 @@
name: Run Tests
on:
workflow_call:
pull_request:
push:
branches:
@@ -38,6 +39,9 @@ jobs:
telegram-tests:
name: Telegram Channel Tests
if: >
github.event_name == 'push' ||
(github.event_name == 'pull_request' && github.base_ref != 'staging')
runs-on: ubuntu-latest
steps:
- name: Checkout repository
@@ -50,6 +54,9 @@ jobs:
windows-build:
name: Windows Build (${{ matrix.name }})
if: >
github.event_name == 'push' ||
(github.event_name == 'pull_request' && github.base_ref != 'staging')
runs-on: windows-latest
strategy:
fail-fast: false
@@ -74,6 +81,9 @@ jobs:
wasm-wit-compat:
name: WASM WIT Compatibility
if: >
github.event_name == 'push' ||
(github.event_name == 'pull_request' && github.base_ref != 'staging')
runs-on: ubuntu-latest
steps:
- name: Checkout repository
@@ -94,6 +104,9 @@ jobs:
docker-build:
name: Docker Build
if: >
github.event_name == 'push' ||
(github.event_name == 'pull_request' && github.base_ref != 'staging')
runs-on: ubuntu-latest
steps:
- name: Checkout repository
@@ -123,12 +136,22 @@ jobs:
needs: [tests, telegram-tests, wasm-wit-compat, docker-build, windows-build, version-check]
steps:
- run: |
if [[ "${{ needs.tests.result }}" != "success" || "${{ needs.telegram-tests.result }}" != "success" || "${{ needs.wasm-wit-compat.result }}" != "success" || "${{ needs.docker-build.result }}" != "success" || "${{ needs.windows-build.result }}" != "success" ]]; then
echo "One or more jobs failed"
exit 1
fi
# version-check only runs on PRs, so skip/success are both acceptable
if [[ "${{ needs.version-check.result }}" == "failure" ]]; then
echo "Version bump check failed"
# Unit tests must always pass
if [[ "${{ needs.tests.result }}" != "success" ]]; then
echo "Unit tests failed"
exit 1
fi
# Gated jobs: must pass on promotion PRs / push, skipped on developer PRs
for job in telegram-tests wasm-wit-compat docker-build windows-build version-check; do
case "$job" in
telegram-tests) result="${{ needs.telegram-tests.result }}" ;;
wasm-wit-compat) result="${{ needs.wasm-wit-compat.result }}" ;;
docker-build) result="${{ needs.docker-build.result }}" ;;
windows-build) result="${{ needs.windows-build.result }}" ;;
version-check) result="${{ needs.version-check.result }}" ;;
esac
if [[ "$result" == "failure" || "$result" == "cancelled" ]]; then
echo "$job failed"
exit 1
fi
done
+3 -1
View File
@@ -4,8 +4,9 @@
.env.*
!.env.example
# Claude Code worktrees
# Claude Code worktrees and lock files
.claude/worktrees/
.claude/scheduled_tasks.lock
# Sidecar tool data
.sidecar/
@@ -22,4 +23,5 @@ bench-results/
# WASM build artifacts (loaded from disk, not bundled)
*.wasm
# Traces
trace_*.json
+4
View File
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
## [Unreleased]
### Added
- AWS Bedrock LLM provider via native Converse API with IAM and SSO auth support (feature-gated: `--features bedrock`)
## [0.16.1](https://github.com/nearai/ironclaw/compare/v0.16.0...v0.16.1) - 2026-03-06
### Fixed
+10 -1
View File
@@ -484,6 +484,13 @@ SKILLS_AUTO_DISCOVER=true # Scan skill directories on startup
TINFOIL_API_KEY=... # Required when LLM_BACKEND=tinfoil
TINFOIL_MODEL=kimi-k2-5 # Default model
# AWS Bedrock (native Converse API, requires --features bedrock)
# LLM_BACKEND=bedrock
# BEDROCK_REGION=us-east-1 # AWS region
# BEDROCK_MODEL=anthropic.claude-opus-4-6-v1 # Required model ID
# BEDROCK_CROSS_REGION=us # Cross-region prefix (us/eu/apac/global)
# AWS_PROFILE=my-profile # Named profile (SSO/assume-role)
# Tunnel (public internet exposure for webhooks)
TUNNEL_URL=https://abc123.ngrok.io # Static public URL (manual tunnel)
# Or use a managed tunnel provider:
@@ -500,7 +507,9 @@ OBSERVABILITY_BACKEND=none # none/noop (default) or log
### LLM Providers
Backends: `nearai` (default), `openai`, `anthropic`, `ollama`, `openai_compatible`, `tinfoil` — set via `LLM_BACKEND`. See [src/llm/CLAUDE.md](src/llm/CLAUDE.md) for per-provider auth and configuration details.
Backends: `nearai` (default), `openai`, `anthropic`, `ollama`, `openai_compatible`, `tinfoil`, `bedrock` (requires `--features bedrock`) — set via `LLM_BACKEND`. See [src/llm/CLAUDE.md](src/llm/CLAUDE.md) for per-provider auth and configuration details.
**AWS Bedrock** -- Uses the native Converse API via `aws-sdk-bedrockruntime`. Requires `--features bedrock` at build time (not included in default features due to heavy AWS SDK dependencies). Supports standard AWS auth methods: IAM credentials (`AWS_ACCESS_KEY_ID`/`AWS_SECRET_ACCESS_KEY`), SSO profiles (`AWS_PROFILE`), and instance roles. Configure with `BEDROCK_REGION` (default: `us-east-1`), `BEDROCK_MODEL` (required, e.g., `anthropic.claude-opus-4-6-v1`), and `BEDROCK_CROSS_REGION` (optional: `us`, `eu`, `apac`, `global` for cross-region inference profiles). The SDK credential chain resolves auth automatically from the environment.
## Database
Generated
+707 -23
View File
File diff suppressed because it is too large Load Diff
+8
View File
@@ -73,6 +73,8 @@ toml = "0.8"
# Core types
uuid = { version = "1", features = ["v4", "v5", "serde"] }
chrono = { version = "0.4", features = ["serde"] }
chrono-tz = "0.10"
iana-time-zone = "0.1"
rust_decimal = { version = "1", features = ["serde", "serde-with-str", "maths"] }
rust_decimal_macros = "1"
@@ -140,6 +142,11 @@ subtle = "2" # Constant-time comparisons for token validation
# Multi-provider LLM support
rig-core = "0.30"
# AWS Bedrock (native Converse API, opt-in via --features bedrock)
aws-config = { version = "1", features = ["behavior-version-latest"], optional = true }
aws-sdk-bedrockruntime = { version = "1", optional = true }
aws-smithy-types = { version = "1", optional = true }
# Docker sandbox
bollard = "0.18"
@@ -201,6 +208,7 @@ postgres = [
libsql = ["dep:libsql"]
integration = []
html-to-markdown = ["dep:html-to-markdown-rs", "dep:readabilityrs"]
bedrock = ["dep:aws-config", "dep:aws-sdk-bedrockruntime", "dep:aws-smithy-types"]
[[test]]
name = "html_to_markdown"
+2 -2
View File
@@ -39,7 +39,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
| Network modes (loopback/LAN/remote) | ✅ | 🚧 | HTTP only |
| OpenAI-compatible HTTP API | ✅ | ✅ | /v1/chat/completions, per-request `model` override |
| Canvas hosting | ✅ | ❌ | Agent-driven UI |
| Gateway lock (PID-based) | ✅ | | |
| Gateway lock (PID-based) | ✅ | | `fs4` flock-based, acquired in `main.rs` before agent startup |
| launchd/systemd integration | ✅ | ❌ | |
| Bonjour/mDNS discovery | ✅ | ❌ | |
| Tailscale integration | ✅ | ❌ | |
@@ -215,7 +215,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
| NEAR AI | ✅ | ✅ | - | Primary provider |
| Anthropic (Claude) | ✅ | 🚧 | - | Via NEAR AI proxy; Opus 4.5, Sonnet 4, Sonnet 4.6 |
| OpenAI | ✅ | 🚧 | - | Via NEAR AI proxy |
| AWS Bedrock | ✅ | ✅ | P3 | Via `openai_compatible` adapter (e.g. LiteLLM) |
| AWS Bedrock | ✅ | ✅ | - | Native Converse API via aws-sdk-bedrockruntime (requires `--features bedrock`) |
| Google Gemini | ✅ | ✅ | P3 | Via `gemini` adapter |
| io.net | ✅ | ✅ | P3 | Via `ionet` adapter |
| Mistral | ✅ | ✅ | P3 | Via `mistral` adapter |
+5
View File
@@ -14,6 +14,11 @@
<a href="https://www.reddit.com/r/ironclawAI/"><img src="https://img.shields.io/badge/Reddit-r%2FironclawAI-FF4500?style=flat&logo=reddit&logoColor=white" alt="Reddit: r/ironclawAI" /></a>
</p>
<p align="center">
<a href="README.md">English</a> |
<a href="README.zh-CN.md">简体中文</a>
</p>
<p align="center">
<a href="#philosophy">Philosophy</a> •
<a href="#features">Features</a> •
+319
View File
@@ -0,0 +1,319 @@
<p align="center">
<img src="ironclaw.png?v=2" alt="IronClaw" width="200"/>
</p>
<h1 align="center">IronClaw</h1>
<p align="center">
<strong>安全可靠的个人 AI 助手,始终站在你这边</strong>
</p>
<p align="center">
<a href="#license"><img src="https://img.shields.io/badge/license-MIT%20OR%20Apache%202.0-blue.svg" alt="License: MIT OR Apache-2.0" /></a>
<a href="https://t.me/ironclawAI"><img src="https://img.shields.io/badge/Telegram-%40ironclawAI-26A5E4?style=flat&logo=telegram&logoColor=white" alt="Telegram: @ironclawAI" /></a>
<a href="https://www.reddit.com/r/ironclawAI/"><img src="https://img.shields.io/badge/Reddit-r%2FironclawAI-FF4500?style=flat&logo=reddit&logoColor=white" alt="Reddit: r/ironclawAI" /></a>
</p>
<p align="center">
<a href="README.md">English</a> |
<a href="README.zh-CN.md">简体中文</a>
</p>
<p align="center">
<a href="#设计理念">设计理念</a> •
<a href="#功能特性">功能特性</a> •
<a href="#安装">安装</a> •
<a href="#配置">配置</a> •
<a href="#安全机制">安全机制</a> •
<a href="#系统架构">系统架构</a>
</p>
---
## 设计理念
IronClaw 基于一个简单的原则:**你的 AI 助手应该为你服务,而不是与你为敌。**
在 AI 系统对数据处理日益不透明、与企业利益捆绑的今天,IronClaw 选择了一条不同的路:
- **数据归你所有** — 所有信息存储在本地,加密保护,始终在你掌控之下
- **透明至上** — 完全开源,可审计,没有隐藏的遥测或数据收集
- **自主扩展** — 随时构建新工具,无需等待供应商更新
- **纵深防御** — 多层安全机制抵御提示注入和数据泄露
IronClaw 是一个你真正可以信赖的 AI 助手,无论是个人生活还是工作。
## 功能特性
### 安全优先
- **WASM 沙箱** — 不受信任的工具在隔离的 WebAssembly 容器中运行,采用基于能力的权限模型
- **凭据保护** — 密钥永远不会暴露给工具;在宿主边界注入并进行泄露检测
- **提示注入防御** — 模式检测、内容清理和策略执行
- **端点白名单** — HTTP 请求仅限于明确批准的主机和路径
### 随时可用
- **多渠道接入** — REPL、HTTP webhook、WASM 渠道(Telegram、Slack)和 Web 网关
- **Docker 沙箱** — 隔离的容器执行,支持每任务令牌和编排器/工作器模式
- **Web 网关** — 浏览器 UI,支持实时 SSE/WebSocket 流式传输
- **定时任务** — Cron 调度、事件触发器、Webhook 处理器,实现后台自动化
- **心跳系统** — 主动后台执行,用于监控和维护任务
- **并行任务** — 使用隔离上下文同时处理多个请求
- **自修复** — 自动检测并恢复卡住的操作
### 自主扩展
- **动态工具构建** — 描述你的需求,IronClaw 会将其构建为 WASM 工具
- **MCP 协议** — 连接模型上下文协议(Model Context Protocol)服务器以获取额外能力
- **插件架构** — 无需重启即可加载新的 WASM 工具和渠道
### 持久记忆
- **混合搜索** — 全文搜索 + 向量搜索,采用倒数排名融合(Reciprocal Rank Fusion
- **工作空间文件系统** — 灵活的基于路径的存储,用于笔记、日志和上下文
- **身份文件** — 跨会话保持一致的个性和偏好设置
## 安装
### 前置要求
- Rust 1.85+
- PostgreSQL 15+,需安装 [pgvector](https://github.com/pgvector/pgvector) 扩展
- NEAR AI 账户(通过设置向导进行身份验证)
## 下载或编译
访问 [Releases 页面](https://github.com/nearai/ironclaw/releases/) 查看最新版本。
<details>
<summary>通过 Windows 安装程序安装 (Windows)</summary>
下载 [Windows 安装程序](https://github.com/nearai/ironclaw/releases/latest/download/ironclaw-x86_64-pc-windows-msvc.msi) 并运行。
</details>
<details>
<summary>通过 PowerShell 脚本安装 (Windows)</summary>
```sh
irm https://github.com/nearai/ironclaw/releases/latest/download/ironclaw-installer.ps1 | iex
```
</details>
<details>
<summary>通过 Shell 脚本安装 (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>通过 Homebrew 安装 (macOS/Linux)</summary>
```sh
brew install ironclaw
```
</details>
<details>
<summary>从源码编译 (Windows、Linux、macOS 上使用 Cargo)</summary>
确保你已安装 [Rust](https://rustup.rs)。
```bash
# 克隆仓库
git clone https://github.com/nearai/ironclaw.git
cd ironclaw
# 编译
cargo build --release
# 运行测试
cargo test
```
如需进行**完整发布构建**(修改了渠道源码后),先运行 `./scripts/build-all.sh` 重新编译渠道。
</details>
### 数据库设置
```bash
# 创建数据库
createdb ironclaw
# 启用 pgvector 扩展
psql ironclaw -c "CREATE EXTENSION IF NOT EXISTS vector;"
```
## 配置
运行设置向导来配置 IronClaw:
```bash
ironclaw onboard
```
向导将引导你完成数据库连接、NEAR AI 身份验证(通过浏览器 OAuth)和密钥加密(使用系统钥匙串)。设置会保存在数据库中;引导变量(如 `DATABASE_URL``LLM_BACKEND`)写入 `~/.ironclaw/.env`,以便在数据库连接前可用。
### 替代 LLM 提供商
IronClaw 默认使用 NEAR AI,但兼容任何 OpenAI 兼容的端点。
常用选项包括 **OpenRouter**300+ 模型)、**Together AI**、**Fireworks AI**、**Ollama**(本地部署)以及自托管服务器如 **vLLM****LiteLLM**
在向导中选择 *"OpenAI-compatible"*,或直接设置环境变量:
```env
LLM_BACKEND=openai_compatible
LLM_BASE_URL=https://openrouter.ai/api/v1
LLM_API_KEY=sk-or-...
LLM_MODEL=anthropic/claude-sonnet-4
```
详见 [docs/LLM_PROVIDERS.md](docs/LLM_PROVIDERS.md) 获取完整的提供商指南。
## 安全机制
IronClaw 实现了纵深防御策略来保护你的数据并防止滥用。
### WASM 沙箱
所有不受信任的工具都在隔离的 WebAssembly 容器中运行:
- **基于能力的权限** — 明确授权 HTTP、密钥、工具调用等能力
- **端点白名单** — HTTP 请求仅限已批准的主机和路径
- **凭据注入** — 密钥在宿主边界注入,永远不会暴露给 WASM 代码
- **泄露检测** — 扫描请求和响应以防止密钥外泄
- **速率限制** — 每个工具独立的请求限制,防止滥用
- **资源限制** — 内存、CPU 和执行时间约束
```
WASM ──► 白名单 ──► 泄露扫描 ──► 凭据 ──► 执行 ──► 泄露扫描 ──► WASM
验证器 (请求) 注入器 请求 (响应)
```
### 提示注入防御
外部内容需通过多个安全层:
- 基于模式的注入尝试检测
- 内容清理和转义
- 带严重级别的策略规则(阻止/警告/审核/清理)
- 工具输出包装,确保安全的 LLM 上下文注入
### 数据保护
- 所有数据存储在本地 PostgreSQL 数据库中
- 密钥使用 AES-256-GCM 加密
- 无遥测、无分析、无数据共享
- 所有工具执行的完整审计日志
## 系统架构
```
┌────────────────────────────────────────────────────────────────┐
│ 渠道 │
│ ┌──────┐ ┌──────┐ ┌─────────────┐ ┌─────────────┐ │
│ │ REPL │ │ HTTP │ │ WASM 渠道 │ │ Web 网关 │ │
│ └──┬───┘ └──┬───┘ └──────┬──────┘ │ (SSE + WS) │ │
│ │ │ │ └──────┬──────┘ │
│ └─────────┴──────────────┴────────────────┘ │
│ │ │
│ ┌─────────▼─────────┐ │
│ │ 代理循环 │ 意图路由 │
│ └────┬──────────┬───┘ │
│ │ │ │
│ ┌──────────▼────┐ ┌──▼───────────────┐ │
│ │ 调度器 │ │ 定时任务引擎 │ │
│ │ (并行任务) │ │(cron, 事件, wh) │ │
│ └──────┬────────┘ └────────┬─────────┘ │
│ │ │ │
│ ┌─────────────┼────────────────────┘ │
│ │ │ │
│ ┌───▼─────┐ ┌────▼────────────────┐ │
│ │ 本地 │ │ 编排器 │ │
│ │ 工作器 │ │ ┌───────────────┐ │ │
│ │(进程内) │ │ │ Docker 沙箱 │ │ │
│ └───┬─────┘ │ │ 容器 │ │ │
│ │ │ │ ┌───────────┐ │ │ │
│ │ │ │ │工作器/CC │ │ │ │
│ │ │ │ └───────────┘ │ │ │
│ │ │ └───────────────┘ │ │
│ │ └─────────┬───────────┘ │
│ └──────────────────┤ │
│ │ │
│ ┌───────────▼──────────┐ │
│ │ 工具注册表 │ │
│ │ 内置、MCP、WASM │ │
│ └──────────────────────┘ │
└────────────────────────────────────────────────────────────────┘
```
### 核心组件
| 组件 | 用途 |
|------|------|
| **代理循环** | 主消息处理和任务协调 |
| **路由器** | 分类用户意图(命令、查询、任务) |
| **调度器** | 管理带优先级的并行任务执行 |
| **工作器** | 执行包含 LLM 推理和工具调用的任务 |
| **编排器** | 容器生命周期、LLM 代理、每任务认证 |
| **Web 网关** | 浏览器 UI,含聊天、记忆、任务、日志、扩展、定时任务 |
| **定时任务引擎** | 定时(cron)和响应式(事件、webhook)后台任务 |
| **工作空间** | 带混合搜索的持久记忆 |
| **安全层** | 提示注入防御和内容清理 |
## 使用方式
```bash
# 首次设置(配置数据库、认证等)
ironclaw onboard
# 启动交互式 REPL
cargo run
# 启用调试日志
RUST_LOG=ironclaw=debug cargo run
```
## 开发
```bash
# 格式化代码
cargo fmt
# 代码检查
cargo clippy --all --benches --tests --examples --all-features
# 运行测试
createdb ironclaw_test
cargo test
# 运行指定测试
cargo test test_name
```
- **Telegram 渠道**:参见 [docs/TELEGRAM_SETUP.md](docs/TELEGRAM_SETUP.md) 了解设置和私信配对。
- **修改渠道源码**:在 `cargo build` 之前运行 `./channels-src/telegram/build.sh` 以便打包更新后的 WASM。
## OpenClaw 传承
IronClaw 是受 [OpenClaw](https://github.com/openclaw/openclaw) 启发的 Rust 重新实现。参见 [FEATURE_PARITY.md](FEATURE_PARITY.md) 了解完整的功能追踪矩阵。
主要差异:
- **Rust vs TypeScript** — 原生性能、内存安全、单一二进制文件
- **WASM 沙箱 vs Docker** — 轻量级、基于能力的安全机制
- **PostgreSQL vs SQLite** — 生产级持久化存储
- **安全优先设计** — 多层防御、凭据保护
## 许可证
可选择以下任一许可证:
- Apache License, Version 2.0 ([LICENSE-APACHE](LICENSE-APACHE))
- MIT License ([LICENSE-MIT](LICENSE-MIT))
+1 -1
View File
@@ -267,7 +267,7 @@ dependencies = [
[[package]]
name = "slack-channel"
version = "0.1.0"
version = "0.2.1"
dependencies = [
"hex",
"hmac",
+1 -1
View File
@@ -1,6 +1,6 @@
[package]
name = "slack-channel"
version = "0.2.0"
version = "0.2.1"
edition = "2021"
description = "Slack Events API channel for IronClaw"
license = "MIT OR Apache-2.0"
+104
View File
@@ -357,10 +357,108 @@ fn extract_slack_attachments(files: &Option<Vec<SlackFile>>) -> Vec<InboundAttac
.collect()
}
/// Download a file from Slack using the url_private endpoint.
///
/// Slack file downloads require Bearer auth with the bot token, which is
/// injected by the host credential system via `channel_host::http_request`.
fn download_slack_file(url: &str) -> Result<Vec<u8>, String> {
let headers = serde_json::json!({});
let result = channel_host::http_request("GET", url, &headers.to_string(), None, None);
let response = result.map_err(|e| format!("Slack file download failed: {}", e))?;
if response.status != 200 {
let body_str = String::from_utf8_lossy(&response.body);
return Err(format!(
"Slack file download returned {}: {}",
response.status, body_str
));
}
Ok(response.body)
}
/// Download file bytes and store them via the host for processing.
///
/// Downloads all file types (images, documents, etc.) so the host-side
/// middleware can process them (vision pipeline for images, text extraction
/// for documents, transcription for audio, etc.).
/// Maximum file size to download (20 MB). Files larger than this are skipped
/// to avoid excessive memory use and slow downloads in the WASM runtime.
const MAX_DOWNLOAD_SIZE_BYTES: u64 = 20 * 1024 * 1024;
fn download_and_store_slack_files(attachments: &[InboundAttachment]) {
for att in attachments {
let Some(ref url) = att.source_url else {
continue;
};
// Skip files that exceed the size limit
if let Some(size) = att.size_bytes {
if size > MAX_DOWNLOAD_SIZE_BYTES {
channel_host::log(
channel_host::LogLevel::Warn,
&format!(
"Skipping Slack file download: {} bytes exceeds {} MB limit (id={})",
size,
MAX_DOWNLOAD_SIZE_BYTES / (1024 * 1024),
att.id
),
);
continue;
}
}
match download_slack_file(url) {
Ok(bytes) => {
// Post-download size guard: metadata size_bytes is optional,
// so a file with no size info could bypass the pre-download check.
if bytes.len() as u64 > MAX_DOWNLOAD_SIZE_BYTES {
channel_host::log(
channel_host::LogLevel::Warn,
&format!(
"Discarding Slack file after download: {} bytes exceeds {} MB limit (id={})",
bytes.len(),
MAX_DOWNLOAD_SIZE_BYTES / (1024 * 1024),
att.id
),
);
continue;
}
channel_host::log(
channel_host::LogLevel::Info,
&format!(
"Downloaded Slack file: {} bytes, mime={}",
bytes.len(),
att.mime_type
),
);
if let Err(e) = channel_host::store_attachment_data(&att.id, &bytes) {
channel_host::log(
channel_host::LogLevel::Error,
&format!("Failed to store Slack file data: {}", e),
);
}
}
Err(e) => {
channel_host::log(
channel_host::LogLevel::Error,
&format!("Failed to download Slack file: {}", e),
);
}
}
}
}
/// Handle a Slack event and emit message if applicable.
fn handle_slack_event(event: SlackEvent, team_id: Option<String>, _event_id: Option<String>) {
let attachments = extract_slack_attachments(&event.files);
// Download and store file attachments for host-side processing
download_and_store_slack_files(&attachments);
match event.event_type.as_str() {
// Direct mention of the bot (always in a channel, not a DM)
"app_mention" => {
@@ -722,4 +820,10 @@ mod tests {
let event: SlackEvent = serde_json::from_str(json).unwrap();
assert!(event.files.is_none());
}
#[test]
fn test_max_download_size_constant() {
// Verify the constant is 20 MB
assert_eq!(MAX_DOWNLOAD_SIZE_BYTES, 20 * 1024 * 1024);
}
}
+1 -1
View File
@@ -212,7 +212,7 @@ dependencies = [
[[package]]
name = "telegram-channel"
version = "0.2.0"
version = "0.2.1"
dependencies = [
"serde",
"serde_json",
+1 -1
View File
@@ -1,6 +1,6 @@
[package]
name = "telegram-channel"
version = "0.2.0"
version = "0.2.1"
edition = "2021"
description = "Telegram Bot API channel for IronClaw"
license = "MIT OR Apache-2.0"
+57 -5
View File
@@ -878,10 +878,6 @@ fn send_message(
// Voice File Download
// ============================================================================
/// Download a voice file from Telegram by file_id.
///
/// 1. Call getFile to get the file_path.
/// 2. Download the file bytes from /file/bot{TOKEN}/{file_path}.
/// Percent-encode a string for safe use as a URL query parameter value.
fn percent_encode(s: &str) -> String {
let mut out = String::with_capacity(s.len());
@@ -898,6 +894,10 @@ fn percent_encode(s: &str) -> String {
out
}
/// Maximum file size to download (20 MB). Files larger than this are discarded
/// to avoid excessive memory use and slow downloads in the WASM runtime.
const MAX_DOWNLOAD_SIZE_BYTES: u64 = 20 * 1024 * 1024;
fn download_telegram_file(file_id: &str) -> Result<Vec<u8>, String> {
// Reject file_id containing curly braces to prevent credential placeholder injection
if file_id.contains('{') || file_id.contains('}') {
@@ -965,6 +965,16 @@ fn download_telegram_file(file_id: &str) -> Result<Vec<u8>, String> {
));
}
// Post-download size guard: Telegram metadata file_size is optional,
// so enforce the limit on actual downloaded bytes.
if response.body.len() as u64 > MAX_DOWNLOAD_SIZE_BYTES {
return Err(format!(
"Downloaded file exceeds {} MB limit ({} bytes)",
MAX_DOWNLOAD_SIZE_BYTES / (1024 * 1024),
response.body.len()
));
}
Ok(response.body)
}
@@ -1535,6 +1545,39 @@ fn download_and_store_voice(attachments: &[InboundAttachment]) {
}
}
/// Download image file bytes and store them via the host for the vision pipeline.
///
/// Separated from `extract_attachments` so that function stays pure (no host
/// calls) and remains testable in native unit tests.
fn download_and_store_images(attachments: &[InboundAttachment]) {
for att in attachments {
if !att.mime_type.starts_with("image/") {
continue;
}
match download_telegram_file(&att.id) {
Ok(bytes) => {
channel_host::log(
channel_host::LogLevel::Info,
&format!("Downloaded image file: {} bytes", bytes.len()),
);
if let Err(e) = channel_host::store_attachment_data(&att.id, &bytes) {
channel_host::log(
channel_host::LogLevel::Error,
&format!("Failed to store image data: {}", e),
);
}
}
Err(e) => {
channel_host::log(
channel_host::LogLevel::Error,
&format!("Failed to download image file: {}", e),
);
}
}
}
}
/// Returns true if the attachment should be downloaded for document text extraction.
///
/// Excludes voice (handled by transcription), image (vision pipeline),
@@ -1608,6 +1651,9 @@ fn handle_message(message: TelegramMessage) {
// Download and store voice attachments for host-side transcription
download_and_store_voice(&attachments);
// Download and store image attachments for host-side vision pipeline
download_and_store_images(&attachments);
// Download and store document attachments for host-side text extraction
download_and_store_documents(&mut attachments);
@@ -1681,7 +1727,7 @@ fn handle_message(message: TelegramMessage) {
let username_opt = from.username.as_deref();
let is_allowed = allowed.contains(&"*".to_string())
|| allowed.contains(&id_str)
|| username_opt.map_or(false, |u| allowed.contains(&u.to_string()));
|| username_opt.is_some_and(|u| allowed.contains(&u.to_string()));
if !is_allowed {
if is_private && dm_policy == "pairing" {
@@ -2605,4 +2651,10 @@ mod tests {
assert!(!is_downloadable_document(&make("audio/mpeg", Some("song.mp3"))));
assert!(!is_downloadable_document(&make("video/mp4", Some("clip.mp4"))));
}
#[test]
fn test_max_download_size_constant() {
// Verify the constant is 20 MB, matching the Slack channel limit
assert_eq!(MAX_DOWNLOAD_SIZE_BYTES, 20 * 1024 * 1024);
}
}
@@ -1,9 +1,17 @@
{
"version": "0.2.0",
"version": "0.2.2",
"wit_version": "0.3.0",
"type": "channel",
"name": "telegram",
"description": "Telegram Bot API channel for receiving and responding to Telegram messages",
"auth": {
"secret_name": "telegram_bot_token",
"display_name": "Telegram",
"instructions": "Get your bot token from @BotFather on Telegram (https://t.me/BotFather). Send /newbot or /token to get it.",
"setup_url": "https://t.me/BotFather",
"token_hint": "Looks like 123456789:AABBccDDeeFFgg...",
"env_var": "TELEGRAM_BOT_TOKEN"
},
"setup": {
"required_secrets": [
{
+1 -1
View File
@@ -3,7 +3,7 @@ services:
postgres:
image: pgvector/pgvector:pg16
ports:
- "5432:5432"
- "127.0.0.1:5432:5432"
environment:
POSTGRES_DB: ironclaw
POSTGRES_USER: ironclaw
+50 -1
View File
@@ -12,12 +12,12 @@ configurations.
| Anthropic | `anthropic` | `ANTHROPIC_API_KEY` | Claude models |
| OpenAI | `openai` | `OPENAI_API_KEY` | GPT models |
| Google Gemini | `gemini` | `GEMINI_API_KEY` | Gemini models |
| AWS Bedrock | `bedrock` | `BEDROCK_ACCESS_KEY` | Requires OpenAI proxy (e.g. LiteLLM) |
| io.net | `ionet` | `IONET_API_KEY` | Intelligence API |
| Mistral | `mistral` | `MISTRAL_API_KEY` | Mistral models |
| Yandex AI Studio | `yandex` | `YANDEX_API_KEY` | YandexGPT models |
| Cloudflare Workers AI | `cloudflare` | `CLOUDFLARE_API_KEY` | Access to Workers AI |
| Ollama | `ollama` | No | Local inference |
| AWS Bedrock | `bedrock` | AWS credentials | Native Converse API |
| OpenRouter | `openai_compatible` | `LLM_API_KEY` | 300+ models |
| Together AI | `openai_compatible` | `LLM_API_KEY` | Fast inference |
| Fireworks AI | `openai_compatible` | `LLM_API_KEY` | Fast inference |
@@ -74,6 +74,55 @@ Pull a model first: `ollama pull llama3.2`
---
## AWS Bedrock (requires `--features bedrock`)
Uses the native AWS Converse API via `aws-sdk-bedrockruntime`. Supports standard AWS
authentication methods: IAM credentials, SSO profiles, and instance roles.
> **Build prerequisite:** The `aws-lc-sys` crate (transitive dependency via AWS SDK)
> requires **CMake** to compile. Install it before building with `--features bedrock`:
> - macOS: `brew install cmake`
> - Ubuntu/Debian: `sudo apt install cmake`
> - Fedora: `sudo dnf install cmake`
### With AWS credentials (IAM, SSO, instance roles)
```env
LLM_BACKEND=bedrock
BEDROCK_MODEL=anthropic.claude-opus-4-6-v1
BEDROCK_REGION=us-east-1
BEDROCK_CROSS_REGION=us
# AWS_PROFILE=my-sso-profile # optional, for named profiles
```
The AWS SDK credential chain automatically resolves credentials from environment
variables (`AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`), shared credentials file
(`~/.aws/credentials`), SSO profiles, and EC2/ECS instance roles.
### Cross-region inference
Set `BEDROCK_CROSS_REGION` to route requests across AWS regions for capacity:
| Prefix | Routing |
|---|---|
| `us` | US regions (us-east-1, us-east-2, us-west-2) |
| `eu` | European regions |
| `apac` | Asia-Pacific regions |
| `global` | All commercial AWS regions |
| _(unset)_ | Single-region only |
### Popular Bedrock model IDs
| Model | ID |
|---|---|
| Claude Opus 4.6 | `anthropic.claude-opus-4-6-v1` |
| Claude Sonnet 4.5 | `anthropic.claude-sonnet-4-5-20250929-v1:0` |
| Claude Haiku 4.5 | `anthropic.claude-haiku-4-5-20251001-v1:0` |
| Amazon Nova Pro | `amazon.nova-pro-v1:0` |
| Llama 4 Maverick | `meta.llama4-maverick-17b-instruct-v1:0` |
---
## OpenAI-Compatible Endpoints
All providers below use `LLM_BACKEND=openai_compatible`. Set `LLM_BASE_URL` to the
-20
View File
@@ -295,26 +295,6 @@
"can_list_models": true
}
},
{
"id": "bedrock",
"aliases": [
"aws_bedrock",
"aws"
],
"protocol": "open_ai_completions",
"api_key_env": "BEDROCK_ACCESS_KEY",
"api_key_required": false,
"base_url_env": "BEDROCK_BASE_URL",
"model_env": "BEDROCK_MODEL",
"default_model": "anthropic.claude-3-5-sonnet-20241022-v2:0",
"description": "AWS Bedrock (requires LiteLLM or OpenAI-compatible proxy)",
"setup": {
"kind": "open_ai_compatible",
"secret_name": "llm_bedrock_api_key",
"display_name": "AWS Bedrock",
"can_list_models": false
}
},
{
"id": "ionet",
"aliases": [
+1 -1
View File
@@ -2,7 +2,7 @@
"name": "slack",
"display_name": "Slack Channel",
"kind": "channel",
"version": "0.2.0",
"version": "0.2.1",
"wit_version": "0.3.0",
"description": "Talk to your agent in Slack",
"keywords": [
+1 -1
View File
@@ -2,7 +2,7 @@
"name": "telegram",
"display_name": "Telegram Channel",
"kind": "channel",
"version": "0.2.0",
"version": "0.2.2",
"wit_version": "0.3.0",
"description": "Talk to your agent through a Telegram bot",
"keywords": [
+6 -1
View File
@@ -356,6 +356,12 @@ impl Agent {
if let Some(workspace) = self.workspace() {
let mut config = AgentHeartbeatConfig::default()
.with_interval(std::time::Duration::from_secs(hb_config.interval_secs));
config.quiet_hours_start = hb_config.quiet_hours_start;
config.quiet_hours_end = hb_config.quiet_hours_end;
config.timezone = hb_config
.timezone
.clone()
.or_else(|| Some(self.config.default_timezone.clone()));
if let (Some(user), Some(channel)) =
(&hb_config.notify_user, &hb_config.notify_channel)
{
@@ -411,7 +417,6 @@ impl Agent {
hygiene,
workspace.clone(),
self.cheap_llm().clone(),
self.safety().clone(),
Some(notify_tx),
self.store().map(Arc::clone),
))
+48 -7
View File
@@ -345,7 +345,6 @@ impl Agent {
crate::workspace::hygiene::HygieneConfig::default(),
workspace.clone(),
self.llm().clone(),
self.safety().clone(),
);
match runner.check_heartbeat().await {
@@ -406,7 +405,7 @@ impl Agent {
.with_max_tokens(512)
.with_temperature(0.3);
let reasoning = Reasoning::new(self.llm().clone(), self.safety().clone());
let reasoning = Reasoning::new(self.llm().clone());
match reasoning.complete(request).await {
Ok((text, _usage)) => Ok(SubmissionResult::response(format!(
"Thread Summary:\n\n{}",
@@ -454,7 +453,7 @@ impl Agent {
.with_max_tokens(512)
.with_temperature(0.5);
let reasoning = Reasoning::new(self.llm().clone(), self.safety().clone());
let reasoning = Reasoning::new(self.llm().clone());
match reasoning.complete(request).await {
Ok((text, _usage)) => Ok(SubmissionResult::response(format!(
"Suggested Next Steps:\n\n{}",
@@ -663,10 +662,14 @@ impl Agent {
}
match self.llm().set_model(requested) {
Ok(()) => Ok(SubmissionResult::response(format!(
"Switched model to: {}",
requested
))),
Ok(()) => {
// Persist the model choice so it survives restarts.
self.persist_selected_model(requested).await;
Ok(SubmissionResult::response(format!(
"Switched model to: {}",
requested
)))
}
Err(e) => Ok(SubmissionResult::error(format!(
"Failed to switch model: {}",
e
@@ -822,4 +825,42 @@ impl Agent {
_ => Ok(None),
}
}
/// Persist the selected model to the settings store (DB and/or TOML config).
///
/// Best-effort: logs warnings on failure but does not propagate errors,
/// since the in-memory model switch already succeeded.
async fn persist_selected_model(&self, model: &str) {
// 1. Persist to DB if available.
if let Some(store) = self.store() {
let value = serde_json::Value::String(model.to_string());
if let Err(e) = store.set_setting("default", "selected_model", &value).await {
tracing::warn!("Failed to persist model to DB: {}", e);
}
}
// 2. Update TOML config file if it exists (sync I/O in spawn_blocking).
let model_owned = model.to_string();
if let Err(e) = tokio::task::spawn_blocking(move || {
let toml_path = crate::settings::Settings::default_toml_path();
match crate::settings::Settings::load_toml(&toml_path) {
Ok(Some(mut settings)) => {
settings.selected_model = Some(model_owned);
if let Err(e) = settings.save_toml(&toml_path) {
tracing::warn!("Failed to persist model to config.toml: {}", e);
}
}
Ok(None) => {
// No config file on disk; nothing to update.
}
Err(e) => {
tracing::warn!("Failed to load config.toml for model persistence: {}", e);
}
}
})
.await
{
tracing::warn!("Model TOML persistence task failed: {}", e);
}
}
}
+112 -37
View File
@@ -13,7 +13,6 @@ use crate::agent::context_monitor::{CompactionStrategy, ContextBreakdown};
use crate::agent::session::Thread;
use crate::error::Error;
use crate::llm::{ChatMessage, CompletionRequest, LlmProvider, Reasoning};
use crate::safety::SafetyLayer;
use crate::workspace::Workspace;
/// Result of a compaction operation.
@@ -34,13 +33,12 @@ pub struct CompactionResult {
/// Compacts conversation context to stay within limits.
pub struct ContextCompactor {
llm: Arc<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
}
impl ContextCompactor {
/// Create a new context compactor.
pub fn new(llm: Arc<dyn LlmProvider>, safety: Arc<SafetyLayer>) -> Self {
Self { llm, safety }
pub fn new(llm: Arc<dyn LlmProvider>) -> Self {
Self { llm }
}
/// Compact a thread's context using the given strategy.
@@ -105,27 +103,26 @@ impl ContextCompactor {
// Generate summary
let summary = self.generate_summary(&to_summarize).await?;
// Write to workspace if available
let summary_written = if let Some(ws) = workspace {
// Write to workspace if available.
// If archival fails, preserve turns to avoid context loss.
let (summary_written, turns_removed) = if let Some(ws) = workspace {
match self.write_summary_to_workspace(ws, &summary).await {
Ok(()) => true,
Ok(()) => {
thread.truncate_turns(keep_recent);
(true, turns_to_remove)
}
Err(e) => {
tracing::warn!(
"Compaction summary write failed (turns will still be truncated): {}",
e
);
false
tracing::warn!("Compaction summary write failed (turns preserved): {}", e);
(false, 0)
}
}
} else {
false
thread.truncate_turns(keep_recent);
(false, turns_to_remove)
};
// Truncate thread
thread.truncate_turns(keep_recent);
Ok(CompactionPartial {
turns_removed: turns_to_remove,
turns_removed,
summary_written,
summary: Some(summary),
})
@@ -167,23 +164,20 @@ impl ContextCompactor {
// Format turns for storage
let content = format_turns_for_storage(old_turns);
// Write to workspace
let written = match self.write_context_to_workspace(ws, &content).await {
Ok(()) => true,
// Write to workspace. If archival fails, preserve turns.
let (written, turns_removed) = match self.write_context_to_workspace(ws, &content).await {
Ok(()) => {
thread.truncate_turns(keep_recent);
(true, turns_to_remove)
}
Err(e) => {
tracing::warn!(
"Compaction context write failed (turns will still be truncated): {}",
e
);
false
tracing::warn!("Compaction context write failed (turns preserved): {}", e);
(false, 0)
}
};
// Truncate
thread.truncate_turns(keep_recent);
Ok(CompactionPartial {
turns_removed: turns_to_remove,
turns_removed,
summary_written: written,
summary: None,
})
@@ -233,7 +227,7 @@ Be brief but capture all important details. Use bullet points."#,
.with_max_tokens(1024)
.with_temperature(0.3);
let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone());
let reasoning = Reasoning::new(self.llm.clone());
let (text, _) = reasoning.complete(request).await?;
Ok(text)
}
@@ -346,17 +340,11 @@ mod tests {
// === QA Plan - Compaction strategy tests ===
use crate::agent::context_monitor::CompactionStrategy;
use crate::config::SafetyConfig;
use crate::safety::SafetyLayer;
use crate::testing::StubLlm;
/// Helper: build a `ContextCompactor` with the given `StubLlm`.
fn make_compactor(llm: Arc<StubLlm>) -> ContextCompactor {
let safety = Arc::new(SafetyLayer::new(&SafetyConfig {
max_output_length: 100_000,
injection_check_enabled: false,
}));
ContextCompactor::new(llm, safety)
ContextCompactor::new(llm)
}
/// Helper: build a thread with `n` completed turns.
@@ -370,6 +358,19 @@ mod tests {
thread
}
#[cfg(feature = "libsql")]
async fn make_unmigrated_workspace() -> crate::workspace::Workspace {
use crate::db::Database;
use crate::db::libsql::LibSqlBackend;
// Intentionally skip migrations so workspace append operations fail.
let backend = LibSqlBackend::new_memory()
.await
.expect("should create in-memory libsql backend");
let db: Arc<dyn Database> = Arc::new(backend);
crate::workspace::Workspace::new_with_db("compaction-test", db)
}
// ------------------------------------------------------------------
// 1. compact_truncate keeps last N turns
// ------------------------------------------------------------------
@@ -568,6 +569,43 @@ mod tests {
assert_eq!(llm.calls(), 0);
}
#[cfg(feature = "libsql")]
#[tokio::test]
async fn test_compact_with_summary_preserves_turns_when_workspace_write_fails() {
let llm = Arc::new(StubLlm::new("summary"));
let compactor = make_compactor(llm.clone());
let mut thread = make_thread(8);
let original_inputs: Vec<String> =
thread.turns.iter().map(|t| t.user_input.clone()).collect();
let workspace = make_unmigrated_workspace().await;
let result = compactor
.compact(
&mut thread,
CompactionStrategy::Summarize { keep_recent: 3 },
Some(&workspace),
)
.await
.expect("compact should succeed even when workspace write fails");
// On archival failure, no turns should be removed.
assert_eq!(thread.turns.len(), 8);
assert_eq!(
thread
.turns
.iter()
.map(|t| t.user_input.as_str())
.collect::<Vec<_>>(),
original_inputs
.iter()
.map(|s| s.as_str())
.collect::<Vec<_>>()
);
assert_eq!(result.turns_removed, 0);
assert!(!result.summary_written);
assert_eq!(llm.calls(), 1);
}
// ------------------------------------------------------------------
// 7. compact_to_workspace without workspace falls back to truncation
// ------------------------------------------------------------------
@@ -616,6 +654,43 @@ mod tests {
assert_eq!(result.turns_removed, 0);
}
#[cfg(feature = "libsql")]
#[tokio::test]
async fn test_compact_to_workspace_preserves_turns_when_workspace_write_fails() {
let llm = Arc::new(StubLlm::new("unused"));
let compactor = make_compactor(llm.clone());
let mut thread = make_thread(20);
let original_inputs: Vec<String> =
thread.turns.iter().map(|t| t.user_input.clone()).collect();
let workspace = make_unmigrated_workspace().await;
let result = compactor
.compact(
&mut thread,
CompactionStrategy::MoveToWorkspace,
Some(&workspace),
)
.await
.expect("compact should succeed even when workspace write fails");
// On archival failure, no turns should be removed.
assert_eq!(thread.turns.len(), 20);
assert_eq!(
thread
.turns
.iter()
.map(|t| t.user_input.as_str())
.collect::<Vec<_>>(),
original_inputs
.iter()
.map(|s| s.as_str())
.collect::<Vec<_>>()
);
assert_eq!(result.turns_removed, 0);
assert!(!result.summary_written);
assert_eq!(llm.calls(), 0);
}
// ------------------------------------------------------------------
// 9. format_turns_for_storage includes tool calls
// ------------------------------------------------------------------
+128 -31
View File
@@ -50,8 +50,18 @@ impl Agent {
// Load workspace system prompt (identity files: AGENTS.md, SOUL.md, etc.)
// In group chats, MEMORY.md is excluded to prevent leaking personal context.
// Resolve the user's timezone
let user_tz = crate::timezone::resolve_timezone(
message.timezone.as_deref(),
None, // user setting lookup can be added later
&self.config.default_timezone,
);
let system_prompt = if let Some(ws) = self.workspace() {
match ws.system_prompt_for_context(is_group_chat).await {
match ws
.system_prompt_for_context_tz(is_group_chat, user_tz)
.await
{
Ok(prompt) if !prompt.is_empty() => Some(prompt),
Ok(_) => None,
Err(e) => {
@@ -103,7 +113,7 @@ impl Agent {
None
};
let mut reasoning = Reasoning::new(self.llm().clone(), self.safety().clone())
let mut reasoning = Reasoning::new(self.llm().clone())
.with_channel(message.channel.clone())
.with_model_name(self.llm().active_model_name())
.with_group_chat(is_group_chat);
@@ -130,6 +140,7 @@ impl Agent {
let mut job_ctx =
JobContext::with_user(&message.user_id, "chat", "Interactive chat session");
job_ctx.http_interceptor = self.deps.http_interceptor.clone();
job_ctx.user_timezone = user_tz.name().to_string();
// Build system prompts once for this turn. Two variants: with tools
// (normal iterations) and without (force_text final iteration).
@@ -670,8 +681,53 @@ impl Agent {
.into())
});
// Send ToolResult preview
if let Ok(ref output) = tool_result
// Detect image generation sentinel in tool output
// (only from image tools — avoids parsing all tool outputs)
let is_image_sentinel = if let Ok(ref output) = tool_result
&& matches!(tc.name.as_str(), "image_generate" | "image_edit")
{
if let Ok(sentinel) =
serde_json::from_str::<serde_json::Value>(output)
&& sentinel.get("type").and_then(|v| v.as_str())
== Some("image_generated")
{
let data_url = sentinel
.get("data")
.and_then(|v| v.as_str())
.unwrap_or_default()
.to_string();
let path = sentinel
.get("path")
.and_then(|v| v.as_str())
.map(String::from);
// Skip broadcasting if data_url is empty to avoid
// sending a broken ImageGenerated SSE event.
if data_url.is_empty() {
tracing::warn!(
"Image generation sentinel has empty data URL, skipping broadcast"
);
} else {
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::ImageGenerated { data_url, path },
&message.metadata,
)
.await;
}
true
} else {
false
}
} else {
false
};
// Send ToolResult preview (skip for image sentinels to avoid
// broadcasting multi-MB base64 data as a preview)
if !is_image_sentinel
&& let Ok(ref output) = tool_result
&& !output.is_empty()
{
let _ = self
@@ -687,23 +743,6 @@ impl Agent {
.await;
}
// Record result in thread
{
let mut sess = session.lock().await;
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());
}
}
}
}
// Check for auth awaiting — defer the return
// until all results are recorded.
if deferred_auth.is_none()
@@ -743,6 +782,7 @@ impl Agent {
}
// Sanitize and add tool result to context
let is_tool_error = tool_result.is_err();
let result_content = match tool_result {
Ok(output) => {
let sanitized =
@@ -756,6 +796,23 @@ impl Agent {
Err(e) => format!("Tool '{}' failed: {}", tc.name, e),
};
// Record sanitized result in thread so messages()
// and persist_tool_calls() use cleaned content.
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id)
&& let Some(turn) = thread.last_turn_mut()
{
if is_tool_error {
turn.record_tool_error(result_content.clone());
} else {
turn.record_tool_result(serde_json::json!(
result_content
));
}
}
}
context_messages.push(ChatMessage::tool_result(
&tc.id,
&tc.name,
@@ -785,6 +842,7 @@ impl Agent {
tool_call_id: tc.id.clone(),
context_messages: context_messages.clone(),
deferred_tool_calls: tool_calls[approval_idx + 1..].to_vec(),
user_timezone: Some(user_tz.name().to_string()),
};
return Ok(AgenticLoopResult::NeedApproval { pending });
@@ -1146,6 +1204,7 @@ mod tests {
max_actions_per_hour: None,
max_tool_iterations: 50,
auto_approve_tools: false,
default_timezone: "UTC".to_string(),
},
deps,
Arc::new(ChannelManager::new()),
@@ -1248,6 +1307,7 @@ mod tests {
arguments: serde_json::json!({"message": "done"}),
},
],
user_timezone: None,
};
let json = serde_json::to_string(&pending).expect("serialize");
@@ -1595,12 +1655,8 @@ mod tests {
use crate::testing::StubLlm;
let stub = Arc::new(StubLlm::failing_non_transient("ctx-bomb"));
let safety = Arc::new(SafetyLayer::new(&SafetyConfig {
max_output_length: 100_000,
injection_check_enabled: false,
}));
let reasoning = Reasoning::new(stub.clone(), safety);
let reasoning = Reasoning::new(stub.clone());
// Build a fat context with lots of history.
let messages = vec![
@@ -1710,11 +1766,7 @@ mod tests {
use crate::llm::{Reasoning, ReasoningContext, RespondResult, ToolDefinition};
let provider = Arc::new(AlwaysToolCallProvider);
let safety = Arc::new(SafetyLayer::new(&SafetyConfig {
max_output_length: 100_000,
injection_check_enabled: false,
}));
let reasoning = Reasoning::new(provider, safety);
let reasoning = Reasoning::new(provider);
let tool_def = ToolDefinition {
name: "echo".to_string(),
@@ -1900,6 +1952,7 @@ mod tests {
max_actions_per_hour: None,
max_tool_iterations,
auto_approve_tools: true,
default_timezone: "UTC".to_string(),
},
deps,
Arc::new(ChannelManager::new()),
@@ -2015,6 +2068,7 @@ mod tests {
max_actions_per_hour: None,
max_tool_iterations: max_iter,
auto_approve_tools: true,
default_timezone: "UTC".to_string(),
},
deps,
Arc::new(ChannelManager::new()),
@@ -2116,4 +2170,47 @@ mod tests {
"Error should include the underlying reason, got: {formatted}"
);
}
#[test]
fn test_image_sentinel_empty_data_url_should_be_skipped() {
// Regression: unwrap_or_default() on missing "data" field produces an empty
// string. Broadcasting an empty data_url would send a broken SSE event.
let sentinel = serde_json::json!({
"type": "image_generated",
"path": "/tmp/image.png"
// "data" field is missing
});
let data_url = sentinel
.get("data")
.and_then(|v| v.as_str())
.unwrap_or_default()
.to_string();
assert!(
data_url.is_empty(),
"Missing 'data' field should produce empty string"
);
// The fix: empty data_url means we skip broadcasting
}
#[test]
fn test_image_sentinel_present_data_url_is_valid() {
let sentinel = serde_json::json!({
"type": "image_generated",
"data": "data:image/png;base64,abc123",
"path": "/tmp/image.png"
});
let data_url = sentinel
.get("data")
.and_then(|v| v.as_str())
.unwrap_or_default()
.to_string();
assert!(
!data_url.is_empty(),
"Present 'data' field should produce non-empty string"
);
}
}
+114 -8
View File
@@ -31,7 +31,6 @@ use tokio::sync::mpsc;
use crate::channels::OutgoingResponse;
use crate::db::Database;
use crate::llm::{ChatMessage, CompletionRequest, LlmProvider, Reasoning};
use crate::safety::SafetyLayer;
use crate::workspace::Workspace;
use crate::workspace::hygiene::HygieneConfig;
@@ -48,6 +47,12 @@ pub struct HeartbeatConfig {
pub notify_user_id: Option<String>,
/// Channel to notify on heartbeat findings.
pub notify_channel: Option<String>,
/// Hour (0-23) when quiet hours start.
pub quiet_hours_start: Option<u32>,
/// Hour (0-23) when quiet hours end.
pub quiet_hours_end: Option<u32>,
/// Timezone for quiet hours evaluation (IANA name).
pub timezone: Option<String>,
}
impl Default for HeartbeatConfig {
@@ -58,6 +63,9 @@ impl Default for HeartbeatConfig {
max_failures: 3,
notify_user_id: None,
notify_channel: None,
quiet_hours_start: None,
quiet_hours_end: None,
timezone: None,
}
}
}
@@ -75,6 +83,26 @@ impl HeartbeatConfig {
self
}
/// Check whether the current time falls within configured quiet hours.
pub fn is_quiet_hours(&self) -> bool {
use chrono::Timelike;
let (Some(start), Some(end)) = (self.quiet_hours_start, self.quiet_hours_end) else {
return false;
};
let tz = self
.timezone
.as_deref()
.and_then(crate::timezone::parse_timezone)
.unwrap_or(chrono_tz::UTC);
let now_hour = crate::timezone::now_in_tz(tz).hour();
if start <= end {
now_hour >= start && now_hour < end
} else {
// Wraps midnight, e.g. 22..06
now_hour >= start || now_hour < end
}
}
/// Set the notification target.
pub fn with_notify(mut self, user_id: impl Into<String>, channel: impl Into<String>) -> Self {
self.notify_user_id = Some(user_id.into());
@@ -102,7 +130,6 @@ pub struct HeartbeatRunner {
hygiene_config: HygieneConfig,
workspace: Arc<Workspace>,
llm: Arc<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
store: Option<Arc<dyn Database>>,
consecutive_failures: u32,
@@ -115,14 +142,12 @@ impl HeartbeatRunner {
hygiene_config: HygieneConfig,
workspace: Arc<Workspace>,
llm: Arc<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
) -> Self {
Self {
config,
hygiene_config,
workspace,
llm,
safety,
response_tx: None,
store: None,
consecutive_failures: 0,
@@ -162,6 +187,12 @@ impl HeartbeatRunner {
loop {
interval.tick().await;
// Skip during quiet hours
if self.config.is_quiet_hours() {
tracing::debug!("Heartbeat skipped: quiet hours");
continue;
}
// Run memory hygiene in the background so it never delays the
// heartbeat checklist. Failures are logged inside run_if_due.
let hygiene_workspace = Arc::clone(&self.workspace);
@@ -272,7 +303,7 @@ impl HeartbeatRunner {
.with_max_tokens(max_tokens)
.with_temperature(0.3);
let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone());
let reasoning = Reasoning::new(self.llm.clone());
let (content, _usage) = match reasoning.complete(request).await {
Ok(r) => r,
Err(e) => return HeartbeatResult::Failed(format!("LLM call failed: {}", e)),
@@ -386,11 +417,10 @@ pub fn spawn_heartbeat(
hygiene_config: HygieneConfig,
workspace: Arc<Workspace>,
llm: Arc<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
store: Option<Arc<dyn Database>>,
) -> tokio::task::JoinHandle<()> {
let mut runner = HeartbeatRunner::new(config, hygiene_config, workspace, llm, safety);
let mut runner = HeartbeatRunner::new(config, hygiene_config, workspace, llm);
if let Some(tx) = response_tx {
runner = runner.with_response_channel(tx);
}
@@ -532,6 +562,83 @@ mod tests {
assert!(!is_effectively_empty(content));
}
// ==================== quiet hours ====================
#[test]
fn test_quiet_hours_inside() {
use chrono::{Timelike, Utc};
let now_utc = Utc::now();
let hour = now_utc.hour();
let start = hour;
let end = (hour + 1) % 24;
let config = HeartbeatConfig {
quiet_hours_start: Some(start),
quiet_hours_end: Some(end),
timezone: Some("UTC".to_string()),
..HeartbeatConfig::default()
};
// Current UTC hour is inside [start, end) by construction
assert!(config.is_quiet_hours());
}
#[test]
fn test_quiet_hours_outside() {
use chrono::{Timelike, Utc};
let now_utc = Utc::now();
let hour = now_utc.hour();
let start = (hour + 1) % 24;
let end = (hour + 2) % 24;
let config = HeartbeatConfig {
quiet_hours_start: Some(start),
quiet_hours_end: Some(end),
timezone: Some("UTC".to_string()),
..HeartbeatConfig::default()
};
// Current UTC hour is outside [start, end) by construction
assert!(!config.is_quiet_hours());
}
#[test]
fn test_quiet_hours_wraparound_excludes_now() {
use chrono::{Timelike, Utc};
let now_utc = Utc::now();
let hour = now_utc.hour();
// Window covers all hours except the current one
let start = (hour + 1) % 24;
let end = hour;
let config = HeartbeatConfig {
quiet_hours_start: Some(start),
quiet_hours_end: Some(end),
timezone: Some("UTC".to_string()),
..HeartbeatConfig::default()
};
assert!(!config.is_quiet_hours());
}
#[test]
fn test_quiet_hours_none_configured() {
let config = HeartbeatConfig::default();
assert!(!config.is_quiet_hours());
}
#[test]
fn test_quiet_hours_same_start_end() {
let config = HeartbeatConfig {
quiet_hours_start: Some(10),
quiet_hours_end: Some(10),
timezone: Some("UTC".to_string()),
..HeartbeatConfig::default()
};
// start == end means zero-width window, should be false
assert!(!config.is_quiet_hours());
}
#[test]
fn test_spawn_heartbeat_accepts_store_param() {
// Regression: spawn_heartbeat must accept an optional Database store
@@ -543,7 +650,6 @@ mod tests {
HygieneConfig,
Arc<crate::workspace::Workspace>,
Arc<dyn crate::llm::LlmProvider>,
Arc<crate::safety::SafetyLayer>,
Option<tokio::sync::mpsc::Sender<crate::channels::OutgoingResponse>>,
Option<Arc<dyn crate::db::Database>>,
) -> tokio::task::JoinHandle<()> = spawn_heartbeat;
+87 -9
View File
@@ -57,7 +57,11 @@ pub struct Routine {
#[serde(tag = "type", rename_all = "snake_case")]
pub enum Trigger {
/// Fire on a cron schedule (e.g. "0 9 * * MON-FRI" or "every 2h").
Cron { schedule: String },
Cron {
schedule: String,
#[serde(default)]
timezone: Option<String>,
},
/// Fire when a channel message matches a pattern.
Event {
/// Optional channel filter (e.g. "telegram", "slack").
@@ -99,7 +103,21 @@ impl Trigger {
field: "schedule".into(),
})?
.to_string();
Ok(Trigger::Cron { schedule })
let timezone = config
.get("timezone")
.and_then(|v| v.as_str())
.and_then(|tz| {
if crate::timezone::parse_timezone(tz).is_some() {
Some(tz.to_string())
} else {
tracing::warn!(
"Ignoring invalid timezone '{}' from DB for cron trigger",
tz
);
None
}
});
Ok(Trigger::Cron { schedule, timezone })
}
"event" => {
let pattern = config
@@ -137,7 +155,10 @@ impl Trigger {
/// Serialize trigger-specific config to JSON for DB storage.
pub fn to_config_json(&self) -> serde_json::Value {
match self {
Trigger::Cron { schedule } => serde_json::json!({ "schedule": schedule }),
Trigger::Cron { schedule, timezone } => serde_json::json!({
"schedule": schedule,
"timezone": timezone,
}),
Trigger::Event { channel, pattern } => serde_json::json!({
"pattern": pattern,
"channel": channel,
@@ -415,12 +436,25 @@ pub fn content_hash(content: &str) -> u64 {
}
/// Parse a cron expression and compute the next fire time from now.
pub fn next_cron_fire(schedule: &str) -> Result<Option<DateTime<Utc>>, RoutineError> {
///
/// When `timezone` is provided and valid, the schedule is evaluated in that
/// timezone and the result is converted back to UTC. Otherwise UTC is used.
pub fn next_cron_fire(
schedule: &str,
timezone: Option<&str>,
) -> Result<Option<DateTime<Utc>>, RoutineError> {
let cron_schedule =
cron::Schedule::from_str(schedule).map_err(|e| RoutineError::InvalidCron {
reason: e.to_string(),
})?;
Ok(cron_schedule.upcoming(Utc).next())
if let Some(tz) = timezone.and_then(crate::timezone::parse_timezone) {
Ok(cron_schedule
.upcoming(tz)
.next()
.map(|dt| dt.with_timezone(&Utc)))
} else {
Ok(cron_schedule.upcoming(Utc).next())
}
}
#[cfg(test)]
@@ -433,10 +467,11 @@ mod tests {
fn test_trigger_roundtrip() {
let trigger = Trigger::Cron {
schedule: "0 9 * * MON-FRI".to_string(),
timezone: None,
};
let json = trigger.to_config_json();
let parsed = Trigger::from_db("cron", json).expect("parse cron");
assert!(matches!(parsed, Trigger::Cron { schedule } if schedule == "0 9 * * MON-FRI"));
assert!(matches!(parsed, Trigger::Cron { schedule, .. } if schedule == "0 9 * * MON-FRI"));
}
#[test]
@@ -509,16 +544,58 @@ mod tests {
#[test]
fn test_next_cron_fire_valid() {
// Every minute should always have a next fire
let next = next_cron_fire("* * * * * *").expect("valid cron");
let next = next_cron_fire("* * * * * *", None).expect("valid cron");
assert!(next.is_some());
}
#[test]
fn test_next_cron_fire_invalid() {
let result = next_cron_fire("not a cron");
let result = next_cron_fire("not a cron", None);
assert!(result.is_err());
}
#[test]
fn test_trigger_cron_timezone_roundtrip() {
let trigger = Trigger::Cron {
schedule: "0 9 * * MON-FRI".to_string(),
timezone: Some("America/New_York".to_string()),
};
let json = trigger.to_config_json();
let parsed = Trigger::from_db("cron", json).expect("parse cron");
assert!(matches!(parsed, Trigger::Cron { schedule, timezone }
if schedule == "0 9 * * MON-FRI"
&& timezone.as_deref() == Some("America/New_York")));
}
#[test]
fn test_trigger_cron_no_timezone_backward_compat() {
let json = serde_json::json!({"schedule": "0 9 * * *"});
let parsed = Trigger::from_db("cron", json).expect("parse cron");
assert!(matches!(parsed, Trigger::Cron { timezone, .. } if timezone.is_none()));
}
#[test]
fn test_trigger_cron_invalid_timezone_coerced_to_none() {
let json = serde_json::json!({"schedule": "0 9 * * *", "timezone": "Fake/Zone"});
let parsed = Trigger::from_db("cron", json).expect("parse cron");
assert!(
matches!(parsed, Trigger::Cron { timezone, .. } if timezone.is_none()),
"invalid timezone should be coerced to None"
);
}
#[test]
fn test_next_cron_fire_with_timezone() {
let next_utc = next_cron_fire("0 0 9 * * * *", None)
.expect("valid cron")
.expect("has next");
let next_est = next_cron_fire("0 0 9 * * * *", Some("America/New_York"))
.expect("valid cron")
.expect("has next");
// EST is UTC-5 (or EDT UTC-4), so the UTC result should differ
assert_ne!(next_utc, next_est, "timezone should shift the fire time");
}
#[test]
fn test_guardrails_default() {
let g = RoutineGuardrails::default();
@@ -531,7 +608,8 @@ mod tests {
fn test_trigger_type_tag() {
assert_eq!(
Trigger::Cron {
schedule: String::new()
schedule: String::new(),
timezone: None,
}
.type_tag(),
"cron"
+12 -13
View File
@@ -170,7 +170,7 @@ impl RoutineEngine {
continue;
}
let detail = if let Trigger::Cron { ref schedule } = routine.trigger {
let detail = if let Trigger::Cron { ref schedule, .. } = routine.trigger {
Some(schedule.clone())
} else {
None
@@ -380,8 +380,12 @@ async fn execute_routine(ctx: EngineContext, routine: Routine, run: RoutineRun)
// Update routine runtime state
let now = Utc::now();
let next_fire = if let Trigger::Cron { ref schedule } = routine.trigger {
next_cron_fire(schedule).unwrap_or(None)
let next_fire = if let Trigger::Cron {
ref schedule,
ref timezone,
} = routine.trigger
{
next_cron_fire(schedule, timezone.as_deref()).unwrap_or(None)
} else {
None
};
@@ -488,18 +492,13 @@ async fn execute_full_job(
reason: "scheduler not available".to_string(),
})?;
// Set the message tool's default channel/target from the routine's notify config
// so the LLM can send results without triggering cross-channel approval.
// TODO: This mutates shared global state and can race with concurrent jobs.
// Move notify config into JobContext metadata and apply per-job instead.
let mut metadata = serde_json::json!({ "max_iterations": max_iterations });
// Carry the routine's notify config in job metadata so the message tool
// can resolve channel/target per-job without global state mutation.
if let Some(channel) = &routine.notify.channel {
scheduler
.tools()
.set_message_tool_context(Some(channel.clone()), Some(routine.notify.user.clone()))
.await;
metadata["notify_channel"] = serde_json::json!(channel);
}
let metadata = serde_json::json!({ "max_iterations": max_iterations });
metadata["notify_user"] = serde_json::json!(&routine.notify.user);
// Build approval context: UnlessAutoApproved tools are auto-approved for routines;
// Always tools require explicit listing in tool_permissions.
+324 -11
View File
@@ -16,6 +16,7 @@ use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use crate::channels::web::util::truncate_preview;
use crate::llm::{ChatMessage, ToolCall};
/// A session containing one or more threads.
@@ -164,6 +165,10 @@ pub struct PendingApproval {
/// executed yet when approval was requested.
#[serde(default)]
pub deferred_tool_calls: Vec<ToolCall>,
/// User timezone at the time the approval was requested, so it persists
/// through the approval flow even if the approval message lacks timezone.
#[serde(default)]
pub user_timezone: Option<String>,
}
/// A conversation thread within a session.
@@ -316,7 +321,13 @@ impl Thread {
}
}
/// Get all messages for context building.
/// Get all messages for context building, including tool call history.
///
/// Emits the full LLM-compatible message sequence per turn:
/// `user → [assistant_with_tool_calls → tool_result*] → assistant`
///
/// This ensures the LLM sees prior tool executions and won't re-attempt
/// completed actions in subsequent turns.
pub fn messages(&self) -> Vec<ChatMessage> {
let mut messages = Vec::new();
for turn in &self.turns {
@@ -328,6 +339,42 @@ impl Thread {
turn.image_content_parts.clone(),
));
}
if !turn.tool_calls.is_empty() {
// Build ToolCall objects with synthetic stable IDs
let tool_calls: Vec<ToolCall> = turn
.tool_calls
.iter()
.enumerate()
.map(|(i, tc)| ToolCall {
id: format!("turn{}_{}", turn.turn_number, i),
name: tc.name.clone(),
arguments: tc.parameters.clone(),
})
.collect();
// Assistant message declaring the tool calls (no text content)
messages.push(ChatMessage::assistant_with_tool_calls(None, tool_calls));
// Individual tool result messages, truncated to limit context size.
for (i, tc) in turn.tool_calls.iter().enumerate() {
let call_id = format!("turn{}_{}", turn.turn_number, i);
let content = if let Some(ref err) = tc.error {
// .error already contains the full error text;
// pass through without wrapping to avoid double-prefix.
truncate_preview(err, 1000)
} else if let Some(ref res) = tc.result {
let raw = match res {
serde_json::Value::String(s) => s.clone(),
other => other.to_string(),
};
truncate_preview(&raw, 1000)
} else {
"OK".to_string()
};
messages.push(ChatMessage::tool_result(call_id, &tc.name, content));
}
}
if let Some(ref response) = turn.response {
messages.push(ChatMessage::assistant(response));
}
@@ -349,13 +396,16 @@ impl Thread {
/// Restore thread state from a checkpoint's messages.
///
/// Clears existing turns and rebuilds from message pairs.
/// Messages should alternate: user, assistant, user, assistant...
/// Clears existing turns and rebuilds from the message sequence.
/// Handles the full message pattern including tool messages:
/// `user → [assistant_with_tool_calls → tool_result*] → assistant`
///
/// Also supports the legacy pattern (user/assistant pairs only) for
/// backward compatibility with old checkpoint data.
pub fn restore_from_messages(&mut self, messages: Vec<ChatMessage>) {
self.turns.clear();
self.state = ThreadState::Idle;
// Messages alternate: user, assistant, user, assistant...
let mut iter = messages.into_iter().peekable();
let mut turn_number = 0;
@@ -363,18 +413,58 @@ impl Thread {
if msg.role == crate::llm::Role::User {
let mut turn = Turn::new(turn_number, &msg.content);
// Check if next is assistant response
if let Some(next) = iter.peek()
&& next.role == crate::llm::Role::Assistant
{
// iter.next() is guaranteed Some after a successful peek()
if let Some(response) = iter.next() {
turn.complete(&response.content);
// Consume tool call sequences (assistant_with_tool_calls + tool_results).
// A single turn may contain multiple rounds of tool calls, so we
// track the cumulative base index into turn.tool_calls.
while let Some(next) = iter.peek() {
if next.role == crate::llm::Role::Assistant && next.tool_calls.is_some() {
let call_base_idx = turn.tool_calls.len();
if let Some(assistant_msg) = iter.next()
&& let Some(ref tcs) = assistant_msg.tool_calls
{
for tc in tcs {
turn.record_tool_call(&tc.name, tc.arguments.clone());
}
}
// Consume the corresponding tool_result messages,
// indexing relative to this batch's base offset.
let mut pos = 0;
while let Some(tr) = iter.peek() {
if tr.role != crate::llm::Role::Tool {
break;
}
if let Some(tool_msg) = iter.next() {
let idx = call_base_idx + pos;
if idx < turn.tool_calls.len() {
// Store as result — the error/success distinction
// is for the live turn only; restored context just
// needs the content the LLM originally saw.
turn.tool_calls[idx].result =
Some(serde_json::Value::String(tool_msg.content.clone()));
}
}
pos += 1;
}
} else {
break;
}
}
// Check if next is the final assistant response for this turn
let is_final_assistant = iter.peek().is_some_and(|n| {
n.role == crate::llm::Role::Assistant && n.tool_calls.is_none()
});
if is_final_assistant && let Some(response) = iter.next() {
turn.complete(&response.content);
}
self.turns.push(turn);
turn_number += 1;
} else {
// Skip non-user messages that aren't anchored to a turn
continue;
}
}
@@ -976,6 +1066,7 @@ mod tests {
tool_call_id: "call_123".to_string(),
context_messages: vec![ChatMessage::user("do it")],
deferred_tool_calls: vec![],
user_timezone: None,
};
thread.await_approval(approval);
@@ -1001,6 +1092,7 @@ mod tests {
tool_call_id: "call_456".to_string(),
context_messages: vec![],
deferred_tool_calls: vec![],
user_timezone: None,
};
thread.await_approval(approval);
@@ -1029,4 +1121,225 @@ mod tests {
ThreadState::Processing
);
}
// Regression tests for #568: tool call history must survive hydration.
#[test]
fn test_messages_includes_tool_calls() {
let mut thread = Thread::new(Uuid::new_v4());
thread.start_turn("Search for X");
{
let turn = thread.turns.last_mut().unwrap();
turn.record_tool_call("memory_search", serde_json::json!({"query": "X"}));
turn.record_tool_result(serde_json::json!("Found X in doc.md"));
}
thread.complete_turn("I found X in doc.md.");
let messages = thread.messages();
// user + assistant_with_tool_calls + tool_result + assistant = 4
assert_eq!(messages.len(), 4);
assert_eq!(messages[0].role, crate::llm::Role::User);
assert_eq!(messages[0].content, "Search for X");
assert_eq!(messages[1].role, crate::llm::Role::Assistant);
assert!(messages[1].tool_calls.is_some());
let tcs = messages[1].tool_calls.as_ref().unwrap();
assert_eq!(tcs.len(), 1);
assert_eq!(tcs[0].name, "memory_search");
assert_eq!(messages[2].role, crate::llm::Role::Tool);
assert!(messages[2].content.contains("Found X"));
assert_eq!(messages[3].role, crate::llm::Role::Assistant);
assert_eq!(messages[3].content, "I found X in doc.md.");
}
#[test]
fn test_messages_multiple_tool_calls_per_turn() {
let mut thread = Thread::new(Uuid::new_v4());
thread.start_turn("Do two things");
{
let turn = thread.turns.last_mut().unwrap();
turn.record_tool_call("echo", serde_json::json!({"msg": "a"}));
turn.record_tool_result(serde_json::json!("a"));
turn.record_tool_call("time", serde_json::json!({}));
turn.record_tool_error("timeout");
}
thread.complete_turn("Done.");
let messages = thread.messages();
// user + assistant_with_calls(2) + tool_result + tool_result + assistant = 5
assert_eq!(messages.len(), 5);
let tcs = messages[1].tool_calls.as_ref().unwrap();
assert_eq!(tcs.len(), 2);
// First tool: success
assert_eq!(messages[2].content, "a");
// Second tool: error (passed through directly, no wrapping)
assert!(messages[3].content.contains("timeout"));
}
#[test]
fn test_restore_from_messages_with_tool_calls() {
let mut thread = Thread::new(Uuid::new_v4());
// Build a message sequence with tool calls
let tc = ToolCall {
id: "call_0".to_string(),
name: "search".to_string(),
arguments: serde_json::json!({"q": "test"}),
};
let messages = vec![
ChatMessage::user("Find test"),
ChatMessage::assistant_with_tool_calls(None, vec![tc]),
ChatMessage::tool_result("call_0", "search", "result: found"),
ChatMessage::assistant("Found it."),
];
thread.restore_from_messages(messages);
assert_eq!(thread.turns.len(), 1);
let turn = &thread.turns[0];
assert_eq!(turn.user_input, "Find test");
assert_eq!(turn.tool_calls.len(), 1);
assert_eq!(turn.tool_calls[0].name, "search");
assert_eq!(
turn.tool_calls[0].result,
Some(serde_json::Value::String("result: found".to_string()))
);
assert_eq!(turn.response, Some("Found it.".to_string()));
}
#[test]
fn test_restore_from_messages_with_tool_error() {
let mut thread = Thread::new(Uuid::new_v4());
let tc = ToolCall {
id: "call_0".to_string(),
name: "http".to_string(),
arguments: serde_json::json!({}),
};
let messages = vec![
ChatMessage::user("Fetch URL"),
ChatMessage::assistant_with_tool_calls(None, vec![tc]),
ChatMessage::tool_result("call_0", "http", "Error: timeout"),
ChatMessage::assistant("The request timed out."),
];
thread.restore_from_messages(messages);
// restore_from_messages stores all tool content as result (not error),
// because it can't reliably distinguish errors from results that happen
// to start with "Error: ". The content is preserved for LLM context.
let turn = &thread.turns[0];
assert_eq!(
turn.tool_calls[0].result,
Some(serde_json::Value::String("Error: timeout".to_string()))
);
}
#[test]
fn test_messages_round_trip_with_tools() {
// Build a thread with tool calls, get messages(), restore, get messages() again
// The two message sequences should be equivalent.
let mut thread = Thread::new(Uuid::new_v4());
thread.start_turn("Do search");
{
let turn = thread.turns.last_mut().unwrap();
turn.record_tool_call("search", serde_json::json!({"q": "test"}));
turn.record_tool_result(serde_json::json!("found"));
}
thread.complete_turn("Here are results.");
let messages_original = thread.messages();
// Restore into a new thread
let mut thread2 = Thread::new(Uuid::new_v4());
thread2.restore_from_messages(messages_original.clone());
let messages_restored = thread2.messages();
// Same number of messages
assert_eq!(messages_original.len(), messages_restored.len());
// Same roles
for (orig, rest) in messages_original.iter().zip(messages_restored.iter()) {
assert_eq!(orig.role, rest.role);
}
// Same final response
assert_eq!(
messages_original.last().unwrap().content,
messages_restored.last().unwrap().content
);
}
#[test]
fn test_restore_multi_stage_tool_calls() {
let mut thread = Thread::new(Uuid::new_v4());
let tc1 = ToolCall {
id: "call_a".to_string(),
name: "search".to_string(),
arguments: serde_json::json!({"q": "data"}),
};
let tc2 = ToolCall {
id: "call_b".to_string(),
name: "write".to_string(),
arguments: serde_json::json!({"path": "out.txt"}),
};
let messages = vec![
ChatMessage::user("Find and save"),
ChatMessage::assistant_with_tool_calls(None, vec![tc1]),
ChatMessage::tool_result("call_a", "search", "found data"),
ChatMessage::assistant_with_tool_calls(None, vec![tc2]),
ChatMessage::tool_result("call_b", "write", "written"),
ChatMessage::assistant("Done, saved to out.txt"),
];
thread.restore_from_messages(messages);
assert_eq!(thread.turns.len(), 1);
let turn = &thread.turns[0];
assert_eq!(turn.tool_calls.len(), 2);
assert_eq!(turn.tool_calls[0].name, "search");
assert_eq!(turn.tool_calls[1].name, "write");
assert_eq!(
turn.tool_calls[0].result,
Some(serde_json::Value::String("found data".to_string()))
);
assert_eq!(
turn.tool_calls[1].result,
Some(serde_json::Value::String("written".to_string()))
);
assert_eq!(turn.response, Some("Done, saved to out.txt".to_string()));
}
#[test]
fn test_messages_truncates_large_tool_results() {
let mut thread = Thread::new(Uuid::new_v4());
thread.start_turn("Read big file");
{
let turn = thread.turns.last_mut().unwrap();
turn.record_tool_call("read_file", serde_json::json!({"path": "big.txt"}));
let big_result = "x".repeat(2000);
turn.record_tool_result(serde_json::json!(big_result));
}
thread.complete_turn("Here's the file content.");
let messages = thread.messages();
let tool_result_content = &messages[2].content;
assert!(
tool_result_content.len() <= 1010,
"Tool result should be truncated, got {} chars",
tool_result_content.len()
);
assert!(tool_result_content.ends_with("..."));
}
}
+305 -60
View File
@@ -20,7 +20,7 @@ use crate::channels::web::util::truncate_preview;
use crate::channels::{IncomingMessage, StatusUpdate};
use crate::context::JobContext;
use crate::error::Error;
use crate::llm::ChatMessage;
use crate::llm::{ChatMessage, ToolCall};
use crate::tools::redact_params;
impl Agent {
@@ -66,16 +66,7 @@ impl Agent {
.await
.unwrap_or_default();
msg_count = db_messages.len();
chat_messages = db_messages
.iter()
.filter_map(|m| match m.role.as_str() {
"user" => Some(ChatMessage::user(&m.content)),
"assistant" => Some(ChatMessage::assistant(&m.content)),
// tool_calls rows are UI metadata (tool name + preview),
// not part of the LLM conversation context.
_ => None,
})
.collect();
chat_messages = rebuild_chat_messages_from_db(&db_messages);
} else {
msg_count = 0;
}
@@ -230,7 +221,7 @@ impl Agent {
)
.await;
let compactor = ContextCompactor::new(self.llm().clone(), self.safety().clone());
let compactor = ContextCompactor::new(self.llm().clone());
if let Err(e) = compactor
.compact(thread, strategy, self.workspace().map(|w| w.as_ref()))
.await
@@ -340,10 +331,10 @@ impl Agent {
};
thread.complete_turn(&response);
let tool_calls = thread
let (turn_number, tool_calls) = thread
.turns
.last()
.map(|t| t.tool_calls.clone())
.map(|t| (t.turn_number, t.tool_calls.clone()))
.unwrap_or_default();
let _ = self
.channels
@@ -355,7 +346,7 @@ impl Agent {
.await;
// Persist tool calls then assistant response (user message already persisted at turn start)
self.persist_tool_calls(thread_id, &message.user_id, &tool_calls)
self.persist_tool_calls(thread_id, &message.user_id, turn_number, &tool_calls)
.await;
self.persist_assistant_response(thread_id, &message.user_id, &response)
.await;
@@ -464,6 +455,7 @@ impl Agent {
&self,
thread_id: Uuid,
user_id: &str,
turn_number: usize,
tool_calls: &[crate::agent::session::TurnToolCall],
) {
if tool_calls.is_empty() {
@@ -477,14 +469,24 @@ impl Agent {
let summaries: Vec<serde_json::Value> = tool_calls
.iter()
.map(|tc| {
let mut obj = serde_json::json!({ "name": tc.name });
.enumerate()
.map(|(i, tc)| {
let mut obj = serde_json::json!({
"name": tc.name,
"call_id": format!("turn{}_{}", turn_number, i),
});
if let Some(ref result) = tc.result {
let preview = match result {
serde_json::Value::String(s) => truncate_preview(s, 500),
other => truncate_preview(&other.to_string(), 500),
};
obj["result_preview"] = serde_json::Value::String(preview);
// Store full result (truncated to ~1000 chars) for LLM context rebuild
let full_result = match result {
serde_json::Value::String(s) => truncate_preview(s, 1000),
other => truncate_preview(&other.to_string(), 1000),
};
obj["result"] = serde_json::Value::String(full_result);
}
if let Some(ref error) = tc.error {
obj["error"] = serde_json::Value::String(truncate_preview(error, 200));
@@ -627,7 +629,7 @@ impl Agent {
crate::agent::context_monitor::CompactionStrategy::Summarize { keep_recent: 5 },
);
let compactor = ContextCompactor::new(self.llm().clone(), self.safety().clone());
let compactor = ContextCompactor::new(self.llm().clone());
match compactor
.compact(thread, strategy, self.workspace().map(|w| w.as_ref()))
.await
@@ -746,6 +748,16 @@ impl Agent {
let mut job_ctx =
JobContext::with_user(&message.user_id, "chat", "Interactive chat session");
job_ctx.http_interceptor = self.deps.http_interceptor.clone();
// Prefer a valid timezone from the approval message, fall back to the
// resolved timezone stored when the approval was originally requested.
let tz_candidate = message
.timezone
.as_deref()
.filter(|tz| crate::timezone::parse_timezone(tz).is_some())
.or(pending.user_timezone.as_deref());
if let Some(tz) = tz_candidate {
job_ctx.user_timezone = tz.to_string();
}
let _ = self
.channels
@@ -797,19 +809,33 @@ impl Agent {
let mut context_messages = pending.context_messages;
let deferred_tool_calls = pending.deferred_tool_calls;
// Record result in thread
// Sanitize tool result, then record the cleaned version in the
// thread. Must happen before auth intercept check which may return early.
let is_tool_error = tool_result.is_err();
let result_content = match &tool_result {
Ok(output) => {
let sanitized = self
.safety()
.sanitize_tool_output(&pending.tool_name, output);
self.safety().wrap_for_llm(
&pending.tool_name,
&sanitized.content,
sanitized.was_modified,
)
}
Err(e) => format!("Error: {}", e),
};
// Record sanitized result in thread
{
let mut sess = session.lock().await;
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());
}
if is_tool_error {
turn.record_tool_error(result_content.clone());
} else {
turn.record_tool_result(serde_json::json!(result_content));
}
}
}
@@ -831,21 +857,6 @@ impl Agent {
return Ok(SubmissionResult::response(instructions));
}
// Add tool result to context
let result_content = match tool_result {
Ok(output) => {
let sanitized = self
.safety()
.sanitize_tool_output(&pending.tool_name, &output);
self.safety().wrap_for_llm(
&pending.tool_name,
&sanitized.content,
sanitized.was_modified,
)
}
Err(e) => format!("Error: {}", e),
};
context_messages.push(ChatMessage::tool_result(
&pending.tool_call_id,
&pending.tool_name,
@@ -1050,15 +1061,31 @@ impl Agent {
.await;
}
// Record in thread
// Sanitize first, then record the cleaned version in thread.
// Must happen before auth detection which may set deferred_auth.
let is_deferred_error = deferred_result.is_err();
let deferred_content = match &deferred_result {
Ok(output) => {
let sanitized = self.safety().sanitize_tool_output(&tc.name, output);
self.safety().wrap_for_llm(
&tc.name,
&sanitized.content,
sanitized.was_modified,
)
}
Err(e) => format!("Error: {}", e),
};
// Record sanitized result in thread
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id)
&& let Some(turn) = thread.last_turn_mut()
{
match &deferred_result {
Ok(output) => turn.record_tool_result(serde_json::json!(output)),
Err(e) => turn.record_tool_error(e.to_string()),
if is_deferred_error {
turn.record_tool_error(deferred_content.clone());
} else {
turn.record_tool_result(serde_json::json!(deferred_content));
}
}
}
@@ -1080,18 +1107,6 @@ impl Agent {
deferred_auth = Some(instructions);
}
let deferred_content = match deferred_result {
Ok(output) => {
let sanitized = self.safety().sanitize_tool_output(&tc.name, &output);
self.safety().wrap_for_llm(
&tc.name,
&sanitized.content,
sanitized.was_modified,
)
}
Err(e) => format!("Error: {}", e),
};
context_messages.push(ChatMessage::tool_result(&tc.id, &tc.name, deferred_content));
}
@@ -1111,6 +1126,8 @@ impl Agent {
tool_call_id: tc.id.clone(),
context_messages: context_messages.clone(),
deferred_tool_calls: deferred_tool_calls[approval_idx + 1..].to_vec(),
// Carry forward the resolved timezone from the original pending approval
user_timezone: pending.user_timezone.clone(),
};
let request_id = new_pending.request_id;
@@ -1157,13 +1174,13 @@ impl Agent {
match result {
Ok(AgenticLoopResult::Response(response)) => {
thread.complete_turn(&response);
let tool_calls = thread
let (turn_number, tool_calls) = thread
.turns
.last()
.map(|t| t.tool_calls.clone())
.map(|t| (t.turn_number, t.tool_calls.clone()))
.unwrap_or_default();
// User message already persisted at turn start; save tool calls then assistant response
self.persist_tool_calls(thread_id, &message.user_id, &tool_calls)
self.persist_tool_calls(thread_id, &message.user_id, turn_number, &tool_calls)
.await;
self.persist_assistant_response(thread_id, &message.user_id, &response)
.await;
@@ -1478,3 +1495,231 @@ impl Agent {
}
}
}
/// Rebuild full LLM-compatible `ChatMessage` sequence from DB messages.
///
/// Parses `role="tool_calls"` rows to reconstruct `assistant_with_tool_calls`
/// and `tool_result` messages so that the LLM sees the complete tool execution
/// history on thread hydration. Falls back gracefully for legacy rows that
/// lack the enriched fields (`call_id`, `parameters`, `result`).
fn rebuild_chat_messages_from_db(
db_messages: &[crate::history::ConversationMessage],
) -> Vec<ChatMessage> {
let mut result = Vec::new();
for msg in db_messages {
match msg.role.as_str() {
"user" => result.push(ChatMessage::user(&msg.content)),
"assistant" => result.push(ChatMessage::assistant(&msg.content)),
"tool_calls" => {
// Try to parse the enriched JSON and rebuild tool messages.
if let Ok(calls) = serde_json::from_str::<Vec<serde_json::Value>>(&msg.content) {
if calls.is_empty() {
continue;
}
// Check if this is an enriched row (has call_id) or legacy
let has_call_id = calls
.first()
.and_then(|c| c.get("call_id"))
.and_then(|v| v.as_str())
.is_some();
if has_call_id {
// Build assistant_with_tool_calls + tool_result messages
let tool_calls: Vec<ToolCall> = calls
.iter()
.map(|c| ToolCall {
id: c["call_id"].as_str().unwrap_or("call_0").to_string(),
name: c["name"].as_str().unwrap_or("unknown").to_string(),
arguments: c
.get("parameters")
.cloned()
.unwrap_or(serde_json::json!({})),
})
.collect();
// The assistant text for tool_calls is always None here;
// the final assistant response comes as a separate
// "assistant" row after this tool_calls row.
result.push(ChatMessage::assistant_with_tool_calls(None, tool_calls));
// Emit tool_result messages for each call
for c in &calls {
let call_id = c["call_id"].as_str().unwrap_or("call_0").to_string();
let name = c["name"].as_str().unwrap_or("unknown").to_string();
let content = if let Some(err) = c.get("error").and_then(|v| v.as_str())
{
format!("Error: {}", err)
} else if let Some(res) = c.get("result").and_then(|v| v.as_str()) {
res.to_string()
} else if let Some(preview) =
c.get("result_preview").and_then(|v| v.as_str())
{
preview.to_string()
} else {
"OK".to_string()
};
result.push(ChatMessage::tool_result(call_id, name, content));
}
}
// Legacy rows without call_id: skip (will appear as
// simple user/assistant pairs, same as before this fix).
}
}
_ => {} // Skip unknown roles
}
}
result
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_rebuild_chat_messages_user_assistant_only() {
let messages = vec![
make_db_msg("user", "Hello"),
make_db_msg("assistant", "Hi there!"),
];
let result = rebuild_chat_messages_from_db(&messages);
assert_eq!(result.len(), 2);
assert_eq!(result[0].role, crate::llm::Role::User);
assert_eq!(result[1].role, crate::llm::Role::Assistant);
}
#[test]
fn test_rebuild_chat_messages_with_enriched_tool_calls() {
let tool_json = serde_json::json!([
{
"name": "memory_search",
"call_id": "call_0",
"parameters": {"query": "test"},
"result": "Found 3 results",
"result_preview": "Found 3 re..."
},
{
"name": "echo",
"call_id": "call_1",
"parameters": {"message": "hi"},
"error": "timeout"
}
]);
let messages = vec![
make_db_msg("user", "Search for test"),
make_db_msg("tool_calls", &tool_json.to_string()),
make_db_msg("assistant", "I found some results."),
];
let result = rebuild_chat_messages_from_db(&messages);
// user + assistant_with_tool_calls + tool_result*2 + assistant
assert_eq!(result.len(), 5);
// user
assert_eq!(result[0].role, crate::llm::Role::User);
// assistant with tool_calls
assert_eq!(result[1].role, crate::llm::Role::Assistant);
assert!(result[1].tool_calls.is_some());
let tcs = result[1].tool_calls.as_ref().unwrap();
assert_eq!(tcs.len(), 2);
assert_eq!(tcs[0].name, "memory_search");
assert_eq!(tcs[0].id, "call_0");
assert_eq!(tcs[1].name, "echo");
// tool results
assert_eq!(result[2].role, crate::llm::Role::Tool);
assert_eq!(result[2].tool_call_id, Some("call_0".to_string()));
assert!(result[2].content.contains("Found 3 results"));
assert_eq!(result[3].role, crate::llm::Role::Tool);
assert_eq!(result[3].tool_call_id, Some("call_1".to_string()));
assert!(result[3].content.contains("Error: timeout"));
// final assistant
assert_eq!(result[4].role, crate::llm::Role::Assistant);
assert_eq!(result[4].content, "I found some results.");
}
#[test]
fn test_rebuild_chat_messages_legacy_tool_calls_skipped() {
// Legacy format: no call_id field
let tool_json = serde_json::json!([
{"name": "echo", "result_preview": "hello"}
]);
let messages = vec![
make_db_msg("user", "Hi"),
make_db_msg("tool_calls", &tool_json.to_string()),
make_db_msg("assistant", "Done"),
];
let result = rebuild_chat_messages_from_db(&messages);
// Legacy rows are skipped, only user + assistant
assert_eq!(result.len(), 2);
assert_eq!(result[0].role, crate::llm::Role::User);
assert_eq!(result[1].role, crate::llm::Role::Assistant);
}
#[test]
fn test_rebuild_chat_messages_empty() {
let result = rebuild_chat_messages_from_db(&[]);
assert!(result.is_empty());
}
#[test]
fn test_rebuild_chat_messages_malformed_tool_calls_json() {
let messages = vec![
make_db_msg("user", "Hi"),
make_db_msg("tool_calls", "not valid json"),
make_db_msg("assistant", "Done"),
];
let result = rebuild_chat_messages_from_db(&messages);
// Malformed JSON is silently skipped
assert_eq!(result.len(), 2);
}
#[test]
fn test_rebuild_chat_messages_multi_turn_with_tools() {
let tool_json_1 = serde_json::json!([
{"name": "search", "call_id": "call_0", "parameters": {}, "result": "found it"}
]);
let tool_json_2 = serde_json::json!([
{"name": "write", "call_id": "call_0", "parameters": {"path": "a.txt"}, "result": "ok"}
]);
let messages = vec![
make_db_msg("user", "Find X"),
make_db_msg("tool_calls", &tool_json_1.to_string()),
make_db_msg("assistant", "Found X"),
make_db_msg("user", "Write it"),
make_db_msg("tool_calls", &tool_json_2.to_string()),
make_db_msg("assistant", "Written"),
];
let result = rebuild_chat_messages_from_db(&messages);
// Turn 1: user + assistant_with_calls + tool_result + assistant = 4
// Turn 2: user + assistant_with_calls + tool_result + assistant = 4
assert_eq!(result.len(), 8);
// Verify turn boundaries
assert_eq!(result[0].content, "Find X");
assert!(result[1].tool_calls.is_some());
assert_eq!(result[2].role, crate::llm::Role::Tool);
assert_eq!(result[3].content, "Found X");
assert_eq!(result[4].content, "Write it");
assert!(result[5].tool_calls.is_some());
assert_eq!(result[6].role, crate::llm::Role::Tool);
assert_eq!(result[7].content, "Written");
}
fn make_db_msg(role: &str, content: &str) -> crate::history::ConversationMessage {
crate::history::ConversationMessage {
id: uuid::Uuid::new_v4(),
role: role.to_string(),
content: content.to_string(),
created_at: chrono::Utc::now(),
}
}
}
+1 -1
View File
@@ -212,7 +212,7 @@ impl Worker {
let job_ctx = self.context_manager().get_context(self.job_id).await?;
// Create reasoning engine
let reasoning = Reasoning::new(self.llm().clone(), self.safety().clone());
let reasoning = Reasoning::new(self.llm().clone());
// Build initial reasoning context (tool definitions refreshed each iteration in execution_loop)
let mut reason_ctx = ReasoningContext::new().with_job(&job_ctx.description);
+228 -102
View File
@@ -21,7 +21,7 @@ use crate::secrets::SecretsStore;
use crate::skills::SkillRegistry;
use crate::skills::catalog::SkillCatalog;
use crate::tools::ToolRegistry;
use crate::tools::mcp::McpSessionManager;
use crate::tools::mcp::{McpProcessManager, McpSessionManager};
use crate::tools::wasm::SharedCredentialRegistry;
use crate::tools::wasm::WasmToolRuntime;
use crate::workspace::{EmbeddingProvider, Workspace};
@@ -41,6 +41,7 @@ pub struct AppComponents {
pub workspace: Option<Arc<Workspace>>,
pub extension_manager: Option<Arc<ExtensionManager>>,
pub mcp_session_manager: Arc<McpSessionManager>,
pub mcp_process_manager: Arc<McpProcessManager>,
pub wasm_tool_runtime: Option<Arc<WasmToolRuntime>>,
pub log_broadcaster: Arc<LogBroadcaster>,
pub context_manager: Arc<ContextManager>,
@@ -255,15 +256,18 @@ impl AppBuilder {
self.libsql_db.take();
}
// Re-resolve config with OS credentials
if let Some(ref db) = self.db {
let toml_path = self.toml_path.as_deref();
if let Ok(refreshed) =
Config::from_db_with_toml(db.as_ref(), "default", toml_path).await
{
self.config = refreshed;
tracing::debug!("LlmConfig re-resolved after OS credential injection");
}
// Re-resolve only the LLM config with OS credentials.
let store: Option<&(dyn crate::db::SettingsStore + Sync)> =
self.db.as_ref().map(|db| db.as_ref() as _);
let toml_path = self.toml_path.as_deref();
if let Err(e) = self
.config
.re_resolve_llm(store, "default", toml_path)
.await
{
tracing::warn!(
"Failed to re-resolve LLM config after OS credential injection: {e}"
);
}
return Ok(());
@@ -308,18 +312,16 @@ impl AppBuilder {
// Inject LLM API keys from encrypted storage
crate::config::inject_llm_keys_from_secrets(secrets.as_ref(), "default").await;
// Re-resolve config with newly available keys
if let Some(ref db) = self.db {
let toml_path = self.toml_path.as_deref();
match Config::from_db_with_toml(db.as_ref(), "default", toml_path).await {
Ok(refreshed) => {
self.config = refreshed;
tracing::debug!("LlmConfig re-resolved after secret injection");
}
Err(e) => {
tracing::warn!("Failed to re-resolve config after secret injection: {}", e);
}
}
// Re-resolve only the LLM config with newly available keys.
let store: Option<&(dyn crate::db::SettingsStore + Sync)> =
self.db.as_ref().map(|db| db.as_ref() as _);
let toml_path = self.toml_path.as_deref();
if let Err(e) = self
.config
.re_resolve_llm(store, "default", toml_path)
.await
{
tracing::warn!("Failed to re-resolve LLM config after secret injection: {e}");
}
}
@@ -332,7 +334,7 @@ impl AppBuilder {
/// Delegates to `build_provider_chain` which applies all decorators
/// (retry, smart routing, failover, circuit breaker, response cache).
#[allow(clippy::type_complexity)]
pub fn init_llm(
pub async fn init_llm(
&self,
) -> Result<
(
@@ -343,7 +345,7 @@ impl AppBuilder {
anyhow::Error,
> {
let (llm, cheap_llm, recording_handle) =
crate::llm::build_provider_chain(&self.config.llm, self.session.clone())?;
crate::llm::build_provider_chain(&self.config.llm, self.session.clone()).await?;
Ok((llm, cheap_llm, recording_handle))
}
@@ -398,16 +400,55 @@ impl AppBuilder {
None
};
// Register image/vision tools if we have a workspace and LLM API credentials
if workspace.is_some() {
let (api_base, api_key_opt) = if let Some(ref provider) = self.config.llm.provider {
(
provider.base_url.clone(),
provider.api_key.as_ref().map(|s| {
use secrecy::ExposeSecret;
s.expose_secret().to_string()
}),
)
} else {
(
self.config.llm.nearai.base_url.clone(),
self.config.llm.nearai.api_key.as_ref().map(|s| {
use secrecy::ExposeSecret;
s.expose_secret().to_string()
}),
)
};
if let Some(api_key) = api_key_opt {
// Check for image generation models
let model_name = self
.config
.llm
.provider
.as_ref()
.map(|p| p.model.clone())
.unwrap_or_else(|| self.config.llm.nearai.model.clone());
let models = vec![model_name.clone()];
let gen_model = crate::llm::image_models::suggest_image_model(&models)
.unwrap_or("flux-1.1-pro")
.to_string();
tools.register_image_tools(api_base.clone(), api_key.clone(), gen_model, None);
// Check for vision models
let vision_model = crate::llm::vision_models::suggest_vision_model(&models)
.unwrap_or(&model_name)
.to_string();
tools.register_vision_tools(api_base, api_key, vision_model, None);
}
}
// Register builder tool if enabled
if self.config.builder.enabled
&& (self.config.agent.allow_local_tools || !self.config.sandbox.enabled)
{
tools
.register_builder_tool(
llm.clone(),
safety.clone(),
Some(self.config.builder.to_builder_config()),
)
.register_builder_tool(llm.clone(), Some(self.config.builder.to_builder_config()))
.await;
tracing::info!("Builder mode enabled");
}
@@ -423,6 +464,7 @@ impl AppBuilder {
) -> Result<
(
Arc<McpSessionManager>,
Arc<McpProcessManager>,
Option<Arc<WasmToolRuntime>>,
Option<Arc<ExtensionManager>>,
Vec<crate::extensions::RegistryEntry>,
@@ -430,10 +472,13 @@ impl AppBuilder {
),
anyhow::Error,
> {
use crate::tools::mcp::{McpClient, config::load_mcp_servers_from_db, is_authenticated};
use crate::tools::mcp::{
McpClient, McpTransport, config::load_mcp_servers_from_db, is_authenticated,
};
use crate::tools::wasm::{WasmToolLoader, load_dev_tools};
let mcp_session_manager = Arc::new(McpSessionManager::new());
let mcp_process_manager = Arc::new(McpProcessManager::new());
// Create WASM tool runtime eagerly so extensions installed after startup
// (e.g. via the web UI) can still be activated. The tools directory is only
@@ -509,97 +554,175 @@ impl AppBuilder {
let db = self.db.clone();
let tools = Arc::clone(tools);
let mcp_sm = Arc::clone(&mcp_session_manager);
let pm = Arc::clone(&mcp_process_manager);
async move {
if let Some(ref secrets) = secrets_store {
let servers_result = if let Some(ref d) = db {
load_mcp_servers_from_db(d.as_ref(), "default").await
} else {
crate::tools::mcp::config::load_mcp_servers().await
};
match servers_result {
Ok(servers) => {
let enabled: Vec<_> = servers.enabled_servers().cloned().collect();
if !enabled.is_empty() {
tracing::info!(
"Loading {} configured MCP server(s)...",
enabled.len()
);
}
let servers_result = if let Some(ref d) = db {
load_mcp_servers_from_db(d.as_ref(), "default").await
} else {
crate::tools::mcp::config::load_mcp_servers().await
};
match servers_result {
Ok(servers) => {
let enabled: Vec<_> = servers.enabled_servers().cloned().collect();
if !enabled.is_empty() {
tracing::info!("Loading {} configured MCP server(s)...", enabled.len());
}
let mut join_set = tokio::task::JoinSet::new();
for server in enabled {
let mcp_sm = Arc::clone(&mcp_sm);
let secrets = Arc::clone(secrets);
let tools = Arc::clone(&tools);
let mut join_set = tokio::task::JoinSet::new();
for server in enabled {
let mcp_sm = Arc::clone(&mcp_sm);
let secrets = secrets_store.clone();
let tools = Arc::clone(&tools);
let pm = Arc::clone(&pm);
join_set.spawn(async move {
let server_name = server.name.clone();
let has_tokens =
is_authenticated(&server, &secrets, "default").await;
join_set.spawn(async move {
let server_name = server.name.clone();
let client = if has_tokens || server.requires_auth() {
McpClient::new_authenticated(
server, mcp_sm, secrets, "default",
)
} else {
McpClient::new_with_name(&server_name, &server.url)
};
match client.list_tools().await {
Ok(mcp_tools) => {
let tool_count = mcp_tools.len();
match client.create_tools().await {
Ok(tool_impls) => {
for tool in tool_impls {
tools.register(tool).await;
}
tracing::info!(
"Loaded {} tools from MCP server '{}'",
tool_count,
server_name
);
}
Err(e) => {
tracing::warn!(
"Failed to create tools from MCP server '{}': {}",
server_name,
e
);
}
let client: McpClient = match server.effective_transport() {
crate::tools::mcp::config::EffectiveTransport::Stdio {
command,
args,
env,
} => {
match pm
.spawn_stdio(
&server_name,
command,
args.to_vec(),
env.clone(),
)
.await
{
Ok(transport) => McpClient::new_with_transport(
&server_name,
transport as Arc<dyn McpTransport>,
None,
secrets,
"default",
Some(server),
),
Err(e) => {
tracing::warn!(
"Failed to spawn stdio MCP server '{}': {}",
server_name,
e
);
return;
}
}
Err(e) => {
let err_str = e.to_string();
if err_str.contains("401")
|| err_str.contains("authentication")
{
}
#[cfg(unix)]
crate::tools::mcp::config::EffectiveTransport::Unix {
socket_path,
} => {
match crate::tools::mcp::unix_transport::UnixMcpTransport::connect(
&server_name,
socket_path,
)
.await
{
Ok(transport) => McpClient::new_with_transport(
&server_name,
Arc::new(transport) as Arc<dyn McpTransport>,
None,
secrets,
"default",
Some(server),
),
Err(e) => {
tracing::warn!(
"MCP server '{}' requires authentication. \
Run: ironclaw mcp auth {}",
"Failed to connect to Unix MCP server '{}': {}",
server_name,
e
);
return;
}
}
}
#[cfg(not(unix))]
crate::tools::mcp::config::EffectiveTransport::Unix { .. } => {
tracing::warn!(
"Unix socket transport is not supported on this platform (server '{}')",
server_name
);
return;
}
crate::tools::mcp::config::EffectiveTransport::Http => {
if let Some(ref secrets) = secrets {
let has_tokens =
is_authenticated(&server, secrets, "default")
.await;
if has_tokens || server.requires_auth() {
McpClient::new_authenticated(
server,
Arc::clone(&mcp_sm),
Arc::clone(secrets),
"default",
)
} else {
McpClient::new_with_config(server)
}
} else {
McpClient::new_with_config(server)
}
}
};
match client.list_tools().await {
Ok(mcp_tools) => {
let tool_count = mcp_tools.len();
match client.create_tools().await {
Ok(tool_impls) => {
for tool in tool_impls {
tools.register(tool).await;
}
tracing::info!(
"Loaded {} tools from MCP server '{}'",
tool_count,
server_name
);
} else {
}
Err(e) => {
tracing::warn!(
"Failed to connect to MCP server '{}': {}",
"Failed to create tools from MCP server '{}': {}",
server_name,
e
);
}
}
}
});
}
while let Some(result) = join_set.join_next().await {
if let Err(e) = result {
tracing::warn!("MCP server loading task panicked: {}", e);
Err(e) => {
let err_str = e.to_string();
if err_str.contains("401")
|| err_str.contains("authentication")
{
tracing::warn!(
"MCP server '{}' requires authentication. \
Run: ironclaw mcp auth {}",
server_name,
server_name
);
} else {
tracing::warn!(
"Failed to connect to MCP server '{}': {}",
server_name,
e
);
}
}
}
});
}
while let Some(result) = join_set.join_next().await {
if let Err(e) = result {
tracing::warn!("MCP server loading task panicked: {}", e);
}
}
Err(e) => {
tracing::debug!("No MCP servers configured ({})", e);
}
}
Err(e) => {
tracing::debug!("No MCP servers configured ({})", e);
}
}
}
@@ -670,6 +793,7 @@ impl AppBuilder {
Ok((
mcp_session_manager,
mcp_process_manager,
wasm_tool_runtime,
extension_manager,
catalog_entries,
@@ -696,7 +820,7 @@ impl AppBuilder {
let (llm, cheap_llm, recording_handle) = if let Some(llm) = self.llm_override.take() {
(llm, None, None)
} else {
self.init_llm()?
self.init_llm().await?
};
let (safety, tools, embeddings, workspace) = self.init_tools(&llm).await?;
@@ -705,6 +829,7 @@ impl AppBuilder {
let (
mcp_session_manager,
mcp_process_manager,
wasm_tool_runtime,
extension_manager,
catalog_entries,
@@ -802,6 +927,7 @@ impl AppBuilder {
workspace,
extension_manager,
mcp_session_manager,
mcp_process_manager,
wasm_tool_runtime,
log_broadcaster: self.log_broadcaster,
context_manager,
+251
View File
@@ -414,10 +414,103 @@ pub enum MigrationError {
Io(String),
}
// ── PID Lock ──────────────────────────────────────────────────────────────
/// Path to the PID lock file: `~/.ironclaw/ironclaw.pid`.
pub fn pid_lock_path() -> PathBuf {
ironclaw_base_dir().join("ironclaw.pid")
}
/// A PID-based lock that prevents multiple IronClaw instances from running
/// simultaneously.
///
/// Uses `fs4::try_lock_exclusive()` for atomic locking (no TOCTOU race),
/// then writes the current PID into the locked file for diagnostics.
/// The OS-level lock is held for the lifetime of this struct and
/// automatically released on drop (along with the PID file cleanup).
#[derive(Debug)]
pub struct PidLock {
path: PathBuf,
/// Held open to maintain the OS-level exclusive lock.
_file: std::fs::File,
}
/// Errors from PID lock acquisition.
#[derive(Debug, thiserror::Error)]
pub enum PidLockError {
#[error("Another IronClaw instance is already running (PID {pid})")]
AlreadyRunning { pid: u32 },
#[error("Failed to acquire PID lock: {0}")]
Io(#[from] std::io::Error),
}
impl PidLock {
/// Try to acquire the PID lock.
///
/// Uses an exclusive file lock (`flock`/`LockFileEx`) so that two
/// concurrent processes cannot both acquire the lock — no TOCTOU race.
/// If the lock file exists but the holding process is gone (stale),
/// the lock is reclaimed automatically by the OS.
pub fn acquire() -> Result<Self, PidLockError> {
Self::acquire_at(pid_lock_path())
}
/// Acquire at a specific path (for testing).
fn acquire_at(path: PathBuf) -> Result<Self, PidLockError> {
use fs4::FileExt;
use std::fs::OpenOptions;
use std::io::Write;
// Ensure parent directory exists
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
// Open (or create) the lock file
let mut file = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(&path)?;
// Try non-blocking exclusive lock — if another process holds it,
// this fails immediately instead of blocking.
if let Err(e) = file.try_lock_exclusive() {
if e.kind() == std::io::ErrorKind::WouldBlock {
// Lock held by another process — read its PID for the error message
let pid = std::fs::read_to_string(&path)
.ok()
.and_then(|s| s.trim().parse::<u32>().ok())
.unwrap_or(0);
return Err(PidLockError::AlreadyRunning { pid });
}
// Other errors (permissions, unsupported filesystem, etc.)
return Err(PidLockError::Io(e));
}
// We hold the exclusive lock — write our PID
file.set_len(0)?; // truncate
write!(file, "{}", std::process::id())?;
Ok(PidLock { path, _file: file })
}
}
impl Drop for PidLock {
fn drop(&mut self) {
// Remove the PID file; the OS-level lock is released when _file is dropped.
let _ = std::fs::remove_file(&self.path);
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::process::Command;
use std::sync::Mutex;
use std::thread;
use std::time::{Duration, Instant};
use tempfile::tempdir;
static ENV_MUTEX: Mutex<()> = Mutex::new(());
@@ -986,4 +1079,162 @@ INJECTED="pwned"#;
unsafe { std::env::remove_var("IRONCLAW_BASE_DIR") };
}
}
// ── PID Lock tests ───────────────────────────────────────────────
#[test]
fn test_pid_lock_acquire_and_drop() {
let dir = tempdir().unwrap();
let pid_path = dir.path().join("ironclaw.pid");
// Acquire lock
let lock = PidLock::acquire_at(pid_path.clone()).unwrap();
assert!(pid_path.exists());
// PID file should contain our PID
let contents = std::fs::read_to_string(&pid_path).unwrap();
assert_eq!(contents.trim().parse::<u32>().unwrap(), std::process::id());
// Drop should remove the file
drop(lock);
assert!(!pid_path.exists());
}
#[test]
fn test_pid_lock_rejects_second_acquire() {
let dir = tempdir().unwrap();
let pid_path = dir.path().join("ironclaw.pid");
// First lock succeeds
let _lock1 = PidLock::acquire_at(pid_path.clone()).unwrap();
// Second acquire on same file must fail (exclusive flock held)
let result = PidLock::acquire_at(pid_path.clone());
assert!(result.is_err());
match result.unwrap_err() {
PidLockError::AlreadyRunning { pid } => {
assert_eq!(pid, std::process::id());
}
other => panic!("expected AlreadyRunning, got: {}", other),
}
}
#[test]
fn test_pid_lock_reclaims_after_drop() {
let dir = tempdir().unwrap();
let pid_path = dir.path().join("ironclaw.pid");
// Acquire and release
let lock = PidLock::acquire_at(pid_path.clone()).unwrap();
drop(lock);
// Should succeed — OS lock was released on drop
let lock2 = PidLock::acquire_at(pid_path).unwrap();
drop(lock2);
}
#[test]
fn test_pid_lock_reclaims_stale_file_without_flock() {
let dir = tempdir().unwrap();
let pid_path = dir.path().join("ironclaw.pid");
// Write a stale PID file manually (no flock held)
std::fs::write(&pid_path, "4294967294").unwrap();
// Should succeed because no OS lock is held on the file
let lock = PidLock::acquire_at(pid_path.clone()).unwrap();
let contents = std::fs::read_to_string(&pid_path).unwrap();
assert_eq!(contents.trim().parse::<u32>().unwrap(), std::process::id());
drop(lock);
}
#[test]
fn test_pid_lock_handles_corrupt_pid_file() {
let dir = tempdir().unwrap();
let pid_path = dir.path().join("ironclaw.pid");
// Write garbage (no flock held)
std::fs::write(&pid_path, "not-a-number").unwrap();
// Should succeed — no OS lock held, file is reclaimed
let lock = PidLock::acquire_at(pid_path).unwrap();
drop(lock);
}
#[test]
fn test_pid_lock_creates_parent_dirs() {
let dir = tempdir().unwrap();
let pid_path = dir.path().join("nested").join("deep").join("ironclaw.pid");
let lock = PidLock::acquire_at(pid_path.clone()).unwrap();
assert!(pid_path.exists());
drop(lock);
}
#[test]
fn test_pid_lock_child_helper_holds_lock() {
if std::env::var("IRONCLAW_PID_LOCK_CHILD").ok().as_deref() != Some("1") {
return;
}
let pid_path = PathBuf::from(
std::env::var("IRONCLAW_PID_LOCK_PATH").expect("IRONCLAW_PID_LOCK_PATH missing"),
);
let hold_ms = std::env::var("IRONCLAW_PID_LOCK_HOLD_MS")
.ok()
.and_then(|s| s.parse::<u64>().ok())
.unwrap_or(3000);
let _lock = PidLock::acquire_at(pid_path).expect("child failed to acquire pid lock");
thread::sleep(Duration::from_millis(hold_ms));
}
#[test]
fn test_pid_lock_rejects_lock_held_by_other_process() {
let dir = tempdir().unwrap();
let pid_path = dir.path().join("ironclaw.pid");
let current_exe = std::env::current_exe().unwrap();
let mut child = Command::new(current_exe)
.args([
"--exact",
"bootstrap::tests::test_pid_lock_child_helper_holds_lock",
"--nocapture",
"--test-threads=1",
])
.env("IRONCLAW_PID_LOCK_CHILD", "1")
.env("IRONCLAW_PID_LOCK_PATH", pid_path.display().to_string())
.env("IRONCLAW_PID_LOCK_HOLD_MS", "3000")
.spawn()
.unwrap();
let started = Instant::now();
while started.elapsed() < Duration::from_secs(2) {
if pid_path.exists() {
break;
}
if let Some(status) = child.try_wait().unwrap() {
panic!("child exited before acquiring lock: {}", status);
}
thread::sleep(Duration::from_millis(20));
}
assert!(
pid_path.exists(),
"child did not create lock file in time: {}",
pid_path.display()
);
let result = PidLock::acquire_at(pid_path.clone());
match result.unwrap_err() {
PidLockError::AlreadyRunning { .. } => {}
other => panic!("expected AlreadyRunning, got: {}", other),
}
let status = child.wait().unwrap();
assert!(status.success(), "child process failed: {}", status);
// After the child exits, lock should be released and reacquirable.
let lock = PidLock::acquire_at(pid_path).unwrap();
drop(lock);
}
}
+22
View File
@@ -79,6 +79,8 @@ pub struct IncomingMessage {
pub received_at: DateTime<Utc>,
/// Channel-specific metadata.
pub metadata: serde_json::Value,
/// IANA timezone string from the client (e.g. "America/New_York").
pub timezone: Option<String>,
/// File or media attachments on this message.
pub attachments: Vec<IncomingAttachment>,
}
@@ -99,6 +101,7 @@ impl IncomingMessage {
thread_id: None,
received_at: Utc::now(),
metadata: serde_json::Value::Null,
timezone: None,
attachments: Vec::new(),
}
}
@@ -121,6 +124,12 @@ impl IncomingMessage {
self
}
/// Set the client timezone.
pub fn with_timezone(mut self, tz: impl Into<String>) -> Self {
self.timezone = Some(tz.into());
self
}
/// Set attachments.
pub fn with_attachments(mut self, attachments: Vec<IncomingAttachment>) -> Self {
self.attachments = attachments;
@@ -222,6 +231,13 @@ pub enum StatusUpdate {
success: bool,
message: String,
},
/// An image was generated by a tool.
ImageGenerated {
/// Base64 data URL of the generated image.
data_url: String,
/// Optional workspace path where the image was saved.
path: Option<String>,
},
}
impl StatusUpdate {
@@ -454,4 +470,10 @@ mod tests {
panic!("expected ToolCompleted variant");
}
}
#[test]
fn test_incoming_message_with_timezone() {
let msg = IncomingMessage::new("test", "user1", "hello").with_timezone("America/New_York");
assert_eq!(msg.timezone.as_deref(), Some("America/New_York"));
}
}
+126 -6
View File
@@ -17,7 +17,9 @@ use tokio::sync::{RwLock, mpsc, oneshot};
use tokio_stream::wrappers::ReceiverStream;
use uuid::Uuid;
use crate::channels::{Channel, IncomingMessage, MessageStream, OutgoingResponse};
use crate::channels::{
AttachmentKind, Channel, IncomingAttachment, IncomingMessage, MessageStream, OutgoingResponse,
};
use crate::config::HttpConfig;
use crate::error::ChannelError;
@@ -46,8 +48,9 @@ struct RateLimitState {
request_count: u32,
}
/// Maximum JSON body size for webhook requests (64 KB).
const MAX_BODY_BYTES: usize = 64 * 1024;
/// Maximum JSON body size for webhook requests (15 MB, to support base64 image attachments
/// with ~33% overhead from base64 encoding).
const MAX_BODY_BYTES: usize = 15 * 1024 * 1024;
/// Maximum number of pending wait-for-response requests.
const MAX_PENDING_RESPONSES: usize = 100;
@@ -115,8 +118,34 @@ struct WebhookRequest {
/// Whether to wait for a synchronous response.
#[serde(default)]
wait_for_response: bool,
/// Optional file attachments (base64-encoded).
#[serde(default)]
attachments: Vec<AttachmentData>,
}
/// A file attachment in a webhook request.
#[derive(Debug, Deserialize)]
struct AttachmentData {
/// MIME type (e.g. "image/png", "application/pdf").
mime_type: String,
/// Optional filename.
#[serde(default)]
filename: Option<String>,
/// Base64-encoded file data.
#[serde(default)]
data_base64: Option<String>,
/// URL to fetch the file from (not downloaded server-side for SSRF prevention).
#[serde(default)]
url: Option<String>,
}
/// Maximum size per attachment (5 MB decoded).
const MAX_ATTACHMENT_BYTES: usize = 5 * 1024 * 1024;
/// Maximum total attachment size (10 MB decoded).
const MAX_TOTAL_ATTACHMENT_BYTES: usize = 10 * 1024 * 1024;
/// Maximum number of attachments per request.
const MAX_ATTACHMENTS: usize = 5;
#[derive(Debug, Serialize)]
struct WebhookResponse {
/// Message ID assigned to this request.
@@ -211,15 +240,106 @@ async fn webhook_handler(
);
}
let msg = IncomingMessage::new("http", &state.user_id, &req.content).with_metadata(
// Validate and decode attachments
let attachments = if !req.attachments.is_empty() {
if req.attachments.len() > MAX_ATTACHMENTS {
return (
StatusCode::BAD_REQUEST,
Json(WebhookResponse {
message_id: Uuid::nil(),
status: "error".to_string(),
response: Some(format!("Too many attachments (max {})", MAX_ATTACHMENTS)),
}),
);
}
let mut decoded_attachments = Vec::new();
let mut total_bytes: usize = 0;
for att in &req.attachments {
if let Some(ref b64) = att.data_base64 {
use base64::Engine;
let data = match base64::engine::general_purpose::STANDARD.decode(b64) {
Ok(d) => d,
Err(_) => {
return (
StatusCode::BAD_REQUEST,
Json(WebhookResponse {
message_id: Uuid::nil(),
status: "error".to_string(),
response: Some("Invalid base64 in attachment".to_string()),
}),
);
}
};
if data.len() > MAX_ATTACHMENT_BYTES {
return (
StatusCode::PAYLOAD_TOO_LARGE,
Json(WebhookResponse {
message_id: Uuid::nil(),
status: "error".to_string(),
response: Some(format!(
"Attachment too large (max {} bytes)",
MAX_ATTACHMENT_BYTES
)),
}),
);
}
total_bytes += data.len();
if total_bytes > MAX_TOTAL_ATTACHMENT_BYTES {
return (
StatusCode::PAYLOAD_TOO_LARGE,
Json(WebhookResponse {
message_id: Uuid::nil(),
status: "error".to_string(),
response: Some("Total attachment size exceeds limit".to_string()),
}),
);
}
decoded_attachments.push(IncomingAttachment {
id: Uuid::new_v4().to_string(),
kind: AttachmentKind::from_mime_type(&att.mime_type),
mime_type: att.mime_type.clone(),
filename: att.filename.clone(),
size_bytes: Some(data.len() as u64),
source_url: None,
storage_key: None,
extracted_text: None,
data,
duration_secs: None,
});
} else if let Some(ref url) = att.url {
// URL-only attachment: set source_url but don't download (SSRF prevention)
decoded_attachments.push(IncomingAttachment {
id: Uuid::new_v4().to_string(),
kind: AttachmentKind::from_mime_type(&att.mime_type),
mime_type: att.mime_type.clone(),
filename: att.filename.clone(),
size_bytes: None,
source_url: Some(url.clone()),
storage_key: None,
extracted_text: None,
data: Vec::new(),
duration_secs: None,
});
}
}
decoded_attachments
} else {
Vec::new()
};
let mut msg = IncomingMessage::new("http", &state.user_id, &req.content).with_metadata(
serde_json::json!({
"wait_for_response": req.wait_for_response,
}),
);
if !attachments.is_empty() {
msg = msg.with_attachments(attachments);
}
if let Some(thread_id) = &req.thread_id {
let msg = msg.with_thread(thread_id);
return process_message(state, msg, req.wait_for_response).await;
msg = msg.with_thread(thread_id);
}
process_message(state, msg, req.wait_for_response).await
+57 -9
View File
@@ -18,7 +18,7 @@
//! - `Esc` - Interrupt current operation
use std::borrow::Cow;
use std::io::{self, Write};
use std::io::{self, IsTerminal, Write};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
@@ -297,10 +297,15 @@ impl Channel for ReplChannel {
let esc_interrupt_triggered_for_thread = Arc::new(AtomicBool::new(false));
std::thread::spawn(move || {
let sys_tz = crate::timezone::detect_system_timezone().name().to_string();
// Single message mode: send it and return
if let Some(msg) = single_message {
let incoming = IncomingMessage::new("repl", "default", &msg);
let incoming = IncomingMessage::new("repl", "default", &msg).with_timezone(&sys_tz);
let _ = tx.blocking_send(incoming);
// Ensure the agent exits after handling exactly one turn in -m mode,
// even when other channels (gateway/http) are enabled.
let _ = tx.blocking_send(IncomingMessage::new("repl", "default", "/quit"));
return;
}
@@ -361,7 +366,8 @@ impl Channel for ReplChannel {
"/quit" | "/exit" => {
// Forward shutdown command so the agent loop exits even
// when other channels (e.g. web gateway) are still active.
let msg = IncomingMessage::new("repl", "default", "/quit");
let msg = IncomingMessage::new("repl", "default", "/quit")
.with_timezone(&sys_tz);
let _ = tx.blocking_send(msg);
break;
}
@@ -382,7 +388,8 @@ impl Channel for ReplChannel {
_ => {}
}
let msg = IncomingMessage::new("repl", "default", line);
let msg =
IncomingMessage::new("repl", "default", line).with_timezone(&sys_tz);
if tx.blocking_send(msg).is_err() {
break;
}
@@ -390,21 +397,29 @@ impl Channel for ReplChannel {
Err(ReadlineError::Interrupted) => {
if esc_interrupt_triggered_for_thread.swap(false, Ordering::Relaxed) {
// Esc: interrupt current operation and keep REPL open.
let msg = IncomingMessage::new("repl", "default", "/interrupt");
let msg = IncomingMessage::new("repl", "default", "/interrupt")
.with_timezone(&sys_tz);
if tx.blocking_send(msg).is_err() {
break;
}
} else {
// Ctrl+C (VINTR): request graceful shutdown.
let msg = IncomingMessage::new("repl", "default", "/quit");
let msg = IncomingMessage::new("repl", "default", "/quit")
.with_timezone(&sys_tz);
let _ = tx.blocking_send(msg);
break;
}
}
Err(ReadlineError::Eof) => {
// Ctrl+D: send /quit so the agent loop runs graceful shutdown
let msg = IncomingMessage::new("repl", "default", "/quit");
let _ = tx.blocking_send(msg);
// Ctrl+D in interactive mode: graceful shutdown.
// In daemon mode (stdin = /dev/null, no TTY), EOF arrives
// immediately — just drop the REPL thread silently so other
// channels (gateway, telegram, …) keep running.
if std::io::stdin().is_terminal() {
let msg = IncomingMessage::new("repl", "default", "/quit")
.with_timezone(&sys_tz);
let _ = tx.blocking_send(msg);
}
break;
}
Err(e) => {
@@ -585,6 +600,13 @@ impl Channel for ReplChannel {
eprintln!("\x1b[31m {extension_name}: {message}\x1b[0m");
}
}
StatusUpdate::ImageGenerated { path, .. } => {
if let Some(ref p) = path {
eprintln!("\x1b[36m [image] {p}\x1b[0m");
} else {
eprintln!("\x1b[36m [image generated]\x1b[0m");
}
}
}
Ok(())
}
@@ -614,3 +636,29 @@ impl Channel for ReplChannel {
Ok(())
}
}
#[cfg(test)]
mod tests {
use futures::StreamExt;
use super::*;
#[tokio::test]
async fn single_message_mode_sends_message_then_quit() {
let repl = ReplChannel::with_message("hi".to_string());
let mut stream = repl.start().await.expect("repl start should succeed");
let first = stream.next().await.expect("first message missing");
assert_eq!(first.channel, "repl");
assert_eq!(first.content, "hi");
let second = stream.next().await.expect("quit message missing");
assert_eq!(second.channel, "repl");
assert_eq!(second.content, "/quit");
assert!(
stream.next().await.is_none(),
"stream should end after /quit"
);
}
}
+8
View File
@@ -2809,6 +2809,14 @@ fn status_to_wit(status: &StatusUpdate, metadata: &serde_json::Value) -> wit_cha
),
metadata_json,
},
StatusUpdate::ImageGenerated { path, .. } => wit_channel::StatusUpdate {
status: wit_channel::StatusType::Status,
message: match path {
Some(p) => format!("[image] {}", p),
None => "[image generated]".to_string(),
},
metadata_json,
},
}
}
+1 -1
View File
@@ -264,7 +264,7 @@ pub async fn routines_runs_handler(
/// Convert a Routine to the trimmed RoutineInfo for list display.
fn routine_to_info(r: &crate::agent::routine::Routine) -> RoutineInfo {
let (trigger_type, trigger_summary) = match &r.trigger {
crate::agent::routine::Trigger::Cron { schedule } => {
crate::agent::routine::Trigger::Cron { schedule, .. } => {
("cron".to_string(), format!("cron: {}", schedule))
}
crate::agent::routine::Trigger::Event {
+5
View File
@@ -386,6 +386,11 @@ impl Channel for GatewayChannel {
success,
message,
},
StatusUpdate::ImageGenerated { data_url, path } => SseEvent::ImageGenerated {
data_url,
path,
thread_id,
},
};
self.state.sse.broadcast(event);
+70 -4
View File
@@ -351,7 +351,7 @@ pub async fn start_server(
.merge(statics)
.merge(projects)
.merge(protected)
.layer(DefaultBodyLimit::max(1024 * 1024)) // 1 MB max request body
.layer(DefaultBodyLimit::max(10 * 1024 * 1024)) // 10 MB max request body (image uploads)
.layer(cors)
.layer(SetResponseHeaderLayer::if_not_present(
header::X_CONTENT_TYPE_OPTIONS,
@@ -608,8 +608,59 @@ async fn oauth_callback_handler(
// --- Chat handlers ---
/// Convert web gateway `ImageData` to `IncomingAttachment` objects.
pub(crate) fn images_to_attachments(
images: &[ImageData],
) -> Vec<crate::channels::IncomingAttachment> {
use base64::Engine;
images
.iter()
.enumerate()
.filter_map(|(i, img)| {
if !img.media_type.starts_with("image/") {
tracing::warn!(
"Skipping image {i}: invalid media type '{}' (must start with 'image/')",
img.media_type
);
return None;
}
let data = match base64::engine::general_purpose::STANDARD.decode(&img.data) {
Ok(d) => d,
Err(e) => {
tracing::warn!("Skipping image {i}: invalid base64 data: {e}");
return None;
}
};
Some(crate::channels::IncomingAttachment {
id: format!("web-image-{i}"),
kind: crate::channels::AttachmentKind::Image,
mime_type: img.media_type.clone(),
filename: Some(format!("image-{i}.{}", mime_to_ext(&img.media_type))),
size_bytes: Some(data.len() as u64),
source_url: None,
storage_key: None,
extracted_text: None,
data,
duration_secs: None,
})
})
.collect()
}
/// Map MIME type to file extension.
fn mime_to_ext(mime: &str) -> &str {
match mime {
"image/png" => "png",
"image/gif" => "gif",
"image/webp" => "webp",
"image/svg+xml" => "svg",
_ => "jpg",
}
}
async fn chat_send_handler(
State(state): State<Arc<GatewayState>>,
headers: axum::http::HeaderMap,
Json(req): Json<SendMessageRequest>,
) -> Result<(StatusCode, Json<SendMessageResponse>), (StatusCode, String)> {
tracing::debug!(
@@ -626,17 +677,32 @@ async fn chat_send_handler(
}
let mut msg = IncomingMessage::new("gateway", &state.user_id, &req.content);
// Prefer timezone from JSON body, fall back to X-Timezone header
let tz = req
.timezone
.as_deref()
.or_else(|| headers.get("X-Timezone").and_then(|v| v.to_str().ok()));
if let Some(tz) = tz {
msg = msg.with_timezone(tz);
}
if let Some(ref thread_id) = req.thread_id {
msg = msg.with_thread(thread_id);
msg = msg.with_metadata(serde_json::json!({"thread_id": thread_id}));
}
// Convert uploaded images to IncomingAttachments
if !req.images.is_empty() {
let attachments = images_to_attachments(&req.images);
msg = msg.with_attachments(attachments);
}
let msg_id = msg.id;
tracing::debug!(
"[chat_send_handler] Created message id={}, content={:?}",
"[chat_send_handler] Created message id={}, content={:?}, images={}",
msg_id,
req.content
req.content,
req.images.len()
);
let tx_guard = state.msg_tx.read().await;
@@ -2115,7 +2181,7 @@ async fn routines_runs_handler(
/// Convert a Routine to the trimmed RoutineInfo for list display.
fn routine_to_info(r: &crate::agent::routine::Routine) -> RoutineInfo {
let (trigger_type, trigger_summary) = match &r.trigger {
crate::agent::routine::Trigger::Cron { schedule } => {
crate::agent::routine::Trigger::Cron { schedule, .. } => {
("cron".to_string(), format!("cron: {}", schedule))
}
crate::agent::routine::Trigger::Event {
+1
View File
@@ -142,6 +142,7 @@ impl SseManager {
SseEvent::JobStatus { .. } => "job_status",
SseEvent::JobResult { .. } => "job_result",
SseEvent::Heartbeat => "heartbeat",
SseEvent::ImageGenerated { .. } => "image_generated",
SseEvent::ExtensionStatus { .. } => "extension_status",
};
Ok(Event::default().event(event_type).data(data))
+150 -32
View File
@@ -18,6 +18,7 @@ let unreadThreads = new Map(); // thread_id -> unread count
let _loadThreadsTimer = null;
const JOB_EVENTS_CAP = 500;
const MEMORY_SEARCH_QUERY_MAX_LENGTH = 100;
let stagedImages = [];
// --- Slash Commands ---
@@ -181,6 +182,7 @@ function confirmRestart() {
body: {
content: '/restart',
thread_id: currentThreadId,
timezone: Intl.DateTimeFormat().resolvedOptions().timeZone,
},
})
.then((response) => {
@@ -223,23 +225,6 @@ function updateRestartButtonVisibility() {
}
}
function startGatewayStatusPolling() {
fetchGatewayStatus();
// Poll every 5 seconds
setInterval(fetchGatewayStatus, 5000);
}
function fetchGatewayStatus() {
apiFetch('/api/gateway/status')
.then((data) => {
restartEnabled = data.restart_enabled || false;
updateRestartButtonVisibility();
})
.catch((err) => {
console.warn('[gateway status] Failed to fetch:', err);
});
}
// --- SSE ---
function connectSSE() {
@@ -388,6 +373,12 @@ function connectSSE() {
if (currentTab === 'extensions') loadExtensions();
});
eventSource.addEventListener('image_generated', (e) => {
const data = JSON.parse(e.data);
if (!isCurrentThread(data.thread_id)) return;
addGeneratedImage(data.data_url, data.path);
});
eventSource.addEventListener('error', (e) => {
if (e.data) {
const data = JSON.parse(e.data);
@@ -445,16 +436,23 @@ function sendMessage() {
return;
}
const content = input.value.trim();
if (!content) return;
if (!content && stagedImages.length === 0) return;
addMessage('user', content);
addMessage('user', content || '(images attached)');
input.value = '';
autoResizeTextarea(input);
input.focus();
const body = { content, thread_id: currentThreadId || undefined, timezone: Intl.DateTimeFormat().resolvedOptions().timeZone };
if (stagedImages.length > 0) {
body.images = stagedImages.map(img => ({ media_type: img.media_type, data: img.data }));
stagedImages = [];
renderImagePreviews();
}
apiFetch('/api/chat/send', {
method: 'POST',
body: { content, thread_id: currentThreadId || undefined },
body: body,
}).catch((err) => {
addMessage('system', 'Failed to send: ' + err.message);
});
@@ -471,6 +469,104 @@ function enableChatInput() {
if (btn) btn.disabled = false;
}
// --- Image Upload ---
function renderImagePreviews() {
const strip = document.getElementById('image-preview-strip');
strip.innerHTML = '';
stagedImages.forEach((img, idx) => {
const container = document.createElement('div');
container.className = 'image-preview-container';
const preview = document.createElement('img');
preview.className = 'image-preview';
preview.src = img.dataUrl;
preview.alt = 'Attached image';
const removeBtn = document.createElement('button');
removeBtn.className = 'image-preview-remove';
removeBtn.textContent = '\u00d7';
removeBtn.addEventListener('click', () => {
stagedImages.splice(idx, 1);
renderImagePreviews();
});
container.appendChild(preview);
container.appendChild(removeBtn);
strip.appendChild(container);
});
}
const MAX_IMAGE_SIZE_BYTES = 5 * 1024 * 1024; // 5 MB per image
const MAX_STAGED_IMAGES = 5;
function handleImageFiles(files) {
Array.from(files).forEach(file => {
if (!file.type.startsWith('image/')) return;
if (file.size > MAX_IMAGE_SIZE_BYTES) {
alert(`Image "${file.name}" exceeds 5 MB limit (${(file.size / 1024 / 1024).toFixed(1)} MB)`);
return;
}
if (stagedImages.length >= MAX_STAGED_IMAGES) {
alert(`Maximum ${MAX_STAGED_IMAGES} images allowed per message`);
return;
}
const reader = new FileReader();
reader.onload = function(e) {
const dataUrl = e.target.result;
const commaIdx = dataUrl.indexOf(',');
const meta = dataUrl.substring(0, commaIdx); // e.g. "data:image/png;base64"
const base64 = dataUrl.substring(commaIdx + 1);
const mediaType = meta.replace('data:', '').replace(';base64', '');
stagedImages.push({ media_type: mediaType, data: base64, dataUrl: dataUrl });
renderImagePreviews();
};
reader.readAsDataURL(file);
});
}
document.getElementById('attach-btn').addEventListener('click', () => {
document.getElementById('image-file-input').click();
});
document.getElementById('image-file-input').addEventListener('change', (e) => {
handleImageFiles(e.target.files);
e.target.value = '';
});
document.getElementById('chat-input').addEventListener('paste', (e) => {
const items = (e.clipboardData || e.originalEvent.clipboardData).items;
for (let i = 0; i < items.length; i++) {
if (items[i].kind === 'file' && items[i].type.startsWith('image/')) {
const file = items[i].getAsFile();
if (file) handleImageFiles([file]);
}
}
});
function addGeneratedImage(dataUrl, path) {
const container = document.getElementById('chat-messages');
const card = document.createElement('div');
card.className = 'generated-image-card';
const img = document.createElement('img');
img.className = 'generated-image';
img.src = dataUrl;
img.alt = 'Generated image';
card.appendChild(img);
if (path) {
const pathLabel = document.createElement('div');
pathLabel.className = 'generated-image-path';
pathLabel.textContent = path;
card.appendChild(pathLabel);
}
container.appendChild(card);
container.scrollTop = container.scrollHeight;
}
// --- Slash Autocomplete ---
function showSlashAutocomplete(matches) {
@@ -563,6 +659,13 @@ function sendApprovalAction(requestId, action) {
function renderMarkdown(text) {
if (typeof marked !== 'undefined') {
// Escape raw HTML error pages instead of rendering them as markup.
// Only triggers when the text *starts with* a doctype or <html> tag
// (after optional whitespace), so normal messages that mention HTML
// tags in prose or code fences are not affected. See #263.
if (/^\s*<!doctype\s/i.test(text) || /^\s*<html[\s>]/i.test(text)) {
return escapeHtml(text);
}
let html = marked.parse(text);
// Sanitize HTML output to prevent XSS from tool output or LLM responses.
html = sanitizeRenderedHtml(html);
@@ -1473,7 +1576,7 @@ chatInput.addEventListener('keydown', (e) => {
}
}
if (e.key === 'Enter' && !e.shiftKey) {
if (e.key === 'Enter' && !e.shiftKey && !e.isComposing) {
e.preventDefault();
hideSlashAutocomplete();
sendMessage();
@@ -3408,6 +3511,10 @@ function shortModelName(model) {
function fetchGatewayStatus() {
apiFetch('/api/gateway/status').then(function(data) {
// Update restart button visibility
restartEnabled = data.restart_enabled || false;
updateRestartButtonVisibility();
var popover = document.getElementById('gateway-popover');
var html = '';
@@ -3473,10 +3580,15 @@ let teeReportCache = null;
let teeReportLoading = false;
function teeApiBase() {
var parts = window.location.hostname.split('.');
if (parts.length < 2) return null;
var domain = parts.slice(1).join('.');
return window.location.protocol + '//api.' + domain;
var hostname = window.location.hostname;
// Skip IP addresses (IPv4 and IPv6) and localhost
if (hostname === "localhost" || /^(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$/.test(hostname) || hostname.indexOf(":") !== -1) {
return null;
}
var parts = hostname.split(".");
if (parts.length < 2) return null;
var domain = parts.slice(1).join(".");
return window.location.protocol + "//api." + domain;
}
function teeInstanceName() {
@@ -3487,13 +3599,19 @@ function checkTeeStatus() {
var base = teeApiBase();
if (!base) return;
var name = teeInstanceName();
fetch(base + '/instances/' + encodeURIComponent(name) + '/attestation').then(function(res) {
if (!res.ok) throw new Error(res.status);
return res.json();
}).then(function(data) {
teeInfo = data;
document.getElementById('tee-shield').style.display = 'flex';
}).catch(function() {});
try {
fetch(base + '/instances/' + encodeURIComponent(name) + '/attestation').then(function(res) {
if (!res.ok) throw new Error(res.status);
return res.json();
}).then(function(data) {
teeInfo = data;
document.getElementById('tee-shield').style.display = 'flex';
}).catch(function(err) {
console.warn('Failed to fetch TEE attestation:', err);
});
} catch (e) {
console.warn("Failed to check TEE status:", e);
}
}
function fetchTeeReport() {
+4 -1
View File
@@ -99,7 +99,7 @@
<span id="sse-status">Connected</span>
<div class="gateway-popover" id="gateway-popover"></div>
</div>
<button class="restart-btn" id="restart-btn" onclick="triggerRestart()" title="Gracefully restart the process">
<button class="restart-btn" id="restart-btn" onclick="triggerRestart()" title="Gracefully restart the process" style="display: none;">
<svg id="restart-icon" width="13" height="13" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
<path d="M23 4v6h-6"></path>
<path d="M1 20v-6h6"></path>
@@ -130,7 +130,10 @@
<div class="chat-messages" id="chat-messages"></div>
<div id="slash-autocomplete" class="slash-autocomplete" style="display:none"></div>
<div class="chat-input">
<div id="image-preview-strip" class="image-preview-strip"></div>
<textarea id="chat-input" placeholder="Message or / for commands..." rows="1"></textarea>
<input type="file" id="image-file-input" accept="image/*" multiple style="display:none">
<button id="attach-btn" class="attach-btn" title="Attach images" aria-label="Attach images">&#x1F4CE;</button>
<button id="send-btn" onclick="sendMessage()">Send</button>
</div>
</div>
+102 -11
View File
@@ -159,7 +159,7 @@ body {
flex-shrink: 0;
}
.tab-bar button:not(.status-logs-btn) {
.tab-bar button:not(.status-logs-btn):not(.restart-btn) {
padding: 10px 20px;
background: none;
border: none;
@@ -171,11 +171,11 @@ body {
transition: color 0.2s, border-color 0.2s;
}
.tab-bar button:not(.status-logs-btn):hover {
.tab-bar button:not(.status-logs-btn):not(.restart-btn):hover {
color: var(--text);
}
.tab-bar button:not(.status-logs-btn).active {
.tab-bar button:not(.status-logs-btn):not(.restart-btn).active {
color: var(--accent);
border-bottom-color: var(--accent);
}
@@ -260,42 +260,42 @@ body {
}
/* Restart Button */
.restart-btn {
.tab-bar .restart-btn {
display: flex;
align-items: center;
gap: 0.375rem;
margin: 0.375rem;
padding: 0.25rem 0.75rem;
border-radius: 0.5rem;
font-size: 0.8rem;
border: 1px solid;
border-color: #00d894;
border: 1px solid #00d894;
color: #00d894;
background-color: transparent;
cursor: pointer;
transition: color 150ms, background-color 150ms, border-color 150ms;
}
.restart-btn:hover:not(:disabled) {
.tab-bar .restart-btn:hover:not(:disabled) {
background-color: rgba(0, 216, 148, 0.1);
}
.restart-btn:disabled {
.tab-bar .restart-btn:disabled {
border-color: #333;
color: #666;
cursor: not-allowed;
}
.restart-btn:disabled:hover {
.tab-bar .restart-btn:disabled:hover {
background-color: transparent;
}
.restart-btn svg {
.tab-bar .restart-btn svg {
flex-shrink: 0;
width: 13px;
height: 13px;
}
.restart-btn svg.spinning {
.tab-bar .restart-btn svg.spinning {
animation: spin-icon 1s linear infinite;
}
@@ -1272,6 +1272,7 @@ body {
/* Chat input */
.chat-input {
display: flex;
flex-wrap: wrap;
padding: 12px 16px max(12px, env(safe-area-inset-bottom)) 16px;
gap: 8px;
background: var(--bg-secondary);
@@ -3761,3 +3762,93 @@ mark {
text-overflow: ellipsis;
white-space: nowrap;
}
/* Image Upload */
.attach-btn {
background: none;
border: none;
cursor: pointer;
font-size: 1.2em;
padding: 8px;
align-self: flex-end;
color: var(--text-secondary);
transition: color 0.2s;
min-height: 40px;
display: flex;
align-items: center;
justify-content: center;
}
.attach-btn:hover {
color: var(--text);
}
.image-preview-strip {
display: flex;
flex-direction: row;
gap: 8px;
padding: 4px;
overflow-x: auto;
min-height: 0;
width: 100%;
}
.image-preview-strip:empty {
display: none;
}
.image-preview-container {
position: relative;
display: inline-block;
flex-shrink: 0;
}
.image-preview {
width: 60px;
height: 60px;
border-radius: 6px;
object-fit: cover;
display: block;
}
.image-preview-remove {
position: absolute;
top: -6px;
right: -6px;
width: 18px;
height: 18px;
border-radius: 50%;
background: var(--danger);
color: #fff;
border: none;
font-size: 12px;
line-height: 18px;
text-align: center;
cursor: pointer;
padding: 0;
}
.image-preview-remove:hover {
background: #c33;
}
/* Generated Image */
.generated-image-card {
max-width: 512px;
margin: 8px 0;
border-radius: 8px;
overflow: hidden;
border: 1px solid var(--border);
}
.generated-image {
max-width: 100%;
display: block;
}
.generated-image-path {
font-size: 12px;
color: var(--text-secondary);
padding: 4px 8px;
background: var(--bg-secondary);
}
+34 -2
View File
@@ -5,10 +5,23 @@ use uuid::Uuid;
// --- Chat ---
/// Base64-encoded image data sent from the web frontend.
#[derive(Debug, Clone, Deserialize)]
pub struct ImageData {
/// MIME type (e.g., "image/png", "image/jpeg").
pub media_type: String,
/// Base64-encoded image data (without data: URL prefix).
pub data: String,
}
#[derive(Debug, Deserialize)]
pub struct SendMessageRequest {
pub content: String,
pub thread_id: Option<String>,
pub timezone: Option<String>,
/// Optional images attached to the message.
#[serde(default)]
pub images: Vec<ImageData>,
}
#[derive(Debug, Serialize)]
@@ -219,6 +232,16 @@ pub enum SseEvent {
session_id: Option<String>,
},
/// An image was generated by a tool.
#[serde(rename = "image_generated")]
ImageGenerated {
data_url: String,
#[serde(skip_serializing_if = "Option::is_none")]
path: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
/// Extension activation status change (WASM channels).
#[serde(rename = "extension_status")]
ExtensionStatus {
@@ -613,6 +636,10 @@ pub enum WsClientMessage {
Message {
content: String,
thread_id: Option<String>,
timezone: Option<String>,
/// Optional images attached to the message.
#[serde(default)]
images: Vec<ImageData>,
},
/// Approve or deny a pending tool execution.
#[serde(rename = "approval")]
@@ -679,6 +706,7 @@ impl WsServerMessage {
SseEvent::JobToolResult { .. } => "job_tool_result",
SseEvent::JobStatus { .. } => "job_status",
SseEvent::JobResult { .. } => "job_result",
SseEvent::ImageGenerated { .. } => "image_generated",
SseEvent::ExtensionStatus { .. } => "extension_status",
};
let data = serde_json::to_value(event).unwrap_or(serde_json::Value::Null);
@@ -798,7 +826,9 @@ mod tests {
let json = r#"{"type":"message","content":"hello","thread_id":"t1"}"#;
let msg: WsClientMessage = serde_json::from_str(json).unwrap();
match msg {
WsClientMessage::Message { content, thread_id } => {
WsClientMessage::Message {
content, thread_id, ..
} => {
assert_eq!(content, "hello");
assert_eq!(thread_id.as_deref(), Some("t1"));
}
@@ -811,7 +841,9 @@ mod tests {
let json = r#"{"type":"message","content":"hi"}"#;
let msg: WsClientMessage = serde_json::from_str(json).unwrap();
match msg {
WsClientMessage::Message { content, thread_id } => {
WsClientMessage::Message {
content, thread_id, ..
} => {
assert_eq!(content, "hi");
assert!(thread_id.is_none());
}
+39 -1
View File
@@ -3,6 +3,10 @@
use crate::channels::web::types::{ToolCallInfo, TurnInfo};
/// Truncate a string to at most `max_bytes` bytes at a char boundary, appending "...".
///
/// If the input is wrapped in `<tool_output …>…</tool_output>` and truncation
/// removes the closing tag, the tag is re-appended so downstream XML parsers
/// never see an unclosed element.
pub fn truncate_preview(s: &str, max_bytes: usize) -> String {
if s.len() <= max_bytes {
return s.to_string();
@@ -12,7 +16,14 @@ pub fn truncate_preview(s: &str, max_bytes: usize) -> String {
while end > 0 && !s.is_char_boundary(end) {
end -= 1;
}
format!("{}...", &s[..end])
let mut result = format!("{}...", &s[..end]);
// Re-close <tool_output> if truncation cut through the closing tag.
if s.starts_with("<tool_output") && !result.ends_with("</tool_output>") {
result.push_str("\n</tool_output>");
}
result
}
/// Build TurnInfo pairs from flat DB messages (user/tool_calls/assistant triples).
@@ -162,6 +173,33 @@ mod tests {
assert_eq!(truncate_preview("hello", 0), "...");
}
#[test]
fn test_truncate_preview_closes_tool_output_tag() {
let s = "<tool_output name=\"search\" sanitized=\"true\">\nSome very long content here\n</tool_output>";
// Truncate so it cuts before the closing tag
let result = truncate_preview(s, 60);
assert!(result.ends_with("</tool_output>"));
assert!(result.contains("..."));
}
#[test]
fn test_truncate_preview_no_extra_close_when_intact() {
let s = "<tool_output name=\"echo\" sanitized=\"false\">\nshort\n</tool_output>";
// The string is short enough not to be truncated
let result = truncate_preview(s, 500);
assert_eq!(result, s);
// Should not have a duplicate closing tag
assert_eq!(result.matches("</tool_output>").count(), 1);
}
#[test]
fn test_truncate_preview_non_xml_unaffected() {
let s = "Just a plain long string that gets truncated";
let result = truncate_preview(s, 10);
assert_eq!(result, "Just a pla...");
assert!(!result.contains("</tool_output>"));
}
// ---- build_turns_from_db_messages tests ----
fn make_msg(role: &str, content: &str, offset_ms: i64) -> crate::history::ConversationMessage {
+19 -1
View File
@@ -156,12 +156,26 @@ async fn handle_client_message(
direct_tx: &mpsc::Sender<WsServerMessage>,
) {
match msg {
WsClientMessage::Message { content, thread_id } => {
WsClientMessage::Message {
content,
thread_id,
timezone,
images,
} => {
let mut incoming = IncomingMessage::new("gateway", user_id, &content);
if let Some(ref tz) = timezone {
incoming = incoming.with_timezone(tz);
}
if let Some(ref tid) = thread_id {
incoming = incoming.with_thread(tid);
}
// Convert uploaded images to IncomingAttachments
if !images.is_empty() {
let attachments = crate::channels::web::server::images_to_attachments(&images);
incoming = incoming.with_attachments(attachments);
}
let tx_guard = state.msg_tx.read().await;
if let Some(ref tx) = *tx_guard {
if tx.send(incoming).await.is_err() {
@@ -349,6 +363,8 @@ mod tests {
WsClientMessage::Message {
content: "hello agent".to_string(),
thread_id: Some("t1".to_string()),
timezone: None,
images: Vec::new(),
},
&state,
"user1",
@@ -373,6 +389,8 @@ mod tests {
WsClientMessage::Message {
content: "hello".to_string(),
thread_id: None,
timezone: None,
images: Vec::new(),
},
&state,
"user1",
+243 -121
View File
@@ -2,52 +2,79 @@
//!
//! Commands for adding, removing, authenticating, and testing MCP servers.
use std::collections::HashMap;
use std::io::Write;
use std::sync::Arc;
use clap::Subcommand;
use clap::{Args, Subcommand};
use crate::config::Config;
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},
config::{self, McpServersFile},
config::{self, EffectiveTransport, McpServersFile},
};
/// Arguments for the `mcp add` subcommand.
#[derive(Args, Debug, Clone)]
pub struct McpAddArgs {
/// Server name (e.g., "notion", "github")
pub name: String,
/// Server URL (e.g., "https://mcp.notion.com") -- required for http transport
pub url: Option<String>,
/// Transport type: http (default), stdio, unix
#[arg(long, default_value = "http")]
pub transport: String,
/// Command to run (stdio transport)
#[arg(long)]
pub command: Option<String>,
/// Command arguments (stdio transport, can be repeated)
#[arg(long = "arg", num_args = 1..)]
pub cmd_args: Vec<String>,
/// Environment variables (stdio transport, KEY=VALUE format, can be repeated)
#[arg(long = "env", value_parser = parse_env_var)]
pub env: Vec<(String, String)>,
/// Unix socket path (unix transport)
#[arg(long)]
pub socket: Option<String>,
/// Custom HTTP headers (KEY:VALUE format, can be repeated)
#[arg(long = "header", value_parser = parse_header)]
pub headers: Vec<(String, String)>,
/// OAuth client ID (if authentication is required)
#[arg(long)]
pub client_id: Option<String>,
/// OAuth authorization URL (optional, can be discovered)
#[arg(long)]
pub auth_url: Option<String>,
/// OAuth token URL (optional, can be discovered)
#[arg(long)]
pub token_url: Option<String>,
/// Scopes to request (comma-separated)
#[arg(long)]
pub scopes: Option<String>,
/// Server description
#[arg(long)]
pub description: Option<String>,
}
#[derive(Subcommand, Debug, Clone)]
pub enum McpCommand {
/// Add an MCP server
Add {
/// Server name (e.g., "notion", "github")
name: String,
/// Server URL (e.g., "https://mcp.notion.com")
url: String,
/// OAuth client ID (if authentication is required)
#[arg(long)]
client_id: Option<String>,
/// OAuth authorization URL (optional, can be discovered)
#[arg(long)]
auth_url: Option<String>,
/// OAuth token URL (optional, can be discovered)
#[arg(long)]
token_url: Option<String>,
/// Scopes to request (comma-separated)
#[arg(long)]
scopes: Option<String>,
/// Server description
#[arg(long)]
description: Option<String>,
},
Add(Box<McpAddArgs>),
/// Remove an MCP server
Remove {
@@ -97,29 +124,24 @@ pub enum McpCommand {
},
}
fn parse_header(s: &str) -> Result<(String, String), String> {
let pos = s
.find(':')
.ok_or_else(|| format!("invalid header format '{}', expected KEY:VALUE", s))?;
Ok((s[..pos].trim().to_string(), s[pos + 1..].trim().to_string()))
}
fn parse_env_var(s: &str) -> Result<(String, String), String> {
let pos = s
.find('=')
.ok_or_else(|| format!("invalid env var format '{}', expected KEY=VALUE", s))?;
Ok((s[..pos].to_string(), s[pos + 1..].to_string()))
}
/// Run an MCP command.
pub async fn run_mcp_command(cmd: McpCommand) -> anyhow::Result<()> {
match cmd {
McpCommand::Add {
name,
url,
client_id,
auth_url,
token_url,
scopes,
description,
} => {
add_server(
name,
url,
client_id,
auth_url,
token_url,
scopes,
description,
)
.await
}
McpCommand::Add(args) => add_server(*args).await,
McpCommand::Remove { name } => remove_server(name).await,
McpCommand::List { verbose } => list_servers(verbose).await,
McpCommand::Auth { name, user } => auth_server(name, user).await,
@@ -133,16 +155,58 @@ pub async fn run_mcp_command(cmd: McpCommand) -> anyhow::Result<()> {
}
/// Add a new MCP server.
async fn add_server(
name: String,
url: String,
client_id: Option<String>,
auth_url: Option<String>,
token_url: Option<String>,
scopes: Option<String>,
description: Option<String>,
) -> anyhow::Result<()> {
let mut config = McpServerConfig::new(&name, &url);
async fn add_server(args: McpAddArgs) -> anyhow::Result<()> {
let McpAddArgs {
name,
url,
transport,
command,
cmd_args,
env,
socket,
headers,
client_id,
auth_url,
token_url,
scopes,
description,
} = args;
let transport_lower = transport.to_lowercase();
let mut config = match transport_lower.as_str() {
"stdio" => {
let cmd = command
.clone()
.ok_or_else(|| anyhow::anyhow!("--command is required for stdio transport"))?;
let env_map: HashMap<String, String> = env.into_iter().collect();
McpServerConfig::new_stdio(&name, &cmd, cmd_args.clone(), env_map)
}
"unix" => {
let socket_path = socket
.clone()
.ok_or_else(|| anyhow::anyhow!("--socket is required for unix transport"))?;
McpServerConfig::new_unix(&name, &socket_path)
}
"http" => {
let url_val = url
.as_deref()
.ok_or_else(|| anyhow::anyhow!("URL is required for http transport"))?;
McpServerConfig::new(&name, url_val)
}
other => {
anyhow::bail!(
"Unknown transport type '{}'. Supported: http, stdio, unix",
other
);
}
};
// Apply headers if any
if !headers.is_empty() {
let headers_map: HashMap<String, String> = headers.into_iter().collect();
config = config.with_headers(headers_map);
}
if let Some(desc) = description {
config = config.with_description(desc);
@@ -151,8 +215,12 @@ async fn add_server(
// Track if auth is required
let requires_auth = client_id.is_some();
// Set up OAuth if client_id is provided
// Set up OAuth if client_id is provided (HTTP transport only)
if let Some(client_id) = client_id {
if transport_lower != "http" {
anyhow::bail!("OAuth authentication is only supported with http transport");
}
let mut oauth = OAuthConfig::new(client_id);
if let (Some(auth), Some(token)) = (auth_url, token_url) {
@@ -181,7 +249,24 @@ async fn add_server(
println!();
println!(" ✓ Added MCP server '{}'", name);
println!(" URL: {}", url);
match transport_lower.as_str() {
"stdio" => {
println!(
" Transport: stdio (command: {})",
command.as_deref().unwrap_or("")
);
}
"unix" => {
println!(
" Transport: unix (socket: {})",
socket.as_deref().unwrap_or("")
);
}
_ => {
println!(" URL: {}", url.as_deref().unwrap_or(""));
}
}
if requires_auth {
println!();
@@ -236,9 +321,40 @@ async fn list_servers(verbose: bool) -> anyhow::Result<()> {
""
};
let effective = server.effective_transport();
let transport_label = match &effective {
EffectiveTransport::Http => "http".to_string(),
EffectiveTransport::Stdio { command, .. } => {
format!("stdio ({})", command)
}
EffectiveTransport::Unix { socket_path } => {
format!("unix ({})", socket_path)
}
};
if verbose {
println!(" {} {}{}", status, server.name, auth_status);
println!(" URL: {}", server.url);
println!(" Transport: {}", transport_label);
match &effective {
EffectiveTransport::Http => {
println!(" URL: {}", server.url);
}
EffectiveTransport::Stdio { command, args, env } => {
println!(" Command: {}", command);
if !args.is_empty() {
println!(" Args: {}", args.join(", "));
}
if !env.is_empty() {
// Only print env var names, not values (may contain secrets).
let env_keys: Vec<&str> = env.keys().map(|k| k.as_str()).collect();
println!(" Env: {}", env_keys.join(", "));
}
}
EffectiveTransport::Unix { socket_path } => {
println!(" Socket: {}", socket_path);
}
}
if let Some(ref desc) = server.description {
println!(" Description: {}", desc);
}
@@ -248,11 +364,27 @@ async fn list_servers(verbose: bool) -> anyhow::Result<()> {
println!(" Scopes: {}", oauth.scopes.join(", "));
}
}
if !server.headers.is_empty() {
let header_keys: Vec<&String> = server.headers.keys().collect();
println!(
" Headers: {}",
header_keys
.iter()
.map(|k| k.as_str())
.collect::<Vec<_>>()
.join(", ")
);
}
println!();
} else {
let display = match &effective {
EffectiveTransport::Http => server.url.clone(),
EffectiveTransport::Stdio { command, .. } => command.to_string(),
EffectiveTransport::Unix { socket_path } => socket_path.to_string(),
};
println!(
" {} {} - {}{}",
status, server.name, server.url, auth_status
" {} {} - {} [{}]{}",
status, server.name, display, transport_label, auth_status
);
}
}
@@ -374,7 +506,7 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> {
return Ok(());
} else {
// No OAuth and no tokens - try unauthenticated
McpClient::new_with_name(&server.name, &server.url)
McpClient::new_with_config(server.clone())
};
// Test connection
@@ -504,61 +636,9 @@ async fn get_secrets_store() -> anyhow::Result<Arc<dyn SecretsStore + Send + Syn
)
})?;
let crypto = SecretsCrypto::new(master_key.clone())?;
let crypto = Arc::new(SecretsCrypto::new(master_key.clone())?);
#[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::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))?;
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."
);
}
Ok(crate::db::create_secrets_store(&config.database, crypto).await?)
}
#[cfg(test)]
@@ -579,4 +659,46 @@ mod tests {
TestCli::command().debug_assert();
}
#[test]
fn test_parse_header_valid() {
let result = parse_header("Authorization: Bearer token123").unwrap();
assert_eq!(result.0, "Authorization");
assert_eq!(result.1, "Bearer token123");
}
#[test]
fn test_parse_header_no_spaces() {
let result = parse_header("X-Api-Key:abc123").unwrap();
assert_eq!(result.0, "X-Api-Key");
assert_eq!(result.1, "abc123");
}
#[test]
fn test_parse_header_invalid() {
let result = parse_header("no-colon-here");
assert!(result.is_err());
assert!(result.unwrap_err().contains("invalid header format"));
}
#[test]
fn test_parse_env_var_valid() {
let result = parse_env_var("NODE_ENV=production").unwrap();
assert_eq!(result.0, "NODE_ENV");
assert_eq!(result.1, "production");
}
#[test]
fn test_parse_env_var_with_equals_in_value() {
let result = parse_env_var("KEY=value=with=equals").unwrap();
assert_eq!(result.0, "KEY");
assert_eq!(result.1, "value=with=equals");
}
#[test]
fn test_parse_env_var_invalid() {
let result = parse_env_var("no-equals-here");
assert!(result.is_err());
assert!(result.unwrap_err().contains("invalid env var format"));
}
}
+1 -1
View File
@@ -132,7 +132,7 @@ pub enum Command {
about = "Manage MCP servers",
long_about = "Add, auth, list, or test MCP servers.\nExample: ironclaw mcp add notion https://mcp.notion.com"
)]
Mcp(McpCommand),
Mcp(Box<McpCommand>),
/// Query and manage workspace memory
#[command(
+123 -1
View File
@@ -8,9 +8,35 @@ use std::path::PathBuf;
use crate::bootstrap::ironclaw_base_dir;
use crate::settings::Settings;
/// Load settings from JSON and TOML config files, matching the runtime
/// priority: TOML overlay > settings.json > defaults.
///
/// This mirrors the loading chain in `Config::from_env_with_toml()` but
/// without resolving the full `Config` (which requires async + secrets).
fn load_settings() -> Settings {
load_settings_from(&Settings::default_path(), &Settings::default_toml_path())
}
/// Inner implementation with injectable paths (testable).
fn load_settings_from(json_path: &std::path::Path, toml_path: &std::path::Path) -> Settings {
let mut settings = Settings::load_from(json_path);
match Settings::load_toml(toml_path) {
Ok(Some(toml_settings)) => {
settings.merge_from(&toml_settings);
}
Ok(None) => {} // File not found — fine for default path
Err(e) => {
eprintln!("Warning: failed to parse {}: {}", toml_path.display(), e);
}
}
settings
}
/// Run the status command, printing system health info.
pub async fn run_status_command() -> anyhow::Result<()> {
let settings = Settings::default();
let settings = load_settings();
println!("IronClaw Status");
println!("===============\n");
@@ -209,3 +235,99 @@ fn default_tools_dir() -> PathBuf {
fn default_channels_dir() -> PathBuf {
ironclaw_base_dir().join("channels")
}
#[cfg(test)]
mod tests {
use super::load_settings_from;
/// Regression test for #354: load_settings_from must read config.toml.
#[test]
fn reads_toml_heartbeat_enabled() {
let dir = tempfile::tempdir().expect("tempdir");
let json_path = dir.path().join("settings.json");
let toml_path = dir.path().join("config.toml");
// No JSON file — only TOML
std::fs::write(
&toml_path,
"[heartbeat]\nenabled = true\ninterval_secs = 600",
)
.expect("write toml");
let settings = load_settings_from(&json_path, &toml_path);
assert!(settings.heartbeat.enabled);
assert_eq!(settings.heartbeat.interval_secs, 600);
}
/// Without any config files, defaults are returned.
#[test]
fn defaults_without_config_files() {
let dir = tempfile::tempdir().expect("tempdir");
let settings = load_settings_from(
&dir.path().join("nonexistent.json"),
&dir.path().join("nonexistent.toml"),
);
assert!(!settings.heartbeat.enabled);
}
/// settings.json is respected.
#[test]
fn reads_json_heartbeat_enabled() {
let dir = tempfile::tempdir().expect("tempdir");
let json_path = dir.path().join("settings.json");
let toml_path = dir.path().join("nonexistent.toml");
std::fs::write(
&json_path,
r#"{"heartbeat":{"enabled":true,"interval_secs":900}}"#,
)
.expect("write json");
let settings = load_settings_from(&json_path, &toml_path);
assert!(settings.heartbeat.enabled);
assert_eq!(settings.heartbeat.interval_secs, 900);
}
/// TOML overlay wins over JSON settings.
#[test]
fn toml_overlay_wins_over_json() {
let dir = tempfile::tempdir().expect("tempdir");
let json_path = dir.path().join("settings.json");
let toml_path = dir.path().join("config.toml");
std::fs::write(
&json_path,
r#"{"heartbeat":{"enabled":false,"interval_secs":100}}"#,
)
.expect("write json");
std::fs::write(
&toml_path,
"[heartbeat]\nenabled = true\ninterval_secs = 200",
)
.expect("write toml");
let settings = load_settings_from(&json_path, &toml_path);
assert!(settings.heartbeat.enabled);
assert_eq!(settings.heartbeat.interval_secs, 200);
}
/// Invalid TOML is warned but doesn't crash; falls back to JSON/defaults.
#[test]
fn invalid_toml_falls_back_gracefully() {
let dir = tempfile::tempdir().expect("tempdir");
let json_path = dir.path().join("settings.json");
let toml_path = dir.path().join("config.toml");
std::fs::write(
&json_path,
r#"{"heartbeat":{"enabled":true,"interval_secs":500}}"#,
)
.expect("write json");
std::fs::write(&toml_path, "this is not valid toml [[[").expect("write bad toml");
let settings = load_settings_from(&json_path, &toml_path);
// Should fall back to JSON values, not crash
assert!(settings.heartbeat.enabled);
assert_eq!(settings.heartbeat.interval_secs, 500);
}
}
+2 -56
View File
@@ -11,10 +11,6 @@ use tokio::fs;
use crate::bootstrap::ironclaw_base_dir;
use crate::config::Config;
#[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};
@@ -563,59 +559,9 @@ async fn init_secrets_store() -> anyhow::Result<Arc<dyn SecretsStore + Send + Sy
)
})?;
let crypto = SecretsCrypto::new(master_key.clone())?;
let crypto = Arc::new(SecretsCrypto::new(master_key.clone())?);
let 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::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."
);
}
};
Ok(store)
Ok(crate::db::create_secrets_store(&config.database, crypto).await?)
}
/// Configure authentication for a tool.
+37
View File
@@ -27,6 +27,8 @@ pub struct AgentConfig {
pub max_tool_iterations: usize,
/// When true, skip tool approval checks entirely. For benchmarks/CI.
pub auto_approve_tools: bool,
/// Default timezone for new sessions (IANA name, e.g. "America/New_York").
pub default_timezone: String,
}
impl AgentConfig {
@@ -47,6 +49,7 @@ impl AgentConfig {
max_actions_per_hour: None,
max_tool_iterations: 10,
auto_approve_tools: true,
default_timezone: "UTC".to_string(),
}
}
@@ -89,6 +92,40 @@ impl AgentConfig {
"AGENT_AUTO_APPROVE_TOOLS",
settings.agent.auto_approve_tools,
)?,
default_timezone: {
let tz: String = parse_optional_env(
"DEFAULT_TIMEZONE",
settings.agent.default_timezone.clone(),
)?;
if crate::timezone::parse_timezone(&tz).is_none() {
return Err(ConfigError::InvalidValue {
key: "DEFAULT_TIMEZONE".into(),
message: format!("invalid IANA timezone: '{tz}'"),
});
}
tz
},
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_timezone_rejects_invalid() {
let mut settings = Settings::default();
settings.agent.default_timezone = "Fake/Zone".to_string();
let result = AgentConfig::resolve(&settings);
assert!(result.is_err(), "invalid IANA timezone should be rejected");
}
#[test]
fn test_default_timezone_accepts_valid() {
let settings = Settings::default(); // default is "UTC"
let config = AgentConfig::resolve(&settings).expect("resolve");
assert_eq!(config.default_timezone, "UTC");
}
}
+102 -1
View File
@@ -1,4 +1,4 @@
use crate::config::helpers::{optional_env, parse_bool_env, parse_optional_env};
use crate::config::helpers::{optional_env, parse_bool_env, parse_option_env, parse_optional_env};
use crate::error::ConfigError;
use crate::settings::Settings;
@@ -13,6 +13,12 @@ pub struct HeartbeatConfig {
pub notify_channel: Option<String>,
/// User ID to notify on heartbeat findings.
pub notify_user: Option<String>,
/// Hour (0-23) when quiet hours start.
pub quiet_hours_start: Option<u32>,
/// Hour (0-23) when quiet hours end.
pub quiet_hours_end: Option<u32>,
/// Timezone for quiet hours evaluation (IANA name).
pub timezone: Option<String>,
}
impl Default for HeartbeatConfig {
@@ -22,6 +28,9 @@ impl Default for HeartbeatConfig {
interval_secs: 1800, // 30 minutes
notify_channel: None,
notify_user: None,
quiet_hours_start: None,
quiet_hours_end: None,
timezone: None,
}
}
}
@@ -38,6 +47,98 @@ impl HeartbeatConfig {
.or_else(|| settings.heartbeat.notify_channel.clone()),
notify_user: optional_env("HEARTBEAT_NOTIFY_USER")?
.or_else(|| settings.heartbeat.notify_user.clone()),
quiet_hours_start: parse_option_env::<u32>("HEARTBEAT_QUIET_START")?
.or(settings.heartbeat.quiet_hours_start)
.map(|h| {
if h > 23 {
return Err(ConfigError::InvalidValue {
key: "HEARTBEAT_QUIET_START".into(),
message: "must be 0-23".into(),
});
}
Ok(h)
})
.transpose()?,
quiet_hours_end: parse_option_env::<u32>("HEARTBEAT_QUIET_END")?
.or(settings.heartbeat.quiet_hours_end)
.map(|h| {
if h > 23 {
return Err(ConfigError::InvalidValue {
key: "HEARTBEAT_QUIET_END".into(),
message: "must be 0-23".into(),
});
}
Ok(h)
})
.transpose()?,
timezone: {
let tz = optional_env("HEARTBEAT_TIMEZONE")?
.or_else(|| settings.heartbeat.timezone.clone());
if let Some(ref tz_str) = tz
&& crate::timezone::parse_timezone(tz_str).is_none()
{
return Err(ConfigError::InvalidValue {
key: "HEARTBEAT_TIMEZONE".into(),
message: format!("invalid IANA timezone: '{tz_str}'"),
});
}
tz
},
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_quiet_hours_settings_fallback() {
// When env vars are not set, settings values should be used
let mut settings = Settings::default();
settings.heartbeat.quiet_hours_start = Some(22);
settings.heartbeat.quiet_hours_end = Some(6);
let config = HeartbeatConfig::resolve(&settings).expect("resolve");
assert_eq!(config.quiet_hours_start, Some(22));
assert_eq!(config.quiet_hours_end, Some(6));
}
#[test]
fn test_quiet_hours_rejects_invalid_hour() {
let mut settings = Settings::default();
settings.heartbeat.quiet_hours_start = Some(24);
let result = HeartbeatConfig::resolve(&settings);
assert!(result.is_err());
}
#[test]
fn test_quiet_hours_accepts_boundary_values() {
let mut settings = Settings::default();
settings.heartbeat.quiet_hours_start = Some(0);
settings.heartbeat.quiet_hours_end = Some(23);
let config = HeartbeatConfig::resolve(&settings).expect("resolve");
assert_eq!(config.quiet_hours_start, Some(0));
assert_eq!(config.quiet_hours_end, Some(23));
}
#[test]
fn test_heartbeat_timezone_rejects_invalid() {
let mut settings = Settings::default();
settings.heartbeat.timezone = Some("Fake/Zone".to_string());
let result = HeartbeatConfig::resolve(&settings);
assert!(result.is_err(), "invalid IANA timezone should be rejected");
}
#[test]
fn test_heartbeat_timezone_accepts_valid() {
let mut settings = Settings::default();
settings.heartbeat.timezone = Some("America/New_York".to_string());
let config = HeartbeatConfig::resolve(&settings).expect("resolve");
assert_eq!(config.timezone.as_deref(), Some("America/New_York"));
}
}
+96 -4
View File
@@ -86,6 +86,19 @@ pub struct RegistryProviderConfig {
pub oauth_token: Option<SecretString>,
}
/// Configuration for AWS Bedrock (native Converse API).
#[derive(Debug, Clone)]
pub struct BedrockConfig {
/// AWS region (e.g. "us-east-1").
pub region: String,
/// Bedrock model ID (e.g. "anthropic.claude-opus-4-6-v1").
pub model: String,
/// Cross-region inference prefix: "us", "eu", "apac", "global", or None.
pub cross_region: Option<String>,
/// AWS named profile (for SSO / assume-role workflows).
pub profile: Option<String>,
}
/// LLM provider configuration.
///
/// NearAI remains the default backend with its own config struct (session auth).
@@ -101,8 +114,14 @@ pub struct LlmConfig {
/// NEAR AI config (always populated, also used for embeddings).
pub nearai: NearAiConfig,
/// Resolved provider config for registry-based providers.
/// `None` when backend is "nearai".
/// `None` when backend is "nearai" or "bedrock".
pub provider: Option<RegistryProviderConfig>,
/// AWS Bedrock config (populated when backend=bedrock, requires --features bedrock).
pub bedrock: Option<BedrockConfig>,
/// HTTP request timeout in seconds for LLM API calls.
/// Default: 120. Increase for local LLMs (Ollama, vLLM, LM Studio) that
/// need more time for prompt evaluation on consumer hardware.
pub request_timeout_secs: u64,
}
/// NEAR AI configuration.
@@ -165,6 +184,8 @@ impl LlmConfig {
smart_routing_cascade: false,
},
provider: None,
bedrock: None,
request_timeout_secs: 120,
}
}
@@ -195,8 +216,10 @@ impl LlmConfig {
let backend_lower = backend.to_lowercase();
let is_nearai =
backend_lower == "nearai" || backend_lower == "near_ai" || backend_lower == "near";
let is_bedrock =
backend_lower == "bedrock" || backend_lower == "aws_bedrock" || backend_lower == "aws";
if !is_nearai && registry.find(&backend_lower).is_none() {
if !is_nearai && !is_bedrock && registry.find(&backend_lower).is_none() {
tracing::warn!(
"Unknown LLM backend '{}'. Will attempt as openai_compatible fallback.",
backend
@@ -243,8 +266,8 @@ impl LlmConfig {
smart_routing_cascade: parse_optional_env("SMART_ROUTING_CASCADE", true)?,
};
// Resolve registry provider config (for non-NearAI backends)
let provider = if is_nearai {
// Resolve registry provider config (for non-NearAI, non-Bedrock backends)
let provider = if is_nearai || is_bedrock {
None
} else {
Some(Self::resolve_registry_provider(
@@ -254,9 +277,50 @@ impl LlmConfig {
)?)
};
let bedrock = if is_bedrock {
let explicit_region =
optional_env("BEDROCK_REGION")?.or_else(|| settings.bedrock_region.clone());
if explicit_region.is_none() {
tracing::info!("BEDROCK_REGION not set, defaulting to us-east-1");
}
let region = explicit_region.unwrap_or_else(|| "us-east-1".to_string());
let model = optional_env("BEDROCK_MODEL")?
.or_else(|| settings.selected_model.clone())
.ok_or_else(|| ConfigError::MissingRequired {
key: "BEDROCK_MODEL".to_string(),
hint: "Set BEDROCK_MODEL when LLM_BACKEND=bedrock".to_string(),
})?;
let cross_region = optional_env("BEDROCK_CROSS_REGION")?
.or_else(|| settings.bedrock_cross_region.clone());
if let Some(ref cr) = cross_region
&& !matches!(cr.as_str(), "us" | "eu" | "apac" | "global")
{
return Err(ConfigError::InvalidValue {
key: "BEDROCK_CROSS_REGION".to_string(),
message: format!(
"'{}' is not valid, expected one of: us, eu, apac, global",
cr
),
});
}
let profile = optional_env("AWS_PROFILE")?.or_else(|| settings.bedrock_profile.clone());
Some(BedrockConfig {
region,
model,
cross_region,
profile,
})
} else {
None
};
let request_timeout_secs = parse_optional_env("LLM_REQUEST_TIMEOUT_SECS", 120)?;
Ok(Self {
backend: if is_nearai {
"nearai".to_string()
} else if is_bedrock {
"bedrock".to_string()
} else if let Some(ref p) = provider {
p.provider_id.clone()
} else {
@@ -265,6 +329,8 @@ impl LlmConfig {
session,
nearai,
provider,
bedrock,
request_timeout_secs,
})
}
@@ -1016,4 +1082,30 @@ mod tests {
assert_eq!(parsed, variant, "round-trip failed for {s}");
}
}
#[test]
fn test_request_timeout_defaults_to_120() {
let _guard = ENV_MUTEX.lock().expect("env mutex poisoned");
// SAFETY: Under ENV_MUTEX.
unsafe {
std::env::remove_var("LLM_REQUEST_TIMEOUT_SECS");
}
let config = LlmConfig::resolve(&Settings::default()).expect("resolve");
assert_eq!(config.request_timeout_secs, 120);
}
#[test]
fn test_request_timeout_configurable() {
let _guard = ENV_MUTEX.lock().expect("env mutex poisoned");
// SAFETY: Under ENV_MUTEX.
unsafe {
std::env::set_var("LLM_REQUEST_TIMEOUT_SECS", "300");
}
let config = LlmConfig::resolve(&Settings::default()).expect("resolve");
assert_eq!(config.request_timeout_secs, 300);
// SAFETY: Cleanup
unsafe {
std::env::remove_var("LLM_REQUEST_TIMEOUT_SECS");
}
}
}
+29 -1
View File
@@ -37,7 +37,9 @@ pub use self::database::{DatabaseBackend, DatabaseConfig, SslMode, default_libsq
pub use self::embeddings::EmbeddingsConfig;
pub use self::heartbeat::HeartbeatConfig;
pub use self::hygiene::HygieneConfig;
pub use self::llm::{CacheRetention, LlmConfig, NearAiConfig, RegistryProviderConfig};
pub use self::llm::{
BedrockConfig, CacheRetention, LlmConfig, NearAiConfig, RegistryProviderConfig,
};
pub use self::routines::RoutineConfig;
pub use self::safety::SafetyConfig;
pub use self::sandbox::{ClaudeCodeConfig, SandboxModeConfig};
@@ -257,6 +259,32 @@ impl Config {
Ok(())
}
/// Re-resolve only the LLM config after credential injection.
///
/// Called by `AppBuilder::init_secrets()` after injecting API keys into
/// the env overlay. Only rebuilds `self.llm` — all other config fields
/// are unaffected, preserving values from the initial config load (or
/// from `Config::for_testing()` in test mode).
pub async fn re_resolve_llm(
&mut self,
store: Option<&(dyn crate::db::SettingsStore + Sync)>,
user_id: &str,
toml_path: Option<&std::path::Path>,
) -> Result<(), ConfigError> {
let settings = if let Some(store) = store {
let mut s = match store.get_all_settings(user_id).await {
Ok(map) => Settings::from_db_map(&map),
Err(_) => Settings::default(),
};
Self::apply_toml_overlay(&mut s, toml_path)?;
s
} else {
Settings::default()
};
self.llm = LlmConfig::resolve(&settings)?;
Ok(())
}
/// Build config from settings (shared by from_env and from_db).
async fn build(settings: &Settings) -> Result<Self, ConfigError> {
Ok(Self {
+9
View File
@@ -164,6 +164,8 @@ pub struct JobContext {
/// previous results by ID via `$tool_call_id` parameter syntax.
#[serde(skip)]
pub tool_output_stash: Arc<tokio::sync::RwLock<HashMap<String, String>>>,
/// User's preferred timezone (IANA name, e.g. "America/New_York"). Defaults to "UTC".
pub user_timezone: String,
}
impl JobContext {
@@ -203,9 +205,16 @@ impl JobContext {
http_interceptor: None,
metadata: serde_json::Value::Null,
tool_output_stash: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
user_timezone: "UTC".to_string(),
}
}
/// Set the user timezone on this context.
pub fn with_timezone(mut self, tz: impl Into<String>) -> Self {
self.user_timezone = tz.into();
self
}
/// Transition to a new state.
pub fn transition_to(
&mut self,
+3
View File
@@ -121,6 +121,9 @@ impl JobStore for LibSqlBackend {
tool_output_stash: std::sync::Arc::new(tokio::sync::RwLock::new(
std::collections::HashMap::new(),
)),
// TODO(#661): persist user_timezone in agent_jobs table so
// background/routine jobs retain the session's timezone context.
user_timezone: "UTC".to_string(),
}))
}
None => Ok(None),
+44 -99
View File
@@ -170,16 +170,16 @@ pub(crate) fn parse_timestamp(s: &str) -> Result<DateTime<Utc>, String> {
// Naive with fractional seconds (legacy or SQLite datetime() output)
if let Ok(ndt) = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f") {
tracing::warn!(
timestamp = s,
"parsing naive timestamp without timezone; assuming UTC — consider re-running migrations"
timestamp = %s,
"parsed naive timestamp without timezone; assuming UTC for backward compatibility"
);
return Ok(ndt.and_utc());
}
// Naive without fractional seconds (legacy format)
if let Ok(ndt) = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S") {
tracing::warn!(
timestamp = s,
"parsing naive timestamp without timezone; assuming UTC — consider re-running migrations"
timestamp = %s,
"parsed naive timestamp without timezone; assuming UTC for backward compatibility"
);
return Ok(ndt.and_utc());
}
@@ -410,8 +410,47 @@ pub(crate) fn row_to_routine_run_libsql(row: &libsql::Row) -> Result<RoutineRun,
#[cfg(test)]
mod tests {
use chrono::{TimeZone, Utc};
use crate::db::Database;
use crate::db::libsql::LibSqlBackend;
use crate::db::libsql::{LibSqlBackend, parse_timestamp};
#[test]
fn test_parse_timestamp_accepts_rfc3339_and_legacy_naive_formats() {
let expected = Utc.with_ymd_and_hms(2026, 3, 7, 12, 34, 56).unwrap();
let with_millis = parse_timestamp("2026-03-07T12:34:56.789Z").unwrap();
assert_eq!(with_millis, expected + chrono::Duration::milliseconds(789));
let naive_with_millis = parse_timestamp("2026-03-07 12:34:56.789").unwrap();
assert_eq!(
naive_with_millis,
expected + chrono::Duration::milliseconds(789)
);
let naive_without_millis = parse_timestamp("2026-03-07 12:34:56").unwrap();
assert_eq!(naive_without_millis, expected);
}
#[tokio::test]
async fn test_libsql_now_format_is_rfc3339_and_parseable() {
let backend = LibSqlBackend::new_memory().await.unwrap();
backend.run_migrations().await.unwrap();
let conn = backend.connect().await.unwrap();
let mut rows = conn
.query("SELECT strftime('%Y-%m-%dT%H:%M:%fZ', 'now')", ())
.await
.unwrap();
let row = rows.next().await.unwrap().unwrap();
let ts: String = row.get(0).unwrap();
let parsed = parse_timestamp(&ts).unwrap();
assert_eq!(
ts,
parsed.to_rfc3339_opts(chrono::SecondsFormat::Millis, true)
);
}
#[tokio::test]
async fn test_wal_mode_after_migrations() {
@@ -518,98 +557,4 @@ mod tests {
);
}
}
#[test]
fn test_parse_timestamp_rfc3339() {
use super::parse_timestamp;
// Standard RFC 3339 with Z suffix
let dt = parse_timestamp("2024-01-15T10:30:00.123Z").unwrap();
assert_eq!(
dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
"2024-01-15T10:30:00.123Z"
);
// RFC 3339 with +00:00 offset
let dt = parse_timestamp("2024-01-15T10:30:00.000+00:00").unwrap();
assert_eq!(
dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
"2024-01-15T10:30:00.000Z"
);
}
#[test]
fn test_parse_timestamp_naive_fallback() {
use super::parse_timestamp;
// Naive with fractional seconds (legacy datetime('now') output)
let dt = parse_timestamp("2024-01-15 10:30:00.123").unwrap();
assert_eq!(
dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
"2024-01-15T10:30:00.123Z"
);
// Naive without fractional seconds
let dt = parse_timestamp("2024-01-15 10:30:00").unwrap();
assert_eq!(
dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
"2024-01-15T10:30:00.000Z"
);
}
#[test]
fn test_parse_timestamp_invalid() {
use super::parse_timestamp;
assert!(parse_timestamp("not-a-timestamp").is_err());
assert!(parse_timestamp("").is_err());
}
#[tokio::test]
async fn test_default_timestamps_are_rfc3339() {
// Verify that DEFAULT column values produce RFC 3339 timestamps
// after the migration change from datetime('now') to strftime.
// Use file-based DB because in-memory doesn't share schema across connections.
let dir = tempfile::tempdir().unwrap();
let db_path = dir.path().join("test_ts.db");
let backend = LibSqlBackend::new_local(&db_path).await.unwrap();
backend.run_migrations().await.unwrap();
let conn = backend.connect().await.unwrap();
let id = uuid::Uuid::new_v4().to_string();
conn.execute(
"INSERT INTO conversations (id, channel, user_id) VALUES (?1, ?2, ?3)",
libsql::params![id.clone(), "test", "user1"],
)
.await
.unwrap();
let mut rows = conn
.query(
"SELECT started_at, last_activity FROM conversations WHERE id = ?1",
libsql::params![id],
)
.await
.unwrap();
let row = rows.next().await.unwrap().unwrap();
let started_at: String = row.get(0).unwrap();
let last_activity: String = row.get(1).unwrap();
// Must end with 'Z' (RFC 3339 UTC) and contain 'T' separator
assert!(
started_at.ends_with('Z') && started_at.contains('T'),
"started_at should be RFC 3339, got: {started_at}"
);
assert!(
last_activity.ends_with('Z') && last_activity.contains('T'),
"last_activity should be RFC 3339, got: {last_activity}"
);
// Must be parseable by the RFC 3339 parser directly (not just naive fallback)
use chrono::DateTime;
assert!(
DateTime::parse_from_rfc3339(&started_at).is_ok(),
"started_at not valid RFC 3339: {started_at}"
);
}
}
+101
View File
@@ -91,6 +91,64 @@ pub async fn connect_from_config(
}
}
/// Create a secrets store from database and secrets configuration.
///
/// This is the shared factory for CLI commands and other call sites that need
/// a `SecretsStore` without going through the full `AppBuilder`. Mirrors the
/// pattern of [`connect_from_config`] but returns a secrets-specific store.
pub async fn create_secrets_store(
config: &crate::config::DatabaseConfig,
crypto: Arc<crate::secrets::SecretsCrypto>,
) -> Result<Arc<dyn crate::secrets::SecretsStore + Send + Sync>, 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::LibSqlBackend::new_remote_replica(db_path, url, token.expose_secret())
.await
.map_err(|e| DatabaseError::Pool(e.to_string()))?
} else {
libsql::LibSqlBackend::new_local(db_path)
.await
.map_err(|e| DatabaseError::Pool(e.to_string()))?
};
backend.run_migrations().await?;
Ok(Arc::new(crate::secrets::LibSqlSecretsStore::new(
backend.shared_db(),
crypto,
)))
}
#[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(crate::secrets::PostgresSecretsStore::new(
pg.pool(),
crypto,
)))
}
#[cfg(not(feature = "postgres"))]
_ => Err(DatabaseError::Pool(
"No database backend available for secrets. Enable 'postgres' or 'libsql' feature."
.to_string(),
)),
}
}
// ==================== Sub-traits ====================
//
// Each sub-trait groups related persistence methods. The `Database` supertrait
@@ -435,3 +493,46 @@ pub trait Database:
/// Run schema migrations for this backend.
async fn run_migrations(&self) -> Result<(), DatabaseError>;
}
#[cfg(test)]
mod tests {
use super::*;
/// Regression test: `create_secrets_store` selects the correct backend at
/// runtime based on `DatabaseConfig`, not at compile time. Previously the
/// CLI duplicated this logic with compile-time `#[cfg]` gates that always
/// chose postgres when both features were enabled (PR #209).
#[cfg(feature = "libsql")]
#[tokio::test]
async fn test_create_secrets_store_libsql_backend() {
use secrecy::SecretString;
let tmp = tempfile::tempdir().unwrap();
let db_path = tmp.path().join("test.db");
let config = crate::config::DatabaseConfig {
backend: crate::config::DatabaseBackend::LibSql,
libsql_path: Some(db_path),
libsql_url: None,
libsql_auth_token: None,
url: SecretString::from("unused://libsql".to_string()),
pool_size: 1,
ssl_mode: crate::config::SslMode::default(),
};
let master_key = SecretString::from("a]".repeat(16));
let crypto = Arc::new(crate::secrets::SecretsCrypto::new(master_key).unwrap());
let store = create_secrets_store(&config, crypto).await;
assert!(
store.is_ok(),
"create_secrets_store should succeed for libsql backend"
);
// Verify basic operation works
let store = store.unwrap();
let exists = store.exists("test_user", "nonexistent_secret").await;
assert!(exists.is_ok());
assert!(!exists.unwrap());
}
}
+123 -227
View File
@@ -1,13 +1,10 @@
//! Success evaluation for jobs.
use std::sync::Arc;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use crate::context::{ActionRecord, JobContext};
use crate::error::EvaluationError;
use crate::llm::LlmProvider;
/// Result of evaluating job success.
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -64,233 +61,132 @@ pub trait SuccessEvaluator: Send + Sync {
) -> Result<EvaluationResult, EvaluationError>;
}
/// Rule-based success evaluator.
pub struct RuleBasedEvaluator {
/// Minimum success rate for actions.
min_action_success_rate: f64,
/// Maximum allowed failures.
max_failures: u32,
}
impl RuleBasedEvaluator {
/// Create a new rule-based evaluator.
pub fn new() -> Self {
Self {
min_action_success_rate: 0.8,
max_failures: 3,
}
}
/// Set minimum action success rate.
#[allow(dead_code)] // Public API for configuring evaluation threshold
pub fn with_min_success_rate(mut self, rate: f64) -> Self {
self.min_action_success_rate = rate;
self
}
/// Set maximum failures.
#[allow(dead_code)] // Public API for configuring failure tolerance
pub fn with_max_failures(mut self, max: u32) -> Self {
self.max_failures = max;
self
}
}
impl Default for RuleBasedEvaluator {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl SuccessEvaluator for RuleBasedEvaluator {
async fn evaluate(
&self,
job: &JobContext,
actions: &[ActionRecord],
_output: Option<&str>,
) -> Result<EvaluationResult, EvaluationError> {
let mut issues = Vec::new();
// Check if there were any actions
if actions.is_empty() {
return Ok(EvaluationResult::failure(
"No actions were taken",
vec!["No actions recorded".to_string()],
));
}
// Calculate action success rate
let successful = actions.iter().filter(|a| a.success).count();
let total = actions.len();
let success_rate = successful as f64 / total as f64;
if success_rate < self.min_action_success_rate {
issues.push(format!(
"Action success rate {:.1}% below threshold {:.1}%",
success_rate * 100.0,
self.min_action_success_rate * 100.0
));
}
// Count failures
let failures = actions.iter().filter(|a| !a.success).count() as u32;
if failures > self.max_failures {
issues.push(format!(
"Too many failures: {} (max {})",
failures, self.max_failures
));
}
// Check for critical errors
for action in actions.iter().filter(|a| !a.success) {
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));
}
}
// Check job state
if job.state != crate::context::JobState::Completed
&& job.state != crate::context::JobState::Submitted
{
issues.push(format!("Job not in completed state: {:?}", job.state));
}
// Calculate quality score
let quality_score = if issues.is_empty() {
let base_score = (success_rate * 80.0) as u32;
let completion_bonus = if job.state == crate::context::JobState::Completed {
20
} else {
0
};
(base_score + completion_bonus).min(100)
} else {
((success_rate * 50.0) as u32).min(50)
};
if issues.is_empty() {
Ok(EvaluationResult::success(
format!(
"Job completed successfully with {}/{} actions succeeding ({:.1}%)",
successful,
total,
success_rate * 100.0
),
quality_score,
))
} else {
Ok(EvaluationResult {
success: false,
confidence: 0.85,
reasoning: format!("Job had {} issues", issues.len()),
issues,
suggestions: vec![
"Review failed actions for common patterns".to_string(),
"Consider adjusting retry logic".to_string(),
],
quality_score,
})
}
}
}
/// LLM-based success evaluator for more nuanced evaluation.
pub struct LlmEvaluator {
llm: Arc<dyn LlmProvider>,
}
impl LlmEvaluator {
/// Create a new LLM-based evaluator.
#[allow(dead_code)] // Public API for LLM-based evaluation
pub fn new(llm: Arc<dyn LlmProvider>) -> Self {
Self { llm }
}
}
#[async_trait]
impl SuccessEvaluator for LlmEvaluator {
async fn evaluate(
&self,
job: &JobContext,
actions: &[ActionRecord],
output: Option<&str>,
) -> Result<EvaluationResult, EvaluationError> {
// Build evaluation prompt
let actions_summary: Vec<String> = actions
.iter()
.map(|a| {
format!(
"- {}: {} ({})",
a.tool_name,
if a.success { "success" } else { "failed" },
a.error.as_deref().unwrap_or("ok")
)
})
.collect();
let prompt = format!(
r#"Evaluate if this job was completed successfully.
Job: {}
Description: {}
State: {:?}
Actions taken:
{}
{}
Respond in JSON format:
{{
"success": true/false,
"confidence": 0.0-1.0,
"reasoning": "...",
"issues": ["..."],
"suggestions": ["..."],
"quality_score": 0-100
}}"#,
job.title,
job.description,
job.state,
actions_summary.join("\n"),
output
.map(|o| format!("Output:\n{}", o))
.unwrap_or_default()
);
let request =
crate::llm::CompletionRequest::new(vec![crate::llm::ChatMessage::user(prompt)])
.with_max_tokens(1024)
.with_temperature(0.1);
let response = self
.llm
.complete(request)
.await
.map_err(|e| EvaluationError::Failed {
job_id: job.job_id,
reason: e.to_string(),
})?;
// Parse the response
let result: EvaluationResult =
serde_json::from_str(&response.content).map_err(|e| EvaluationError::Failed {
job_id: job.job_id,
reason: format!("Failed to parse LLM evaluation: {}", e),
})?;
Ok(result)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::context::JobContext;
use crate::context::{ActionRecord, JobContext};
use crate::error::EvaluationError;
/// Rule-based success evaluator (test-only; no production callers).
struct RuleBasedEvaluator {
min_action_success_rate: f64,
max_failures: u32,
}
impl RuleBasedEvaluator {
fn new() -> Self {
Self {
min_action_success_rate: 0.8,
max_failures: 3,
}
}
fn with_min_success_rate(mut self, rate: f64) -> Self {
self.min_action_success_rate = rate;
self
}
fn with_max_failures(mut self, max: u32) -> Self {
self.max_failures = max;
self
}
}
impl Default for RuleBasedEvaluator {
fn default() -> Self {
Self::new()
}
}
#[async_trait::async_trait]
impl SuccessEvaluator for RuleBasedEvaluator {
async fn evaluate(
&self,
job: &JobContext,
actions: &[ActionRecord],
_output: Option<&str>,
) -> Result<EvaluationResult, EvaluationError> {
let mut issues = Vec::new();
if actions.is_empty() {
return Ok(EvaluationResult::failure(
"No actions were taken",
vec!["No actions recorded".to_string()],
));
}
let successful = actions.iter().filter(|a| a.success).count();
let total = actions.len();
let success_rate = successful as f64 / total as f64;
if success_rate < self.min_action_success_rate {
issues.push(format!(
"Action success rate {:.1}% below threshold {:.1}%",
success_rate * 100.0,
self.min_action_success_rate * 100.0
));
}
let failures = actions.iter().filter(|a| !a.success).count() as u32;
if failures > self.max_failures {
issues.push(format!(
"Too many failures: {} (max {})",
failures, self.max_failures
));
}
for action in actions.iter().filter(|a| !a.success) {
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));
}
}
if job.state != crate::context::JobState::Completed
&& job.state != crate::context::JobState::Submitted
{
issues.push(format!("Job not in completed state: {:?}", job.state));
}
let quality_score = if issues.is_empty() {
let base_score = (success_rate * 80.0) as u32;
let completion_bonus = if job.state == crate::context::JobState::Completed {
20
} else {
0
};
(base_score + completion_bonus).min(100)
} else {
((success_rate * 50.0) as u32).min(50)
};
if issues.is_empty() {
Ok(EvaluationResult::success(
format!(
"Job completed successfully with {}/{} actions succeeding ({:.1}%)",
successful,
total,
success_rate * 100.0
),
quality_score,
))
} else {
Ok(EvaluationResult {
success: false,
confidence: 0.85,
reasoning: format!("Job had {} issues", issues.len()),
issues,
suggestions: vec![
"Review failed actions for common patterns".to_string(),
"Consider adjusting retry logic".to_string(),
],
quality_score,
})
}
}
}
#[tokio::test]
async fn test_rule_based_evaluator_success() {
+2 -33
View File
@@ -1405,38 +1405,6 @@ impl ExtensionManager {
Ok(())
}
#[allow(dead_code)] // Used by upcoming hot-activation flow
async fn install_bundled_channel_from_artifacts(
&self,
name: &str,
) -> Result<InstallResult, ExtensionError> {
// Check if already installed
let channel_wasm = self.wasm_channels_dir.join(format!("{}.wasm", name));
if channel_wasm.exists() {
return Err(ExtensionError::AlreadyInstalled(name.to_string()));
}
crate::channels::wasm::install_bundled_channel(name, &self.wasm_channels_dir, false)
.await
.map_err(ExtensionError::InstallFailed)?;
tracing::info!(
"Installed bundled channel '{}' to {}",
name,
self.wasm_channels_dir.display()
);
Ok(InstallResult {
name: name.to_string(),
kind: ExtensionKind::WasmChannel,
message: format!(
"Channel '{}' installed. \
Run tool_auth('{}') to configure authentication, then activate.",
name, name,
),
})
}
/// Install a WASM extension from local build artifacts (WasmBuildable source).
///
/// Resolves the build directory (relative to `CARGO_MANIFEST_DIR` or absolute),
@@ -1628,6 +1596,7 @@ impl ExtensionManager {
&metadata.scopes_supported,
Some(&pkce),
&std::collections::HashMap::new(),
None,
);
// Store pending auth for later callback handling
@@ -2508,7 +2477,7 @@ impl ExtensionManager {
&self.user_id,
)
} else {
McpClient::new_with_name(&server.name, &server.url)
McpClient::new_with_config(server.clone())
};
// Try to list and create tools
+3
View File
@@ -241,6 +241,9 @@ impl Store {
tool_output_stash: std::sync::Arc::new(tokio::sync::RwLock::new(
std::collections::HashMap::new(),
)),
// TODO(#661): persist user_timezone in agent_jobs table so
// background/routine jobs retain the session's timezone context.
user_timezone: "UTC".to_string(),
}))
}
None => Ok(None),
+1
View File
@@ -66,6 +66,7 @@ pub mod service;
pub mod settings;
pub mod setup;
pub mod skills;
pub mod timezone;
pub mod tools;
pub mod tracing_fmt;
pub mod transcription;
+1148
View File
File diff suppressed because it is too large Load Diff
+95
View File
@@ -0,0 +1,95 @@
//! Image generation model detection utilities.
/// Known image generation model families.
const IMAGE_GEN_PATTERNS: &[&str] = &[
"flux",
"dall-e",
"dalle",
"stable-diffusion",
"sdxl",
"imagen",
"midjourney",
"ideogram",
"playground",
];
/// Check if a model name indicates an image generation model.
pub fn is_image_generation_model(model: &str) -> bool {
let lower = model.to_lowercase();
IMAGE_GEN_PATTERNS.iter().any(|p| lower.contains(p))
}
/// Suggest the best image generation model from a list of available models.
///
/// Priority: FLUX > DALL-E > Stable Diffusion > others.
pub fn suggest_image_model(models: &[String]) -> Option<&str> {
let priorities: &[&str] = &[
"flux",
"dall-e",
"dalle",
"stable-diffusion",
"sdxl",
"imagen",
];
for priority in priorities {
if let Some(model) = models.iter().find(|m| m.to_lowercase().contains(priority)) {
return Some(model);
}
}
// Fall back to any image gen model
models.iter().find_map(|m| {
if is_image_generation_model(m) {
Some(m.as_str())
} else {
None
}
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn detects_flux_models() {
assert!(is_image_generation_model(
"black-forest-labs/FLUX.1-schnell"
));
assert!(is_image_generation_model("flux-pro"));
}
#[test]
fn detects_dalle_models() {
assert!(is_image_generation_model("dall-e-3"));
assert!(is_image_generation_model("dalle-3"));
}
#[test]
fn rejects_non_image_models() {
assert!(!is_image_generation_model("gpt-4o"));
assert!(!is_image_generation_model("claude-3-sonnet"));
assert!(!is_image_generation_model("llama-3.1-70b"));
}
#[test]
fn suggests_flux_first() {
let models = vec![
"gpt-4o".to_string(),
"dall-e-3".to_string(),
"flux-pro".to_string(),
];
assert_eq!(suggest_image_model(&models), Some("flux-pro"));
}
#[test]
fn suggests_dalle_without_flux() {
let models = vec!["gpt-4o".to_string(), "dall-e-3".to_string()];
assert_eq!(suggest_image_model(&models), Some("dall-e-3"));
}
#[test]
fn returns_none_when_no_image_models() {
let models = vec!["gpt-4o".to_string(), "claude-3-sonnet".to_string()];
assert_eq!(suggest_image_model(&models), None);
}
}
+64 -7
View File
@@ -6,8 +6,11 @@
//! - **Anthropic**: Direct API access with your own key
//! - **Ollama**: Local model inference
//! - **OpenAI-compatible**: Any endpoint that speaks the OpenAI API
//! - **AWS Bedrock**: Native Converse API via aws-sdk-bedrockruntime
mod anthropic_oauth;
#[cfg(feature = "bedrock")]
mod bedrock;
pub mod circuit_breaker;
pub mod costs;
pub mod failover;
@@ -22,6 +25,9 @@ mod rig_adapter;
pub mod session;
pub mod smart_routing;
pub mod image_models;
pub mod vision_models;
pub use circuit_breaker::{CircuitBreakerConfig, CircuitBreakerProvider};
pub use failover::{CooldownConfig, FailoverProvider};
pub use nearai_chat::{ModelInfo, NearAiChatProvider};
@@ -54,12 +60,29 @@ use crate::error::LlmError;
///
/// - NearAI backend: Uses session manager for authentication
/// - Registry providers: Looked up by protocol and constructed generically
pub fn create_llm_provider(
pub async fn create_llm_provider(
config: &LlmConfig,
session: Arc<SessionManager>,
) -> Result<Arc<dyn LlmProvider>, LlmError> {
let timeout = config.request_timeout_secs;
if config.backend == "nearai" || config.backend == "near_ai" || config.backend == "near" {
return create_llm_provider_with_config(&config.nearai, session);
return create_llm_provider_with_config(&config.nearai, session, timeout);
}
// Bedrock uses a native AWS SDK, not the rig-core registry
if config.backend == "bedrock" {
#[cfg(feature = "bedrock")]
{
return create_bedrock_provider(config).await;
}
#[cfg(not(feature = "bedrock"))]
{
return Err(LlmError::RequestFailed {
provider: "bedrock".to_string(),
reason: "Bedrock support not compiled. Rebuild with --features bedrock".to_string(),
});
}
}
let reg_config = config
@@ -79,6 +102,7 @@ pub fn create_llm_provider(
pub fn create_llm_provider_with_config(
config: &NearAiConfig,
session: Arc<SessionManager>,
request_timeout_secs: u64,
) -> Result<Arc<dyn LlmProvider>, LlmError> {
let auth_mode = if config.api_key.is_some() {
"API key"
@@ -89,9 +113,14 @@ pub fn create_llm_provider_with_config(
model = %config.model,
base_url = %config.base_url,
auth = auth_mode,
timeout_secs = request_timeout_secs,
"Using NEAR AI (Chat Completions API)"
);
Ok(Arc::new(NearAiChatProvider::new(config.clone(), session)?))
Ok(Arc::new(NearAiChatProvider::new_with_timeout(
config.clone(),
session,
request_timeout_secs,
)?))
}
/// Create a provider from a registry-resolved config.
@@ -109,6 +138,24 @@ fn create_registry_provider(
}
}
#[cfg(feature = "bedrock")]
async fn create_bedrock_provider(config: &LlmConfig) -> Result<Arc<dyn LlmProvider>, LlmError> {
let br = config
.bedrock
.as_ref()
.ok_or_else(|| LlmError::AuthFailed {
provider: "bedrock".to_string(),
})?;
let provider = bedrock::BedrockProvider::new(br).await?;
tracing::info!(
"Using AWS Bedrock (Converse API, region: {}, model: {})",
br.region,
provider.active_model_name(),
);
Ok(Arc::new(provider))
}
fn create_openai_compat_from_registry(
config: &RegistryProviderConfig,
) -> Result<Arc<dyn LlmProvider>, LlmError> {
@@ -333,7 +380,7 @@ pub fn create_cheap_llm_provider(
/// This is the single source of truth for provider chain construction,
/// called by both `main.rs` and `app.rs`.
#[allow(clippy::type_complexity)]
pub fn build_provider_chain(
pub async fn build_provider_chain(
config: &LlmConfig,
session: Arc<SessionManager>,
) -> Result<
@@ -344,7 +391,7 @@ pub fn build_provider_chain(
),
LlmError,
> {
let llm = create_llm_provider(config, session.clone())?;
let llm = create_llm_provider(config, session.clone()).await?;
tracing::info!("LLM provider initialized: {}", llm.model_name());
// 1. Retry
@@ -365,7 +412,11 @@ pub fn build_provider_chain(
let llm: Arc<dyn LlmProvider> = if let Some(ref cheap_model) = config.nearai.cheap_model {
let mut cheap_config = config.nearai.clone();
cheap_config.model = cheap_model.clone();
let cheap = create_llm_provider_with_config(&cheap_config, session.clone())?;
let cheap = create_llm_provider_with_config(
&cheap_config,
session.clone(),
config.request_timeout_secs,
)?;
let cheap: Arc<dyn LlmProvider> = if retry_config.max_retries > 0 {
Arc::new(RetryProvider::new(cheap, retry_config.clone()))
} else {
@@ -397,7 +448,11 @@ pub fn build_provider_chain(
}
let mut fallback_config = config.nearai.clone();
fallback_config.model = fallback_model.clone();
let fallback = create_llm_provider_with_config(&fallback_config, session.clone())?;
let fallback = create_llm_provider_with_config(
&fallback_config,
session.clone(),
config.request_timeout_secs,
)?;
tracing::info!(
primary = %llm.model_name(),
fallback = %fallback.model_name(),
@@ -503,6 +558,8 @@ mod tests {
session: SessionConfig::default(),
nearai: test_nearai_config(),
provider: None,
bedrock: None,
request_timeout_secs: 120,
}
}
+15 -4
View File
@@ -58,17 +58,28 @@ impl NearAiChatProvider {
/// By default this enables tool-message flattening for compatibility with
/// providers that reject `role: "tool"` messages.
pub fn new(config: NearAiConfig, session: Arc<SessionManager>) -> Result<Self, LlmError> {
Self::new_with_flatten(config, session, true)
Self::new_with_options(config, session, true, 120)
}
/// Create a chat completions provider with configurable tool-message flattening.
pub fn new_with_flatten(
/// Create a new provider with a custom request timeout.
pub fn new_with_timeout(
config: NearAiConfig,
session: Arc<SessionManager>,
request_timeout_secs: u64,
) -> Result<Self, LlmError> {
Self::new_with_options(config, session, true, request_timeout_secs)
}
/// Create a chat completions provider with configurable tool-message flattening
/// and request timeout.
pub fn new_with_options(
config: NearAiConfig,
session: Arc<SessionManager>,
flatten_tool_messages: bool,
request_timeout_secs: u64,
) -> Result<Self, LlmError> {
let client = Client::builder()
.timeout(std::time::Duration::from_secs(120))
.timeout(std::time::Duration::from_secs(request_timeout_secs))
.build()
.map_err(|e| LlmError::RequestFailed {
provider: "nearai_chat".to_string(),
+2 -12
View File
@@ -11,7 +11,6 @@ use crate::llm::{
ChatMessage, CompletionRequest, LlmProvider, Role, ToolCall, ToolCompletionRequest,
ToolDefinition,
};
use crate::safety::SafetyLayer;
/// Token the agent returns when it has nothing to say (e.g. in group chats).
/// The dispatcher should check for this and suppress the message.
@@ -343,8 +342,6 @@ pub struct RespondOutput {
/// Reasoning engine for the agent.
pub struct Reasoning {
llm: Arc<dyn LlmProvider>,
#[allow(dead_code)] // Will be used for sanitizing tool outputs
safety: Arc<SafetyLayer>,
/// Optional workspace for loading identity/system prompts.
workspace_system_prompt: Option<String>,
/// Optional skill context block to inject into system prompt.
@@ -362,10 +359,9 @@ pub struct Reasoning {
impl Reasoning {
/// Create a new reasoning engine.
pub fn new(llm: Arc<dyn LlmProvider>, safety: Arc<SafetyLayer>) -> Self {
pub fn new(llm: Arc<dyn LlmProvider>) -> Self {
Self {
llm,
safety,
workspace_system_prompt: None,
skill_context: None,
channel: None,
@@ -2117,15 +2113,9 @@ That's my plan."#;
// ---- System prompt building tests (issue #565) ----
fn make_test_reasoning() -> Reasoning {
use crate::config::SafetyConfig;
use crate::safety::SafetyLayer;
use crate::testing::StubLlm;
let llm = Arc::new(StubLlm::new("test"));
let safety = Arc::new(SafetyLayer::new(&SafetyConfig {
max_output_length: 100_000,
injection_check_enabled: false,
}));
Reasoning::new(llm, safety)
Reasoning::new(llm)
}
#[test]
+2 -1
View File
@@ -200,9 +200,10 @@ impl SessionManager {
let status = response.status();
let body = response.text().await.unwrap_or_default();
let preview = crate::agent::truncate_for_preview(&body, 200);
Err(LlmError::SessionRenewalFailed {
provider: "nearai".to_string(),
reason: format!("Validation failed: HTTP {}: {}", status, body),
reason: format!("Validation failed: HTTP {status}: {preview}"),
})
}
+104
View File
@@ -0,0 +1,104 @@
//! Vision model detection utilities.
/// Known vision-capable model families.
const VISION_PATTERNS: &[&str] = &[
"claude-3",
"claude-4",
"gpt-4o",
"gpt-4-turbo",
"gpt-4-vision",
"gemini-pro-vision",
"gemini-1.5",
"gemini-2",
"llava",
"cogvlm",
"internvl",
"qwen-vl",
"qwen2-vl",
"pixtral",
];
/// Check if a model name indicates vision capabilities.
pub fn is_vision_model(model: &str) -> bool {
let lower = model.to_lowercase();
VISION_PATTERNS.iter().any(|p| lower.contains(p))
}
/// Suggest the best vision model from a list of available models.
///
/// Priority: Claude > GPT-4 > Gemini > others.
pub fn suggest_vision_model(models: &[String]) -> Option<&str> {
let priorities: &[&str] = &[
"claude-3",
"claude-4",
"gpt-4o",
"gpt-4-turbo",
"gpt-4-vision",
"gemini",
"llava",
"pixtral",
];
for priority in priorities {
if let Some(model) = models.iter().find(|m| m.to_lowercase().contains(priority)) {
return Some(model);
}
}
models.iter().find_map(|m| {
if is_vision_model(m) {
Some(m.as_str())
} else {
None
}
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn detects_claude_vision() {
assert!(is_vision_model("claude-3-5-sonnet-20241022"));
assert!(is_vision_model("claude-3-opus"));
assert!(is_vision_model("claude-4-sonnet"));
}
#[test]
fn detects_gpt4_vision() {
assert!(is_vision_model("gpt-4o"));
assert!(is_vision_model("gpt-4-turbo"));
assert!(is_vision_model("gpt-4-vision-preview"));
}
#[test]
fn detects_other_vision_models() {
assert!(is_vision_model("gemini-1.5-pro"));
assert!(is_vision_model("llava-v1.6"));
assert!(is_vision_model("pixtral-12b"));
}
#[test]
fn rejects_non_vision_models() {
assert!(!is_vision_model("gpt-3.5-turbo"));
assert!(!is_vision_model("llama-3.1-70b"));
assert!(!is_vision_model("mistral-7b"));
}
#[test]
fn suggests_claude_first() {
let models = vec![
"gpt-4o".to_string(),
"claude-3-5-sonnet-20241022".to_string(),
];
assert_eq!(
suggest_vision_model(&models),
Some("claude-3-5-sonnet-20241022")
);
}
#[test]
fn returns_none_when_no_vision_models() {
let models = vec!["gpt-3.5-turbo".to_string(), "llama-3.1-70b".to_string()];
assert_eq!(suggest_vision_model(&models), None);
}
}
+27 -7
View File
@@ -75,7 +75,7 @@ async fn async_main() -> anyhow::Result<()> {
}
Some(Command::Mcp(mcp_cmd)) => {
init_cli_tracing();
return run_mcp_command(mcp_cmd.clone()).await;
return run_mcp_command(*mcp_cmd.clone()).await;
}
Some(Command::Memory(mem_cmd)) => {
init_cli_tracing();
@@ -145,6 +145,24 @@ async fn async_main() -> anyhow::Result<()> {
}
}
// ── PID lock (prevent multiple instances) ────────────────────────
let _pid_lock = match ironclaw::bootstrap::PidLock::acquire() {
Ok(lock) => Some(lock),
Err(ironclaw::bootstrap::PidLockError::AlreadyRunning { pid }) => {
anyhow::bail!(
"Another IronClaw instance is already running (PID {}). \
If this is incorrect, remove the stale PID file: {}",
pid,
ironclaw::bootstrap::pid_lock_path().display()
);
}
Err(e) => {
eprintln!("Warning: Could not acquire PID lock: {}", e);
eprintln!("Continuing without PID lock protection.");
None
}
};
// ── Agent startup ──────────────────────────────────────────────────
// Enhanced first-run detection
@@ -166,13 +184,12 @@ async fn async_main() -> anyhow::Result<()> {
let config = match Config::from_env_with_toml(toml_path).await {
Ok(c) => c,
Err(ironclaw::error::ConfigError::MissingRequired { key, hint }) => {
eprintln!("Configuration error: Missing required setting '{}'", key);
eprintln!(" {}", hint);
eprintln!();
eprintln!(
"Run 'ironclaw onboard' to configure, or set the required environment variables."
anyhow::bail!(
"Configuration error: Missing required setting '{}'. {}. \
Run 'ironclaw onboard' to configure, or set the required environment variables.",
key,
hint
);
std::process::exit(1);
}
Err(e) => return Err(e.into()),
};
@@ -706,6 +723,9 @@ async fn async_main() -> anyhow::Result<()> {
// ── Shutdown ────────────────────────────────────────────────────────
// Shut down all stdio MCP server child processes.
components.mcp_process_manager.shutdown_all().await;
// Flush LLM trace recording if enabled
if let Some(ref recorder) = components.recording_handle
&& let Err(e) = recorder.flush().await
+66 -1
View File
@@ -42,8 +42,12 @@ pub struct Settings {
#[serde(default)]
pub secrets_master_key_source: KeySource,
/// Generated master key hex (env var mode only, written to .env by wizard).
#[serde(default, skip_serializing)]
pub secrets_master_key_hex: Option<String>,
// === Step 3: Inference Provider ===
/// LLM backend: "nearai", "anthropic", "openai", "ollama", "openai_compatible".
/// LLM backend: "nearai", "anthropic", "openai", "ollama", "openai_compatible", "tinfoil", "bedrock".
#[serde(default)]
pub llm_backend: Option<String>,
@@ -55,6 +59,18 @@ pub struct Settings {
#[serde(default)]
pub openai_compatible_base_url: Option<String>,
/// Bedrock region (when llm_backend = "bedrock").
#[serde(default)]
pub bedrock_region: Option<String>,
/// Bedrock cross-region inference prefix (when llm_backend = "bedrock").
#[serde(default)]
pub bedrock_cross_region: Option<String>,
/// AWS profile name for Bedrock (when llm_backend = "bedrock").
#[serde(default)]
pub bedrock_profile: Option<String>,
// === Step 4: Model Selection ===
/// Currently selected model.
#[serde(default)]
@@ -291,6 +307,18 @@ pub struct HeartbeatSettings {
/// User ID to notify on heartbeat findings.
#[serde(default)]
pub notify_user: Option<String>,
/// Hour (0-23) when quiet hours start (heartbeat skipped).
#[serde(default)]
pub quiet_hours_start: Option<u32>,
/// Hour (0-23) when quiet hours end (heartbeat resumes).
#[serde(default)]
pub quiet_hours_end: Option<u32>,
/// Timezone for quiet hours evaluation (IANA name, e.g. "America/New_York").
#[serde(default)]
pub timezone: Option<String>,
}
fn default_heartbeat_interval() -> u64 {
@@ -304,6 +332,9 @@ impl Default for HeartbeatSettings {
interval_secs: default_heartbeat_interval(),
notify_channel: None,
notify_user: None,
quiet_hours_start: None,
quiet_hours_end: None,
timezone: None,
}
}
}
@@ -351,6 +382,10 @@ pub struct AgentSettings {
/// When true, skip tool approval checks entirely. For benchmarks/CI.
#[serde(default)]
pub auto_approve_tools: bool,
/// Default timezone for new sessions (IANA name, e.g. "America/New_York").
#[serde(default = "default_timezone")]
pub default_timezone: String,
}
fn default_agent_name() -> String {
@@ -385,6 +420,10 @@ fn default_max_tool_iterations() -> usize {
50
}
fn default_timezone() -> String {
"UTC".to_string()
}
fn default_true() -> bool {
true
}
@@ -402,6 +441,7 @@ impl Default for AgentSettings {
session_idle_timeout_secs: default_session_idle_timeout(),
max_tool_iterations: default_max_tool_iterations(),
auto_approve_tools: false,
default_timezone: default_timezone(),
}
}
}
@@ -1174,6 +1214,31 @@ mod tests {
assert_eq!(loaded.heartbeat.interval_secs, 900);
}
/// Regression test: /model command must persist selected_model to TOML config.
/// Prior to the fix, `set_model()` only changed the in-memory provider and the
/// choice was lost on restart.
#[test]
fn toml_selected_model_update_persists() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("config.toml");
// Start with a config that has a different model.
let settings = Settings {
selected_model: Some("old-model".to_string()),
..Default::default()
};
settings.save_toml(&path).unwrap();
// Simulate what persist_selected_model does: load, update, save.
let mut loaded = Settings::load_toml(&path).unwrap().unwrap();
loaded.selected_model = Some("new-model".to_string());
loaded.save_toml(&path).unwrap();
// Verify the change survived a reload.
let reloaded = Settings::load_toml(&path).unwrap().unwrap();
assert_eq!(reloaded.selected_model, Some("new-model".to_string()));
}
#[test]
fn toml_missing_file_returns_none() {
let result = Settings::load_toml(std::path::Path::new("/tmp/nonexistent_config.toml"));
+2 -1
View File
@@ -174,6 +174,7 @@ env-var mode or skipped secrets.
| Ollama | None | - | - |
| OpenRouter¹ | API key | `llm_compatible_api_key` | `LLM_API_KEY` |
| OpenAI-compatible¹ | Optional API key | `llm_compatible_api_key` | `LLM_API_KEY` |
| AWS Bedrock | AWS credentials (IAM, SSO, instance roles) | - | - |
¹ OpenRouter and OpenAI-compatible share the same secret name and env var because
OpenRouter is stored as `llm_backend = "openai_compatible"` under the hood.
@@ -479,7 +480,7 @@ pub struct Settings {
pub secrets_master_key_source: KeySource, // Keychain | Env | None
// Step 3: Inference
pub llm_backend: Option<String>, // "nearai" | "anthropic" | "openai" | "ollama" | "openai_compatible"
pub llm_backend: Option<String>, // "nearai" | "anthropic" | "openai" | "ollama" | "openai_compatible" | "bedrock"
pub ollama_base_url: Option<String>,
pub openai_compatible_base_url: Option<String>,
+330 -43
View File
@@ -769,13 +769,28 @@ impl SetupWizard {
print_success("Master key generated and stored in OS keychain");
}
1 => {
// Env var mode
print_info("Generate a key and add it to your environment:");
// Env var mode — generate key, init crypto, and persist to .env
let key_hex = crate::secrets::keychain::generate_master_key_hex();
// Initialize crypto so subsequent wizard steps (channel setup,
// API key storage) can encrypt secrets immediately.
self.secrets_crypto = Some(Arc::new(
SecretsCrypto::new(SecretString::from(key_hex.clone()))
.map_err(|e| SetupError::Config(e.to_string()))?,
));
// Make visible to optional_env() for any subsequent config resolution.
crate::config::inject_single_var("SECRETS_MASTER_KEY", &key_hex);
// Store hex for write_bootstrap_env to persist to ~/.ironclaw/.env.
self.settings.secrets_master_key_hex = Some(key_hex.clone());
println!();
println!(" export SECRETS_MASTER_KEY={}", key_hex);
print_info("Master key generated and will be saved to ~/.ironclaw/.env");
println!();
print_info("Add this to your shell profile or .env file.");
println!(" SECRETS_MASTER_KEY={}", key_hex);
println!();
print_info("You can also copy this to another .env file or CI secrets.");
self.settings.secrets_master_key_source = KeySource::Env;
print_success("Configured for environment variable");
@@ -812,9 +827,16 @@ impl SetupWizard {
print_info(&format!("Current provider: {}", display));
println!();
let is_known = current == "nearai" || registry.is_known(&current);
let is_known =
current == "nearai" || current == "bedrock" || registry.is_known(&current);
if is_known && confirm("Keep current provider?", true).map_err(SetupError::Io)? {
if current == "bedrock" {
// Keeping the existing Bedrock config — no need to re-run
// the full setup flow (region, auth, cross-region).
print_info("Keeping existing AWS Bedrock configuration.");
return Ok(());
}
return self.run_provider_setup(&current, &registry).await;
}
@@ -829,10 +851,10 @@ impl SetupWizard {
print_info("Select your inference provider:");
println!();
// Build menu: NearAI first, then all registry providers with setup hints
// Build menu: NearAI first, then all registry providers with setup hints, then Bedrock
let selectable = registry.selectable();
let mut options: Vec<String> = Vec::with_capacity(1 + selectable.len());
let mut provider_ids: Vec<String> = Vec::with_capacity(1 + selectable.len());
let mut options: Vec<String> = Vec::with_capacity(2 + selectable.len());
let mut provider_ids: Vec<String> = Vec::with_capacity(2 + selectable.len());
options.push("NEAR AI - multi-model access via NEAR account".to_string());
provider_ids.push("nearai".to_string());
@@ -850,11 +872,19 @@ impl SetupWizard {
provider_ids.push(def.id.clone());
}
// Bedrock is a special case (native AWS SDK, not registry-based)
options.push("AWS Bedrock - Claude & other models via AWS (IAM, SSO)".to_string());
provider_ids.push("bedrock".to_string());
let option_refs: Vec<&str> = options.iter().map(|s| s.as_str()).collect();
let choice = select_one("Provider:", &option_refs).map_err(SetupError::Io)?;
let selected_id = &provider_ids[choice];
self.run_provider_setup(selected_id, &registry).await?;
if selected_id == "bedrock" {
self.setup_bedrock().await?;
} else {
self.run_provider_setup(selected_id, &registry).await?;
}
Ok(())
}
@@ -1022,10 +1052,11 @@ impl SetupWizard {
/// Anthropic OAuth setup: extract token from `claude login` credentials.
async fn setup_anthropic_oauth(&mut self) -> Result<(), SetupError> {
self.settings.llm_backend = Some("anthropic".to_string());
if self.settings.selected_model.is_some() {
// Clear model only when switching providers (old model may be invalid)
if self.settings.llm_backend.as_deref() != Some("anthropic") {
self.settings.selected_model = None;
}
self.settings.llm_backend = Some("anthropic".to_string());
// Try to extract existing OAuth token from Claude Code credentials
if let Some(token) = crate::config::ClaudeCodeConfig::extract_oauth_token() {
@@ -1119,10 +1150,11 @@ impl SetupWizard {
other => other,
});
self.settings.llm_backend = Some(backend.to_string());
if self.settings.selected_model.is_some() {
// Clear model only when switching providers (old model may be invalid)
if self.settings.llm_backend.as_deref() != Some(backend) {
self.settings.selected_model = None;
}
self.settings.llm_backend = Some(backend.to_string());
// Check env var first
if let Ok(existing) = std::env::var(env_var) {
@@ -1181,10 +1213,11 @@ impl SetupWizard {
&mut self,
def: &crate::llm::ProviderDefinition,
) -> Result<(), SetupError> {
self.settings.llm_backend = Some(def.id.clone());
if self.settings.selected_model.is_some() {
// Clear model only when switching providers (old model may be invalid)
if self.settings.llm_backend.as_deref() != Some(&def.id) {
self.settings.selected_model = None;
}
self.settings.llm_backend = Some(def.id.clone());
let default_url = self
.settings
@@ -1212,6 +1245,95 @@ impl SetupWizard {
Ok(())
}
/// AWS Bedrock provider setup: region, auth, and cross-region config.
async fn setup_bedrock(&mut self) -> Result<(), SetupError> {
if self.settings.llm_backend.as_deref() != Some("bedrock") {
self.settings.selected_model = None;
}
self.settings.llm_backend = Some("bedrock".to_string());
// Region
let default_region = self
.settings
.bedrock_region
.as_deref()
.unwrap_or("us-east-1");
let region_input =
optional_input("AWS region", Some(&format!("default: {}", default_region)))
.map_err(SetupError::Io)?;
let region = region_input.unwrap_or_else(|| default_region.to_string());
self.settings.bedrock_region = Some(region.clone());
// Auth method
print_info("Select authentication method:");
println!();
let auth_options = &[
"AWS default credentials (env vars, ~/.aws/credentials, IAM roles)",
"AWS named profile (SSO / assume-role)",
];
let auth_choice = select_one("Auth:", auth_options).map_err(SetupError::Io)?;
match auth_choice {
0 => {
// Default AWS credentials — clear any stale named profile
self.settings.bedrock_profile = None;
print_info(
"Using default AWS credential chain (env vars, ~/.aws/credentials, IAM roles).",
);
}
1 => {
// Named profile
let profile =
input("AWS profile name (from ~/.aws/config)").map_err(SetupError::Io)?;
if profile.trim().is_empty() {
// Empty input clears any previously configured profile
self.settings.bedrock_profile = None;
print_info("AWS profile cleared; using default AWS credential chain instead.");
} else {
self.settings.bedrock_profile = Some(profile.clone());
print_success(&format!("AWS profile '{}' saved", profile));
}
}
_ => return Err(SetupError::Config("Invalid auth selection".to_string())),
}
self.setup_bedrock_cross_region()
}
/// Bedrock cross-region inference prefix selection (sub-step of setup_bedrock).
fn setup_bedrock_cross_region(&mut self) -> Result<(), SetupError> {
print_info("Cross-region inference routes requests across AWS regions for capacity:");
println!();
let cross_options = &[
"us - route within US regions (recommended for us-east-1)",
"global - route to any AWS region worldwide",
"eu - route within European regions",
"apac - route within Asia-Pacific regions",
"none - single-region only (no cross-region routing)",
];
let cross_choice = select_one("Cross-region:", cross_options).map_err(SetupError::Io)?;
let cross_region = match cross_choice {
0 => Some("us".to_string()),
1 => Some("global".to_string()),
2 => Some("eu".to_string()),
3 => Some("apac".to_string()),
4 => None,
_ => None,
};
self.settings.bedrock_cross_region = cross_region;
let region = self
.settings
.bedrock_region
.as_deref()
.unwrap_or("us-east-1");
print_success(&format!("AWS Bedrock configured (region: {})", region));
Ok(())
}
/// Generic OpenAI-compatible setup: base URL + optional API key.
async fn setup_openai_compatible_generic(
&mut self,
@@ -1219,10 +1341,11 @@ impl SetupWizard {
secret_name: &str,
display_name: &str,
) -> Result<(), SetupError> {
self.settings.llm_backend = Some(backend_id.to_string());
if self.settings.selected_model.is_some() {
// Clear model only when switching providers (old model may be invalid)
if self.settings.llm_backend.as_deref() != Some(backend_id) {
self.settings.selected_model = None;
}
self.settings.llm_backend = Some(backend_id.to_string());
let existing_url = self
.settings
@@ -1393,6 +1516,14 @@ impl SetupWizard {
self.settings.selected_model = Some(model_id.clone());
print_success(&format!("Selected {}", model_id));
}
} else if backend == "bedrock" {
let model_id = input("Bedrock model ID (e.g., anthropic.claude-opus-4-6-v1)")
.map_err(SetupError::Io)?;
if model_id.is_empty() {
return Err(SetupError::Config("Model ID is required".to_string()));
}
self.settings.selected_model = Some(model_id.clone());
print_success(&format!("Selected {}", model_id));
} else {
// Unknown provider, manual entry
let model_id = input("Model name (e.g., meta-llama/Llama-3-8b-chat-hf)")
@@ -1476,9 +1607,11 @@ impl SetupWizard {
smart_routing_cascade: true,
},
provider: None,
bedrock: None,
request_timeout_secs: 120,
};
match create_llm_provider(&config, session) {
match create_llm_provider(&config, session).await {
Ok(provider) => match provider.list_models().await {
Ok(models) => models,
Err(e) => {
@@ -1592,48 +1725,50 @@ impl SetupWizard {
};
// Create backend-appropriate secrets store.
// Respect the user's selected backend when both features are compiled,
// so we don't accidentally use a postgres pool from DATABASE_URL when
// libsql was chosen (or vice versa).
// Use runtime dispatch based on the user's selected backend.
// Default to whichever backend is compiled in. When only libsql is
// available, we must not default to "postgres" or we'd skip store creation.
let default_backend = {
#[cfg(feature = "postgres")]
{
"postgres"
}
#[cfg(not(feature = "postgres"))]
{
"libsql"
}
};
let selected_backend = self
.settings
.database_backend
.as_deref()
.unwrap_or("postgres");
.unwrap_or(default_backend);
#[cfg(all(feature = "libsql", feature = "postgres"))]
{
if selected_backend == "libsql" {
match selected_backend {
#[cfg(feature = "libsql")]
"libsql" | "turso" | "sqlite" => {
if let Some(store) = self.create_libsql_secrets_store(&crypto)? {
return Ok(SecretsContext::from_store(store, "default"));
}
// Fallback to postgres if libsql store creation returned None
#[cfg(feature = "postgres")]
if let Some(store) = self.create_postgres_secrets_store(&crypto).await? {
return Ok(SecretsContext::from_store(store, "default"));
}
} else {
}
#[cfg(feature = "postgres")]
_ => {
if let Some(store) = self.create_postgres_secrets_store(&crypto).await? {
return Ok(SecretsContext::from_store(store, "default"));
}
// Fallback to libsql if postgres store creation returned None
#[cfg(feature = "libsql")]
if let Some(store) = self.create_libsql_secrets_store(&crypto)? {
return Ok(SecretsContext::from_store(store, "default"));
}
}
}
#[cfg(all(feature = "postgres", not(feature = "libsql")))]
{
let _ = selected_backend;
if let Some(store) = self.create_postgres_secrets_store(&crypto).await? {
return Ok(SecretsContext::from_store(store, "default"));
}
}
#[cfg(all(feature = "libsql", not(feature = "postgres")))]
{
let _ = selected_backend;
if let Some(store) = self.create_libsql_secrets_store(&crypto)? {
return Ok(SecretsContext::from_store(store, "default"));
}
#[cfg(not(feature = "postgres"))]
_ => {}
}
Err(SetupError::Config(
@@ -2293,12 +2428,29 @@ impl SetupWizard {
if let Some(ref url) = self.settings.ollama_base_url {
env_vars.push(("OLLAMA_BASE_URL".to_string(), url.clone()));
}
if let Some(ref region) = self.settings.bedrock_region {
env_vars.push(("BEDROCK_REGION".to_string(), region.clone()));
}
if self.settings.llm_backend.as_deref() == Some("bedrock") {
if let Some(ref model) = self.settings.selected_model {
env_vars.push(("BEDROCK_MODEL".to_string(), model.clone()));
}
if let Some(ref cross) = self.settings.bedrock_cross_region {
env_vars.push(("BEDROCK_CROSS_REGION".to_string(), cross.clone()));
}
if let Some(ref profile) = self.settings.bedrock_profile {
env_vars.push(("AWS_PROFILE".to_string(), profile.clone()));
}
}
// Model name: same chicken-and-egg — Config::from_env() resolves the
// model before the DB is connected, so we must persist it to .env.
// Write the backend-specific env var so the correct resolution path
// picks it up (looked up from the provider registry).
if let Some(ref model) = self.settings.selected_model {
// Bedrock model is already written above as BEDROCK_MODEL, skip here.
if self.settings.llm_backend.as_deref() != Some("bedrock")
&& let Some(ref model) = self.settings.selected_model
{
let backend_str = self.settings.llm_backend.as_deref().unwrap_or("nearai");
let model_env = registry.model_env_var(backend_str);
env_vars.push((model_env.to_string(), model.clone()));
@@ -2324,6 +2476,12 @@ impl SetupWizard {
env_vars.push(("NEARAI_API_KEY".to_string(), api_key));
}
// Secrets master key (env var mode): write to .env so it's available
// on next startup before the DB is connected.
if let Some(ref key_hex) = self.settings.secrets_master_key_hex {
env_vars.push(("SECRETS_MASTER_KEY".to_string(), key_hex.clone()));
}
// Always write ONBOARD_COMPLETED so that check_onboard_needed()
// (which runs before the DB is connected) knows to skip re-onboarding.
if self.settings.onboard_completed {
@@ -2577,6 +2735,7 @@ impl SetupWizard {
"openai" => "OpenAI",
"ollama" => "Ollama",
"openai_compatible" => "OpenAI-compatible",
"bedrock" => "AWS Bedrock",
other => other,
};
println!(" Provider: {}", display);
@@ -3499,6 +3658,108 @@ mod tests {
}
}
/// Regression test for #600: re-running provider setup for the same backend
/// must NOT clear selected_model. Only switching to a different backend should.
#[test]
fn test_same_provider_preserves_selected_model() {
let mut wizard = SetupWizard::new();
wizard.settings.llm_backend = Some("ollama".to_string());
wizard.settings.selected_model = Some("llama3".to_string());
// Simulate re-entering the same provider -- model should survive
// (This is the check that each setup_* function now performs)
if wizard.settings.llm_backend.as_deref() != Some("ollama") {
wizard.settings.selected_model = None;
}
wizard.settings.llm_backend = Some("ollama".to_string());
assert_eq!(
wizard.settings.selected_model.as_deref(),
Some("llama3"),
"model should be preserved when re-selecting the same provider"
);
}
/// Regression test for #600: switching to a different provider must clear
/// selected_model since the old model may not be valid for the new backend.
#[test]
fn test_different_provider_clears_selected_model() {
let mut wizard = SetupWizard::new();
wizard.settings.llm_backend = Some("ollama".to_string());
wizard.settings.selected_model = Some("llama3".to_string());
// Simulate switching to a different provider -- model should be cleared
if wizard.settings.llm_backend.as_deref() != Some("openai") {
wizard.settings.selected_model = None;
}
wizard.settings.llm_backend = Some("openai".to_string());
assert!(
wizard.settings.selected_model.is_none(),
"model should be cleared when switching providers"
);
}
/// Regression: Bedrock setup_bedrock() should preserve selected_model
/// when re-entering the same provider (matches pattern from #600).
#[test]
fn test_bedrock_same_provider_preserves_model() {
let mut wizard = SetupWizard::new();
wizard.settings.llm_backend = Some("bedrock".to_string());
wizard.settings.selected_model = Some("anthropic.claude-opus-4-6-v1".to_string());
// Simulate the conditional clearing logic from setup_bedrock()
if wizard.settings.llm_backend.as_deref() != Some("bedrock") {
wizard.settings.selected_model = None;
}
wizard.settings.llm_backend = Some("bedrock".to_string());
assert_eq!(
wizard.settings.selected_model.as_deref(),
Some("anthropic.claude-opus-4-6-v1"),
"bedrock model should be preserved when re-selecting bedrock"
);
}
/// Regression: switching from another provider to bedrock must clear
/// selected_model, and choosing "default credentials" must clear
/// bedrock_profile.
#[test]
fn test_bedrock_clears_stale_profile_on_default_creds() {
let mut wizard = SetupWizard::new();
wizard.settings.llm_backend = Some("bedrock".to_string());
wizard.settings.bedrock_profile = Some("old-sso-profile".to_string());
// Simulate auth_choice == 0 (default credentials) clearing the profile
wizard.settings.bedrock_profile = None;
assert!(
wizard.settings.bedrock_profile.is_none(),
"bedrock_profile should be cleared when selecting default credentials"
);
}
/// Regression: empty profile input in named-profile auth should clear
/// any previously configured profile instead of leaving it stale.
#[test]
fn test_bedrock_empty_profile_clears_existing() {
let mut wizard = SetupWizard::new();
wizard.settings.bedrock_profile = Some("old-profile".to_string());
// Simulate auth_choice == 1 with empty input
let profile = "".to_string();
if profile.trim().is_empty() {
wizard.settings.bedrock_profile = None;
} else {
wizard.settings.bedrock_profile = Some(profile);
}
assert!(
wizard.settings.bedrock_profile.is_none(),
"empty profile input should clear existing bedrock_profile"
);
}
#[tokio::test]
async fn test_run_provider_setup_no_setup_hint() {
// A provider with setup: None should not error. It should set the
@@ -3536,4 +3797,30 @@ mod tests {
"backend should be set even without setup hint"
);
}
/// Regression test for #666: env-var security option must initialize
/// secrets_crypto so subsequent steps can encrypt API keys.
#[test]
fn test_env_var_security_initializes_crypto() {
use crate::secrets::SecretsCrypto;
use secrecy::SecretString;
// Simulate what option 1 in step_security() does after the fix:
let key_hex = crate::secrets::keychain::generate_master_key_hex();
// The fix: create SecretsCrypto from the generated key.
// Before the fix, this was skipped, leaving secrets_crypto = None.
let crypto = SecretsCrypto::new(SecretString::from(key_hex.clone()));
assert!(
crypto.is_ok(),
"generated key hex must produce valid SecretsCrypto"
);
// Verify the key is stored for bootstrap env persistence.
let settings = Settings {
secrets_master_key_hex: Some(key_hex),
..Settings::default()
};
assert!(settings.secrets_master_key_hex.is_some());
}
}
+1
View File
@@ -142,6 +142,7 @@ mod tests {
content_hash: "sha256:000".to_string(),
compiled_patterns: vec![],
lowercased_keywords: vec![],
lowercased_exclude_keywords: vec![],
lowercased_tags: vec![],
}
}
+11
View File
@@ -98,6 +98,10 @@ pub struct ActivationCriteria {
/// Capped at `MAX_KEYWORDS_PER_SKILL` during loading.
#[serde(default)]
pub keywords: Vec<String>,
/// Keywords that veto this skill — if any match, score is 0 regardless of
/// keyword/pattern matches. Prevents cross-skill interference.
#[serde(default)]
pub exclude_keywords: Vec<String>,
/// Regex patterns for more complex matching.
/// Capped at `MAX_PATTERNS_PER_SKILL` during loading.
#[serde(default)]
@@ -118,6 +122,9 @@ impl ActivationCriteria {
pub fn enforce_limits(&mut self) {
self.keywords.retain(|k| k.len() >= MIN_KEYWORD_TAG_LENGTH);
self.keywords.truncate(MAX_KEYWORDS_PER_SKILL);
self.exclude_keywords
.retain(|k| k.len() >= MIN_KEYWORD_TAG_LENGTH);
self.exclude_keywords.truncate(MAX_KEYWORDS_PER_SKILL);
self.patterns.truncate(MAX_PATTERNS_PER_SKILL);
self.tags.retain(|t| t.len() >= MIN_KEYWORD_TAG_LENGTH);
self.tags.truncate(MAX_TAGS_PER_SKILL);
@@ -199,6 +206,9 @@ pub struct LoadedSkill {
/// Pre-computed lowercased keywords for scoring (avoids per-message allocation).
/// Derived from `manifest.activation.keywords` at load time — do not mutate independently.
pub lowercased_keywords: Vec<String>,
/// Pre-computed lowercased exclude keywords for veto scoring.
/// Derived from `manifest.activation.exclude_keywords` at load time.
pub lowercased_exclude_keywords: Vec<String>,
/// Pre-computed lowercased tags for scoring (avoids per-message allocation).
/// Derived from `manifest.activation.tags` at load time — do not mutate independently.
pub lowercased_tags: Vec<String>,
@@ -513,6 +523,7 @@ metadata:
content_hash: "sha256:000".to_string(),
compiled_patterns: vec![],
lowercased_keywords: vec![],
lowercased_exclude_keywords: vec![],
lowercased_tags: vec![],
};
assert_eq!(skill.name(), "test");
+8 -12
View File
@@ -24,6 +24,10 @@ use crate::skills::{
/// Prevents resource exhaustion from a directory with thousands of entries.
const MAX_DISCOVERED_SKILLS: usize = 100;
fn to_lowercase_vec(items: &[String]) -> Vec<String> {
items.iter().map(|s| s.to_lowercase()).collect()
}
/// Error type for skill registry operations.
#[derive(Debug, thiserror::Error)]
pub enum SkillRegistryError {
@@ -582,18 +586,9 @@ async fn load_and_validate_skill(
let compiled_patterns = LoadedSkill::compile_patterns(&manifest.activation.patterns);
// Pre-compute lowercased keywords and tags for efficient scoring
let lowercased_keywords = manifest
.activation
.keywords
.iter()
.map(|k| k.to_lowercase())
.collect();
let lowercased_tags = manifest
.activation
.tags
.iter()
.map(|t| t.to_lowercase())
.collect();
let lowercased_keywords = to_lowercase_vec(&manifest.activation.keywords);
let lowercased_exclude_keywords = to_lowercase_vec(&manifest.activation.exclude_keywords);
let lowercased_tags = to_lowercase_vec(&manifest.activation.tags);
let name = manifest.name.clone();
let skill = LoadedSkill {
@@ -604,6 +599,7 @@ async fn load_and_validate_skill(
content_hash,
compiled_patterns,
lowercased_keywords,
lowercased_exclude_keywords,
lowercased_tags,
};
+118
View File
@@ -99,6 +99,15 @@ pub fn prefilter_skills<'a>(
/// Score a skill against a user message.
fn score_skill(skill: &LoadedSkill, message_lower: &str, message_original: &str) -> u32 {
// Exclusion veto: if any exclude_keyword is present in the message, score 0
if skill
.lowercased_exclude_keywords
.iter()
.any(|excl| message_lower.contains(excl.as_str()))
{
return 0;
}
let mut score: u32 = 0;
// Keyword scoring with cap to prevent gaming via keyword stuffing
@@ -158,6 +167,7 @@ mod tests {
description: format!("{} skill", name),
activation: ActivationCriteria {
keywords: kw_vec,
exclude_keywords: vec![],
patterns: pattern_strings,
tags: tag_vec,
max_context_tokens: 1000,
@@ -170,6 +180,7 @@ mod tests {
content_hash: "sha256:000".to_string(),
compiled_patterns: compiled,
lowercased_keywords,
lowercased_exclude_keywords: vec![],
lowercased_tags,
}
}
@@ -368,4 +379,111 @@ mod tests {
let result = prefilter_skills("test", &skills, 5, 1);
assert_eq!(result.len(), 1);
}
fn make_skill_with_excludes(
name: &str,
keywords: &[&str],
exclude_keywords: &[&str],
tags: &[&str],
patterns: &[&str],
) -> LoadedSkill {
let mut skill = make_skill(name, keywords, tags, patterns);
let excl_vec: Vec<String> = exclude_keywords.iter().map(|s| s.to_string()).collect();
skill.lowercased_exclude_keywords = excl_vec.iter().map(|k| k.to_lowercase()).collect();
skill.manifest.activation.exclude_keywords = excl_vec;
skill
}
// --- exclude_keywords tests ---
#[test]
fn test_exclude_keyword_vetos_match() {
// Skill matches on "write" but exclude_keywords: ["route"] — message contains "route"
// so the skill should score 0 and be excluded.
let skills = vec![make_skill_with_excludes(
"writer",
&["write"],
&["route"],
&[],
&[],
)];
let result = prefilter_skills(
"route this write request to another agent",
&skills,
3,
MAX_SKILL_CONTEXT_TOKENS,
);
assert!(
result.is_empty(),
"skill with matching exclude_keyword should score 0"
);
}
#[test]
fn test_exclude_keyword_absent_does_not_block() {
// Same skill, message does NOT contain the exclude keyword — should activate normally.
let skills = vec![make_skill_with_excludes(
"writer",
&["write"],
&["route"],
&[],
&[],
)];
let result = prefilter_skills(
"help me write an email",
&skills,
3,
MAX_SKILL_CONTEXT_TOKENS,
);
assert_eq!(
result.len(),
1,
"skill should activate when no exclude_keyword is present"
);
}
#[test]
fn test_exclude_keyword_veto_wins_over_positive_match() {
// Both a keyword match AND an exclude_keyword match are present.
// The veto must win regardless of how high the positive score is.
let skills = vec![make_skill_with_excludes(
"writer",
&["write", "draft", "compose"],
&["redirect"],
&[],
&[],
)];
let result = prefilter_skills(
"write and draft and compose — but redirect this somewhere else",
&skills,
3,
MAX_SKILL_CONTEXT_TOKENS,
);
assert!(
result.is_empty(),
"exclude_keyword veto must win even when multiple positive keywords match"
);
}
#[test]
fn test_exclude_keyword_case_insensitive() {
// exclude_keywords are pre-lowercased; the veto must fire regardless of case in the message.
let skills = vec![make_skill_with_excludes(
"writer",
&["write"],
&["Route"],
&[],
&[],
)];
let result = prefilter_skills(
"please ROUTE this write request",
&skills,
3,
MAX_SKILL_CONTEXT_TOKENS,
);
assert!(
result.is_empty(),
"exclude_keyword veto should be case-insensitive"
);
}
}
+1
View File
@@ -1009,6 +1009,7 @@ mod tests {
enabled: true,
trigger: Trigger::Cron {
schedule: "0 * * * *".to_string(),
timezone: None,
},
action: RoutineAction::Lightweight {
prompt: "Check status".to_string(),
+110
View File
@@ -0,0 +1,110 @@
//! Timezone resolution and utilities.
use chrono::{DateTime, NaiveDate, Utc};
use chrono_tz::Tz;
/// Resolve the effective timezone from a priority chain.
///
/// Priority: client_tz > user_setting > config_default > UTC
pub fn resolve_timezone(
client_tz: Option<&str>,
user_setting: Option<&str>,
config_default: &str,
) -> Tz {
// Try each in priority order, skipping invalid values
for candidate in [client_tz, user_setting, Some(config_default)] {
if let Some(tz) = candidate.and_then(parse_timezone) {
return tz;
}
}
Tz::UTC
}
/// Parse a timezone string (IANA name) into a `Tz`.
pub fn parse_timezone(s: &str) -> Option<Tz> {
s.parse::<Tz>().ok()
}
/// Get today's date in the given timezone.
pub fn today_in_tz(tz: Tz) -> NaiveDate {
Utc::now().with_timezone(&tz).date_naive()
}
/// Get the current time in the given timezone.
pub fn now_in_tz(tz: Tz) -> DateTime<Tz> {
Utc::now().with_timezone(&tz)
}
/// Detect the system's timezone, falling back to UTC.
pub fn detect_system_timezone() -> Tz {
iana_time_zone::get_timezone()
.ok()
.and_then(|s| parse_timezone(&s))
.unwrap_or(Tz::UTC)
}
#[cfg(test)]
mod tests {
use chrono::Datelike;
use super::*;
#[test]
fn test_resolve_client_wins() {
let tz = resolve_timezone(Some("America/New_York"), Some("Europe/London"), "UTC");
assert_eq!(tz, chrono_tz::America::New_York);
}
#[test]
fn test_resolve_user_setting_fallback() {
let tz = resolve_timezone(None, Some("Europe/London"), "UTC");
assert_eq!(tz, chrono_tz::Europe::London);
}
#[test]
fn test_resolve_config_fallback() {
let tz = resolve_timezone(None, None, "Asia/Tokyo");
assert_eq!(tz, chrono_tz::Asia::Tokyo);
}
#[test]
fn test_resolve_all_none_utc() {
let tz = resolve_timezone(None, None, "UTC");
assert_eq!(tz, Tz::UTC);
}
#[test]
fn test_resolve_invalid_client_skipped() {
let tz = resolve_timezone(Some("Fake/Zone"), Some("Europe/London"), "UTC");
assert_eq!(tz, chrono_tz::Europe::London);
}
#[test]
fn test_parse_valid() {
assert_eq!(
parse_timezone("America/Chicago"),
Some(chrono_tz::America::Chicago)
);
}
#[test]
fn test_parse_invalid() {
assert_eq!(parse_timezone("Fake/Zone"), None);
}
#[test]
fn test_detect_system_tz() {
// Should always return a valid Tz (at minimum UTC)
let tz = detect_system_timezone();
let _ = now_in_tz(tz); // Should not panic
}
#[test]
fn test_today_in_tz_returns_valid_date() {
let date = today_in_tz(Tz::UTC);
// Verify it returns a valid date (year, month, day are all positive)
assert!(date.year() > 0);
assert!((1..=12).contains(&date.month()));
assert!((1..=31).contains(&date.day()));
}
}
+4 -16
View File
@@ -43,7 +43,6 @@ use crate::error::ToolError as AgentToolError;
use crate::llm::{
ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult, ToolDefinition,
};
use crate::safety::SafetyLayer;
use crate::tools::ToolRegistry;
use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput};
@@ -251,29 +250,18 @@ pub trait SoftwareBuilder: Send + Sync {
pub struct LlmSoftwareBuilder {
config: BuilderConfig,
llm: Arc<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
tools: Arc<ToolRegistry>,
}
impl LlmSoftwareBuilder {
/// Create a new LLM-based software builder.
pub fn new(
config: BuilderConfig,
llm: Arc<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
tools: Arc<ToolRegistry>,
) -> Self {
pub fn new(config: BuilderConfig, llm: Arc<dyn LlmProvider>, tools: Arc<ToolRegistry>) -> Self {
// Ensure build directory exists
if let Err(e) = std::fs::create_dir_all(&config.build_dir) {
tracing::warn!("Failed to create build directory: {}", e);
}
Self {
config,
llm,
safety,
tools,
}
Self { config, llm, tools }
}
/// Get the build tools available for the build loop.
@@ -521,7 +509,7 @@ Create alongside the .wasm file to grant capabilities:
let mut iteration = 0;
// Create reasoning engine
let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone());
let reasoning = Reasoning::new(self.llm.clone());
// Build initial context
let tool_defs = self.get_build_tools().await;
@@ -822,7 +810,7 @@ Create alongside the .wasm file to grant capabilities:
impl SoftwareBuilder for LlmSoftwareBuilder {
async fn analyze(&self, description: &str) -> Result<BuildRequirement, AgentToolError> {
// Use LLM to parse the description
let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone());
let reasoning = Reasoning::new(self.llm.clone());
let prompt = format!(
r#"Analyze this software requirement and extract structured information.
+250
View File
@@ -0,0 +1,250 @@
//! Image analysis tool using vision-capable LLM models.
use std::path::PathBuf;
use async_trait::async_trait;
use base64::Engine;
use secrecy::{ExposeSecret, SecretString};
use crate::context::JobContext;
use crate::tools::builtin::path_utils::validate_path;
use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput};
/// Tool for analyzing images using a vision-capable model.
pub struct ImageAnalyzeTool {
/// API base URL.
api_base_url: String,
/// Bearer token for API auth.
api_key: SecretString,
/// Vision-capable model name.
model: String,
/// HTTP client.
client: reqwest::Client,
/// Optional base directory for resolving relative image paths.
base_dir: Option<PathBuf>,
}
impl ImageAnalyzeTool {
/// Create a new image analysis tool.
pub fn new(
api_base_url: String,
api_key: String,
model: String,
base_dir: Option<PathBuf>,
) -> Self {
let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(120))
.build()
.unwrap_or_default();
Self {
api_base_url,
api_key: SecretString::from(api_key),
model,
client,
base_dir,
}
}
/// Read binary image bytes from filesystem.
///
/// Validates the path against the base directory sandbox to prevent
/// path traversal attacks, then reads the file bytes.
async fn read_image_bytes(&self, image_path: &str) -> Result<Vec<u8>, ToolError> {
let resolved = validate_path(image_path, self.base_dir.as_deref())?;
tokio::fs::read(&resolved)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Failed to read image file: {e}")))
}
}
#[async_trait]
impl Tool for ImageAnalyzeTool {
fn name(&self) -> &str {
"image_analyze"
}
fn description(&self) -> &str {
"Analyze an image using a vision-capable AI model. Provide a workspace path to the image and an optional analysis question."
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"image_path": {
"type": "string",
"description": "Path to the image file in the workspace (e.g., 'images/photo.jpg')"
},
"question": {
"type": "string",
"description": "Specific question to answer about the image. Defaults to general analysis.",
"default": "Describe this image in detail."
}
},
"required": ["image_path"]
})
}
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
ApprovalRequirement::UnlessAutoApproved
}
fn requires_sanitization(&self) -> bool {
true
}
async fn execute(
&self,
params: serde_json::Value,
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let start = std::time::Instant::now();
let image_path = params
.get("image_path")
.and_then(|v| v.as_str())
.ok_or_else(|| {
ToolError::InvalidParameters("Missing required 'image_path' parameter".to_string())
})?;
let question = params
.get("question")
.and_then(|v| v.as_str())
.unwrap_or("Describe this image in detail.");
// Read binary image bytes directly from filesystem
let image_bytes = self.read_image_bytes(image_path).await?;
if image_bytes.is_empty() {
return Err(ToolError::ExecutionFailed(
"Image file is empty".to_string(),
));
}
let media_type = super::media_type_from_path(image_path);
let b64 = base64::engine::general_purpose::STANDARD.encode(&image_bytes);
let data_url = format!("data:{media_type};base64,{b64}");
// Call vision model via chat completions API
let url = format!(
"{}/v1/chat/completions",
self.api_base_url.trim_end_matches('/')
);
let request_body = serde_json::json!({
"model": &self.model,
"messages": [{
"role": "user",
"content": [
{
"type": "text",
"text": question
},
{
"type": "image_url",
"image_url": {
"url": data_url
}
}
]
}],
"max_tokens": 2048
});
let response = self
.client
.post(&url)
.bearer_auth(self.api_key.expose_secret())
.json(&request_body)
.send()
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Vision API request failed: {e}")))?;
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
return Err(ToolError::ExecutionFailed(format!(
"Vision API returned {status}: {body}"
)));
}
let resp: serde_json::Value = response.json().await.map_err(|e| {
ToolError::ExecutionFailed(format!("Failed to parse vision API response: {e}"))
})?;
let analysis = resp
.pointer("/choices/0/message/content")
.and_then(|v| v.as_str())
.unwrap_or("No analysis available.");
Ok(ToolOutput::text(analysis, start.elapsed()))
}
}
#[cfg(test)]
mod tests {
use super::super::media_type_from_path;
use super::*;
use tempfile::TempDir;
#[test]
fn test_media_type_detection() {
assert_eq!(media_type_from_path("photo.png"), "image/png");
assert_eq!(media_type_from_path("photo.jpg"), "image/jpeg");
assert_eq!(media_type_from_path("photo.jpeg"), "image/jpeg");
assert_eq!(media_type_from_path("photo.gif"), "image/gif");
assert_eq!(media_type_from_path("photo.webp"), "image/webp");
assert_eq!(media_type_from_path("photo.bmp"), "image/bmp");
assert_eq!(media_type_from_path("photo.svg"), "image/svg+xml");
}
#[test]
fn test_requires_approval_returns_unless_auto_approved() {
let tool = ImageAnalyzeTool::new(
"https://api.example.com".to_string(),
"test-key".to_string(),
"gpt-4o".to_string(),
None,
);
assert_eq!(
tool.requires_approval(&serde_json::json!({})),
ApprovalRequirement::UnlessAutoApproved
);
}
#[tokio::test]
async fn test_read_image_bytes_rejects_path_traversal() {
let dir = TempDir::new().unwrap();
let tool = ImageAnalyzeTool::new(
"https://api.example.com".to_string(),
"test-key".to_string(),
"gpt-4o".to_string(),
Some(dir.path().to_path_buf()),
);
let result = tool.read_image_bytes("../../etc/passwd").await;
assert!(
result.is_err(),
"Should reject path traversal, got: {:?}",
result
);
}
#[tokio::test]
async fn test_read_image_bytes_rejects_absolute_path_outside_sandbox() {
let dir = TempDir::new().unwrap();
let tool = ImageAnalyzeTool::new(
"https://api.example.com".to_string(),
"test-key".to_string(),
"gpt-4o".to_string(),
Some(dir.path().to_path_buf()),
);
let result = tool.read_image_bytes("/etc/passwd").await;
assert!(
result.is_err(),
"Should reject absolute path outside sandbox, got: {:?}",
result
);
}
}
+322
View File
@@ -0,0 +1,322 @@
//! Image editing tool using cloud API.
use std::path::PathBuf;
use async_trait::async_trait;
use secrecy::{ExposeSecret, SecretString};
use crate::context::JobContext;
use crate::tools::builtin::path_utils::validate_path;
use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput};
/// Tool for editing images using an AI image editing API.
pub struct ImageEditTool {
/// API base URL.
api_base_url: String,
/// Bearer token for API auth.
api_key: SecretString,
/// Model to use.
model: String,
/// HTTP client.
client: reqwest::Client,
/// Optional base directory for resolving relative image paths.
base_dir: Option<PathBuf>,
}
impl ImageEditTool {
/// Create a new image edit tool.
pub fn new(
api_base_url: String,
api_key: String,
model: String,
base_dir: Option<PathBuf>,
) -> Self {
let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(180))
.build()
.unwrap_or_default();
Self {
api_base_url,
api_key: SecretString::from(api_key),
model,
client,
base_dir,
}
}
/// Read binary image bytes from filesystem.
///
/// Validates the path against the base directory sandbox to prevent
/// path traversal attacks, then reads the file bytes.
async fn read_image_bytes(&self, image_path: &str) -> Result<Vec<u8>, ToolError> {
let resolved = validate_path(image_path, self.base_dir.as_deref())?;
tokio::fs::read(&resolved)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Failed to read image file: {e}")))
}
}
#[async_trait]
impl Tool for ImageEditTool {
fn name(&self) -> &str {
"image_edit"
}
fn description(&self) -> &str {
"Edit an existing image using an AI model. Provide the workspace path to the source image and a text prompt describing the desired edits."
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"prompt": {
"type": "string",
"description": "Text description of the edits to apply to the image",
"maxLength": 4000
},
"image_path": {
"type": "string",
"description": "Path to the source image in the workspace (e.g., 'images/photo.jpg')"
}
},
"required": ["prompt", "image_path"]
})
}
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
ApprovalRequirement::UnlessAutoApproved
}
fn requires_sanitization(&self) -> bool {
false
}
async fn execute(
&self,
params: serde_json::Value,
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let start = std::time::Instant::now();
let prompt = params
.get("prompt")
.and_then(|v| v.as_str())
.ok_or_else(|| {
ToolError::InvalidParameters("Missing required 'prompt' parameter".to_string())
})?;
let image_path = params
.get("image_path")
.and_then(|v| v.as_str())
.ok_or_else(|| {
ToolError::InvalidParameters("Missing required 'image_path' parameter".to_string())
})?;
if prompt.len() > 4000 {
return Err(ToolError::InvalidParameters(
"Prompt exceeds 4000 character limit".to_string(),
));
}
// Read binary image bytes directly from filesystem
let image_bytes = self.read_image_bytes(image_path).await?;
if image_bytes.is_empty() {
return Err(ToolError::ExecutionFailed(
"Source image file is empty".to_string(),
));
}
let media_type = super::media_type_from_path(image_path);
// Use multipart form for image edit API
let url = format!(
"{}/v1/images/edits",
self.api_base_url.trim_end_matches('/')
);
let form = reqwest::multipart::Form::new()
.text("model", self.model.clone())
.text("prompt", prompt.to_string())
.text("response_format", "b64_json")
.part(
"image",
reqwest::multipart::Part::bytes(image_bytes)
.mime_str(&media_type)
.map_err(|e| ToolError::ExecutionFailed(format!("Invalid media type: {e}")))?
.file_name("image"),
);
let response = self
.client
.post(&url)
.bearer_auth(self.api_key.expose_secret())
.multipart(form)
.send()
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Image edit request failed: {e}")))?;
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
// Fall back to generation if edits endpoint not available
if status.as_u16() == 404 {
tracing::warn!(
"Image edit endpoint returned 404, falling back to generation API. \
Note: the source image will NOT be used a new image will be generated from the prompt alone."
);
return self.fallback_generate(prompt, start).await;
}
return Err(ToolError::ExecutionFailed(format!(
"Image edit API returned {status}: {body}"
)));
}
let resp: serde_json::Value = response.json().await.map_err(|e| {
ToolError::ExecutionFailed(format!("Failed to parse image edit response: {e}"))
})?;
let edited_data = resp
.pointer("/data/0/b64_json")
.and_then(|v| v.as_str())
.ok_or_else(|| {
ToolError::ExecutionFailed("No image data in edit response".to_string())
})?;
let sentinel = serde_json::json!({
"type": "image_generated",
"data": format!("data:image/png;base64,{}", edited_data),
"media_type": "image/png",
"prompt": prompt,
"source_path": image_path
});
Ok(ToolOutput::text(sentinel.to_string(), start.elapsed()))
}
}
impl ImageEditTool {
/// Fallback: generate a new image from the prompt when the edit endpoint is unavailable.
///
/// The source image is NOT used — this generates a completely new image.
/// The response includes a `note` field warning the user.
async fn fallback_generate(
&self,
prompt: &str,
start: std::time::Instant,
) -> Result<ToolOutput, ToolError> {
let url = format!(
"{}/v1/images/generations",
self.api_base_url.trim_end_matches('/')
);
let request_body = serde_json::json!({
"model": &self.model,
"prompt": prompt,
"size": "1024x1024",
"response_format": "b64_json",
"n": 1
});
let response = self
.client
.post(&url)
.bearer_auth(self.api_key.expose_secret())
.json(&request_body)
.send()
.await
.map_err(|e| {
ToolError::ExecutionFailed(format!("Fallback image generation failed: {e}"))
})?;
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
return Err(ToolError::ExecutionFailed(format!(
"Fallback generation API returned {status}: {body}"
)));
}
let resp: serde_json::Value = response.json().await.map_err(|e| {
ToolError::ExecutionFailed(format!("Failed to parse fallback response: {e}"))
})?;
let image_data = resp
.pointer("/data/0/b64_json")
.and_then(|v| v.as_str())
.ok_or_else(|| {
ToolError::ExecutionFailed("No image data in fallback response".to_string())
})?;
let sentinel = serde_json::json!({
"type": "image_generated",
"data": format!("data:image/png;base64,{}", image_data),
"media_type": "image/png",
"prompt": prompt,
"note": "Generated new image (edit endpoint unavailable — source image was NOT used)"
});
Ok(ToolOutput::text(sentinel.to_string(), start.elapsed()))
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[test]
fn test_tool_metadata() {
let tool = ImageEditTool::new(
"https://api.example.com".to_string(),
"test-key".to_string(),
"flux-1".to_string(),
None,
);
assert_eq!(tool.name(), "image_edit");
assert!(!tool.requires_sanitization());
assert_eq!(
tool.requires_approval(&serde_json::json!({})),
ApprovalRequirement::UnlessAutoApproved
);
}
#[tokio::test]
async fn test_read_image_bytes_rejects_path_traversal() {
let dir = TempDir::new().unwrap();
let tool = ImageEditTool::new(
"https://api.example.com".to_string(),
"test-key".to_string(),
"flux-1".to_string(),
Some(dir.path().to_path_buf()),
);
let result = tool.read_image_bytes("../../etc/passwd").await;
assert!(
result.is_err(),
"Should reject path traversal, got: {:?}",
result
);
}
#[tokio::test]
async fn test_read_image_bytes_rejects_absolute_path_outside_sandbox() {
let dir = TempDir::new().unwrap();
let tool = ImageEditTool::new(
"https://api.example.com".to_string(),
"test-key".to_string(),
"flux-1".to_string(),
Some(dir.path().to_path_buf()),
);
let result = tool.read_image_bytes("/etc/passwd").await;
assert!(
result.is_err(),
"Should reject absolute path outside sandbox, got: {:?}",
result
);
}
}
+251
View File
@@ -0,0 +1,251 @@
//! Image generation tool using cloud API.
use async_trait::async_trait;
use secrecy::{ExposeSecret, SecretString};
use serde::{Deserialize, Serialize};
use crate::context::JobContext;
use crate::tools::tool::ApprovalRequirement;
use crate::tools::{Tool, ToolError, ToolOutput};
/// Tool for generating images using FLUX or compatible image generation APIs.
pub struct ImageGenerateTool {
/// API base URL (e.g., "https://cloud-api.near.ai").
api_base_url: String,
/// Bearer token for API auth.
api_key: SecretString,
/// Model to use (e.g., "black-forest-labs/FLUX.1-schnell").
model: String,
/// HTTP client.
client: reqwest::Client,
}
#[derive(Debug, Serialize)]
struct ImageGenRequest {
model: String,
prompt: String,
size: String,
response_format: String,
n: u32,
}
#[derive(Debug, Deserialize)]
struct ImageGenResponse {
data: Vec<ImageGenData>,
}
#[derive(Debug, Deserialize)]
#[allow(dead_code)]
struct ImageGenData {
b64_json: Option<String>,
url: Option<String>,
}
impl ImageGenerateTool {
/// Create a new image generation tool.
pub fn new(api_base_url: String, api_key: String, model: String) -> Self {
let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(180))
.build()
.unwrap_or_default();
Self {
api_base_url,
api_key: SecretString::from(api_key),
model,
client,
}
}
}
#[async_trait]
impl Tool for ImageGenerateTool {
fn name(&self) -> &str {
"image_generate"
}
fn description(&self) -> &str {
"Generate an image from a text prompt using an AI image generation model (e.g., FLUX). Returns the generated image data."
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"prompt": {
"type": "string",
"description": "Text description of the image to generate (max 4000 chars)",
"maxLength": 4000
},
"size": {
"type": "string",
"description": "Image dimensions",
"enum": ["1024x1024", "1792x1024", "1024x1792"],
"default": "1024x1024"
}
},
"required": ["prompt"]
})
}
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
ApprovalRequirement::UnlessAutoApproved
}
fn requires_sanitization(&self) -> bool {
false
}
async fn execute(
&self,
params: serde_json::Value,
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let start = std::time::Instant::now();
let prompt = params
.get("prompt")
.and_then(|v| v.as_str())
.ok_or_else(|| {
ToolError::InvalidParameters("Missing required 'prompt' parameter".to_string())
})?;
if prompt.len() > 4000 {
return Err(ToolError::InvalidParameters(
"Prompt exceeds 4000 character limit".to_string(),
));
}
let size = params
.get("size")
.and_then(|v| v.as_str())
.unwrap_or("1024x1024");
// Validate size
if !["1024x1024", "1792x1024", "1024x1792"].contains(&size) {
return Err(ToolError::InvalidParameters(format!(
"Invalid size '{}'. Must be 1024x1024, 1792x1024, or 1024x1792",
size
)));
}
let url = format!(
"{}/v1/images/generations",
self.api_base_url.trim_end_matches('/')
);
let request_body = ImageGenRequest {
model: self.model.clone(),
prompt: prompt.to_string(),
size: size.to_string(),
response_format: "b64_json".to_string(),
n: 1,
};
let response = self
.client
.post(&url)
.bearer_auth(self.api_key.expose_secret())
.json(&request_body)
.send()
.await
.map_err(|e| {
ToolError::ExecutionFailed(format!("Image generation request failed: {e}"))
})?;
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
return Err(ToolError::ExecutionFailed(format!(
"Image generation API returned {status}: {body}"
)));
}
let gen_response: ImageGenResponse = response.json().await.map_err(|e| {
ToolError::ExecutionFailed(format!("Failed to parse image generation response: {e}"))
})?;
let image_data = gen_response
.data
.first()
.and_then(|d| d.b64_json.as_deref())
.ok_or_else(|| ToolError::ExecutionFailed("No image data in response".to_string()))?;
// Return sentinel JSON for image display
let sentinel = serde_json::json!({
"type": "image_generated",
"data": format!("data:image/png;base64,{}", image_data),
"media_type": "image/png",
"prompt": prompt,
"size": size
});
Ok(ToolOutput::text(sentinel.to_string(), start.elapsed()))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_tool_metadata() {
let tool = ImageGenerateTool::new(
"https://api.example.com".to_string(),
"test-key".to_string(),
"flux-1".to_string(),
);
assert_eq!(tool.name(), "image_generate");
assert_eq!(
tool.requires_approval(&serde_json::json!({})),
ApprovalRequirement::UnlessAutoApproved
);
let schema = tool.parameters_schema();
assert!(schema["properties"]["prompt"].is_object());
assert!(schema["properties"]["size"].is_object());
}
#[tokio::test]
async fn test_missing_prompt() {
let tool = ImageGenerateTool::new(
"https://api.example.com".to_string(),
"test-key".to_string(),
"flux-1".to_string(),
);
let ctx = JobContext::default();
let result = tool.execute(serde_json::json!({}), &ctx).await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_invalid_size() {
let tool = ImageGenerateTool::new(
"https://api.example.com".to_string(),
"test-key".to_string(),
"flux-1".to_string(),
);
let ctx = JobContext::default();
let result = tool
.execute(
serde_json::json!({"prompt": "a cat", "size": "999x999"}),
&ctx,
)
.await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_prompt_too_long() {
let tool = ImageGenerateTool::new(
"https://api.example.com".to_string(),
"test-key".to_string(),
"flux-1".to_string(),
);
let ctx = JobContext::default();
let long_prompt = "x".repeat(4001);
let result = tool
.execute(serde_json::json!({"prompt": long_prompt}), &ctx)
.await;
assert!(result.is_err());
}
}
+179
View File
@@ -1416,6 +1416,185 @@ mod tests {
);
}
#[tokio::test]
async fn test_create_job_params() {
let manager = Arc::new(ContextManager::new(5));
let tool = CreateJobTool::new(manager);
let ctx = JobContext::default();
let missing_title = tool
.execute(serde_json::json!({ "description": "A test job" }), &ctx)
.await;
assert!(missing_title.is_err());
assert!(
missing_title
.unwrap_err()
.to_string()
.contains("missing 'title' parameter")
);
let missing_description = tool
.execute(serde_json::json!({ "title": "Test Job" }), &ctx)
.await;
assert!(missing_description.is_err());
assert!(
missing_description
.unwrap_err()
.to_string()
.contains("missing 'description' parameter")
);
}
#[tokio::test]
async fn test_list_jobs_formatting() {
let manager = Arc::new(ContextManager::new(10));
let pending_id = manager
.create_job_for_user("default", "Pending Job", "Todo")
.await
.unwrap();
let completed_id = manager
.create_job_for_user("default", "Completed Job", "Done")
.await
.unwrap();
let failed_id = manager
.create_job_for_user("default", "Failed Job", "Oops")
.await
.unwrap();
manager
.create_job_for_user("other-user", "Other User Job", "Ignore")
.await
.unwrap();
manager
.update_context(completed_id, |ctx| {
ctx.transition_to(JobState::InProgress, None)?;
ctx.transition_to(JobState::Completed, Some("done".to_string()))
})
.await
.unwrap()
.unwrap();
manager
.update_context(failed_id, |ctx| {
ctx.transition_to(JobState::InProgress, None)?;
ctx.transition_to(JobState::Failed, Some("boom".to_string()))
})
.await
.unwrap()
.unwrap();
let tool = ListJobsTool::new(Arc::clone(&manager));
let ctx = JobContext::default();
let result = tool.execute(serde_json::json!({}), &ctx).await.unwrap();
let jobs = result.result.get("jobs").unwrap().as_array().unwrap();
assert_eq!(jobs.len(), 3);
assert!(jobs.iter().any(|job| {
job.get("job_id").and_then(|v| v.as_str()) == Some(&pending_id.to_string())
&& job.get("status").and_then(|v| v.as_str()) == Some("Pending")
}));
assert!(jobs.iter().any(|job| {
job.get("job_id").and_then(|v| v.as_str()) == Some(&completed_id.to_string())
&& job.get("status").and_then(|v| v.as_str()) == Some("Completed")
}));
assert!(jobs.iter().any(|job| {
job.get("job_id").and_then(|v| v.as_str()) == Some(&failed_id.to_string())
&& job.get("status").and_then(|v| v.as_str()) == Some("Failed")
}));
let summary = result.result.get("summary").unwrap();
assert_eq!(summary.get("total").and_then(|v| v.as_u64()), Some(3));
assert_eq!(summary.get("pending").and_then(|v| v.as_u64()), Some(1));
assert_eq!(summary.get("completed").and_then(|v| v.as_u64()), Some(1));
assert_eq!(summary.get("failed").and_then(|v| v.as_u64()), Some(1));
}
#[tokio::test]
async fn test_job_status_transitions() {
let manager = Arc::new(ContextManager::new(5));
let job_id = manager
.create_job_for_user("default", "Transition Job", "Track me")
.await
.unwrap();
manager
.update_context(job_id, |ctx| {
ctx.transition_to(JobState::InProgress, Some("started".to_string()))?;
ctx.transition_to(JobState::Completed, Some("finished".to_string()))
})
.await
.unwrap()
.unwrap();
let tool = JobStatusTool::new(Arc::clone(&manager));
let ctx = JobContext::default();
let result = tool
.execute(serde_json::json!({ "job_id": job_id.to_string() }), &ctx)
.await
.unwrap();
assert_eq!(
result.result.get("status").and_then(|v| v.as_str()),
Some("Completed")
);
assert!(result.result.get("started_at").unwrap().is_string());
assert!(result.result.get("completed_at").unwrap().is_string());
}
#[tokio::test]
async fn test_cancel_job_running() {
let manager = Arc::new(ContextManager::new(5));
let job_id = manager
.create_job_for_user("default", "Running Job", "In progress")
.await
.unwrap();
manager
.update_context(job_id, |ctx| ctx.transition_to(JobState::InProgress, None))
.await
.unwrap()
.unwrap();
let tool = CancelJobTool::new(Arc::clone(&manager));
let ctx = JobContext::default();
let result = tool
.execute(serde_json::json!({ "job_id": job_id.to_string() }), &ctx)
.await
.unwrap();
assert_eq!(
result.result.get("status").and_then(|v| v.as_str()),
Some("cancelled")
);
let updated = manager.get_context(job_id).await.unwrap();
assert_eq!(updated.state, JobState::Cancelled);
}
#[tokio::test]
async fn test_cancel_job_completed() {
let manager = Arc::new(ContextManager::new(5));
let job_id = manager
.create_job_for_user("default", "Completed Job", "Already done")
.await
.unwrap();
manager
.update_context(job_id, |ctx| {
ctx.transition_to(JobState::InProgress, None)?;
ctx.transition_to(JobState::Completed, Some("done".to_string()))
})
.await
.unwrap()
.unwrap();
let tool = CancelJobTool::new(Arc::clone(&manager));
let ctx = JobContext::default();
let result = tool
.execute(serde_json::json!({ "job_id": job_id.to_string() }), &ctx)
.await
.unwrap();
let error = result.result.get("error").and_then(|v| v.as_str()).unwrap();
assert!(error.contains("Cannot cancel job"));
assert!(error.contains("completed"));
}
#[test]
fn test_resolve_project_dir_auto() {
let project_id = Uuid::new_v4();
+5 -4
View File
@@ -172,7 +172,7 @@ impl Tool for MemoryWriteTool {
async fn execute(
&self,
params: serde_json::Value,
_ctx: &JobContext,
ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let start = std::time::Instant::now();
@@ -239,11 +239,12 @@ impl Tool for MemoryWriteTool {
paths::MEMORY.to_string()
}
"daily_log" => {
let tz = crate::timezone::parse_timezone(&ctx.user_timezone)
.unwrap_or(chrono_tz::Tz::UTC);
self.workspace
.append_daily_log(content)
.append_daily_log_tz(content, tz)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?;
format!("daily/{}.md", chrono::Utc::now().format("%Y-%m-%d"))
.map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?
}
"heartbeat" => {
if append {
+184 -50
View File
@@ -105,42 +105,47 @@ impl Tool for MessageTool {
async fn execute(
&self,
params: serde_json::Value,
_ctx: &JobContext,
ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let start = std::time::Instant::now();
let content = require_str(&params, "content")?;
// Get channel: use param or fall back to default
let channel = if let Some(c) = params.get("channel").and_then(|v| v.as_str()) {
c.to_string()
} else {
self.default_channel
// Get channel: use param → conversation default → job metadata → None (broadcast all)
let channel: Option<String> =
if let Some(c) = params.get("channel").and_then(|v| v.as_str()) {
Some(c.to_string())
} else if let Some(c) = self
.default_channel
.read()
.unwrap_or_else(|e| e.into_inner())
.clone()
.ok_or_else(|| {
ToolError::ExecutionFailed(
"No channel specified and no active conversation. Provide channel parameter."
.to_string(),
)
})?
};
{
Some(c)
} else {
ctx.metadata
.get("notify_channel")
.and_then(|v| v.as_str())
.map(|c| c.to_string())
};
// Get target: use param or fall back to default
// Get target: use param → conversation default → job metadata
let target = if let Some(t) = params.get("target").and_then(|v| v.as_str()) {
t.to_string()
} else if let Some(t) = self
.default_target
.read()
.unwrap_or_else(|e| e.into_inner())
.clone()
{
t
} else if let Some(t) = ctx.metadata.get("notify_user").and_then(|v| v.as_str()) {
t.to_string()
} else {
self.default_target
.read()
.unwrap_or_else(|e| e.into_inner())
.clone()
.ok_or_else(|| {
ToolError::ExecutionFailed(
"No target specified and no active conversation. Provide target parameter."
.to_string(),
)
})?
return Err(ToolError::ExecutionFailed(
"No target specified and no active conversation. Provide target parameter."
.to_string(),
));
};
let attachments: Vec<String> = match params.get("attachments") {
@@ -181,36 +186,79 @@ impl Tool for MessageTool {
response = response.with_attachments(attachments);
}
match self
.channel_manager
.broadcast(&channel, &target, response)
.await
{
Ok(()) => {
tracing::info!(
message_sent = true,
channel = %channel,
target = %target,
attachments = attachment_count,
"Message sent via message tool"
);
let msg = format!("Sent message to {}:{}", channel, target);
Ok(ToolOutput::text(msg, start.elapsed()))
if let Some(ref channel) = channel {
// Send to a specific channel
match self
.channel_manager
.broadcast(channel, &target, response)
.await
{
Ok(()) => {
tracing::info!(
message_sent = true,
channel = %channel,
target = %target,
attachments = attachment_count,
"Message sent via message tool"
);
let msg = format!("Sent message to {}:{}", channel, target);
Ok(ToolOutput::text(msg, start.elapsed()))
}
Err(e) => {
let available = self.channel_manager.channel_names().await.join(", ");
let err_msg = if available.is_empty() {
format!(
"Failed to send to {}:{}: {}. No channels connected.",
channel, target, e
)
} else {
format!(
"Failed to send to {}:{}. Available channels: {}. Error: {}",
channel, target, available, e
)
};
Err(ToolError::ExecutionFailed(err_msg))
}
}
Err(e) => {
let available = self.channel_manager.channel_names().await.join(", ");
let err_msg = if available.is_empty() {
format!(
"Failed to send to {}:{}: {}. No channels connected.",
channel, target, e
)
} else {
// No channel specified — broadcast to all channels (routine with notify.channel = None)
let results = self.channel_manager.broadcast_all(&target, response).await;
let mut succeeded = Vec::new();
let mut failed: Vec<&str> = Vec::new();
for (ch, result) in &results {
match result {
Ok(()) => succeeded.push(ch.as_str()),
Err(e) => {
tracing::warn!(
channel = %ch,
target = %target,
"broadcast_all: channel failed: {}", e
);
failed.push(ch.as_str());
}
}
}
if succeeded.is_empty() {
let err_msg = if failed.is_empty() {
"No channels connected.".to_string()
} else {
format!(
"Failed to send to {}:{}. Available channels: {}. Error: {}",
channel, target, available, e
)
format!("All channels failed: {}", failed.join(", "))
};
Err(ToolError::ExecutionFailed(err_msg))
} else {
tracing::info!(
message_sent = true,
channels = ?succeeded,
target = %target,
attachments = attachment_count,
"Message broadcast via message tool"
);
let msg = format!(
"Broadcast message to {} (target: {})",
succeeded.join(", "),
target
);
Ok(ToolOutput::text(msg, start.elapsed()))
}
}
}
@@ -576,4 +624,90 @@ mod tests {
ApprovalRequirement::Never,
);
}
#[tokio::test]
async fn message_tool_falls_back_to_job_metadata() {
// Regression: when no conversation context is set (e.g. routine full-job),
// the message tool should fall back to notify_channel/notify_user from
// JobContext metadata instead of returning "No target specified".
let tool = MessageTool::new(Arc::new(ChannelManager::new()));
let mut ctx = crate::context::JobContext::new("routine-job", "price alert");
ctx.metadata = serde_json::json!({
"notify_channel": "telegram",
"notify_user": "123456789",
});
// No set_context called — simulates a routine full-job worker
let result = tool
.execute(serde_json::json!({"content": "NEAR price is $5"}), &ctx)
.await;
// Should fail at channel broadcast (no real channel), NOT at
// "No target specified and no active conversation"
assert!(result.is_err());
let err = result.unwrap_err().to_string();
assert!(
!err.contains("No target specified"),
"Should not get 'No target specified' when metadata has notify_user, got: {}",
err
);
assert!(
!err.contains("No channel specified"),
"Should not get 'No channel specified' when metadata has notify_channel, got: {}",
err
);
}
#[tokio::test]
async fn message_tool_no_metadata_still_errors() {
// When neither conversation context nor metadata is set, should still
// return a clear error (target resolution fails).
let tool = MessageTool::new(Arc::new(ChannelManager::new()));
let ctx = crate::context::JobContext::new("orphan-job", "no notify config");
let result = tool
.execute(serde_json::json!({"content": "hello"}), &ctx)
.await;
assert!(result.is_err());
let err = result.unwrap_err().to_string();
assert!(
err.contains("No target specified"),
"Expected 'No target specified' error, got: {}",
err
);
}
#[tokio::test]
async fn message_tool_broadcasts_all_when_no_channel() {
// Regression: when notify.channel is None but notify_user is set,
// the message tool should attempt broadcast_all instead of erroring
// with "No channel specified".
let tool = MessageTool::new(Arc::new(ChannelManager::new()));
let mut ctx = crate::context::JobContext::new("routine-job", "price alert");
ctx.metadata = serde_json::json!({
"notify_user": "123456789",
});
let result = tool
.execute(serde_json::json!({"content": "NEAR price is $5"}), &ctx)
.await;
// Should fail because no channels are registered (empty ChannelManager),
// NOT because "No channel specified".
assert!(result.is_err());
let err = result.unwrap_err().to_string();
assert!(
!err.contains("No channel specified"),
"Should not get 'No channel specified' when broadcasting, got: {}",
err
);
assert!(
err.contains("No channels connected") || err.contains("All channels failed"),
"Expected channel delivery error, got: {}",
err
);
}
}
+16
View File
@@ -40,5 +40,21 @@ pub use shell::ShellTool;
pub use skill_tools::{SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool};
pub use time::TimeTool;
mod html_converter;
pub mod image_analyze;
pub mod image_edit;
pub mod image_gen;
pub use html_converter::convert_html_to_markdown;
pub use image_analyze::ImageAnalyzeTool;
pub use image_edit::ImageEditTool;
pub use image_gen::ImageGenerateTool;
/// Detect image media type from file extension via `mime_guess`.
/// Falls back to `image/jpeg` for unrecognized or non-image extensions.
pub(crate) fn media_type_from_path(path: &str) -> String {
mime_guess::from_path(path)
.first_raw()
.filter(|m| m.starts_with("image/"))
.unwrap_or("image/jpeg")
.to_string()
}
+68 -10
View File
@@ -107,6 +107,10 @@ impl Tool for RoutineCreateTool {
"notify_user": {
"type": "string",
"description": "User/target to notify (e.g. username, chat ID). Defaults to 'default'."
},
"timezone": {
"type": "string",
"description": "IANA timezone for cron schedule evaluation (e.g. 'America/New_York'). Defaults to UTC."
}
},
"required": ["name", "trigger_type", "prompt"]
@@ -143,12 +147,26 @@ impl Tool for RoutineCreateTool {
"cron trigger requires 'schedule'".to_string(),
)
})?;
let timezone = params
.get("timezone")
.and_then(|v| v.as_str())
.map(|tz| {
crate::timezone::parse_timezone(tz)
.map(|_| tz.to_string())
.ok_or_else(|| {
ToolError::InvalidParameters(format!(
"invalid IANA timezone: '{tz}'"
))
})
})
.transpose()?;
// Validate cron expression
next_cron_fire(schedule).map_err(|e| {
next_cron_fire(schedule, timezone.as_deref()).map_err(|e| {
ToolError::InvalidParameters(format!("invalid cron schedule: {e}"))
})?;
Trigger::Cron {
schedule: schedule.to_string(),
timezone,
}
}
"event" => {
@@ -228,8 +246,12 @@ impl Tool for RoutineCreateTool {
.unwrap_or(300);
// Compute next fire time for cron
let next_fire = if let Trigger::Cron { ref schedule } = trigger {
next_cron_fire(schedule).unwrap_or(None)
let next_fire = if let Trigger::Cron {
ref schedule,
ref timezone,
} = trigger
{
next_cron_fire(schedule, timezone.as_deref()).unwrap_or(None)
} else {
None
};
@@ -412,6 +434,10 @@ impl Tool for RoutineUpdateTool {
"type": "string",
"description": "New cron schedule (for cron triggers)"
},
"timezone": {
"type": "string",
"description": "IANA timezone for cron schedule (e.g. 'America/New_York'). Only valid for cron triggers."
},
"description": {
"type": "string",
"description": "New description"
@@ -453,15 +479,47 @@ impl Tool for RoutineUpdateTool {
}
}
if let Some(schedule) = params.get("schedule").and_then(|v| v.as_str()) {
// Validate
next_cron_fire(schedule)
.map_err(|e| ToolError::InvalidParameters(format!("invalid cron schedule: {e}")))?;
// Validate timezone param if provided
let new_timezone = params
.get("timezone")
.and_then(|v| v.as_str())
.map(|tz| {
crate::timezone::parse_timezone(tz)
.map(|_| tz.to_string())
.ok_or_else(|| {
ToolError::InvalidParameters(format!("invalid IANA timezone: '{tz}'"))
})
})
.transpose()?;
routine.trigger = Trigger::Cron {
schedule: schedule.to_string(),
let new_schedule = params.get("schedule").and_then(|v| v.as_str());
if new_schedule.is_some() || new_timezone.is_some() {
// Extract existing cron fields (cloned to avoid borrow conflict)
let existing_cron = match &routine.trigger {
Trigger::Cron { schedule, timezone } => Some((schedule.clone(), timezone.clone())),
_ => None,
};
routine.next_fire_at = next_cron_fire(schedule).unwrap_or(None);
if let Some((old_schedule, old_tz)) = existing_cron {
let effective_schedule = new_schedule.unwrap_or(&old_schedule);
let effective_tz = new_timezone.or(old_tz);
// Validate
next_cron_fire(effective_schedule, effective_tz.as_deref()).map_err(|e| {
ToolError::InvalidParameters(format!("invalid cron schedule: {e}"))
})?;
routine.trigger = Trigger::Cron {
schedule: effective_schedule.to_string(),
timezone: effective_tz.clone(),
};
routine.next_fire_at =
next_cron_fire(effective_schedule, effective_tz.as_deref()).unwrap_or(None);
} else {
return Err(ToolError::InvalidParameters(
"Cannot update schedule or timezone on a non-cron routine.".to_string(),
));
}
}
self.store
+472 -49
View File
@@ -1,7 +1,8 @@
//! Time utility tool.
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use chrono::{DateTime, LocalResult, NaiveDate, NaiveDateTime, TimeZone, Utc};
use chrono_tz::Tz;
use crate::context::JobContext;
use crate::tools::tool::{Tool, ToolError, ToolOutput, require_str};
@@ -16,7 +17,7 @@ impl Tool for TimeTool {
}
fn description(&self) -> &str {
"Get current time, convert timezones, or calculate time differences."
"Get current time, parse or format timestamps, convert timezones, or calculate time differences."
}
fn parameters_schema(&self) -> serde_json::Value {
@@ -25,20 +26,40 @@ impl Tool for TimeTool {
"properties": {
"operation": {
"type": "string",
"enum": ["now", "parse", "format", "diff"],
"enum": ["now", "parse", "convert", "format", "diff"],
"description": "The time operation to perform"
},
"input": {
"type": "string",
"description": "Input timestamp. Accepts RFC 3339, or a naive timestamp when timezone/from_timezone is provided."
},
"timestamp": {
"type": "string",
"description": "ISO 8601 timestamp (for parse/format/diff operations)"
"description": "Alias for input (kept for backward compatibility)."
},
"timezone": {
"type": "string",
"description": "IANA timezone name (e.g. 'America/New_York'). Used by now/format, and can also interpret naive timestamps."
},
"from_timezone": {
"type": "string",
"description": "Source IANA timezone for naive input timestamps during convert/format/diff."
},
"to_timezone": {
"type": "string",
"description": "Target IANA timezone for convert."
},
"format": {
"type": "string",
"description": "Output format string (for format operation)"
"description": "strftime format string for format (kept for backward compatibility)."
},
"format_string": {
"type": "string",
"description": "strftime format string for format."
},
"timestamp2": {
"type": "string",
"description": "Second timestamp (for diff operation)"
"description": "Second timestamp for diff."
}
},
"required": ["operation"]
@@ -48,55 +69,18 @@ impl Tool for TimeTool {
async fn execute(
&self,
params: serde_json::Value,
_ctx: &JobContext,
ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let start = std::time::Instant::now();
let operation = require_str(&params, "operation")?;
let result = match operation {
"now" => {
let now = Utc::now();
serde_json::json!({
"iso": now.to_rfc3339(),
"unix": now.timestamp(),
"unix_millis": now.timestamp_millis()
})
}
"parse" => {
let timestamp = require_str(&params, "timestamp")?;
let dt: DateTime<Utc> = timestamp.parse().map_err(|e| {
ToolError::InvalidParameters(format!("invalid timestamp: {}", e))
})?;
serde_json::json!({
"iso": dt.to_rfc3339(),
"unix": dt.timestamp(),
"unix_millis": dt.timestamp_millis()
})
}
"diff" => {
let ts1 = require_str(&params, "timestamp")?;
let ts2 = require_str(&params, "timestamp2")?;
let dt1: DateTime<Utc> = ts1.parse().map_err(|e| {
ToolError::InvalidParameters(format!("invalid timestamp: {}", e))
})?;
let dt2: DateTime<Utc> = ts2.parse().map_err(|e| {
ToolError::InvalidParameters(format!("invalid timestamp2: {}", e))
})?;
let diff = dt2.signed_duration_since(dt1);
serde_json::json!({
"seconds": diff.num_seconds(),
"minutes": diff.num_minutes(),
"hours": diff.num_hours(),
"days": diff.num_days()
})
}
"now" => execute_now(&params, ctx)?,
"parse" => execute_parse(&params, ctx)?,
"convert" => execute_convert(&params, ctx)?,
"format" => execute_format(&params, ctx)?,
"diff" => execute_diff(&params, ctx)?,
_ => {
return Err(ToolError::InvalidParameters(format!(
"unknown operation: {}",
@@ -112,3 +96,442 @@ impl Tool for TimeTool {
false // Internal tool, no external data
}
}
fn execute_now(
params: &serde_json::Value,
ctx: &JobContext,
) -> Result<serde_json::Value, ToolError> {
let now = Utc::now();
let mut result = serde_json::json!({
"iso": now.to_rfc3339(),
"utc_iso": now.to_rfc3339(),
"unix": now.timestamp(),
"unix_millis": now.timestamp_millis()
});
if let Some((tz, tz_name)) = resolve_timezone_for_output(params, ctx)? {
let local = now.with_timezone(&tz);
result["local_iso"] = serde_json::Value::String(local.to_rfc3339());
result["timezone"] = serde_json::Value::String(tz_name);
}
Ok(result)
}
fn execute_parse(
params: &serde_json::Value,
ctx: &JobContext,
) -> Result<serde_json::Value, ToolError> {
let input = require_input(params)?;
let parse_tz = resolve_parse_timezone(params, ctx)?;
let dt = parse_timestamp(input, parse_tz.as_ref())?;
Ok(serde_json::json!({
"iso": dt.to_rfc3339(),
"unix": dt.timestamp(),
"unix_millis": dt.timestamp_millis()
}))
}
fn execute_convert(
params: &serde_json::Value,
ctx: &JobContext,
) -> Result<serde_json::Value, ToolError> {
let input = require_input(params)?;
let source_tz = optional_timezone(params, &["from_timezone", "timezone"])?;
let dt = parse_timestamp(input, source_tz.as_ref())?;
let target_name = params
.get("to_timezone")
.and_then(|v| v.as_str())
.ok_or_else(|| {
ToolError::InvalidParameters("convert operation requires 'to_timezone'".to_string())
})?;
let target_tz = parse_timezone(target_name)?;
let converted = dt.with_timezone(&target_tz);
let mut result = serde_json::json!({
"input": input,
"utc_iso": dt.to_rfc3339(),
"output": converted.to_rfc3339(),
"timezone": target_tz.to_string()
});
if let Some((ctx_tz, ctx_tz_name)) = context_timezone(ctx)? {
result["context_timezone"] = serde_json::Value::String(ctx_tz_name);
result["context_iso"] = serde_json::Value::String(dt.with_timezone(&ctx_tz).to_rfc3339());
}
Ok(result)
}
fn execute_format(
params: &serde_json::Value,
ctx: &JobContext,
) -> Result<serde_json::Value, ToolError> {
let input = require_input(params)?;
let output_tz = resolve_timezone_for_output(params, ctx)?;
let source_tz = optional_timezone(params, &["from_timezone"])?
.or_else(|| output_tz.as_ref().map(|(tz, _)| *tz));
let dt = parse_timestamp(input, source_tz.as_ref())?;
let format_string = params
.get("format_string")
.and_then(|v| v.as_str())
.or_else(|| params.get("format").and_then(|v| v.as_str()))
.unwrap_or("%Y-%m-%d %H:%M:%S %Z");
let mut result = if let Some((tz, tz_name)) = output_tz {
serde_json::json!({
"formatted": dt.with_timezone(&tz).format(format_string).to_string(),
"timezone": tz_name
})
} else {
serde_json::json!({
"formatted": dt.format(format_string).to_string()
})
};
result["utc_iso"] = serde_json::Value::String(dt.to_rfc3339());
Ok(result)
}
fn execute_diff(
params: &serde_json::Value,
ctx: &JobContext,
) -> Result<serde_json::Value, ToolError> {
let parse_tz = resolve_parse_timezone(params, ctx)?;
let ts1 = require_input(params)?;
let ts2 = params
.get("timestamp2")
.and_then(|v| v.as_str())
.ok_or_else(|| {
ToolError::InvalidParameters("diff operation requires 'timestamp2'".to_string())
})?;
let dt1 = parse_timestamp(ts1, parse_tz.as_ref())?;
let dt2 = parse_timestamp(ts2, parse_tz.as_ref())?;
let diff = dt2.signed_duration_since(dt1);
Ok(serde_json::json!({
"seconds": diff.num_seconds(),
"minutes": diff.num_minutes(),
"hours": diff.num_hours(),
"days": diff.num_days()
}))
}
fn require_input(params: &serde_json::Value) -> Result<&str, ToolError> {
params
.get("input")
.and_then(|v| v.as_str())
.or_else(|| params.get("timestamp").and_then(|v| v.as_str()))
.ok_or_else(|| {
ToolError::InvalidParameters(
"missing 'input' (or legacy 'timestamp') parameter".to_string(),
)
})
}
fn resolve_parse_timezone(
params: &serde_json::Value,
ctx: &JobContext,
) -> Result<Option<Tz>, ToolError> {
if let Some(tz) = optional_timezone(params, &["from_timezone", "timezone"])? {
return Ok(Some(tz));
}
Ok(context_timezone(ctx)?.map(|(tz, _)| tz))
}
fn resolve_timezone_for_output(
params: &serde_json::Value,
ctx: &JobContext,
) -> Result<Option<(Tz, String)>, ToolError> {
if let Some(name) = params.get("timezone").and_then(|v| v.as_str()) {
let tz = parse_timezone(name)?;
return Ok(Some((tz, tz.to_string())));
}
context_timezone(ctx)
}
/// Resolve the user's timezone from the JobContext.
///
/// Uses `ctx.user_timezone` (set from main's timezone resolution) as the
/// primary source. Falls back to metadata fields for backward compatibility.
fn context_timezone(ctx: &JobContext) -> Result<Option<(Tz, String)>, ToolError> {
// Primary: use the dedicated user_timezone field from JobContext
if ctx.user_timezone != "UTC"
&& !ctx.user_timezone.is_empty()
&& let Some(tz) = crate::timezone::parse_timezone(&ctx.user_timezone)
{
return Ok(Some((tz, tz.to_string())));
}
// Fallback: check metadata for backward compatibility
let tz_name = ctx
.metadata
.get("user_timezone")
.and_then(|v| v.as_str())
.or_else(|| ctx.metadata.get("timezone").and_then(|v| v.as_str()));
match tz_name {
Some(name) => {
let tz = parse_timezone(name)?;
Ok(Some((tz, tz.to_string())))
}
None => Ok(None),
}
}
fn optional_timezone(params: &serde_json::Value, keys: &[&str]) -> Result<Option<Tz>, ToolError> {
for key in keys {
if let Some(value) = params.get(*key).and_then(|v| v.as_str()) {
return parse_timezone(value).map(Some);
}
}
Ok(None)
}
fn parse_timezone(value: &str) -> Result<Tz, ToolError> {
value.parse::<Tz>().map_err(|_| {
ToolError::InvalidParameters(format!(
"Unknown timezone '{}'. Use IANA names like 'America/New_York' or 'Europe/London'.",
value
))
})
}
fn parse_timestamp(input: &str, fallback_tz: Option<&Tz>) -> Result<DateTime<Utc>, ToolError> {
if let Ok(dt) = DateTime::parse_from_rfc3339(input) {
return Ok(dt.with_timezone(&Utc));
}
if let Some(naive) = parse_naive_datetime(input) {
return localize_naive_datetime(naive, fallback_tz, input);
}
Err(ToolError::InvalidParameters(format!(
"invalid timestamp '{}': expected RFC 3339 or a naive timestamp with timezone/from_timezone",
input
)))
}
fn parse_naive_datetime(input: &str) -> Option<NaiveDateTime> {
const DATETIME_FORMATS: &[&str] = &[
"%Y-%m-%d %H:%M:%S%.f",
"%Y-%m-%dT%H:%M:%S%.f",
"%Y-%m-%d %H:%M",
"%Y-%m-%dT%H:%M",
];
const DATE_FORMATS: &[&str] = &["%Y-%m-%d"];
for format in DATETIME_FORMATS {
if let Ok(value) = NaiveDateTime::parse_from_str(input, format) {
return Some(value);
}
}
for format in DATE_FORMATS {
if let Ok(date) = NaiveDate::parse_from_str(input, format) {
return date.and_hms_opt(0, 0, 0);
}
}
None
}
fn localize_naive_datetime(
naive: NaiveDateTime,
fallback_tz: Option<&Tz>,
original_input: &str,
) -> Result<DateTime<Utc>, ToolError> {
let tz = fallback_tz.ok_or_else(|| {
ToolError::InvalidParameters(format!(
"timestamp '{}' has no UTC offset; provide 'timezone' or 'from_timezone'",
original_input
))
})?;
match tz.from_local_datetime(&naive) {
LocalResult::Single(dt) => Ok(dt.with_timezone(&Utc)),
LocalResult::Ambiguous(_, _) => Err(ToolError::InvalidParameters(format!(
"timestamp '{}' is ambiguous in timezone '{}'; include an explicit UTC offset instead",
original_input, tz
))),
LocalResult::None => Err(ToolError::InvalidParameters(format!(
"timestamp '{}' does not exist in timezone '{}'",
original_input, tz
))),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_now_accepts_explicit_timezone() {
let tool = TimeTool;
let ctx = JobContext::with_user("test", "chat", "test");
let output = tool
.execute(
serde_json::json!({
"operation": "now",
"timezone": "America/New_York"
}),
&ctx,
)
.await
.expect("execute");
assert_eq!(output.result["timezone"].as_str(), Some("America/New_York"));
assert!(
output.result.get("utc_iso").is_some(),
"should have utc_iso"
);
assert!(
output.result.get("local_iso").is_some(),
"should have local_iso"
);
}
#[tokio::test]
async fn test_now_includes_local_time_when_user_timezone_set() {
let tool = TimeTool;
let mut ctx = JobContext::with_user("test", "chat", "test");
ctx.user_timezone = "America/New_York".to_string();
let output = tool
.execute(serde_json::json!({"operation": "now"}), &ctx)
.await
.expect("execute");
assert!(
output.result.get("local_iso").is_some(),
"should have local_iso"
);
assert_eq!(
output.result["timezone"].as_str(),
Some("America/New_York"),
"should report timezone"
);
}
#[tokio::test]
async fn test_now_uses_context_metadata_timezone_fallback() {
let tool = TimeTool;
let mut ctx = JobContext::with_user("test", "chat", "test");
ctx.metadata = serde_json::json!({
"user_timezone": "America/Los_Angeles"
});
let output = tool
.execute(serde_json::json!({"operation": "now"}), &ctx)
.await
.expect("execute");
assert_eq!(
output.result["timezone"].as_str(),
Some("America/Los_Angeles")
);
assert!(
output.result.get("local_iso").is_some(),
"should have local_iso"
);
}
#[tokio::test]
async fn test_now_returns_utc_by_default() {
let tool = TimeTool;
let ctx = JobContext::with_user("test", "chat", "test");
// Default user_timezone is "UTC" -- context_timezone skips UTC so no
// local_iso is added, but iso and utc_iso are always present.
let output = tool
.execute(serde_json::json!({"operation": "now"}), &ctx)
.await
.expect("execute");
assert!(output.result.get("iso").is_some(), "should have iso");
}
#[tokio::test]
async fn test_convert_across_dst_boundary() {
let tool = TimeTool;
let ctx = JobContext::with_user("test", "chat", "test");
let output = tool
.execute(
serde_json::json!({
"operation": "convert",
"input": "2026-03-08T07:30:00Z",
"to_timezone": "America/New_York"
}),
&ctx,
)
.await
.expect("execute");
assert_eq!(output.result["timezone"].as_str(), Some("America/New_York"));
assert_eq!(
output.result["output"].as_str(),
Some("2026-03-08T03:30:00-04:00")
);
}
#[tokio::test]
async fn test_format_with_timezone() {
let tool = TimeTool;
let ctx = JobContext::with_user("test", "chat", "test");
let output = tool
.execute(
serde_json::json!({
"operation": "format",
"input": "2026-03-08T07:30:00Z",
"timezone": "America/New_York",
"format_string": "%Y-%m-%d %H:%M:%S %Z"
}),
&ctx,
)
.await
.expect("execute");
assert_eq!(output.result["timezone"].as_str(), Some("America/New_York"));
assert_eq!(
output.result["formatted"].as_str(),
Some("2026-03-08 03:30:00 EDT")
);
}
#[tokio::test]
async fn test_invalid_timezone_returns_clear_error() {
let tool = TimeTool;
let ctx = JobContext::with_user("test", "chat", "test");
let err = tool
.execute(
serde_json::json!({
"operation": "now",
"timezone": "Mars/Olympus"
}),
&ctx,
)
.await
.expect_err("expected invalid timezone error");
match err {
ToolError::InvalidParameters(message) => {
assert!(message.contains("Unknown timezone 'Mars/Olympus'"));
}
other => panic!("unexpected error: {other:?}"),
}
}
#[test]
fn test_parse_naive_timestamp_with_timezone() {
let dt = parse_timestamp("2026-03-08 03:30:00", Some(&chrono_tz::America::New_York))
.expect("parse timestamp");
assert_eq!(dt.to_rfc3339(), "2026-03-08T07:30:00+00:00");
}
}
+570 -23
View File
@@ -4,6 +4,7 @@
//! See: https://spec.modelcontextprotocol.io/specification/2025-03-26/basic/authorization/
use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::Arc;
use std::time::Duration;
@@ -199,23 +200,285 @@ impl PkceChallenge {
}
}
// ---------------------------------------------------------------------------
// Well-known URI construction (RFC 8414 / RFC 9728)
// ---------------------------------------------------------------------------
/// Build a well-known URI according to RFC 8414 / RFC 9728.
///
/// The path component of the base URL is placed *after* the well-known suffix:
/// ```text
/// https://example.com/path + oauth-authorization-server
/// -> https://example.com/.well-known/oauth-authorization-server/path
/// ```
pub fn build_well_known_uri(base_url: &str, suffix: &str) -> Result<String, AuthError> {
let parsed = reqwest::Url::parse(base_url)
.map_err(|e| AuthError::DiscoveryFailed(format!("Invalid URL: {}", e)))?;
let origin = parsed.origin().ascii_serialization();
let path = parsed.path().trim_end_matches('/');
Ok(format!("{}/.well-known/{}{}", origin, suffix, path))
}
// ---------------------------------------------------------------------------
// RFC 8707 resource parameter
// ---------------------------------------------------------------------------
/// Compute the canonical resource URI for RFC 8707.
///
/// Strips fragments and trailing slashes from the server URL.
pub fn canonical_resource_uri(server_url: &str) -> String {
match reqwest::Url::parse(server_url) {
Ok(mut parsed) => {
parsed.set_fragment(None);
let s = parsed.to_string();
s.trim_end_matches('/').to_string()
}
Err(_) => server_url.trim_end_matches('/').to_string(),
}
}
// ---------------------------------------------------------------------------
// SSRF protection
// ---------------------------------------------------------------------------
/// Check if an IP address is dangerous (loopback, link-local, private, etc.)
fn is_dangerous_ip(ip: IpAddr) -> bool {
match ip {
IpAddr::V4(v4) => {
v4.is_loopback()
|| v4.is_private()
|| v4.is_link_local()
|| v4.is_broadcast()
|| v4.is_unspecified()
|| (v4.octets()[0] == 169 && v4.octets()[1] == 254) // link-local
|| (v4.octets()[0] == 100 && (v4.octets()[1] & 0xC0) == 64) // CGNAT 100.64/10
}
IpAddr::V6(v6) => {
let segs = v6.segments();
v6.is_loopback()
|| v6.is_unspecified()
// Link-local (fe80::/10)
|| (segs[0] & 0xffc0) == 0xfe80
// Site-local / deprecated (fec0::/10)
|| (segs[0] & 0xffc0) == 0xfec0
// Unique local (fc00::/7)
|| (segs[0] & 0xfe00) == 0xfc00
// Documentation (2001:db8::/32)
|| (segs[0] == 0x2001 && segs[1] == 0x0db8)
// Check for IPv4-mapped IPv6 (::ffff:x.x.x.x)
|| v6
.to_ipv4_mapped()
.is_some_and(|v4| is_dangerous_ip(IpAddr::V4(v4)))
}
}
}
/// Validate that a URL is safe for server-side requests (SSRF protection).
async fn validate_url_safe(url: &str) -> Result<(), AuthError> {
let parsed = reqwest::Url::parse(url)
.map_err(|e| AuthError::DiscoveryFailed(format!("Invalid URL: {}", e)))?;
// Must be HTTPS. HTTP is only allowed for localhost/loopback (dev scenarios).
let scheme = parsed.scheme();
if scheme != "https" && scheme != "http" {
return Err(AuthError::DiscoveryFailed(format!(
"Unsupported scheme: {}",
scheme
)));
}
if scheme == "http" {
let host = parsed.host_str().unwrap_or("");
let is_localhost =
host == "localhost" || host == "127.0.0.1" || host == "::1" || host == "[::1]";
if !is_localhost {
return Err(AuthError::DiscoveryFailed(format!(
"HTTP is only allowed for localhost; use HTTPS for '{}'",
host
)));
}
// Localhost HTTP is allowed for dev — skip SSRF checks since we've
// already validated the host is localhost/loopback.
return Ok(());
}
let host = parsed
.host_str()
.ok_or_else(|| AuthError::DiscoveryFailed("URL has no host".to_string()))?;
// For IP literals, parse directly and check.
if let Ok(ip) = host.parse::<IpAddr>()
&& is_dangerous_ip(ip)
{
return Err(AuthError::DiscoveryFailed(format!(
"URL points to a restricted IP address: {}",
host
)));
}
// For hostnames, resolve DNS and check each resolved address.
// This prevents DNS-based SSRF where a hostname resolves to an internal IP
// (e.g., 169.254.169.254 for cloud metadata endpoints).
if host.parse::<IpAddr>().is_err() {
let addr = format!("{}:{}", host, parsed.port_or_known_default().unwrap_or(443));
match tokio::net::lookup_host(&addr).await {
Ok(addrs) => {
for socket_addr in addrs {
if is_dangerous_ip(socket_addr.ip()) {
return Err(AuthError::DiscoveryFailed(format!(
"URL hostname '{}' resolves to restricted IP address: {}",
host,
socket_addr.ip()
)));
}
}
}
Err(e) => {
// DNS failure = fail closed (do not allow the request)
return Err(AuthError::DiscoveryFailed(format!(
"DNS resolution failed for '{}': {}",
host, e
)));
}
}
}
Ok(())
}
// ---------------------------------------------------------------------------
// Multi-strategy OAuth discovery helpers
// ---------------------------------------------------------------------------
/// Parse the resource_metadata URL from a WWW-Authenticate header value.
fn parse_resource_metadata_url(www_authenticate: &str) -> Option<String> {
// Try comma-separated parameters first
for part in www_authenticate.split(',') {
let part = part.trim();
if let Some(rest) = part.strip_prefix("resource_metadata=\"") {
return rest.strip_suffix('"').map(|s| s.to_string());
}
if let Some(rest) = part.strip_prefix("resource_metadata=") {
let val = rest.trim_matches('"');
return Some(val.to_string());
}
}
// Also try whitespace-separated tokens (e.g. Bearer resource_metadata="url")
for part in www_authenticate.split_whitespace() {
if let Some(rest) = part.strip_prefix("resource_metadata=\"") {
return rest
.trim_end_matches(',')
.strip_suffix('"')
.map(|s| s.to_string());
}
if let Some(rest) = part.strip_prefix("resource_metadata=") {
let val = rest.trim_matches('"').trim_end_matches(',');
return Some(val.to_string());
}
}
None
}
/// Fetch protected resource metadata from a URL.
async fn fetch_resource_metadata(url: &str) -> Result<ProtectedResourceMetadata, AuthError> {
validate_url_safe(url).await?;
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(10))
.redirect(reqwest::redirect::Policy::none())
.build()
.map_err(|e| AuthError::Http(e.to_string()))?;
let response = client
.get(url)
.send()
.await
.map_err(|e| AuthError::DiscoveryFailed(e.to_string()))?;
if !response.status().is_success() {
return Err(AuthError::DiscoveryFailed(format!(
"HTTP {}",
response.status()
)));
}
response
.json()
.await
.map_err(|e| AuthError::DiscoveryFailed(format!("Invalid metadata: {}", e)))
}
/// Try to discover OAuth metadata via 401 challenge response.
async fn discover_via_401(server_url: &str) -> Result<AuthorizationServerMetadata, AuthError> {
validate_url_safe(server_url).await?;
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(10))
.redirect(reqwest::redirect::Policy::none())
.build()
.map_err(|e| AuthError::Http(e.to_string()))?;
let response = client
.post(server_url)
.header("Content-Type", "application/json")
.body("{}")
.send()
.await
.map_err(|e| AuthError::DiscoveryFailed(e.to_string()))?;
if response.status().as_u16() != 401 {
return Err(AuthError::DiscoveryFailed(format!(
"Expected 401, got {}",
response.status()
)));
}
let www_auth = response
.headers()
.get("WWW-Authenticate")
.and_then(|v| v.to_str().ok())
.ok_or_else(|| {
AuthError::DiscoveryFailed("No WWW-Authenticate header in 401 response".to_string())
})?;
let resource_metadata_url = parse_resource_metadata_url(www_auth).ok_or_else(|| {
AuthError::DiscoveryFailed(
"No resource_metadata URL in WWW-Authenticate header".to_string(),
)
})?;
let resource_meta = fetch_resource_metadata(&resource_metadata_url).await?;
try_discover_from_auth_servers(&resource_meta).await
}
/// Try to discover auth server metadata from resource metadata's authorization_servers list.
async fn try_discover_from_auth_servers(
resource_meta: &ProtectedResourceMetadata,
) -> Result<AuthorizationServerMetadata, AuthError> {
let auth_server_url = resource_meta
.authorization_servers
.first()
.ok_or_else(|| AuthError::DiscoveryFailed("No authorization servers listed".to_string()))?;
discover_authorization_server(auth_server_url).await
}
// ---------------------------------------------------------------------------
// Discovery functions
// ---------------------------------------------------------------------------
/// Discover protected resource metadata from an MCP server.
pub async fn discover_protected_resource(
server_url: &str,
) -> Result<ProtectedResourceMetadata, AuthError> {
validate_url_safe(server_url).await?;
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(10))
.redirect(reqwest::redirect::Policy::none())
.build()
.map_err(|e| AuthError::Http(e.to_string()))?;
// Parse the server URL to extract the origin (scheme + host + port)
// The .well-known endpoints are always at the root of the origin, not under any path
let parsed = reqwest::Url::parse(server_url)
.map_err(|e| AuthError::DiscoveryFailed(format!("Invalid server URL: {}", e)))?;
let origin = parsed.origin().ascii_serialization();
// Try the well-known endpoint at the origin root
let well_known_url = format!("{}/.well-known/oauth-protected-resource", origin);
let well_known_url = build_well_known_uri(server_url, "oauth-protected-resource")?;
let response = client
.get(&well_known_url)
@@ -237,13 +500,15 @@ pub async fn discover_protected_resource(
pub async fn discover_authorization_server(
auth_server_url: &str,
) -> Result<AuthorizationServerMetadata, AuthError> {
validate_url_safe(auth_server_url).await?;
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(10))
.redirect(reqwest::redirect::Policy::none())
.build()
.map_err(|e| AuthError::Http(e.to_string()))?;
let base_url = auth_server_url.trim_end_matches('/');
let well_known_url = format!("{}/.well-known/oauth-authorization-server", base_url);
let well_known_url = build_well_known_uri(auth_server_url, "oauth-authorization-server")?;
let response = client
.get(&well_known_url)
@@ -298,20 +563,27 @@ pub async fn discover_oauth_endpoints(
/// Discover full OAuth metadata including DCR support.
///
/// Returns authorization server metadata which includes registration_endpoint if DCR is supported.
/// Uses a 3-strategy discovery chain:
/// 1. **401-based**: POST to MCP server, parse WWW-Authenticate header for resource_metadata URL
/// 2. **RFC 9728**: Discover protected resource metadata, then authorization server from it
/// 3. **Direct**: Treat MCP server as its own auth server
pub async fn discover_full_oauth_metadata(
server_url: &str,
) -> Result<AuthorizationServerMetadata, AuthError> {
// Try to discover from the server
let resource_meta = discover_protected_resource(server_url).await?;
// Strategy 1: 401-based discovery
if let Ok(meta) = discover_via_401(server_url).await {
return Ok(meta);
}
// Get the first authorization server
let auth_server_url = resource_meta
.authorization_servers
.first()
.ok_or_else(|| AuthError::DiscoveryFailed("No authorization servers listed".to_string()))?;
// Strategy 2: RFC 9728 protected resource discovery
if let Ok(resource_meta) = discover_protected_resource(server_url).await
&& let Ok(meta) = try_discover_from_auth_servers(&resource_meta).await
{
return Ok(meta);
}
// Discover the authorization server metadata
discover_authorization_server(auth_server_url).await
// Strategy 3: Direct - treat MCP server as its own auth server
discover_authorization_server(server_url).await
}
/// Perform Dynamic Client Registration with an authorization server.
@@ -321,8 +593,11 @@ pub async fn register_client(
registration_endpoint: &str,
redirect_uri: &str,
) -> Result<ClientRegistrationResponse, AuthError> {
validate_url_safe(registration_endpoint).await?;
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(30))
.redirect(reqwest::redirect::Policy::none())
.build()
.map_err(|e| AuthError::Http(e.to_string()))?;
@@ -417,7 +692,7 @@ pub async fn authorize_mcp_server(
println!(" Registering client dynamically...");
let registration = register_client(&registration_endpoint, &redirect_uri).await?;
println!(" Client registered: {}", registration.client_id);
println!(" Client registered: {}", registration.client_id);
(
registration.client_id,
@@ -436,6 +711,15 @@ pub async fn authorize_mcp_server(
None
};
// Compute canonical resource URI for RFC 8707
let resource = canonical_resource_uri(&server_config.url);
// Validate the discovered authorization URL to prevent a malicious MCP server
// from redirecting the user to a phishing page or non-HTTPS endpoint.
validate_url_safe(&authorization_url)
.await
.map_err(|e| AuthError::DiscoveryFailed(format!("Unsafe authorization endpoint: {}", e)))?;
// Build authorization URL
let auth_url = build_authorization_url(
&authorization_url,
@@ -444,6 +728,7 @@ pub async fn authorize_mcp_server(
&scopes,
pkce.as_ref(),
&extra_params,
Some(&resource),
);
// Open browser
@@ -462,9 +747,15 @@ pub async fn authorize_mcp_server(
println!(" Exchanging code for token...");
// Exchange code for token
let token =
exchange_code_for_token(&token_url, &client_id, &code, &redirect_uri, pkce.as_ref())
.await?;
let token = exchange_code_for_token(
&token_url,
&client_id,
&code,
&redirect_uri,
pkce.as_ref(),
Some(&resource),
)
.await?;
// Store the tokens
store_tokens(secrets, user_id, server_config, &token).await?;
@@ -493,6 +784,7 @@ pub fn build_authorization_url(
scopes: &[String],
pkce: Option<&PkceChallenge>,
extra_params: &HashMap<String, String>,
resource: Option<&str>,
) -> String {
let mut url = format!(
"{}?client_id={}&response_type=code&redirect_uri={}",
@@ -523,6 +815,10 @@ pub fn build_authorization_url(
));
}
if let Some(resource) = resource {
url.push_str(&format!("&resource={}", urlencoding::encode(resource)));
}
url
}
@@ -553,9 +849,13 @@ pub async fn exchange_code_for_token(
code: &str,
redirect_uri: &str,
pkce: Option<&PkceChallenge>,
resource: Option<&str>,
) -> Result<AccessToken, AuthError> {
validate_url_safe(token_url).await?;
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(30))
.redirect(reqwest::redirect::Policy::none())
.build()
.map_err(|e| AuthError::Http(e.to_string()))?;
@@ -570,6 +870,10 @@ pub async fn exchange_code_for_token(
params.push(("code_verifier", pkce.verifier.clone()));
}
if let Some(resource) = resource {
params.push(("resource", resource.to_string()));
}
let response = client
.post(token_url)
.form(&params)
@@ -738,15 +1042,22 @@ pub async fn refresh_access_token(
auth_meta.token_endpoint
};
validate_url_safe(&token_url).await?;
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(30))
.redirect(reqwest::redirect::Policy::none())
.build()
.map_err(|e| AuthError::Http(e.to_string()))?;
// Compute canonical resource URI for RFC 8707
let resource = canonical_resource_uri(&server_config.url);
let params = vec![
("grant_type", "refresh_token".to_string()),
("refresh_token", refresh_token.expose().to_string()),
("client_id", client_id),
("resource", resource),
];
let response = client
@@ -815,6 +1126,7 @@ mod tests {
&["read".to_string(), "write".to_string()],
None,
&HashMap::new(),
None,
);
assert!(url.starts_with("https://auth.example.com/authorize?"));
@@ -834,6 +1146,7 @@ mod tests {
&[],
Some(&pkce),
&HashMap::new(),
None,
);
assert!(url.contains(&format!("code_challenge={}", pkce.challenge)));
@@ -853,6 +1166,7 @@ mod tests {
&[],
None,
&extra,
None,
);
assert!(url.contains("owner=user"));
@@ -880,6 +1194,7 @@ mod tests {
&[],
None,
&HashMap::new(),
None,
);
// With no scopes, the URL must not contain a scope parameter at all.
@@ -895,6 +1210,7 @@ mod tests {
&[],
None,
&HashMap::new(),
None,
);
// Spaces and ampersands in client_id must be percent-encoded.
@@ -1164,4 +1480,235 @@ mod tests {
);
}
}
// --- New tests for well-known URI construction ---
#[test]
fn test_build_well_known_uri_no_path() {
let uri =
build_well_known_uri("https://example.com", "oauth-authorization-server").unwrap();
assert_eq!(
uri,
"https://example.com/.well-known/oauth-authorization-server"
);
}
#[test]
fn test_build_well_known_uri_with_path() {
let uri =
build_well_known_uri("https://example.com/path", "oauth-authorization-server").unwrap();
assert_eq!(
uri,
"https://example.com/.well-known/oauth-authorization-server/path"
);
}
#[test]
fn test_build_well_known_uri_with_trailing_slash() {
let uri =
build_well_known_uri("https://example.com/path/", "oauth-protected-resource").unwrap();
assert_eq!(
uri,
"https://example.com/.well-known/oauth-protected-resource/path"
);
}
#[test]
fn test_build_well_known_uri_root_trailing_slash() {
let uri =
build_well_known_uri("https://example.com/", "oauth-authorization-server").unwrap();
assert_eq!(
uri,
"https://example.com/.well-known/oauth-authorization-server"
);
}
// --- New tests for canonical_resource_uri ---
#[test]
fn test_canonical_resource_uri_strips_fragment() {
assert_eq!(
canonical_resource_uri("https://mcp.example.com/v1#section"),
"https://mcp.example.com/v1"
);
}
#[test]
fn test_canonical_resource_uri_strips_trailing_slash() {
assert_eq!(
canonical_resource_uri("https://mcp.example.com/v1/"),
"https://mcp.example.com/v1"
);
}
#[test]
fn test_canonical_resource_uri_no_changes_needed() {
assert_eq!(
canonical_resource_uri("https://mcp.example.com/v1"),
"https://mcp.example.com/v1"
);
}
// --- New tests for SSRF protection ---
#[test]
fn test_is_dangerous_ip_loopback_v4() {
assert!(is_dangerous_ip("127.0.0.1".parse().unwrap()));
assert!(is_dangerous_ip("127.0.0.2".parse().unwrap()));
}
#[test]
fn test_is_dangerous_ip_private_v4() {
assert!(is_dangerous_ip("10.0.0.1".parse().unwrap()));
assert!(is_dangerous_ip("172.16.0.1".parse().unwrap()));
assert!(is_dangerous_ip("192.168.1.1".parse().unwrap()));
}
#[test]
fn test_is_dangerous_ip_link_local_v4() {
assert!(is_dangerous_ip("169.254.169.254".parse().unwrap()));
}
#[test]
fn test_is_dangerous_ip_cgnat() {
assert!(is_dangerous_ip("100.64.0.1".parse().unwrap()));
assert!(is_dangerous_ip("100.127.255.254".parse().unwrap()));
}
#[test]
fn test_is_dangerous_ip_safe_v4() {
assert!(!is_dangerous_ip("8.8.8.8".parse().unwrap()));
assert!(!is_dangerous_ip("1.1.1.1".parse().unwrap()));
}
#[test]
fn test_is_dangerous_ip_ipv4_mapped_v6_loopback() {
// ::ffff:127.0.0.1 must be blocked
let ip: IpAddr = "::ffff:127.0.0.1".parse().unwrap();
assert!(is_dangerous_ip(ip));
}
#[test]
fn test_is_dangerous_ip_ipv4_mapped_v6_link_local() {
// ::ffff:169.254.169.254 must be blocked
let ip: IpAddr = "::ffff:169.254.169.254".parse().unwrap();
assert!(is_dangerous_ip(ip));
}
#[test]
fn test_is_dangerous_ip_unspecified() {
assert!(is_dangerous_ip("0.0.0.0".parse().unwrap()));
assert!(is_dangerous_ip("::".parse().unwrap()));
}
#[test]
fn test_is_dangerous_ip_v6_loopback() {
assert!(is_dangerous_ip("::1".parse().unwrap()));
}
#[tokio::test]
async fn test_validate_url_safe_https() {
assert!(validate_url_safe("https://example.com/path").await.is_ok());
}
#[tokio::test]
async fn test_validate_url_safe_http_localhost_allowed() {
// HTTP is only allowed for localhost dev scenarios
assert!(validate_url_safe("http://localhost/path").await.is_ok());
assert!(
validate_url_safe("http://localhost:8080/path")
.await
.is_ok()
);
}
#[tokio::test]
async fn test_validate_url_safe_http_non_localhost_rejected() {
// HTTP to non-localhost hosts must be rejected (plaintext credential risk)
assert!(validate_url_safe("http://example.com/path").await.is_err());
}
#[tokio::test]
async fn test_validate_url_safe_bad_scheme() {
assert!(validate_url_safe("ftp://example.com/path").await.is_err());
assert!(validate_url_safe("file:///etc/passwd").await.is_err());
}
#[tokio::test]
async fn test_validate_url_safe_private_ip() {
// 127.0.0.1 over HTTP is allowed (localhost dev scenario)
assert!(validate_url_safe("http://127.0.0.1/path").await.is_ok());
// Private/link-local IPs over HTTPS are blocked (SSRF protection)
assert!(validate_url_safe("https://10.0.0.1/path").await.is_err());
assert!(
validate_url_safe("https://169.254.169.254/latest/meta-data")
.await
.is_err()
);
// Private IPs over HTTP (non-localhost) are blocked
assert!(validate_url_safe("http://10.0.0.1/path").await.is_err());
}
#[tokio::test]
async fn test_validate_url_safe_public_ip() {
assert!(validate_url_safe("https://8.8.8.8/dns").await.is_ok());
}
// --- New tests for parse_resource_metadata_url ---
#[test]
fn test_parse_resource_metadata_url_bearer() {
let header = r#"Bearer resource_metadata="https://res.example.com/.well-known/oauth-protected-resource""#;
let url = parse_resource_metadata_url(header);
assert_eq!(
url.as_deref(),
Some("https://res.example.com/.well-known/oauth-protected-resource")
);
}
#[test]
fn test_parse_resource_metadata_url_with_other_params() {
let header = r#"Bearer realm="example", resource_metadata="https://res.example.com/meta""#;
let url = parse_resource_metadata_url(header);
assert_eq!(url.as_deref(), Some("https://res.example.com/meta"));
}
#[test]
fn test_parse_resource_metadata_url_missing() {
let header = r#"Bearer realm="example""#;
let url = parse_resource_metadata_url(header);
assert!(url.is_none());
}
// --- New tests for resource parameter in authorization URL ---
#[test]
fn test_build_authorization_url_with_resource() {
let url = build_authorization_url(
"https://auth.example.com/authorize",
"client-123",
"http://localhost:9876/callback",
&[],
None,
&HashMap::new(),
Some("https://mcp.example.com/v1"),
);
assert!(url.contains("resource=https%3A%2F%2Fmcp.example.com%2Fv1"));
}
#[test]
fn test_build_authorization_url_without_resource() {
let url = build_authorization_url(
"https://auth.example.com/authorize",
"client-123",
"http://localhost:9876/callback",
&[],
None,
&HashMap::new(),
None,
);
assert!(!url.contains("resource="));
}
}
+253 -187
View File
@@ -1,11 +1,11 @@
//! MCP client for connecting to MCP servers.
//!
//! Supports both local (unauthenticated) and hosted (OAuth-authenticated) servers.
//! Uses the Streamable HTTP transport with session management.
//! Uses pluggable transports (HTTP, stdio, Unix) via the `McpTransport` trait.
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use async_trait::async_trait;
use tokio::sync::RwLock;
@@ -14,27 +14,29 @@ use crate::context::JobContext;
use crate::secrets::SecretsStore;
use crate::tools::mcp::auth::refresh_access_token;
use crate::tools::mcp::config::McpServerConfig;
use crate::tools::mcp::http_transport::HttpMcpTransport;
use crate::tools::mcp::protocol::{
CallToolResult, InitializeResult, ListToolsResult, McpRequest, McpResponse, McpTool,
};
use crate::tools::mcp::session::McpSessionManager;
use crate::tools::mcp::transport::McpTransport;
use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput};
/// MCP client for communicating with MCP servers.
///
/// Supports two modes:
/// - Simple: Just a URL, no auth or session management (for local/test servers)
/// - Authenticated: Full OAuth support with session management (for hosted servers)
/// Supports multiple transport types:
/// - HTTP: For remote MCP servers (created via `new`, `new_with_name`, `new_authenticated`)
/// - Stdio/Unix: Via `new_with_transport` with a custom `McpTransport` implementation
pub struct McpClient {
/// Server URL (for HTTP transport).
/// Transport for sending requests.
transport: Arc<dyn McpTransport>,
/// Server URL (kept for accessor compatibility).
server_url: String,
/// Server name (for logging and session management).
server_name: String,
/// HTTP client.
http_client: reqwest::Client,
/// Request ID counter.
next_id: AtomicU64,
@@ -52,6 +54,9 @@ pub struct McpClient {
/// Server configuration (for token secret name lookup).
server_config: Option<McpServerConfig>,
/// Custom headers to include in every request.
custom_headers: HashMap<String, String>,
}
impl McpClient {
@@ -59,22 +64,21 @@ impl McpClient {
///
/// Use this for local development servers or servers that don't require auth.
pub fn new(server_url: impl Into<String>) -> Self {
let url = server_url.into();
let url: String = server_url.into();
let name = extract_server_name(&url);
let transport = Arc::new(HttpMcpTransport::new(url.clone(), name.clone()));
Self {
transport,
server_url: url,
server_name: name,
http_client: reqwest::Client::builder()
.timeout(Duration::from_secs(30))
.build()
.expect("Failed to create HTTP client"),
next_id: AtomicU64::new(1),
tools_cache: RwLock::new(None),
session_manager: None,
secrets: None,
user_id: "default".to_string(),
server_config: None,
custom_headers: HashMap::new(),
}
}
@@ -82,19 +86,52 @@ impl McpClient {
///
/// Use this when you have a configured server name but no authentication.
pub fn new_with_name(server_name: impl Into<String>, server_url: impl Into<String>) -> Self {
let name: String = server_name.into();
let url: String = server_url.into();
let transport = Arc::new(HttpMcpTransport::new(url.clone(), name.clone()));
Self {
server_url: server_url.into(),
server_name: server_name.into(),
http_client: reqwest::Client::builder()
.timeout(Duration::from_secs(30))
.build()
.expect("Failed to create HTTP client"),
transport,
server_url: url,
server_name: name,
next_id: AtomicU64::new(1),
tools_cache: RwLock::new(None),
session_manager: None,
secrets: None,
user_id: "default".to_string(),
server_config: None,
custom_headers: HashMap::new(),
}
}
/// Create a new simple MCP client from an HTTP server configuration (no authentication).
///
/// Use this when you have an `McpServerConfig` with custom headers but no OAuth.
/// The config must use HTTP transport (the default); for stdio/UDS use `new_with_transport`.
pub fn new_with_config(config: McpServerConfig) -> Self {
assert!(
matches!(
config.effective_transport(),
crate::tools::mcp::config::EffectiveTransport::Http
),
"new_with_config only supports HTTP transport; use new_with_transport for stdio/UDS"
);
let transport = Arc::new(HttpMcpTransport::new(
config.url.clone(),
config.name.clone(),
));
Self {
transport,
server_url: config.url.clone(),
server_name: config.name.clone(),
next_id: AtomicU64::new(1),
tools_cache: RwLock::new(None),
session_manager: None,
secrets: None,
user_id: "default".to_string(),
custom_headers: config.headers.clone(),
server_config: Some(config),
}
}
@@ -107,19 +144,59 @@ impl McpClient {
secrets: Arc<dyn SecretsStore + Send + Sync>,
user_id: impl Into<String>,
) -> Self {
let transport = Arc::new(
HttpMcpTransport::new(config.url.clone(), config.name.clone())
.with_session_manager(session_manager.clone()),
);
let custom_headers = config.headers.clone();
Self {
transport,
server_url: config.url.clone(),
server_name: config.name.clone(),
http_client: reqwest::Client::builder()
.timeout(Duration::from_secs(30))
.build()
.expect("Failed to create HTTP client"),
next_id: AtomicU64::new(1),
tools_cache: RwLock::new(None),
session_manager: Some(session_manager),
secrets: Some(secrets),
user_id: user_id.into(),
server_config: Some(config),
custom_headers,
}
}
/// Create a new MCP client with a custom transport.
///
/// Use this for stdio, UDS, or other non-HTTP transports.
pub fn new_with_transport(
server_name: impl Into<String>,
transport: Arc<dyn McpTransport>,
session_manager: Option<Arc<McpSessionManager>>,
secrets: Option<Arc<dyn SecretsStore + Send + Sync>>,
user_id: impl Into<String>,
server_config: Option<McpServerConfig>,
) -> Self {
let name: String = server_name.into();
let url = server_config
.as_ref()
.map(|c| c.url.clone())
.unwrap_or_default();
let custom_headers = server_config
.as_ref()
.map(|c| c.headers.clone())
.unwrap_or_default();
Self {
transport,
server_url: url,
server_name: name,
next_id: AtomicU64::new(1),
tools_cache: RwLock::new(None),
session_manager,
secrets,
user_id: user_id.into(),
server_config,
custom_headers,
}
}
@@ -139,19 +216,13 @@ impl McpClient {
}
/// Get the access token for this server (if authenticated).
///
/// Returns the stored token regardless of whether OAuth was pre-configured
/// or obtained via Dynamic Client Registration.
async fn get_access_token(&self) -> Result<Option<String>, ToolError> {
let Some(ref secrets) = self.secrets else {
return Ok(None);
};
let Some(ref config) = self.server_config else {
return Ok(None);
};
// Try to get stored token (from either pre-configured OAuth or DCR)
match secrets
.get_decrypted(&self.user_id, &config.token_secret_name())
.await
@@ -165,46 +236,41 @@ impl McpClient {
}
}
/// Build the headers map for a request (auth, session-id, custom headers).
async fn build_request_headers(&self) -> Result<HashMap<String, String>, ToolError> {
let mut headers = self.custom_headers.clone();
if let Some(token) = self.get_access_token().await? {
headers.insert("Authorization".to_string(), format!("Bearer {}", token));
}
if let Some(ref session_manager) = self.session_manager
&& let Some(session_id) = session_manager.get_session_id(&self.server_name).await
{
headers.insert("Mcp-Session-Id".to_string(), session_id);
}
Ok(headers)
}
/// Send a request to the MCP server with auth and session headers.
/// Automatically attempts token refresh on 401 errors.
/// Automatically attempts token refresh on 401 errors (HTTP transports only).
async fn send_request(&self, request: McpRequest) -> Result<McpResponse, ToolError> {
// Try up to 2 times: first attempt, then retry after token refresh
// For non-HTTP transports, just send directly without retry logic
if !self.transport.supports_http_features() {
let headers = self.build_request_headers().await?;
return self.transport.send(&request, &headers).await;
}
// HTTP transport: try up to 2 times (first attempt, then retry after token refresh)
for attempt in 0..2 {
// Request both JSON and SSE as per MCP spec
let mut req_builder = self
.http_client
.post(&self.server_url)
.header("Accept", "application/json, text/event-stream")
.header("Content-Type", "application/json")
.json(&request);
let headers = self.build_request_headers().await?;
let result = self.transport.send(&request, &headers).await;
// Add Authorization header if we have a token
if let Some(token) = self.get_access_token().await? {
req_builder = req_builder.header("Authorization", format!("Bearer {}", token));
}
// Add Mcp-Session-Id header if we have a session
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| {
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
match result {
Ok(response) => return Ok(response),
Err(ToolError::ExternalService(ref msg))
if msg.contains("401") || msg.contains("Unauthorized") =>
{
if attempt == 0
&& let Some(ref secrets) = self.secrets
&& let Some(ref config) = self.server_config
{
tracing::debug!(
@@ -214,7 +280,6 @@ impl McpClient {
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) => {
@@ -223,108 +288,30 @@ impl McpClient {
self.server_name,
e
);
// Fall through to return auth error
}
}
}
return Err(ToolError::ExternalService(format!(
"MCP server '{}' requires authentication. Run: ironclaw mcp auth {}",
self.server_name, self.server_name
)));
}
return Err(ToolError::ExternalService(format!(
"MCP server '{}' requires authentication. Run: ironclaw mcp auth {}",
self.server_name, self.server_name
)));
Err(e) => return Err(e),
}
// Success path - return the parsed response
return self.parse_response(response).await;
}
// Should not reach here, but just in case
Err(ToolError::ExternalService(
"MCP request failed after retry".to_string(),
))
}
/// 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
&& 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;
}
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
return Err(ToolError::ExternalService(format!(
"MCP server returned status: {} - {}",
status, body
)));
}
// Check content type to handle SSE vs JSON responses
let content_type = response
.headers()
.get("content-type")
.and_then(|v| v.to_str().ok())
.unwrap_or("")
.to_string();
if content_type.contains("text/event-stream") {
// SSE response - read chunks until we get a complete JSON message
use futures::StreamExt;
let mut stream = response.bytes_stream();
let mut buffer = String::new();
while let Some(chunk) = stream.next().await {
let chunk = chunk.map_err(|e| {
ToolError::ExternalService(format!("Failed to read SSE chunk: {}", e))
})?;
buffer.push_str(&String::from_utf8_lossy(&chunk));
// Look for complete SSE data lines
for line in buffer.lines() {
if let Some(json_str) = line.strip_prefix("data: ") {
// Try to parse - if valid JSON, we're done
if let Ok(response) = serde_json::from_str::<McpResponse>(json_str) {
return Ok(response);
}
}
}
}
Err(ToolError::ExternalService(format!(
"No valid data in SSE response: {}",
buffer
)))
} else {
// JSON response
response.json().await.map_err(|e| {
ToolError::ExternalService(format!("Failed to parse MCP response: {}", e))
})
}
}
/// Initialize the connection to the MCP server.
///
/// 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
&& session_manager.is_initialized(&self.server_name).await
{
// Return cached/default capabilities
return Ok(InitializeResult::default());
}
// Ensure we have a session
if let Some(ref session_manager) = self.session_manager {
session_manager
.get_or_create(&self.server_name, &self.server_url)
@@ -352,14 +339,11 @@ impl McpClient {
})
})?;
// Mark session as initialized
if let Some(ref session_manager) = self.session_manager {
session_manager.mark_initialized(&self.server_name).await;
}
// Send initialized notification
let notification = McpRequest::initialized_notification();
// Fire and forget - notifications don't have responses
let _ = self.send_request(notification).await;
Ok(result)
@@ -367,12 +351,9 @@ impl McpClient {
/// List available tools from the MCP server.
pub async fn list_tools(&self) -> Result<Vec<McpTool>, ToolError> {
// Check cache first
if let Some(tools) = self.tools_cache.read().await.as_ref() {
return Ok(tools.clone());
}
// Ensure initialized for authenticated sessions
if self.session_manager.is_some() {
self.initialize().await?;
}
@@ -395,9 +376,7 @@ impl McpClient {
.map_err(|e| ToolError::ExternalService(format!("Invalid tools list: {}", e)))
})?;
// Cache the tools
*self.tools_cache.write().await = Some(result.tools.clone());
Ok(result.tools)
}
@@ -407,7 +386,6 @@ impl McpClient {
name: &str,
arguments: serde_json::Value,
) -> Result<CallToolResult, ToolError> {
// Ensure initialized for authenticated sessions
if self.session_manager.is_some() {
self.initialize().await?;
}
@@ -440,7 +418,6 @@ impl McpClient {
pub async fn create_tools(&self) -> Result<Vec<Arc<dyn Tool>>, ToolError> {
let mcp_tools = self.list_tools().await?;
let client = Arc::new(self.clone());
Ok(mcp_tools
.into_iter()
.map(|t| {
@@ -465,15 +442,16 @@ impl McpClient {
impl Clone for McpClient {
fn clone(&self) -> Self {
Self {
transport: self.transport.clone(),
server_url: self.server_url.clone(),
server_name: self.server_name.clone(),
http_client: self.http_client.clone(),
next_id: AtomicU64::new(self.next_id.load(Ordering::SeqCst)),
tools_cache: RwLock::new(None),
session_manager: self.session_manager.clone(),
secrets: self.secrets.clone(),
user_id: self.user_id.clone(),
server_config: self.server_config.clone(),
custom_headers: self.custom_headers.clone(),
}
}
}
@@ -490,7 +468,6 @@ fn extract_server_name(url: &str) -> String {
/// Wrapper that implements Tool for an MCP tool.
struct McpToolWrapper {
tool: McpTool,
/// Prefixed name (server_name_tool_name) for unique identification.
prefixed_name: String,
client: Arc<McpClient>,
}
@@ -500,11 +477,9 @@ impl Tool for McpToolWrapper {
fn name(&self) -> &str {
&self.prefixed_name
}
fn description(&self) -> &str {
&self.tool.description
}
fn parameters_schema(&self) -> serde_json::Value {
self.tool.input_schema.clone()
}
@@ -515,31 +490,24 @@ impl Tool for McpToolWrapper {
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let start = std::time::Instant::now();
// Use the original tool name (without prefix) for the actual call
let result = self.client.call_tool(&self.tool.name, params).await?;
// Convert content blocks to a single result
let content: String = result
.content
.iter()
.filter_map(|block| block.as_text())
.filter_map(|b| b.as_text())
.collect::<Vec<_>>()
.join("\n");
if result.is_error {
return Err(ToolError::ExecutionFailed(content));
}
Ok(ToolOutput::text(content, start.elapsed()))
}
fn requires_sanitization(&self) -> bool {
true // MCP tools are external, always sanitize
true
}
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
// Delegate to the MCP protocol type's own requires_approval() bool method
if self.tool.requires_approval() {
ApprovalRequirement::UnlessAutoApproved
} else {
@@ -548,6 +516,9 @@ impl Tool for McpToolWrapper {
}
}
/// Sanitize an HTTP error response body for safe display.
///
/// Detects full HTML error pages (containing `<html` or `<!DOCTYPE`) and
#[cfg(test)]
mod tests {
use super::*;
@@ -556,7 +527,7 @@ mod tests {
fn test_mcp_request_list_tools() {
let req = McpRequest::list_tools(1);
assert_eq!(req.method, "tools/list");
assert_eq!(req.id, 1);
assert_eq!(req.id, Some(1));
}
#[test]
@@ -660,42 +631,61 @@ mod tests {
#[test]
fn test_clone_preserves_fields() {
let client = McpClient::new_with_name("cloned-server", "http://localhost:5555");
// Bump the request ID a few times
client.next_request_id();
client.next_request_id();
let cloned = client.clone();
assert_eq!(cloned.server_url(), "http://localhost:5555");
assert_eq!(cloned.server_name(), "cloned-server");
assert_eq!(cloned.user_id, "default");
// The atomic counter value is copied
assert_eq!(cloned.next_id.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn test_clone_resets_tools_cache() {
let client = McpClient::new("http://localhost:5555");
// The clone implementation resets tools_cache to None
let cloned = client.clone();
let cache = cloned.tools_cache.read().await;
assert!(cache.is_none());
}
#[test]
fn test_new_with_config_carries_custom_headers() {
let mut headers = HashMap::new();
headers.insert("X-API-Key".to_string(), "secret".to_string());
headers.insert("X-Custom".to_string(), "value".to_string());
let config = McpServerConfig::new("test", "http://localhost:8080").with_headers(headers);
let client = McpClient::new_with_config(config.clone());
assert_eq!(client.server_name(), "test");
assert_eq!(client.server_url(), "http://localhost:8080");
assert_eq!(client.custom_headers.len(), 2);
assert_eq!(client.custom_headers.get("X-API-Key").unwrap(), "secret");
assert!(client.server_config.is_some());
}
#[test]
fn test_new_with_config_no_headers() {
let config = McpServerConfig::new("bare", "http://localhost:9090");
let client = McpClient::new_with_config(config);
assert_eq!(client.server_name(), "bare");
assert!(client.custom_headers.is_empty());
assert!(client.secrets.is_none());
assert!(client.session_manager.is_none());
}
#[test]
fn test_next_request_id_monotonically_increasing() {
let client = McpClient::new("http://localhost:1234");
let id1 = client.next_request_id();
let id2 = client.next_request_id();
let id3 = client.next_request_id();
assert_eq!(id1, 1);
assert_eq!(id2, 2);
assert_eq!(id3, 3);
assert_eq!(client.next_request_id(), 1);
assert_eq!(client.next_request_id(), 2);
assert_eq!(client.next_request_id(), 3);
}
#[test]
fn test_mcp_tool_requires_approval_destructive() {
use crate::tools::mcp::protocol::{McpTool, McpToolAnnotations};
let tool = McpTool {
name: "delete_all".to_string(),
description: "Deletes everything".to_string(),
@@ -713,7 +703,6 @@ mod tests {
#[test]
fn test_mcp_tool_no_approval_when_not_destructive() {
use crate::tools::mcp::protocol::{McpTool, McpToolAnnotations};
let tool = McpTool {
name: "read_data".to_string(),
description: "Reads data".to_string(),
@@ -731,7 +720,6 @@ mod tests {
#[test]
fn test_mcp_tool_no_approval_when_no_annotations() {
use crate::tools::mcp::protocol::McpTool;
let tool = McpTool {
name: "simple_tool".to_string(),
description: "A simple tool".to_string(),
@@ -740,4 +728,82 @@ mod tests {
};
assert!(!tool.requires_approval());
}
/// Mock transport for testing transport abstraction behavior.
struct MockTransport {
supports_http: bool,
responses: std::sync::Mutex<Vec<McpResponse>>,
recorded_headers: std::sync::Mutex<Vec<HashMap<String, String>>>,
}
impl MockTransport {
fn new(supports_http: bool, responses: Vec<McpResponse>) -> Self {
Self {
supports_http,
responses: std::sync::Mutex::new(responses),
recorded_headers: std::sync::Mutex::new(Vec::new()),
}
}
fn recorded_headers(&self) -> Vec<HashMap<String, String>> {
self.recorded_headers.lock().unwrap().clone()
}
}
#[async_trait]
impl McpTransport for MockTransport {
async fn send(
&self,
_request: &McpRequest,
headers: &HashMap<String, String>,
) -> Result<McpResponse, ToolError> {
self.recorded_headers.lock().unwrap().push(headers.clone());
let mut responses = self.responses.lock().unwrap();
if responses.is_empty() {
return Err(ToolError::ExternalService(
"No more mock responses".to_string(),
));
}
Ok(responses.remove(0))
}
async fn shutdown(&self) -> Result<(), ToolError> {
Ok(())
}
fn supports_http_features(&self) -> bool {
self.supports_http
}
}
#[tokio::test]
async fn test_non_http_transport_skips_401_retry() {
let response = McpResponse {
jsonrpc: "2.0".to_string(),
id: Some(1),
result: Some(serde_json::json!({"tools": []})),
error: None,
};
let transport = Arc::new(MockTransport::new(false, vec![response]));
let client = McpClient::new_with_transport(
"test-stdio",
transport.clone(),
None,
None,
"default",
None,
);
let result = client.list_tools().await;
assert!(result.is_ok());
assert_eq!(result.unwrap().len(), 0);
let headers = transport.recorded_headers();
assert_eq!(headers.len(), 1);
assert!(!headers[0].contains_key("Authorization"));
assert!(!headers[0].contains_key("Mcp-Session-Id"));
}
#[tokio::test]
async fn test_transport_supports_http_features_accessor() {
let http_transport = HttpMcpTransport::new("http://localhost:8080", "test");
assert!(http_transport.supports_http_features());
let mock_non_http = MockTransport::new(false, vec![]);
assert!(!mock_non_http.supports_http_features());
}
}
+382 -12
View File
@@ -12,6 +12,24 @@ use tokio::fs;
use crate::bootstrap::ironclaw_base_dir;
use crate::tools::tool::ToolError;
/// Transport configuration for an MCP server.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "transport", rename_all = "lowercase")]
pub enum McpTransportConfig {
/// HTTP/HTTPS transport (uses the `url` field on McpServerConfig).
Http,
/// Stdio transport — spawns a child process.
Stdio {
command: String,
#[serde(default)]
args: Vec<String>,
#[serde(default)]
env: HashMap<String, String>,
},
/// Unix domain socket transport.
Unix { socket_path: String },
}
/// Configuration for connecting to a remote MCP server.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpServerConfig {
@@ -21,6 +39,14 @@ pub struct McpServerConfig {
/// Server URL (must be HTTPS for remote servers).
pub url: String,
/// Transport configuration. If `None`, defaults to Http using `url`.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub transport: Option<McpTransportConfig>,
/// Custom headers to include in every HTTP request.
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub headers: HashMap<String, String>,
/// OAuth configuration (if server requires authentication).
#[serde(skip_serializing_if = "Option::is_none")]
pub oauth: Option<OAuthConfig>,
@@ -44,6 +70,45 @@ impl McpServerConfig {
Self {
name: name.into(),
url: url.into(),
transport: None,
headers: HashMap::new(),
oauth: None,
enabled: true,
description: None,
}
}
/// Create a new stdio transport MCP server configuration.
pub fn new_stdio(
name: impl Into<String>,
command: impl Into<String>,
args: Vec<String>,
env: HashMap<String, String>,
) -> Self {
Self {
name: name.into(),
url: String::new(),
transport: Some(McpTransportConfig::Stdio {
command: command.into(),
args,
env,
}),
headers: HashMap::new(),
oauth: None,
enabled: true,
description: None,
}
}
/// Create a new Unix socket transport MCP server configuration.
pub fn new_unix(name: impl Into<String>, socket_path: impl Into<String>) -> Self {
Self {
name: name.into(),
url: String::new(),
transport: Some(McpTransportConfig::Unix {
socket_path: socket_path.into(),
}),
headers: HashMap::new(),
oauth: None,
enabled: true,
description: None,
@@ -62,6 +127,25 @@ impl McpServerConfig {
self
}
/// Set custom headers.
pub fn with_headers(mut self, headers: HashMap<String, String>) -> Self {
self.headers = headers;
self
}
/// Get the effective transport type.
pub fn effective_transport(&self) -> EffectiveTransport<'_> {
match &self.transport {
Some(McpTransportConfig::Http) | None => EffectiveTransport::Http,
Some(McpTransportConfig::Stdio { command, args, env }) => {
EffectiveTransport::Stdio { command, args, env }
}
Some(McpTransportConfig::Unix { socket_path }) => {
EffectiveTransport::Unix { socket_path }
}
}
}
/// Validate the server configuration.
pub fn validate(&self) -> Result<(), ConfigError> {
if self.name.is_empty() {
@@ -70,19 +154,38 @@ impl McpServerConfig {
});
}
if self.url.is_empty() {
return Err(ConfigError::InvalidConfig {
reason: "Server URL cannot be empty".to_string(),
});
}
match self.effective_transport() {
EffectiveTransport::Http => {
if self.url.is_empty() {
return Err(ConfigError::InvalidConfig {
reason: "Server URL cannot be empty".to_string(),
});
}
// Remote servers must use HTTPS (localhost is allowed for development)
let url_lower = self.url.to_lowercase();
let is_localhost = url_lower.contains("localhost") || url_lower.contains("127.0.0.1");
if !is_localhost && !url_lower.starts_with("https://") {
return Err(ConfigError::InvalidConfig {
reason: "Remote MCP servers must use HTTPS".to_string(),
});
// Remote servers must use HTTPS (localhost is allowed for development)
let url_lower = self.url.to_lowercase();
let is_localhost =
url_lower.contains("localhost") || url_lower.contains("127.0.0.1");
if !is_localhost && !url_lower.starts_with("https://") {
return Err(ConfigError::InvalidConfig {
reason: "Remote MCP servers must use HTTPS".to_string(),
});
}
}
EffectiveTransport::Stdio { command, .. } => {
if command.is_empty() {
return Err(ConfigError::InvalidConfig {
reason: "Stdio transport command cannot be empty".to_string(),
});
}
}
EffectiveTransport::Unix { socket_path } => {
if socket_path.is_empty() {
return Err(ConfigError::InvalidConfig {
reason: "Unix socket path cannot be empty".to_string(),
});
}
}
}
Ok(())
@@ -92,7 +195,14 @@ impl McpServerConfig {
///
/// 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).
///
/// Non-HTTP transports (stdio, unix) never require auth.
pub fn requires_auth(&self) -> bool {
// Non-HTTP transports don't use HTTP auth
if !matches!(self.effective_transport(), EffectiveTransport::Http) {
return false;
}
if self.oauth.is_some() {
return true;
}
@@ -426,6 +536,20 @@ fn is_localhost_url(url: &str) -> bool {
}
}
/// Resolved transport type (borrows from config).
#[derive(Debug)]
pub enum EffectiveTransport<'a> {
Http,
Stdio {
command: &'a str,
args: &'a [String],
env: &'a HashMap<String, String>,
},
Unix {
socket_path: &'a str,
},
}
#[cfg(test)]
mod tests {
use super::*;
@@ -593,4 +717,250 @@ mod tests {
let config = McpServerConfig::new("bad", "http://mcp.example.com");
assert!(!config.requires_auth());
}
#[test]
fn test_stdio_config_creation() {
let env = HashMap::from([("PATH".to_string(), "/usr/bin".to_string())]);
let config = McpServerConfig::new_stdio(
"my-server",
"npx",
vec!["-y".to_string(), "@modelcontextprotocol/server".to_string()],
env.clone(),
);
assert_eq!(config.name, "my-server");
assert!(config.url.is_empty());
assert!(config.enabled);
assert!(config.oauth.is_none());
assert!(config.headers.is_empty());
match &config.transport {
Some(McpTransportConfig::Stdio {
command,
args,
env: e,
}) => {
assert_eq!(command, "npx");
assert_eq!(
args,
&["-y".to_string(), "@modelcontextprotocol/server".to_string()]
);
assert_eq!(e, &env);
}
other => panic!("Expected Stdio transport, got {:?}", other),
}
}
#[test]
fn test_unix_config_creation() {
let config = McpServerConfig::new_unix("local-server", "/tmp/mcp.sock");
assert_eq!(config.name, "local-server");
assert!(config.url.is_empty());
assert!(config.enabled);
match &config.transport {
Some(McpTransportConfig::Unix { socket_path }) => {
assert_eq!(socket_path, "/tmp/mcp.sock");
}
other => panic!("Expected Unix transport, got {:?}", other),
}
}
#[test]
fn test_stdio_validation() {
// Valid stdio config
let config = McpServerConfig::new_stdio("server", "npx", vec![], HashMap::new());
assert!(config.validate().is_ok());
// Invalid: empty command
let config = McpServerConfig::new_stdio("server", "", vec![], HashMap::new());
assert!(config.validate().is_err());
let err = config.validate().unwrap_err().to_string();
assert!(
err.contains("command"),
"Error should mention command: {}",
err
);
// Invalid: empty name
let config = McpServerConfig::new_stdio("", "npx", vec![], HashMap::new());
assert!(config.validate().is_err());
}
#[test]
fn test_unix_validation() {
// Valid unix config
let config = McpServerConfig::new_unix("server", "/tmp/mcp.sock");
assert!(config.validate().is_ok());
// Invalid: empty socket path
let config = McpServerConfig::new_unix("server", "");
assert!(config.validate().is_err());
let err = config.validate().unwrap_err().to_string();
assert!(
err.contains("socket"),
"Error should mention socket: {}",
err
);
// Invalid: empty name
let config = McpServerConfig::new_unix("", "/tmp/mcp.sock");
assert!(config.validate().is_err());
}
#[test]
fn test_requires_auth_stdio_never() {
// Stdio transport should never require auth, even with OAuth configured
let mut config = McpServerConfig::new_stdio("server", "npx", vec![], HashMap::new());
assert!(!config.requires_auth());
// Even if OAuth is set, stdio doesn't use HTTP auth
config.oauth = Some(OAuthConfig::new("client-123"));
assert!(!config.requires_auth());
}
#[test]
fn test_requires_auth_unix_never() {
// Unix transport should never require auth
let mut config = McpServerConfig::new_unix("server", "/tmp/mcp.sock");
assert!(!config.requires_auth());
config.oauth = Some(OAuthConfig::new("client-123"));
assert!(!config.requires_auth());
}
#[test]
fn test_custom_headers() {
let headers = HashMap::from([
("X-Api-Key".to_string(), "secret".to_string()),
("Authorization".to_string(), "Bearer token".to_string()),
]);
let config =
McpServerConfig::new("server", "https://mcp.example.com").with_headers(headers.clone());
assert_eq!(config.headers, headers);
assert_eq!(config.headers.get("X-Api-Key").unwrap(), "secret");
}
#[test]
fn test_transport_config_serde_http() {
let transport = McpTransportConfig::Http;
let json = serde_json::to_string(&transport).unwrap();
assert!(json.contains("\"transport\":\"http\""));
let parsed: McpTransportConfig = serde_json::from_str(&json).unwrap();
assert!(matches!(parsed, McpTransportConfig::Http));
}
#[test]
fn test_transport_config_serde_stdio() {
let transport = McpTransportConfig::Stdio {
command: "npx".to_string(),
args: vec!["-y".to_string(), "server".to_string()],
env: HashMap::from([("KEY".to_string(), "val".to_string())]),
};
let json = serde_json::to_string(&transport).unwrap();
assert!(json.contains("\"transport\":\"stdio\""));
assert!(json.contains("\"command\":\"npx\""));
let parsed: McpTransportConfig = serde_json::from_str(&json).unwrap();
match parsed {
McpTransportConfig::Stdio { command, args, env } => {
assert_eq!(command, "npx");
assert_eq!(args, vec!["-y".to_string(), "server".to_string()]);
assert_eq!(env.get("KEY").unwrap(), "val");
}
other => panic!("Expected Stdio, got {:?}", other),
}
}
#[test]
fn test_transport_config_serde_unix() {
let transport = McpTransportConfig::Unix {
socket_path: "/tmp/mcp.sock".to_string(),
};
let json = serde_json::to_string(&transport).unwrap();
assert!(json.contains("\"transport\":\"unix\""));
assert!(json.contains("\"socket_path\":\"/tmp/mcp.sock\""));
let parsed: McpTransportConfig = serde_json::from_str(&json).unwrap();
match parsed {
McpTransportConfig::Unix { socket_path } => {
assert_eq!(socket_path, "/tmp/mcp.sock");
}
other => panic!("Expected Unix, got {:?}", other),
}
}
#[test]
fn test_backward_compat_no_transport_field() {
// Existing configs without transport field should still deserialize
let json = r#"{
"name": "notion",
"url": "https://mcp.notion.com",
"enabled": true
}"#;
let config: McpServerConfig = serde_json::from_str(json).unwrap();
assert_eq!(config.name, "notion");
assert_eq!(config.url, "https://mcp.notion.com");
assert!(config.transport.is_none());
assert!(config.headers.is_empty());
assert!(matches!(
config.effective_transport(),
EffectiveTransport::Http
));
}
#[test]
fn test_config_roundtrip_with_transport() {
// Test full roundtrip with stdio transport
let config = McpServerConfig::new_stdio(
"test-server",
"node",
vec!["server.js".to_string()],
HashMap::from([("NODE_ENV".to_string(), "production".to_string())]),
)
.with_description("A test server");
let json = serde_json::to_string_pretty(&config).unwrap();
let parsed: McpServerConfig = serde_json::from_str(&json).unwrap();
assert_eq!(parsed.name, "test-server");
assert!(parsed.url.is_empty());
assert_eq!(parsed.description.as_deref(), Some("A test server"));
match &parsed.transport {
Some(McpTransportConfig::Stdio { command, args, env }) => {
assert_eq!(command, "node");
assert_eq!(args, &["server.js".to_string()]);
assert_eq!(env.get("NODE_ENV").unwrap(), "production");
}
other => panic!("Expected Stdio transport, got {:?}", other),
}
// Test full roundtrip with unix transport
let config = McpServerConfig::new_unix("unix-server", "/var/run/mcp.sock");
let json = serde_json::to_string_pretty(&config).unwrap();
let parsed: McpServerConfig = serde_json::from_str(&json).unwrap();
assert_eq!(parsed.name, "unix-server");
match &parsed.transport {
Some(McpTransportConfig::Unix { socket_path }) => {
assert_eq!(socket_path, "/var/run/mcp.sock");
}
other => panic!("Expected Unix transport, got {:?}", other),
}
// Test roundtrip with HTTP + headers
let headers = HashMap::from([("X-Custom".to_string(), "value".to_string())]);
let config =
McpServerConfig::new("http-server", "https://mcp.example.com").with_headers(headers);
let json = serde_json::to_string_pretty(&config).unwrap();
let parsed: McpServerConfig = serde_json::from_str(&json).unwrap();
assert_eq!(parsed.name, "http-server");
assert!(parsed.transport.is_none());
assert_eq!(parsed.headers.get("X-Custom").unwrap(), "value");
}
}
+386
View File
@@ -0,0 +1,386 @@
//! HTTP transport for MCP servers.
//!
//! Implements the Streamable HTTP transport, communicating with MCP servers
//! over HTTP POST with JSON and SSE response support.
use std::collections::HashMap;
use std::sync::Arc;
use async_trait::async_trait;
use crate::tools::mcp::protocol::{McpRequest, McpResponse};
use crate::tools::mcp::session::McpSessionManager;
use crate::tools::mcp::transport::McpTransport;
use crate::tools::tool::ToolError;
/// MCP transport that communicates with a server over HTTP.
///
/// Sends JSON-RPC requests as HTTP POST with `Content-Type: application/json`
/// and accepts either JSON or SSE (`text/event-stream`) responses. Optionally
/// manages session IDs via [`McpSessionManager`] and supports custom headers.
pub struct HttpMcpTransport {
server_url: String,
server_name: String,
http_client: reqwest::Client,
session_manager: Option<Arc<McpSessionManager>>,
custom_headers: HashMap<String, String>,
}
impl HttpMcpTransport {
/// Create a new HTTP transport for the given server URL.
pub fn new(server_url: impl Into<String>, server_name: impl Into<String>) -> Self {
Self {
server_url: server_url.into(),
server_name: server_name.into(),
// reqwest::Client::builder().build() only fails if the TLS backend
// cannot initialize, which does not happen with the default rustls
// feature set. Panic is acceptable here (same as reqwest's own
// `Client::new()`).
http_client: reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(30))
.build()
.expect("Failed to create HTTP client"),
session_manager: None,
custom_headers: HashMap::new(),
}
}
/// Attach a session manager for Mcp-Session-Id tracking.
pub fn with_session_manager(mut self, session_manager: Arc<McpSessionManager>) -> Self {
self.session_manager = Some(session_manager);
self
}
/// Set custom headers that will be sent with every request.
#[cfg(test)]
pub fn with_custom_headers(mut self, headers: HashMap<String, String>) -> Self {
self.custom_headers = headers;
self
}
/// Get the server URL.
#[cfg(test)]
pub(crate) fn server_url(&self) -> &str {
&self.server_url
}
/// Get the session manager, if one is configured.
#[cfg(test)]
pub(crate) fn session_manager(&self) -> Option<&Arc<McpSessionManager>> {
self.session_manager.as_ref()
}
}
#[async_trait]
impl McpTransport for HttpMcpTransport {
async fn send(
&self,
request: &McpRequest,
headers: &HashMap<String, String>,
) -> Result<McpResponse, ToolError> {
// Build the HTTP request.
let mut req_builder = self
.http_client
.post(&self.server_url)
.header("Content-Type", "application/json")
.header("Accept", "application/json, text/event-stream")
.json(request);
// Apply custom headers configured on the transport.
for (key, value) in &self.custom_headers {
req_builder = req_builder.header(key.as_str(), value.as_str());
}
// Apply per-request headers (e.g. Authorization, Mcp-Session-Id).
for (key, value) in headers {
req_builder = req_builder.header(key.as_str(), value.as_str());
}
// Send the request.
let response = req_builder.send().await.map_err(|e| {
let mut chain = format!("[{}] MCP HTTP request failed: {}", self.server_name, 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)
})?;
// Extract session ID from response headers before consuming the body.
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;
}
// Handle error status codes.
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
let sanitized = sanitize_error_body(&body);
return Err(ToolError::ExternalService(format!(
"[{}] MCP server returned status: {} - {}",
self.server_name, status, sanitized
)));
}
// Determine response format from Content-Type.
let content_type = response
.headers()
.get("content-type")
.and_then(|v| v.to_str().ok())
.unwrap_or("")
.to_string();
if content_type.contains("text/event-stream") {
self.parse_sse_response(response).await
} else {
response.json().await.map_err(|e| {
ToolError::ExternalService(format!(
"[{}] Failed to parse MCP response: {}",
self.server_name, e
))
})
}
}
async fn shutdown(&self) -> Result<(), ToolError> {
// HTTP transport is stateless; nothing to shut down.
Ok(())
}
fn supports_http_features(&self) -> bool {
true
}
}
impl HttpMcpTransport {
/// Parse a Server-Sent Events response, returning the first valid JSON-RPC
/// `data:` line as an [`McpResponse`].
async fn parse_sse_response(
&self,
response: reqwest::Response,
) -> Result<McpResponse, ToolError> {
use futures::StreamExt;
const MAX_SSE_BUFFER: usize = 10 * 1024 * 1024; // 10 MB
let mut stream = response.bytes_stream();
let mut buffer = String::new();
while let Some(chunk) = stream.next().await {
let chunk = chunk.map_err(|e| {
ToolError::ExternalService(format!(
"[{}] Failed to read SSE chunk: {}",
self.server_name, e
))
})?;
buffer.push_str(&String::from_utf8_lossy(&chunk));
if buffer.len() > MAX_SSE_BUFFER {
return Err(ToolError::ExternalService(format!(
"[{}] SSE response exceeded {} byte limit",
self.server_name, MAX_SSE_BUFFER
)));
}
// Process only complete lines (terminated by \n). The last
// element of split('\n') may be an incomplete line; keep it
// in the buffer for the next chunk.
let mut remaining_start = 0;
let bytes = buffer.as_bytes();
for (i, &b) in bytes.iter().enumerate() {
if b == b'\n' {
let line = &buffer[remaining_start..i];
remaining_start = i + 1;
if let Some(json_str) = line.strip_prefix("data: ")
&& let Ok(response) = serde_json::from_str::<McpResponse>(json_str)
{
return Ok(response);
}
}
}
// Keep only the unprocessed trailing fragment.
if remaining_start > 0 {
buffer = buffer[remaining_start..].to_string();
}
}
// Process any remaining data without a trailing newline.
if let Some(json_str) = buffer.strip_prefix("data: ")
&& let Ok(response) = serde_json::from_str::<McpResponse>(json_str.trim())
{
return Ok(response);
}
Err(ToolError::ExternalService(format!(
"[{}] No valid data in SSE response: {}",
self.server_name, buffer
)))
}
}
/// Sanitize an HTTP error body for safe inclusion in error messages.
///
/// When the body looks like a full HTML document (`<html` or `<!doctype`),
/// strips all tags, collapsing whitespace. Non-HTML bodies are left
/// intact. In both cases the result is truncated to 200 *characters*
/// (char-boundary safe) so that large payloads don't bloat error messages.
///
/// See #263 — raw HTML error pages were propagating through the error
/// chain into the web UI, causing a white screen.
pub(crate) fn sanitize_error_body(body: &str) -> String {
const MAX_CHARS: usize = 200;
// Only strip tags when the body looks like a full HTML document.
// Plain text that happens to contain `<` / `>` (e.g. log lines,
// comparison expressions) is left untouched.
let lower = body.to_ascii_lowercase();
let is_html_document = lower.contains("<html") || lower.contains("<!doctype");
let text = if is_html_document {
let stripped = body
.chars()
.fold((String::new(), false), |(mut out, in_tag), c| {
if c == '<' {
(out, true)
} else if c == '>' {
(out, false)
} else if !in_tag {
out.push(c);
(out, false)
} else {
(out, true)
}
})
.0;
stripped.split_whitespace().collect::<Vec<_>>().join(" ")
} else {
body.to_string()
};
// Truncate at a char boundary (safe for multi-byte UTF-8).
if text.chars().count() > MAX_CHARS {
let byte_offset = text
.char_indices()
.nth(MAX_CHARS)
.map(|(i, _)| i)
.unwrap_or(text.len());
format!("{}... ({} bytes total)", &text[..byte_offset], body.len())
} else {
text
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_sanitize_error_body_strips_html_tags() {
let html =
r#"<!DOCTYPE html><html><body><h1>422 Error</h1><p>Invalid token</p></body></html>"#;
let result = sanitize_error_body(html);
assert!(!result.contains('<'), "HTML tags must be stripped");
assert!(!result.contains('>'), "HTML tags must be stripped");
assert!(result.contains("422 Error"));
assert!(result.contains("Invalid token"));
}
#[test]
fn test_sanitize_error_body_truncates_large_html_page() {
let html = format!(
"<html><body><p>{}</p></body></html>",
"error detail ".repeat(50)
);
let result = sanitize_error_body(&html);
assert!(result.contains("..."));
assert!(result.contains("bytes total)"));
assert!(!result.contains('<'));
}
#[test]
fn test_sanitize_error_body_passes_short_plain_text() {
assert_eq!(sanitize_error_body("Not Found"), "Not Found");
}
#[test]
fn test_sanitize_error_body_truncates_long_plain_text() {
let long = "x".repeat(300);
let result = sanitize_error_body(&long);
assert!(result.contains("..."));
assert!(result.contains("300 bytes total)"));
}
#[test]
fn test_sanitize_error_body_multibyte_no_panic() {
// 300 CJK characters = 900 bytes; truncation must land on a
// char boundary, not in the middle of a multi-byte sequence.
let cjk = "错误".repeat(150);
let result = sanitize_error_body(&cjk);
assert!(result.contains("..."));
// Must be valid UTF-8 (would have panicked otherwise).
assert!(result.is_char_boundary(result.len()));
}
#[test]
fn test_sanitize_error_body_strips_uppercase_html() {
let html = "<HTML><BODY><H1>500 Internal Server Error</H1></BODY></HTML>";
let result = sanitize_error_body(html);
assert!(
!result.contains('<'),
"uppercase HTML tags must be stripped"
);
assert!(result.contains("500 Internal Server Error"));
}
#[test]
fn test_sanitize_error_body_preserves_angle_brackets_in_non_html() {
let text = "value < 10 and value > 0";
assert_eq!(sanitize_error_body(text), text);
}
#[test]
fn test_sanitize_error_body_empty_string() {
assert_eq!(sanitize_error_body(""), "");
}
#[test]
fn test_new_creates_transport() {
let transport = HttpMcpTransport::new("http://localhost:8080", "test");
assert_eq!(transport.server_url(), "http://localhost:8080");
assert!(transport.session_manager().is_none());
assert!(transport.custom_headers.is_empty());
}
#[test]
fn test_supports_http_features() {
let http_transport = HttpMcpTransport::new("http://localhost:8080", "test");
assert!(http_transport.supports_http_features());
}
#[test]
fn test_with_session_manager() {
let session_manager = Arc::new(McpSessionManager::new());
let transport = HttpMcpTransport::new("http://localhost:8080", "test")
.with_session_manager(session_manager.clone());
assert!(transport.session_manager().is_some());
}
#[test]
fn test_with_custom_headers() {
let mut headers = HashMap::new();
headers.insert("X-Custom".to_string(), "value".to_string());
let transport =
HttpMcpTransport::new("http://localhost:8080", "test").with_custom_headers(headers);
assert_eq!(transport.custom_headers.get("X-Custom").unwrap(), "value");
}
}
+10
View File
@@ -4,6 +4,8 @@
//! additional capabilities through a standardized protocol.
//!
//! Supports both local (unauthenticated) and hosted (OAuth-authenticated) servers.
//! Transport options include HTTP (Streamable HTTP / SSE), stdio (subprocess),
//! and Unix domain sockets.
//!
//! ## Usage
//!
@@ -29,11 +31,19 @@
pub mod auth;
mod client;
pub mod config;
pub(crate) mod http_transport;
pub(crate) mod process;
mod protocol;
pub mod session;
pub(crate) mod stdio_transport;
pub(crate) mod transport;
#[cfg(unix)]
pub(crate) mod unix_transport;
pub use auth::{is_authenticated, refresh_access_token};
pub use client::McpClient;
pub use config::{McpServerConfig, McpServersFile, OAuthConfig};
pub use process::McpProcessManager;
pub use protocol::{InitializeResult, McpRequest, McpResponse, McpTool};
pub use session::McpSessionManager;
pub use transport::McpTransport;
+206
View File
@@ -0,0 +1,206 @@
//! MCP stdio process manager.
//!
//! Manages the lifecycle of MCP servers running as child processes.
//! Handles spawning, shutdown, and crash recovery with exponential backoff.
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use crate::tools::mcp::stdio_transport::StdioMcpTransport;
use crate::tools::mcp::transport::McpTransport;
use crate::tools::tool::ToolError;
/// Configuration for spawning a stdio MCP server.
#[derive(Debug, Clone)]
pub struct StdioSpawnConfig {
pub command: String,
pub args: Vec<String>,
pub env: HashMap<String, String>,
}
/// Manages stdio MCP server processes.
///
/// Handles spawning, tracking, and shutdown of child processes.
pub struct McpProcessManager {
transports: RwLock<HashMap<String, Arc<StdioMcpTransport>>>,
configs: RwLock<HashMap<String, StdioSpawnConfig>>,
}
impl McpProcessManager {
pub fn new() -> Self {
Self {
transports: RwLock::new(HashMap::new()),
configs: RwLock::new(HashMap::new()),
}
}
/// Spawn a new stdio MCP server process.
pub async fn spawn_stdio(
&self,
name: impl Into<String>,
command: impl Into<String>,
args: Vec<String>,
env: HashMap<String, String>,
) -> Result<Arc<StdioMcpTransport>, ToolError> {
let name = name.into();
let command = command.into();
// Store config for potential restart
self.configs.write().await.insert(
name.clone(),
StdioSpawnConfig {
command: command.clone(),
args: args.clone(),
env: env.clone(),
},
);
let transport = Arc::new(StdioMcpTransport::spawn(&name, &command, args, env).await?);
self.transports
.write()
.await
.insert(name, Arc::clone(&transport));
Ok(transport)
}
/// Get a transport by server name.
pub async fn get(&self, name: &str) -> Option<Arc<StdioMcpTransport>> {
self.transports.read().await.get(name).cloned()
}
/// Shut down all managed transports.
pub async fn shutdown_all(&self) {
let transports: Vec<(String, Arc<StdioMcpTransport>)> = {
let mut map = self.transports.write().await;
map.drain().collect()
};
for (name, transport) in transports {
if let Err(e) = transport.shutdown().await {
tracing::warn!("Failed to shut down MCP stdio server '{}': {}", name, e);
}
}
}
/// Shut down a specific transport by name.
pub async fn shutdown(&self, name: &str) -> Result<(), ToolError> {
let transport = self.transports.write().await.remove(name);
if let Some(transport) = transport {
transport.shutdown().await?;
}
self.configs.write().await.remove(name);
Ok(())
}
/// Attempt to restart a crashed transport with exponential backoff.
///
/// Tries up to 5 times with delays of 1s, 2s, 4s, 8s, 16s (total: 31s max wait).
pub async fn try_restart(&self, name: &str) -> Result<Arc<StdioMcpTransport>, ToolError> {
let config = self
.configs
.read()
.await
.get(name)
.cloned()
.ok_or_else(|| {
ToolError::ExternalService(format!(
"No spawn config for MCP server '{}', cannot restart",
name
))
})?;
// Shut down and remove old transport to avoid orphaning a wedged process.
if let Some(old_transport) = self.transports.write().await.remove(name) {
let _ = old_transport.shutdown().await;
}
let max_retries = 5;
let mut last_err = None;
for attempt in 0..max_retries {
let delay = Duration::from_secs(1 << attempt);
tokio::time::sleep(delay).await;
match StdioMcpTransport::spawn(
name,
&config.command,
config.args.clone(),
config.env.clone(),
)
.await
{
Ok(transport) => {
let transport = Arc::new(transport);
self.transports
.write()
.await
.insert(name.to_string(), Arc::clone(&transport));
tracing::info!(
"MCP stdio server '{}' restarted after {} attempt(s)",
name,
attempt + 1
);
return Ok(transport);
}
Err(e) => {
tracing::warn!(
"Restart attempt {}/{} for MCP server '{}' failed: {}",
attempt + 1,
max_retries,
name,
e
);
last_err = Some(e);
}
}
}
Err(last_err.unwrap_or_else(|| {
ToolError::ExternalService(format!(
"Failed to restart MCP server '{}' after {} attempts",
name, max_retries
))
}))
}
/// Get names of all managed transports.
pub async fn managed_servers(&self) -> Vec<String> {
self.transports.read().await.keys().cloned().collect()
}
}
impl Default for McpProcessManager {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_new_creates_empty_manager() {
let _manager = McpProcessManager::new();
}
#[tokio::test]
async fn test_managed_servers_returns_empty_list_initially() {
let manager = McpProcessManager::new();
let servers = manager.managed_servers().await;
assert!(servers.is_empty());
}
#[tokio::test]
async fn test_shutdown_all_on_empty_manager_does_not_panic() {
let manager = McpProcessManager::new();
manager.shutdown_all().await;
}
}

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