mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
Compare commits
11
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0d02c313ff | ||
|
|
d55b302b39 | ||
|
|
517be42ccc | ||
|
|
09198c68ab | ||
|
|
115b7f38fe | ||
|
|
bb228f6315 | ||
|
|
45f547c711 | ||
|
|
23de75d75b | ||
|
|
ced83d5b4d | ||
|
|
810ba58fd2 | ||
|
|
202665a55c |
@@ -0,0 +1,8 @@
|
||||
target/
|
||||
.git/
|
||||
.env
|
||||
.env.*
|
||||
*.md
|
||||
!CLAUDE.md
|
||||
node_modules/
|
||||
tools-src/
|
||||
+3
-3
@@ -1,15 +1,15 @@
|
||||
# Database Configuration
|
||||
DATABASE_URL=postgres://ironclaw:password@localhost:5432/ironclaw
|
||||
DATABASE_URL=postgres://localhost/ironclaw
|
||||
DATABASE_POOL_SIZE=10
|
||||
|
||||
# LLM Provider (NEAR AI)
|
||||
# NEAR AI provides a unified interface to all models with user authentication
|
||||
# Session token is stored in ~/.near-agent/session.json and managed automatically.
|
||||
# Session token is stored in ~/.ironclaw/session.json and managed automatically.
|
||||
# On first run, the agent will open a browser for OAuth authentication.
|
||||
NEARAI_MODEL=claude-3-5-sonnet-20241022
|
||||
NEARAI_BASE_URL=https://cloud-api.near.ai
|
||||
NEARAI_AUTH_URL=https://private.near.ai
|
||||
# NEARAI_SESSION_PATH=~/.near-agent/session.json # optional, default shown
|
||||
# NEARAI_SESSION_PATH=~/.ironclaw/session.json # optional, default shown
|
||||
|
||||
# Channel Configuration
|
||||
# CLI is always enabled
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
name: Code Style
|
||||
on:
|
||||
pull_request:
|
||||
|
||||
jobs:
|
||||
codestyle:
|
||||
name: Code Style (fmt + clippy)
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
- name: Install Rust
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
with:
|
||||
profile: minimal
|
||||
components: rustfmt, clippy
|
||||
- name: Check formatting
|
||||
run: |
|
||||
cargo fmt --all -- --check
|
||||
- name: Check lints (cargo clippy)
|
||||
run: cargo clippy -- -D warnings
|
||||
@@ -0,0 +1,55 @@
|
||||
name: Release-plz
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
|
||||
jobs:
|
||||
|
||||
# Release unpublished packages.
|
||||
release-plz-release:
|
||||
if: ${{ github.repository_owner == 'nearai' }}
|
||||
name: Release-plz release
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: write
|
||||
steps:
|
||||
- &checkout
|
||||
name: Checkout repository
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
persist-credentials: false
|
||||
- &install-rust
|
||||
name: Install Rust toolchain
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
- name: Run release-plz
|
||||
uses: release-plz/[email protected]
|
||||
with:
|
||||
command: release
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }}
|
||||
|
||||
# Create a PR with the new versions and changelog, preparing the next release.
|
||||
release-plz-pr:
|
||||
if: ${{ github.repository_owner == 'nearai' }}
|
||||
name: Release-plz PR
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: write
|
||||
pull-requests: write
|
||||
concurrency:
|
||||
group: release-plz-${{ github.ref }}
|
||||
cancel-in-progress: false
|
||||
steps:
|
||||
- *checkout
|
||||
- *install-rust
|
||||
- name: Run release-plz
|
||||
uses: release-plz/[email protected]
|
||||
with:
|
||||
command: release-pr
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }}
|
||||
@@ -0,0 +1,324 @@
|
||||
# This file was autogenerated by dist: https://opensource.axo.dev/cargo-dist/
|
||||
#
|
||||
# Copyright 2022-2024, axodotdev
|
||||
# SPDX-License-Identifier: MIT or Apache-2.0
|
||||
#
|
||||
# CI that:
|
||||
#
|
||||
# * checks for a Git Tag that looks like a release
|
||||
# * builds artifacts with dist (archives, installers, hashes)
|
||||
# * uploads those artifacts to temporary workflow zip
|
||||
# * on success, uploads the artifacts to a GitHub Release
|
||||
#
|
||||
# Note that the GitHub Release will be created with a generated
|
||||
# title/body based on your changelogs.
|
||||
|
||||
name: Release
|
||||
permissions:
|
||||
"contents": "write"
|
||||
|
||||
# This task will run whenever you push a git tag that looks like a version
|
||||
# like "1.0.0", "v0.1.0-prerelease.1", "my-app/0.1.0", "releases/v1.0.0", etc.
|
||||
# Various formats will be parsed into a VERSION and an optional PACKAGE_NAME, where
|
||||
# PACKAGE_NAME must be the name of a Cargo package in your workspace, and VERSION
|
||||
# must be a Cargo-style SemVer Version (must have at least major.minor.patch).
|
||||
#
|
||||
# If PACKAGE_NAME is specified, then the announcement will be for that
|
||||
# package (erroring out if it doesn't have the given version or isn't dist-able).
|
||||
#
|
||||
# If PACKAGE_NAME isn't specified, then the announcement will be for all
|
||||
# (dist-able) packages in the workspace with that version (this mode is
|
||||
# intended for workspaces with only one dist-able package, or with all dist-able
|
||||
# packages versioned/released in lockstep).
|
||||
#
|
||||
# If you push multiple tags at once, separate instances of this workflow will
|
||||
# spin up, creating an independent announcement for each one. However, GitHub
|
||||
# will hard limit this to 3 tags per commit, as it will assume more tags is a
|
||||
# mistake.
|
||||
#
|
||||
# If there's a prerelease-style suffix to the version, then the release(s)
|
||||
# will be marked as a prerelease.
|
||||
on:
|
||||
pull_request:
|
||||
push:
|
||||
tags:
|
||||
- '**[0-9]+.[0-9]+.[0-9]+*'
|
||||
|
||||
jobs:
|
||||
# Run 'dist plan' (or host) to determine what tasks we need to do
|
||||
plan:
|
||||
runs-on: "ubuntu-22.04"
|
||||
outputs:
|
||||
val: ${{ steps.plan.outputs.manifest }}
|
||||
tag: ${{ !github.event.pull_request && github.ref_name || '' }}
|
||||
tag-flag: ${{ !github.event.pull_request && format('--tag={0}', github.ref_name) || '' }}
|
||||
publishing: ${{ !github.event.pull_request }}
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
submodules: recursive
|
||||
- name: Install dist
|
||||
# we specify bash to get pipefail; it guards against the `curl` command
|
||||
# failing. otherwise `sh` won't catch that `curl` returned non-0
|
||||
shell: bash
|
||||
run: "curl --proto '=https' --tlsv1.2 -LsSf https://github.com/axodotdev/cargo-dist/releases/download/v0.30.3/cargo-dist-installer.sh | sh"
|
||||
- name: Cache dist
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: cargo-dist-cache
|
||||
path: ~/.cargo/bin/dist
|
||||
# sure would be cool if github gave us proper conditionals...
|
||||
# so here's a doubly-nested ternary-via-truthiness to try to provide the best possible
|
||||
# functionality based on whether this is a pull_request, and whether it's from a fork.
|
||||
# (PRs run on the *source* but secrets are usually on the *target* -- that's *good*
|
||||
# but also really annoying to build CI around when it needs secrets to work right.)
|
||||
- id: plan
|
||||
run: |
|
||||
dist ${{ (!github.event.pull_request && format('host --steps=create --tag={0}', github.ref_name)) || 'plan' }} --output-format=json > plan-dist-manifest.json
|
||||
echo "dist ran successfully"
|
||||
cat plan-dist-manifest.json
|
||||
echo "manifest=$(jq -c "." plan-dist-manifest.json)" >> "$GITHUB_OUTPUT"
|
||||
- name: "Upload dist-manifest.json"
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: artifacts-plan-dist-manifest
|
||||
path: plan-dist-manifest.json
|
||||
|
||||
# Build and packages all the platform-specific things
|
||||
build-local-artifacts:
|
||||
name: build-local-artifacts (${{ join(matrix.targets, ', ') }})
|
||||
# Let the initial task tell us to not run (currently very blunt)
|
||||
needs:
|
||||
- plan
|
||||
if: ${{ fromJson(needs.plan.outputs.val).ci.github.artifacts_matrix.include != null && (needs.plan.outputs.publishing == 'true' || fromJson(needs.plan.outputs.val).ci.github.pr_run_mode == 'upload') }}
|
||||
strategy:
|
||||
fail-fast: false
|
||||
# Target platforms/runners are computed by dist in create-release.
|
||||
# Each member of the matrix has the following arguments:
|
||||
#
|
||||
# - runner: the github runner
|
||||
# - dist-args: cli flags to pass to dist
|
||||
# - install-dist: expression to run to install dist on the runner
|
||||
#
|
||||
# Typically there will be:
|
||||
# - 1 "global" task that builds universal installers
|
||||
# - N "local" tasks that build each platform's binaries and platform-specific installers
|
||||
matrix: ${{ fromJson(needs.plan.outputs.val).ci.github.artifacts_matrix }}
|
||||
runs-on: ${{ matrix.runner }}
|
||||
container: ${{ matrix.container && matrix.container.image || null }}
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
BUILD_MANIFEST_NAME: target/distrib/${{ join(matrix.targets, '-') }}-dist-manifest.json
|
||||
steps:
|
||||
- name: enable windows longpaths
|
||||
run: |
|
||||
git config --global core.longpaths true
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
submodules: recursive
|
||||
- name: Install Rust non-interactively if not already installed
|
||||
if: ${{ matrix.container }}
|
||||
run: |
|
||||
if ! command -v cargo > /dev/null 2>&1; then
|
||||
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
|
||||
echo "$HOME/.cargo/bin" >> $GITHUB_PATH
|
||||
fi
|
||||
- uses: swatinem/rust-cache@v2
|
||||
with:
|
||||
key: ${{ join(matrix.targets, '-') }}
|
||||
cache-provider: ${{ matrix.cache_provider }}
|
||||
- name: Install dist
|
||||
run: ${{ matrix.install_dist.run }}
|
||||
# Get the dist-manifest
|
||||
- name: Fetch local artifacts
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
pattern: artifacts-*
|
||||
path: target/distrib/
|
||||
merge-multiple: true
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
${{ matrix.packages_install }}
|
||||
- name: Build artifacts
|
||||
run: |
|
||||
# Actually do builds and make zips and whatnot
|
||||
dist build ${{ needs.plan.outputs.tag-flag }} --print=linkage --output-format=json ${{ matrix.dist_args }} > dist-manifest.json
|
||||
echo "dist ran successfully"
|
||||
- id: cargo-dist
|
||||
name: Post-build
|
||||
# We force bash here just because github makes it really hard to get values up
|
||||
# to "real" actions without writing to env-vars, and writing to env-vars has
|
||||
# inconsistent syntax between shell and powershell.
|
||||
shell: bash
|
||||
run: |
|
||||
# Parse out what we just built and upload it to scratch storage
|
||||
echo "paths<<EOF" >> "$GITHUB_OUTPUT"
|
||||
dist print-upload-files-from-manifest --manifest dist-manifest.json >> "$GITHUB_OUTPUT"
|
||||
echo "EOF" >> "$GITHUB_OUTPUT"
|
||||
|
||||
cp dist-manifest.json "$BUILD_MANIFEST_NAME"
|
||||
- name: "Upload artifacts"
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: artifacts-build-local-${{ join(matrix.targets, '_') }}
|
||||
path: |
|
||||
${{ steps.cargo-dist.outputs.paths }}
|
||||
${{ env.BUILD_MANIFEST_NAME }}
|
||||
|
||||
# Build and package all the platform-agnostic(ish) things
|
||||
build-global-artifacts:
|
||||
needs:
|
||||
- plan
|
||||
- build-local-artifacts
|
||||
runs-on: "ubuntu-22.04"
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
BUILD_MANIFEST_NAME: target/distrib/global-dist-manifest.json
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
submodules: recursive
|
||||
- name: Install cached dist
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: cargo-dist-cache
|
||||
path: ~/.cargo/bin/
|
||||
- run: chmod +x ~/.cargo/bin/dist
|
||||
# Get all the local artifacts for the global tasks to use (for e.g. checksums)
|
||||
- name: Fetch local artifacts
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
pattern: artifacts-*
|
||||
path: target/distrib/
|
||||
merge-multiple: true
|
||||
- id: cargo-dist
|
||||
shell: bash
|
||||
run: |
|
||||
dist build ${{ needs.plan.outputs.tag-flag }} --output-format=json "--artifacts=global" > dist-manifest.json
|
||||
echo "dist ran successfully"
|
||||
|
||||
# Parse out what we just built and upload it to scratch storage
|
||||
echo "paths<<EOF" >> "$GITHUB_OUTPUT"
|
||||
jq --raw-output ".upload_files[]" dist-manifest.json >> "$GITHUB_OUTPUT"
|
||||
echo "EOF" >> "$GITHUB_OUTPUT"
|
||||
|
||||
cp dist-manifest.json "$BUILD_MANIFEST_NAME"
|
||||
- name: "Upload artifacts"
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: artifacts-build-global
|
||||
path: |
|
||||
${{ steps.cargo-dist.outputs.paths }}
|
||||
${{ env.BUILD_MANIFEST_NAME }}
|
||||
# Determines if we should publish/announce
|
||||
host:
|
||||
needs:
|
||||
- plan
|
||||
- build-local-artifacts
|
||||
- build-global-artifacts
|
||||
# Only run if we're "publishing", and only if local and global didn't fail (skipped is fine)
|
||||
if: ${{ always() && needs.plan.outputs.publishing == 'true' && (needs.build-global-artifacts.result == 'skipped' || needs.build-global-artifacts.result == 'success') && (needs.build-local-artifacts.result == 'skipped' || needs.build-local-artifacts.result == 'success') }}
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
runs-on: "ubuntu-22.04"
|
||||
outputs:
|
||||
val: ${{ steps.host.outputs.manifest }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
submodules: recursive
|
||||
- name: Install cached dist
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: cargo-dist-cache
|
||||
path: ~/.cargo/bin/
|
||||
- run: chmod +x ~/.cargo/bin/dist
|
||||
# Fetch artifacts from scratch-storage
|
||||
- name: Fetch artifacts
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
pattern: artifacts-*
|
||||
path: target/distrib/
|
||||
merge-multiple: true
|
||||
- id: host
|
||||
shell: bash
|
||||
run: |
|
||||
dist host ${{ needs.plan.outputs.tag-flag }} --steps=upload --steps=release --output-format=json > dist-manifest.json
|
||||
echo "artifacts uploaded and released successfully"
|
||||
cat dist-manifest.json
|
||||
echo "manifest=$(jq -c "." dist-manifest.json)" >> "$GITHUB_OUTPUT"
|
||||
- name: "Upload dist-manifest.json"
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
# Overwrite the previous copy
|
||||
name: artifacts-dist-manifest
|
||||
path: dist-manifest.json
|
||||
# Create a GitHub Release while uploading all files to it
|
||||
- name: "Download GitHub Artifacts"
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
pattern: artifacts-*
|
||||
path: artifacts
|
||||
merge-multiple: true
|
||||
- name: Cleanup
|
||||
run: |
|
||||
# Remove the granular manifests
|
||||
rm -f artifacts/*-dist-manifest.json
|
||||
- name: Create GitHub Release
|
||||
env:
|
||||
PRERELEASE_FLAG: "${{ fromJson(steps.host.outputs.manifest).announcement_is_prerelease && '--prerelease' || '' }}"
|
||||
ANNOUNCEMENT_TITLE: "${{ fromJson(steps.host.outputs.manifest).announcement_title }}"
|
||||
ANNOUNCEMENT_BODY: "${{ fromJson(steps.host.outputs.manifest).announcement_github_body }}"
|
||||
RELEASE_COMMIT: "${{ github.sha }}"
|
||||
run: |
|
||||
# Write and read notes from a file to avoid quoting breaking things
|
||||
echo "$ANNOUNCEMENT_BODY" > $RUNNER_TEMP/notes.txt
|
||||
|
||||
gh release create "${{ needs.plan.outputs.tag }}" --target "$RELEASE_COMMIT" $PRERELEASE_FLAG --title "$ANNOUNCEMENT_TITLE" --notes-file "$RUNNER_TEMP/notes.txt" artifacts/*
|
||||
|
||||
publish-npm:
|
||||
needs:
|
||||
- plan
|
||||
- host
|
||||
runs-on: "ubuntu-22.04"
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
PLAN: ${{ needs.plan.outputs.val }}
|
||||
if: ${{ !fromJson(needs.plan.outputs.val).announcement_is_prerelease || fromJson(needs.plan.outputs.val).publish_prereleases }}
|
||||
steps:
|
||||
- name: Fetch npm packages
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
pattern: artifacts-*
|
||||
path: npm/
|
||||
merge-multiple: true
|
||||
- uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: '20.x'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
- run: |
|
||||
for release in $(echo "$PLAN" | jq --compact-output '.releases[] | select([.artifacts[] | endswith("-npm-package.tar.gz")] | any)'); do
|
||||
pkg=$(echo "$release" | jq '.artifacts[] | select(endswith("-npm-package.tar.gz"))' --raw-output)
|
||||
npm publish --access public "./npm/${pkg}"
|
||||
done
|
||||
env:
|
||||
NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
|
||||
|
||||
announce:
|
||||
needs:
|
||||
- plan
|
||||
- host
|
||||
- publish-npm
|
||||
# use "always() && ..." to allow us to wait for all publish jobs while
|
||||
# still allowing individual publish jobs to skip themselves (for prereleases).
|
||||
# "host" however must run to completion, no skipping allowed!
|
||||
if: ${{ always() && needs.host.result == 'success' && (needs.publish-npm.result == 'skipped' || needs.publish-npm.result == 'success') }}
|
||||
runs-on: "ubuntu-22.04"
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
submodules: recursive
|
||||
@@ -0,0 +1,20 @@
|
||||
name: Run Tests
|
||||
on:
|
||||
pull_request:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
|
||||
jobs:
|
||||
tests:
|
||||
name: Run Tests
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
- name: Install Rust
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
with:
|
||||
profile: minimal
|
||||
- name: Run Tests
|
||||
run: cargo test --all-features -- --nocapture
|
||||
@@ -4,3 +4,6 @@
|
||||
|
||||
target/
|
||||
|
||||
# WASM build artifacts (loaded from disk, not bundled)
|
||||
*.wasm
|
||||
|
||||
|
||||
@@ -0,0 +1,96 @@
|
||||
# Changelog
|
||||
|
||||
All notable changes to this project will be documented in this file.
|
||||
|
||||
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
|
||||
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
## [0.1.0](https://github.com/nearai/ironclaw/releases/tag/v0.1.0) - 2026-02-12
|
||||
|
||||
### Added
|
||||
|
||||
- Add multi-provider LLM support via rig-core adapter ([#36](https://github.com/nearai/ironclaw/pull/36))
|
||||
- Sandbox jobs ([#4](https://github.com/nearai/ironclaw/pull/4))
|
||||
- Add Google Suite & Telegram WASM tools ([#9](https://github.com/nearai/ironclaw/pull/9))
|
||||
- Improve CLI ([#5](https://github.com/nearai/ironclaw/pull/5))
|
||||
|
||||
### Fixed
|
||||
|
||||
- resolve runtime panic in Linux keychain integration ([#32](https://github.com/nearai/ironclaw/pull/32))
|
||||
|
||||
### Other
|
||||
|
||||
- Skip release-plz on forks
|
||||
- Upgraded release-plz CD pipeline
|
||||
- Added CI/CD and release pipelines ([#45](https://github.com/nearai/ironclaw/pull/45))
|
||||
- DM pairing + Telegram channel improvements ([#17](https://github.com/nearai/ironclaw/pull/17))
|
||||
- Fixes build, adds missing sse event and correct command ([#11](https://github.com/nearai/ironclaw/pull/11))
|
||||
- Codex/feature parity pr hook ([#6](https://github.com/nearai/ironclaw/pull/6))
|
||||
- Add WebSocket gateway and control plane ([#8](https://github.com/nearai/ironclaw/pull/8))
|
||||
- select bundled Telegram channel and auto-install ([#3](https://github.com/nearai/ironclaw/pull/3))
|
||||
- Adding skills for reusable work
|
||||
- Fix MCP tool calls, approval loop, shutdown, and improve web UI
|
||||
- Add auth mode, fix MCP token handling, and parallelize startup loading
|
||||
- Merge remote-tracking branch 'origin/main' into ui
|
||||
- Adding web UI
|
||||
- Rename `setup` CLI command to `onboard` for compatibility
|
||||
- Add in-chat extension discovery, auth, and activation system
|
||||
- Add Telegram typing indicator via WIT on-status callback
|
||||
- Add proactivity features: memory CLI, session pruning, self-repair notifications, slash commands, status diagnostics, context warnings
|
||||
- Add hosted MCP server support with OAuth 2.1 and token refresh
|
||||
- Add interactive setup wizard and persistent settings
|
||||
- Rebrand to IronClaw with security-first mission
|
||||
- Fix build_software tool stuck in planning mode loop
|
||||
- Enable sandbox by default
|
||||
- Fix Telegram Markdown formatting and clarify tool/memory distinctions
|
||||
- Simplify Telegram channel config with host-injected tunnel/webhook settings
|
||||
- Apply Telegram channel learnings to WhatsApp implementation
|
||||
- Merge remote-tracking branch 'origin/main'
|
||||
- Docker file for sandbox
|
||||
- Replace hardcoded intent patterns with job tools
|
||||
- Fix router test to match intentional job creation patterns
|
||||
- Add Docker execution sandbox for secure shell command isolation
|
||||
- Move setup wizard credentials to database storage
|
||||
- Add interactive setup wizard for first-run configuration
|
||||
- Add Telegram Bot API channel as WASM module
|
||||
- Add OpenClaw feature parity tracking matrix
|
||||
- Add Chat Completions API support and expand REPL debugging
|
||||
- Implementing channels to be handled in wasm
|
||||
- Support non interactive mode and model selection
|
||||
- Implement tool approval, fix tool definition refresh, and wire embeddings
|
||||
- Tool use
|
||||
- Wiring more
|
||||
- Add heartbeat integration, planning phase, and auto-repair
|
||||
- Login flow
|
||||
- Extend support for session management
|
||||
- Adding builder capability
|
||||
- Load tools at launch
|
||||
- Fix multiline message rendering in TUI
|
||||
- Parse NEAR AI alternative response format with output field
|
||||
- Handle NEAR AI plain text responses
|
||||
- Disable mouse capture to allow text selection in TUI
|
||||
- Add verbose logging to debug empty NEAR AI responses
|
||||
- Improve NEAR AI response parsing for varying response formats
|
||||
- Show status/thinking messages in chat window, debug empty responses
|
||||
- Add timeout and logging to NEAR AI provider
|
||||
- Add status updates to show agent thinking/processing state
|
||||
- Add CLI subcommands for WASM tool management
|
||||
- Fix TUI shutdown: send /shutdown message and handle in agent loop
|
||||
- Remove SimpleCliChannel, add Ctrl+D twice quit, redirect logs to TUI
|
||||
- Fix TuiChannel integration and enable in main.rs
|
||||
- Integrate Codex patterns: task scheduler, TUI, sessions, compaction
|
||||
- Adding LICENSE
|
||||
- Add README with IronClaw branding
|
||||
- Add WASM sandbox secure API extension
|
||||
- Wire database Store into agent loop
|
||||
- Implementing WASM runtime
|
||||
- Add workspace integration tests
|
||||
- Compact memory_tree output format
|
||||
- Replace memory_list with memory_tree tool
|
||||
- Simplify workspace to path-based storage, remove legacy code
|
||||
- Add NEAR AI chat-api as default LLM provider
|
||||
- Add CLAUDE.md project documentation
|
||||
- Add workspace and memory system (OpenClaw-inspired)
|
||||
- Initial implementation of the agent framework
|
||||
@@ -11,8 +11,13 @@
|
||||
- **Always available** - Multi-channel access with proactive background execution
|
||||
|
||||
### Features
|
||||
- **Multi-channel input**: TUI (Ratatui), HTTP webhooks, Telegram, WhatsApp, Slack (WASM channels)
|
||||
- **Multi-channel input**: TUI (Ratatui), HTTP webhooks, WASM channels (Telegram, Slack), web gateway
|
||||
- **Parallel job execution** with state machine and self-repair for stuck jobs
|
||||
- **Sandbox execution**: Docker container isolation with orchestrator/worker pattern
|
||||
- **Claude Code mode**: Delegate jobs to Claude CLI inside containers
|
||||
- **Routines**: Scheduled (cron) and reactive (event, webhook) task execution
|
||||
- **Web gateway**: Browser UI with SSE/WebSocket real-time streaming
|
||||
- **Extension management**: Install, auth, activate MCP/WASM extensions
|
||||
- **Extensible tools**: Built-in tools, WASM sandbox, MCP client, dynamic builder
|
||||
- **Persistent memory**: Workspace with hybrid search (FTS + vector via RRF)
|
||||
- **Prompt injection defense**: Sanitizer, validator, policy rules, leak detection
|
||||
@@ -59,7 +64,9 @@ src/
|
||||
│ ├── context_monitor.rs # Memory pressure detection
|
||||
│ ├── undo.rs # Turn-based undo/redo with checkpoints
|
||||
│ ├── submission.rs # Submission parsing (undo, redo, compact, clear, etc.)
|
||||
│ └── task.rs # Sub-task execution framework
|
||||
│ ├── task.rs # Sub-task execution framework
|
||||
│ ├── routine.rs # Routine types (Trigger, Action, Guardrails)
|
||||
│ └── routine_engine.rs # Routine execution (cron ticker, event matcher)
|
||||
│
|
||||
├── channels/ # Multi-channel input
|
||||
│ ├── channel.rs # Channel trait, IncomingMessage, OutgoingResponse
|
||||
@@ -72,8 +79,33 @@ src/
|
||||
│ │ ├── overlay.rs # Approval overlays
|
||||
│ │ └── composer.rs # Message composition
|
||||
│ ├── http.rs # HTTP webhook (axum) with secret validation
|
||||
│ ├── slack.rs # Stub
|
||||
│ └── telegram.rs # Stub
|
||||
│ ├── repl.rs # Simple REPL (for testing)
|
||||
│ ├── web/ # Web gateway (browser UI)
|
||||
│ │ ├── mod.rs # Gateway builder, startup
|
||||
│ │ ├── server.rs # Axum router, 40+ API endpoints
|
||||
│ │ ├── sse.rs # SSE broadcast manager
|
||||
│ │ ├── ws.rs # WebSocket gateway + connection tracking
|
||||
│ │ ├── types.rs # Request/response types, SseEvent enum
|
||||
│ │ ├── auth.rs # Bearer token auth middleware
|
||||
│ │ ├── log_layer.rs # Tracing layer for log streaming
|
||||
│ │ └── static/ # HTML, CSS, JS (single-page app)
|
||||
│ └── wasm/ # WASM channel runtime
|
||||
│ ├── mod.rs
|
||||
│ ├── bundled.rs # Bundled channel discovery
|
||||
│ └── wrapper.rs # Channel trait wrapper for WASM modules
|
||||
│
|
||||
├── orchestrator/ # Internal HTTP API for sandbox containers
|
||||
│ ├── mod.rs
|
||||
│ ├── api.rs # Axum endpoints (LLM proxy, events, prompts)
|
||||
│ ├── auth.rs # Per-job bearer token store
|
||||
│ └── job_manager.rs # Container lifecycle (create, stop, cleanup)
|
||||
│
|
||||
├── worker/ # Runs inside Docker containers
|
||||
│ ├── mod.rs
|
||||
│ ├── runtime.rs # Worker execution loop (tool calls, LLM)
|
||||
│ ├── claude_bridge.rs # Claude Code bridge (spawns claude CLI)
|
||||
│ ├── api.rs # HTTP client to orchestrator
|
||||
│ └── proxy_llm.rs # LlmProvider that proxies through orchestrator
|
||||
│
|
||||
├── safety/ # Prompt injection defense
|
||||
│ ├── sanitizer.rs # Pattern detection, content escaping
|
||||
@@ -96,6 +128,9 @@ src/
|
||||
│ │ ├── file.rs # ReadFile, WriteFile, ListDir, ApplyPatch
|
||||
│ │ ├── shell.rs # Shell command execution
|
||||
│ │ ├── memory.rs # Memory tools (search, write, read, tree)
|
||||
│ │ ├── job.rs # CreateJob, ListJobs, JobStatus, CancelJob
|
||||
│ │ ├── routine.rs # routine_create/list/update/delete/history
|
||||
│ │ ├── extension_tools.rs # Extension install/auth/activate/remove
|
||||
│ │ └── marketplace.rs, ecommerce.rs, taskrabbit.rs, restaurant.rs (stubs)
|
||||
│ ├── builder/ # Dynamic tool building
|
||||
│ │ ├── core.rs # BuildRequirement, SoftwareType, Language
|
||||
@@ -236,6 +271,30 @@ HEARTBEAT_ENABLED=true
|
||||
HEARTBEAT_INTERVAL_SECS=1800 # 30 minutes
|
||||
HEARTBEAT_NOTIFY_CHANNEL=tui
|
||||
HEARTBEAT_NOTIFY_USER=default
|
||||
|
||||
# Web gateway
|
||||
GATEWAY_ENABLED=true
|
||||
GATEWAY_HOST=127.0.0.1
|
||||
GATEWAY_PORT=3001
|
||||
GATEWAY_AUTH_TOKEN=changeme # Required for API access
|
||||
GATEWAY_USER_ID=default
|
||||
|
||||
# Docker sandbox
|
||||
SANDBOX_ENABLED=true
|
||||
SANDBOX_IMAGE=ironclaw-worker:latest
|
||||
SANDBOX_MEMORY_LIMIT_MB=512
|
||||
SANDBOX_TIMEOUT_SECS=1800
|
||||
|
||||
# Claude Code mode (runs inside sandbox containers)
|
||||
CLAUDE_CODE_ENABLED=false
|
||||
CLAUDE_CODE_MODEL=claude-sonnet-4-20250514
|
||||
CLAUDE_CODE_MAX_TURNS=50
|
||||
CLAUDE_CODE_CONFIG_DIR=/home/worker/.claude
|
||||
|
||||
# Routines (scheduled/reactive execution)
|
||||
ROUTINES_ENABLED=true
|
||||
ROUTINES_CRON_INTERVAL=60 # Tick interval in seconds
|
||||
ROUTINES_MAX_CONCURRENT=3
|
||||
```
|
||||
|
||||
### NEAR AI Provider
|
||||
@@ -297,13 +356,14 @@ Key test patterns:
|
||||
|
||||
## Current Limitations / TODOs
|
||||
|
||||
1. **Slack/Telegram channels** - Stubs only, need implementation
|
||||
2. **Domain-specific tools** - `marketplace.rs`, `restaurant.rs`, `taskrabbit.rs`, `ecommerce.rs` return placeholder responses; need real API integrations
|
||||
3. **Integration tests** - Need testcontainers setup for PostgreSQL
|
||||
4. **MCP stdio transport** - Only HTTP transport implemented
|
||||
5. **WIT bindgen integration** - Auto-extract tool description/schema from WASM modules (stubbed)
|
||||
6. **Capability granting after tool build** - Built tools get empty capabilities; need UX for granting HTTP/secrets access
|
||||
7. **Tool versioning workflow** - No version tracking or rollback for dynamically built tools
|
||||
1. **Domain-specific tools** - `marketplace.rs`, `restaurant.rs`, `taskrabbit.rs`, `ecommerce.rs` return placeholder responses; need real API integrations
|
||||
2. **Integration tests** - Need testcontainers setup for PostgreSQL
|
||||
3. **MCP stdio transport** - Only HTTP transport implemented
|
||||
4. **WIT bindgen integration** - Auto-extract tool description/schema from WASM modules (stubbed)
|
||||
5. **Capability granting after tool build** - Built tools get empty capabilities; need UX for granting HTTP/secrets access
|
||||
6. **Tool versioning workflow** - No version tracking or rollback for dynamically built tools
|
||||
7. **Webhook trigger endpoint** - Routines webhook trigger not yet exposed in web gateway
|
||||
8. **Full channel status view** - Gateway status widget exists, but no per-channel connection dashboard
|
||||
|
||||
### Completed
|
||||
|
||||
@@ -320,6 +380,13 @@ Key test patterns:
|
||||
- ✅ **Tool approval enforcement** - Tools with `requires_approval()` (shell, http, file write/patch, build_software) now gate execution, track auto-approved tools per session
|
||||
- ✅ **Tool definition refresh** - Tool definitions refreshed each iteration so newly built tools become visible in same session
|
||||
- ✅ **Worker tool call handling** - Uses `respond_with_tools()` to properly execute tool calls when `select_tools()` returns empty
|
||||
- ✅ **Gateway control plane** - Web gateway with 40+ API endpoints, SSE/WebSocket
|
||||
- ✅ **Web Control UI** - Browser-based dashboard with chat, memory, jobs, logs, extensions, routines
|
||||
- ✅ **Slack/Telegram channels** - Implemented as WASM tools
|
||||
- ✅ **Docker sandbox** - Orchestrator/worker containers with per-job auth
|
||||
- ✅ **Claude Code mode** - Delegate jobs to Claude CLI inside containers
|
||||
- ✅ **Routines system** - Cron, event, webhook, and manual triggers with guardrails
|
||||
- ✅ **Extension management** - Install, auth, activate MCP/WASM extensions via CLI and web UI
|
||||
|
||||
## Adding a New Tool
|
||||
|
||||
|
||||
Generated
+390
-28
@@ -182,6 +182,12 @@ version = "0.7.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50"
|
||||
|
||||
[[package]]
|
||||
name = "as-any"
|
||||
version = "0.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b0f477b951e452a0b6b4a10b53ccd569042d1d01729b519e02074a9c0958a063"
|
||||
|
||||
[[package]]
|
||||
name = "async-broadcast"
|
||||
version = "0.7.2"
|
||||
@@ -307,6 +313,28 @@ dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-stream"
|
||||
version = "0.3.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476"
|
||||
dependencies = [
|
||||
"async-stream-impl",
|
||||
"futures-core",
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-stream-impl"
|
||||
version = "0.3.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.114",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-task"
|
||||
version = "4.7.1"
|
||||
@@ -820,6 +848,16 @@ dependencies = [
|
||||
"crossterm 0.29.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "core-foundation"
|
||||
version = "0.9.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "91e195e091a93c46f7102ec7818a2aa394e1e1771c3ab4825963fa03e45afb8f"
|
||||
dependencies = [
|
||||
"core-foundation-sys",
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "core-foundation"
|
||||
version = "0.10.1"
|
||||
@@ -996,6 +1034,17 @@ dependencies = [
|
||||
"syn 2.0.114",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cron"
|
||||
version = "0.13.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "eee8b2b4516038bc0f1d3c9934bcb4a13dd316e04abbc63c96757a6d75978532"
|
||||
dependencies = [
|
||||
"chrono",
|
||||
"nom",
|
||||
"once_cell",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crossbeam"
|
||||
version = "0.8.4"
|
||||
@@ -1468,6 +1517,17 @@ dependencies = [
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "eventsource-stream"
|
||||
version = "0.2.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "74fef4569247a5f429d9156b9d0a2599914385dd189c539334c625d8099d90ab"
|
||||
dependencies = [
|
||||
"futures-core",
|
||||
"nom",
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fallible-iterator"
|
||||
version = "0.2.0"
|
||||
@@ -1526,6 +1586,21 @@ version = "0.1.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2"
|
||||
|
||||
[[package]]
|
||||
name = "foreign-types"
|
||||
version = "0.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1"
|
||||
dependencies = [
|
||||
"foreign-types-shared",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "foreign-types-shared"
|
||||
version = "0.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b"
|
||||
|
||||
[[package]]
|
||||
name = "form_urlencoded"
|
||||
version = "1.2.2"
|
||||
@@ -1546,6 +1621,16 @@ dependencies = [
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fs4"
|
||||
version = "0.6.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2eeb4ed9e12f43b7fa0baae3f9cdda28352770132ef2e09a23760c29cae8bd47"
|
||||
dependencies = [
|
||||
"rustix 0.38.44",
|
||||
"windows-sys 0.48.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "funty"
|
||||
version = "2.0.0"
|
||||
@@ -1636,6 +1721,12 @@ version = "0.3.31"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988"
|
||||
|
||||
[[package]]
|
||||
name = "futures-timer"
|
||||
version = "3.0.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24"
|
||||
|
||||
[[package]]
|
||||
name = "futures-util"
|
||||
version = "0.3.31"
|
||||
@@ -1734,6 +1825,12 @@ dependencies = [
|
||||
"stable_deref_trait",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "glob"
|
||||
version = "0.3.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280"
|
||||
|
||||
[[package]]
|
||||
name = "h2"
|
||||
version = "0.4.13"
|
||||
@@ -1933,6 +2030,22 @@ dependencies = [
|
||||
"webpki-roots",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hyper-tls"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "70206fc6890eaca9fde8a0bf71caa2ddfc9fe045ac9e5c70df101a7dbde866e0"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"http-body-util",
|
||||
"hyper",
|
||||
"hyper-util",
|
||||
"native-tls",
|
||||
"tokio",
|
||||
"tokio-native-tls",
|
||||
"tower-service",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hyper-util"
|
||||
version = "0.1.20"
|
||||
@@ -1951,9 +2064,11 @@ dependencies = [
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"socket2",
|
||||
"system-configuration",
|
||||
"tokio",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
"windows-registry",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2189,15 +2304,18 @@ dependencies = [
|
||||
"bytes",
|
||||
"chrono",
|
||||
"clap",
|
||||
"cron",
|
||||
"crossterm 0.28.1",
|
||||
"deadpool-postgres",
|
||||
"dirs 6.0.0",
|
||||
"dotenvy",
|
||||
"fs4",
|
||||
"futures",
|
||||
"hkdf",
|
||||
"http-body-util",
|
||||
"hyper",
|
||||
"hyper-util",
|
||||
"mime_guess",
|
||||
"open",
|
||||
"pgvector",
|
||||
"postgres-types",
|
||||
@@ -2206,12 +2324,13 @@ dependencies = [
|
||||
"refinery",
|
||||
"regex",
|
||||
"reqwest",
|
||||
"rig-core",
|
||||
"rust_decimal",
|
||||
"rust_decimal_macros",
|
||||
"rustyline",
|
||||
"secrecy",
|
||||
"secret-service",
|
||||
"security-framework",
|
||||
"security-framework 3.5.1",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2",
|
||||
@@ -2236,30 +2355,6 @@ dependencies = [
|
||||
"zbus",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ironclaw-bench"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"async-trait",
|
||||
"chrono",
|
||||
"clap",
|
||||
"futures",
|
||||
"ironclaw",
|
||||
"regex",
|
||||
"rust_decimal",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tempfile",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
"toml",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "is-docker"
|
||||
version = "0.2.0"
|
||||
@@ -2519,6 +2614,16 @@ version = "0.3.17"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
|
||||
|
||||
[[package]]
|
||||
name = "mime_guess"
|
||||
version = "2.0.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f7c44f8e672c00fe5308fa235f821cb4198414e1c77935c1ab6948d3fd78550e"
|
||||
dependencies = [
|
||||
"mime",
|
||||
"unicase",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "minimad"
|
||||
version = "0.14.0"
|
||||
@@ -2528,6 +2633,12 @@ dependencies = [
|
||||
"once_cell",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "minimal-lexical"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a"
|
||||
|
||||
[[package]]
|
||||
name = "mio"
|
||||
version = "1.1.1"
|
||||
@@ -2540,6 +2651,32 @@ dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nanoid"
|
||||
version = "0.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3ffa00dec017b5b1a8b7cf5e2c008bfda1aa7e0697ac1508b491fdf2622fb4d8"
|
||||
dependencies = [
|
||||
"rand 0.8.5",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "native-tls"
|
||||
version = "0.2.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "87de3442987e9dbec73158d5c715e7ad9072fda936bb03d19d7fa10e00520f0e"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"log",
|
||||
"openssl",
|
||||
"openssl-probe 0.1.6",
|
||||
"openssl-sys",
|
||||
"schannel",
|
||||
"security-framework 2.11.1",
|
||||
"security-framework-sys",
|
||||
"tempfile",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nibble_vec"
|
||||
version = "0.1.0"
|
||||
@@ -2574,6 +2711,16 @@ dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nom"
|
||||
version = "7.1.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a"
|
||||
dependencies = [
|
||||
"memchr",
|
||||
"minimal-lexical",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nu-ansi-term"
|
||||
version = "0.50.3"
|
||||
@@ -2722,18 +2869,71 @@ dependencies = [
|
||||
"pathdiff",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "openssl"
|
||||
version = "0.10.75"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "08838db121398ad17ab8531ce9de97b244589089e290a384c900cb9ff7434328"
|
||||
dependencies = [
|
||||
"bitflags 2.10.0",
|
||||
"cfg-if",
|
||||
"foreign-types",
|
||||
"libc",
|
||||
"once_cell",
|
||||
"openssl-macros",
|
||||
"openssl-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "openssl-macros"
|
||||
version = "0.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.114",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "openssl-probe"
|
||||
version = "0.1.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d05e27ee213611ffe7d6348b942e8f942b37114c00cc03cec254295a4a17852e"
|
||||
|
||||
[[package]]
|
||||
name = "openssl-probe"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe"
|
||||
|
||||
[[package]]
|
||||
name = "openssl-sys"
|
||||
version = "0.9.111"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "82cab2d520aa75e3c58898289429321eb788c3106963d0dc886ec7a5f4adc321"
|
||||
dependencies = [
|
||||
"cc",
|
||||
"libc",
|
||||
"pkg-config",
|
||||
"vcpkg",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "option-ext"
|
||||
version = "0.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d"
|
||||
|
||||
[[package]]
|
||||
name = "ordered-float"
|
||||
version = "5.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7f4779c6901a562440c3786d08192c6fbda7c1c2060edd10006b05ee35d10f2d"
|
||||
dependencies = [
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ordered-stream"
|
||||
version = "0.2.0"
|
||||
@@ -2845,6 +3045,26 @@ dependencies = [
|
||||
"siphasher",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pin-project"
|
||||
version = "1.1.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "677f1add503faace112b9f1373e43e9e054bfdd22ff1a63c1bc485eaec6a6a8a"
|
||||
dependencies = [
|
||||
"pin-project-internal",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pin-project-internal"
|
||||
version = "1.1.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6e918e4ff8c4549eb882f14b3a4bc8c8bc93de829416eacf579f1207a8fbf861"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.114",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pin-project-lite"
|
||||
version = "0.2.16"
|
||||
@@ -3377,16 +3597,22 @@ checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"encoding_rs",
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"h2",
|
||||
"http",
|
||||
"http-body",
|
||||
"http-body-util",
|
||||
"hyper",
|
||||
"hyper-rustls",
|
||||
"hyper-tls",
|
||||
"hyper-util",
|
||||
"js-sys",
|
||||
"log",
|
||||
"mime",
|
||||
"mime_guess",
|
||||
"native-tls",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"quinn",
|
||||
@@ -3397,6 +3623,7 @@ dependencies = [
|
||||
"serde_urlencoded",
|
||||
"sync_wrapper",
|
||||
"tokio",
|
||||
"tokio-native-tls",
|
||||
"tokio-rustls",
|
||||
"tokio-util",
|
||||
"tower",
|
||||
@@ -3410,6 +3637,38 @@ dependencies = [
|
||||
"webpki-roots",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rig-core"
|
||||
version = "0.30.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a8f7a3f0c7c00eaced15a68ee16e1bd6bb709ff598d11b9aedac8b628217dc09"
|
||||
dependencies = [
|
||||
"as-any",
|
||||
"async-stream",
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"eventsource-stream",
|
||||
"fastrand",
|
||||
"futures",
|
||||
"futures-timer",
|
||||
"glob",
|
||||
"http",
|
||||
"mime",
|
||||
"mime_guess",
|
||||
"nanoid",
|
||||
"ordered-float",
|
||||
"pin-project-lite",
|
||||
"reqwest",
|
||||
"schemars 1.2.1",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"tracing-futures",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ring"
|
||||
version = "0.17.14"
|
||||
@@ -3557,10 +3816,10 @@ version = "0.8.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "612460d5f7bea540c490b2b6395d8e34a953e52b491accd6c86c8164c5932a63"
|
||||
dependencies = [
|
||||
"openssl-probe",
|
||||
"openssl-probe 0.2.1",
|
||||
"rustls-pki-types",
|
||||
"schannel",
|
||||
"security-framework",
|
||||
"security-framework 3.5.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3677,10 +3936,23 @@ checksum = "a2b42f36aa1cd011945615b92222f6bf73c599a102a300334cd7f8dbeec726cc"
|
||||
dependencies = [
|
||||
"dyn-clone",
|
||||
"ref-cast",
|
||||
"schemars_derive",
|
||||
"serde",
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "schemars_derive"
|
||||
version = "1.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7d115b50f4aaeea07e79c1912f645c7513d81715d0420f8bc77a18c6260b307f"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"serde_derive_internals",
|
||||
"syn 2.0.114",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "scopeguard"
|
||||
version = "1.2.0"
|
||||
@@ -3722,6 +3994,19 @@ dependencies = [
|
||||
"zbus",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "security-framework"
|
||||
version = "2.11.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "897b2245f0b511c87893af39b033e5ca9cce68824c4d7e7630b5a1d339658d02"
|
||||
dependencies = [
|
||||
"bitflags 2.10.0",
|
||||
"core-foundation 0.9.4",
|
||||
"core-foundation-sys",
|
||||
"libc",
|
||||
"security-framework-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "security-framework"
|
||||
version = "3.5.1"
|
||||
@@ -3729,7 +4014,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b3297343eaf830f66ede390ea39da1d462b6b0c1b000f420d0a83f898bbbe6ef"
|
||||
dependencies = [
|
||||
"bitflags 2.10.0",
|
||||
"core-foundation",
|
||||
"core-foundation 0.10.1",
|
||||
"core-foundation-sys",
|
||||
"libc",
|
||||
"security-framework-sys",
|
||||
@@ -3785,6 +4070,17 @@ dependencies = [
|
||||
"syn 2.0.114",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_derive_internals"
|
||||
version = "0.29.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.114",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_json"
|
||||
version = "1.0.149"
|
||||
@@ -4098,6 +4394,27 @@ dependencies = [
|
||||
"syn 2.0.114",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "system-configuration"
|
||||
version = "0.7.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a13f3d0daba03132c0aa9767f98351b3488edc2c100cda2d2ec2b04f3d8d3c8b"
|
||||
dependencies = [
|
||||
"bitflags 2.10.0",
|
||||
"core-foundation 0.9.4",
|
||||
"system-configuration-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "system-configuration-sys"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8e1d1b10ced5ca923a1fcb8d03e96b8d3268065d724548c0211415ff6ac6bac4"
|
||||
dependencies = [
|
||||
"core-foundation-sys",
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "system-interface"
|
||||
version = "0.27.3"
|
||||
@@ -4336,6 +4653,16 @@ dependencies = [
|
||||
"syn 2.0.114",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-native-tls"
|
||||
version = "0.3.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bbae76ab933c85776efabc971569dd6119c580d8f5d448769dec1764bf796ef2"
|
||||
dependencies = [
|
||||
"native-tls",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-postgres"
|
||||
version = "0.7.16"
|
||||
@@ -4598,6 +4925,18 @@ dependencies = [
|
||||
"valuable",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tracing-futures"
|
||||
version = "0.2.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "97d095ae15e245a057c8e8451bab9b3ee1e1f68e9ba2b4fbc18d0ac5237835f2"
|
||||
dependencies = [
|
||||
"futures",
|
||||
"futures-task",
|
||||
"pin-project",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tracing-log"
|
||||
version = "0.2.0"
|
||||
@@ -4697,6 +5036,12 @@ dependencies = [
|
||||
"winapi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "unicase"
|
||||
version = "2.9.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142"
|
||||
|
||||
[[package]]
|
||||
name = "unicode-bidi"
|
||||
version = "0.3.18"
|
||||
@@ -4819,6 +5164,12 @@ version = "0.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65"
|
||||
|
||||
[[package]]
|
||||
name = "vcpkg"
|
||||
version = "0.2.15"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426"
|
||||
|
||||
[[package]]
|
||||
name = "version_check"
|
||||
version = "0.9.5"
|
||||
@@ -5493,6 +5844,17 @@ version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5"
|
||||
|
||||
[[package]]
|
||||
name = "windows-registry"
|
||||
version = "0.6.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "02752bf7fbdcce7f2a27a742f798510f3e5ad88dbe84871e5168e2120c3d5720"
|
||||
dependencies = [
|
||||
"windows-link",
|
||||
"windows-result",
|
||||
"windows-strings",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-result"
|
||||
version = "0.4.1"
|
||||
|
||||
+61
-3
@@ -1,13 +1,19 @@
|
||||
[workspace]
|
||||
members = [".", "benchmarks"]
|
||||
|
||||
[package]
|
||||
name = "ironclaw"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
rust-version = "1.85"
|
||||
description = "Secure personal AI assistant that protects your data and expands its capabilities on the fly"
|
||||
authors = ["NEAR AI <[email protected]>"]
|
||||
license = "MIT OR Apache-2.0"
|
||||
homepage = "https://github.com/nearai/ironclaw"
|
||||
repository = "https://github.com/nearai/ironclaw"
|
||||
|
||||
[package.metadata.wix]
|
||||
upgrade-guid = "D0156E61-BA37-451E-8AB9-1A2ECCCFA48F"
|
||||
path-guid = "F90B6EA6-87F7-499B-BB19-CF55DE1EB339"
|
||||
license = false
|
||||
eula = false
|
||||
|
||||
[dependencies]
|
||||
# Async runtime
|
||||
@@ -61,12 +67,16 @@ axum = { version = "0.8", features = ["ws"] }
|
||||
tower = "0.5"
|
||||
tower-http = { version = "0.6", features = ["trace", "cors"] }
|
||||
|
||||
# Cron scheduling for routines
|
||||
cron = "0.13"
|
||||
|
||||
# Safety/sanitization
|
||||
regex = "1"
|
||||
aho-corasick = "1"
|
||||
|
||||
# Filesystem paths
|
||||
dirs = "6"
|
||||
fs4 = "0.6"
|
||||
|
||||
# Secrecy for sensitive values
|
||||
secrecy = { version = "0.10", features = ["serde"] }
|
||||
@@ -93,6 +103,9 @@ sha2 = "0.10"
|
||||
blake3 = "1"
|
||||
rand = "0.8"
|
||||
|
||||
# Multi-provider LLM support
|
||||
rig-core = "0.30"
|
||||
|
||||
# Docker sandbox
|
||||
bollard = "0.18"
|
||||
|
||||
@@ -102,6 +115,7 @@ hyper-util = { version = "0.1", features = ["server", "tokio", "http1", "http2"]
|
||||
http-body-util = "0.1"
|
||||
bytes = "1"
|
||||
base64 = "0.22.1"
|
||||
mime_guess = "2.0.5"
|
||||
|
||||
# macOS keychain
|
||||
[target.'cfg(target_os = "macos")'.dependencies]
|
||||
@@ -122,3 +136,47 @@ tempfile = "3"
|
||||
[features]
|
||||
default = []
|
||||
integration = []
|
||||
|
||||
# The profile that 'cargo dist' will build with
|
||||
[profile.dist]
|
||||
inherits = "release"
|
||||
lto = "thin"
|
||||
|
||||
# Config for 'dist'
|
||||
[workspace.metadata.dist]
|
||||
# The preferred dist version to use in CI (Cargo.toml SemVer syntax)
|
||||
cargo-dist-version = "0.30.3"
|
||||
allow-dirty = ["ci"]
|
||||
# CI backends to support
|
||||
ci = "github"
|
||||
# The installers to generate for each app
|
||||
installers = ["shell", "powershell", "npm", "msi"]
|
||||
# Publish jobs to run in CI
|
||||
publish-jobs = ["npm"]
|
||||
# Target platforms to build apps for (Rust target-triple syntax)
|
||||
targets = [
|
||||
"aarch64-apple-darwin",
|
||||
"aarch64-unknown-linux-gnu",
|
||||
"aarch64-pc-windows-msvc",
|
||||
"x86_64-apple-darwin",
|
||||
"x86_64-unknown-linux-gnu",
|
||||
"x86_64-pc-windows-msvc",
|
||||
]
|
||||
# The archive format to use for windows builds (defaults .zip)
|
||||
windows-archive = ".tar.gz"
|
||||
# The archive format to use for non-windows builds (defaults .tar.xz)
|
||||
unix-archive = ".tar.gz"
|
||||
# Which actions to run on pull requests
|
||||
pr-run-mode = "upload"
|
||||
# Path that installers should place binaries in
|
||||
install-path = "CARGO_HOME"
|
||||
# Whether to install an updater program
|
||||
install-updater = false
|
||||
|
||||
[workspace.metadata.dist.github-custom-runners]
|
||||
aarch64-unknown-linux-gnu = "ubuntu-24.04-arm"
|
||||
x86_64-unknown-linux-gnu = "ubuntu-22.04"
|
||||
x86_64-pc-windows-msvc = "windows-2022"
|
||||
aarch64-pc-windows-msvc = "windows-2025"
|
||||
x86_64-apple-darwin = "macos-15-intel"
|
||||
aarch64-apple-darwin = "macos-14"
|
||||
|
||||
@@ -0,0 +1,63 @@
|
||||
# Multi-stage Dockerfile for the IronClaw worker container.
|
||||
#
|
||||
# This image runs the ironclaw binary in worker mode inside Docker containers.
|
||||
# The orchestrator creates instances of this image for sandboxed job execution.
|
||||
#
|
||||
# Build:
|
||||
# docker build -f Dockerfile.worker -t ironclaw-worker .
|
||||
#
|
||||
# The image includes common development tools so workers can build software,
|
||||
# run tests, and execute shell commands.
|
||||
|
||||
FROM rust:1.85-bookworm AS builder
|
||||
|
||||
WORKDIR /build
|
||||
COPY . .
|
||||
|
||||
# Build only the ironclaw binary (release mode)
|
||||
RUN cargo build --release --bin ironclaw
|
||||
|
||||
# ---
|
||||
|
||||
FROM debian:bookworm-slim
|
||||
|
||||
# Install common development tools
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
ca-certificates \
|
||||
curl \
|
||||
git \
|
||||
build-essential \
|
||||
pkg-config \
|
||||
libssl-dev \
|
||||
nodejs \
|
||||
npm \
|
||||
python3 \
|
||||
python3-pip \
|
||||
python3-venv \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install Rust toolchain for the sandbox user
|
||||
ENV RUSTUP_HOME=/usr/local/rustup \
|
||||
CARGO_HOME=/usr/local/cargo \
|
||||
PATH=/usr/local/cargo/bin:$PATH
|
||||
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain 1.85.0 \
|
||||
&& chmod -R a+r /usr/local/rustup /usr/local/cargo
|
||||
|
||||
# Install Claude Code CLI (for claude-bridge mode)
|
||||
RUN npm install -g @anthropic-ai/claude-code@latest
|
||||
|
||||
# Copy the binary
|
||||
COPY --from=builder /build/target/release/ironclaw /usr/local/bin/ironclaw
|
||||
|
||||
# Create non-root user (UID 1000 matches the orchestrator's container config)
|
||||
RUN useradd -m -u 1000 -s /bin/bash sandbox \
|
||||
&& mkdir -p /workspace \
|
||||
&& chown sandbox:sandbox /workspace \
|
||||
&& mkdir -p /home/sandbox/.claude \
|
||||
&& chown sandbox:sandbox /home/sandbox/.claude
|
||||
|
||||
USER sandbox
|
||||
WORKDIR /workspace
|
||||
|
||||
# The orchestrator passes the full command via Docker cmd.
|
||||
ENTRYPOINT ["ironclaw"]
|
||||
+44
-34
@@ -16,8 +16,8 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
|
||||
| Feature | OpenClaw | IronClaw | Notes |
|
||||
|---------|----------|----------|-------|
|
||||
| Hub-and-spoke architecture | ✅ | 🚧 | IronClaw has channels but no central gateway |
|
||||
| WebSocket control plane | ✅ | ❌ | Gateway with ws://127.0.0.1:18789 |
|
||||
| Hub-and-spoke architecture | ✅ | ✅ | Web gateway as central hub |
|
||||
| WebSocket control plane | ✅ | ✅ | Gateway with WebSocket + SSE |
|
||||
| Single-user system | ✅ | ✅ | |
|
||||
| Multi-agent routing | ✅ | ❌ | Workspace isolation per-agent |
|
||||
| Session-based messaging | ✅ | ✅ | Per-sender sessions |
|
||||
@@ -31,9 +31,9 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
|
||||
| Feature | OpenClaw | IronClaw | Notes |
|
||||
|---------|----------|----------|-------|
|
||||
| Gateway control plane | ✅ | ❌ | Central WebSocket server |
|
||||
| HTTP endpoints for Control UI | ✅ | ❌ | Web dashboard |
|
||||
| Channel connection lifecycle | ✅ | 🚧 | ChannelManager handles streams |
|
||||
| Gateway control plane | ✅ | ✅ | Web gateway with 40+ API endpoints |
|
||||
| HTTP endpoints for Control UI | ✅ | ✅ | Web dashboard with chat, memory, jobs, logs, extensions |
|
||||
| Channel connection lifecycle | ✅ | ✅ | ChannelManager + WebSocket tracker |
|
||||
| Session management/routing | ✅ | ✅ | SessionManager exists |
|
||||
| Configuration hot-reload | ✅ | ❌ | |
|
||||
| Network modes (loopback/LAN/remote) | ✅ | 🚧 | HTTP only |
|
||||
@@ -43,7 +43,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
| launchd/systemd integration | ✅ | ❌ | |
|
||||
| Bonjour/mDNS discovery | ✅ | ❌ | |
|
||||
| Tailscale integration | ✅ | ❌ | |
|
||||
| Health check endpoints | ✅ | ❌ | |
|
||||
| Health check endpoints | ✅ | ✅ | /api/health + /api/gateway/status |
|
||||
| `doctor` diagnostics | ✅ | ❌ | |
|
||||
|
||||
### Owner: _Unassigned_
|
||||
@@ -59,14 +59,14 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
| REPL (simple) | ✅ | ✅ | - | For testing |
|
||||
| WASM channels | ❌ | ✅ | - | IronClaw innovation |
|
||||
| WhatsApp | ✅ | ❌ | P1 | Baileys (Web) |
|
||||
| Telegram | ✅ | ❌ | P1 | grammY (Bot API) |
|
||||
| Telegram | ✅ | ✅ | - | WASM channel(MTProto), DM pairing, caption, /start, bot_username |
|
||||
| Discord | ✅ | ❌ | P2 | discord.js |
|
||||
| Signal | ✅ | ❌ | P2 | signal-cli |
|
||||
| Slack | ✅ | 🚧 | P1 | Stub exists, needs implementation |
|
||||
| Slack | ✅ | ✅ | - | WASM tool |
|
||||
| iMessage | ✅ | ❌ | P3 | BlueBubbles recommended |
|
||||
| Feishu/Lark | ✅ | ❌ | P3 | |
|
||||
| LINE | ✅ | ❌ | P3 | |
|
||||
| WebChat | ✅ | ❌ | P2 | Browser-based chat |
|
||||
| WebChat | ✅ | ✅ | - | Web gateway chat |
|
||||
| Matrix | ✅ | ❌ | P3 | E2EE support |
|
||||
| Mattermost | ✅ | ❌ | P3 | |
|
||||
| Google Chat | ✅ | ❌ | P3 | |
|
||||
@@ -79,13 +79,13 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
|
||||
| Feature | OpenClaw | IronClaw | Notes |
|
||||
|---------|----------|----------|-------|
|
||||
| DM pairing codes | ✅ | ❌ | Verification for unknown senders |
|
||||
| Allowlist/blocklist | ✅ | ❌ | Per-channel access control |
|
||||
| DM pairing codes | ✅ | ✅ | `ironclaw pairing list/approve`, host APIs |
|
||||
| Allowlist/blocklist | ✅ | 🚧 | allow_from + pairing store |
|
||||
| Self-message bypass | ✅ | ❌ | Own messages skip pairing |
|
||||
| Mention-based activation | ✅ | ❌ | Configurable patterns |
|
||||
| Mention-based activation | ✅ | ✅ | bot_username + respond_to_all_group_messages |
|
||||
| Per-group tool policies | ✅ | ❌ | Allow/deny specific tools |
|
||||
| Thread isolation | ✅ | ✅ | Separate sessions per thread |
|
||||
| Per-channel media limits | ✅ | ❌ | |
|
||||
| Per-channel media limits | ✅ | 🚧 | Caption support for media; no size limits |
|
||||
| Typing indicators | ✅ | 🚧 | TUI shows status |
|
||||
|
||||
### Owner: _Unassigned_
|
||||
@@ -99,17 +99,17 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
| `run` (agent) | ✅ | ✅ | - | Default command |
|
||||
| `tool install/list/remove` | ✅ | ✅ | - | WASM tools |
|
||||
| `gateway start/stop` | ✅ | ❌ | P2 | |
|
||||
| `onboard` (wizard) | ✅ | ❌ | P2 | Interactive setup |
|
||||
| `onboard` (wizard) | ✅ | ✅ | - | Interactive setup |
|
||||
| `tui` | ✅ | ✅ | - | Ratatui TUI |
|
||||
| `config` | ✅ | ❌ | P2 | Read/write config |
|
||||
| `config` | ✅ | ✅ | - | Read/write config |
|
||||
| `channels` | ✅ | ❌ | P2 | Channel management |
|
||||
| `models` | ✅ | 🚧 | - | Model selector in TUI |
|
||||
| `status` | ✅ | ❌ | P2 | System status |
|
||||
| `status` | ✅ | ✅ | - | System status |
|
||||
| `agents` | ✅ | ❌ | P3 | Multi-agent management |
|
||||
| `sessions` | ✅ | ❌ | P3 | Session listing |
|
||||
| `memory` | ✅ | ❌ | P2 | Memory search CLI |
|
||||
| `memory` | ✅ | ✅ | - | Memory search CLI |
|
||||
| `skills` | ✅ | ❌ | P3 | Agent skills |
|
||||
| `pairing` | ✅ | ❌ | P3 | Node pairing |
|
||||
| `pairing` | ✅ | ✅ | - | list/approve for channel DM pairing |
|
||||
| `nodes` | ✅ | ❌ | P3 | Device management |
|
||||
| `plugins` | ✅ | ❌ | P3 | Plugin management |
|
||||
| `hooks` | ✅ | ❌ | P2 | Lifecycle hooks |
|
||||
@@ -132,7 +132,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
| Feature | OpenClaw | IronClaw | Notes |
|
||||
|---------|----------|----------|-------|
|
||||
| Pi agent runtime | ✅ | ➖ | IronClaw uses custom runtime |
|
||||
| RPC-based execution | ✅ | 🚧 | Worker isolation |
|
||||
| RPC-based execution | ✅ | ✅ | Orchestrator/worker pattern |
|
||||
| Multi-provider failover | ✅ | ❌ | Provider fallback chains |
|
||||
| Per-sender sessions | ✅ | ✅ | |
|
||||
| Global sessions | ✅ | ❌ | Optional shared context |
|
||||
@@ -303,13 +303,13 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
|
||||
| Feature | OpenClaw | IronClaw | Priority | Notes |
|
||||
|---------|----------|----------|----------|-------|
|
||||
| Control UI Dashboard | ✅ | ❌ | P2 | Web status/config |
|
||||
| Channel status view | ✅ | ❌ | P2 | |
|
||||
| Control UI Dashboard | ✅ | ✅ | - | Web gateway with chat, memory, jobs, logs, extensions |
|
||||
| Channel status view | ✅ | 🚧 | P2 | Gateway status widget, full channel view pending |
|
||||
| Agent management | ✅ | ❌ | P3 | |
|
||||
| Model selection | ✅ | ✅ | - | TUI only |
|
||||
| Config editing | ✅ | ❌ | P3 | |
|
||||
| Debug/logs viewer | ✅ | ❌ | P3 | |
|
||||
| WebChat interface | ✅ | ❌ | P2 | Browser chat |
|
||||
| Debug/logs viewer | ✅ | ✅ | - | Real-time log streaming with level/target filters |
|
||||
| WebChat interface | ✅ | ✅ | - | Web gateway chat with SSE/WebSocket |
|
||||
| Canvas system (A2UI) | ✅ | ❌ | P3 | Agent-driven UI |
|
||||
|
||||
### Owner: _Unassigned_
|
||||
@@ -320,13 +320,13 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
|
||||
| Feature | OpenClaw | IronClaw | Priority | Notes |
|
||||
|---------|----------|----------|----------|-------|
|
||||
| Cron jobs | ✅ | ❌ | P2 | Schedule-based tasks |
|
||||
| Timezone support | ✅ | ❌ | P2 | |
|
||||
| One-shot/recurring jobs | ✅ | ❌ | P2 | |
|
||||
| Cron jobs | ✅ | ✅ | - | Routines with cron trigger |
|
||||
| Timezone support | ✅ | ✅ | - | Via cron expressions |
|
||||
| One-shot/recurring jobs | ✅ | ✅ | - | Manual + cron triggers |
|
||||
| `beforeInbound` hook | ✅ | ❌ | P2 | |
|
||||
| `beforeOutbound` hook | ✅ | ❌ | P2 | |
|
||||
| `beforeToolCall` hook | ✅ | ❌ | P2 | |
|
||||
| `onMessage` hook | ✅ | ❌ | P2 | |
|
||||
| `onMessage` hook | ✅ | ✅ | - | Routines with event trigger |
|
||||
| `onSessionStart` hook | ✅ | ❌ | P2 | |
|
||||
| `onSessionEnd` hook | ✅ | ❌ | P2 | |
|
||||
| `transcribeAudio` hook | ✅ | ❌ | P3 | |
|
||||
@@ -346,18 +346,18 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
|
||||
| Feature | OpenClaw | IronClaw | Notes |
|
||||
|---------|----------|----------|-------|
|
||||
| Gateway token auth | ✅ | 🚧 | HTTP webhook secret |
|
||||
| Gateway token auth | ✅ | ✅ | Bearer token auth on web gateway |
|
||||
| Device pairing | ✅ | ❌ | |
|
||||
| Tailscale identity | ✅ | ❌ | |
|
||||
| OAuth flows | ✅ | 🚧 | NEAR AI OAuth |
|
||||
| DM pairing verification | ✅ | ❌ | |
|
||||
| Allowlist/blocklist | ✅ | ❌ | |
|
||||
| DM pairing verification | ✅ | ✅ | ironclaw pairing approve, host APIs |
|
||||
| Allowlist/blocklist | ✅ | 🚧 | allow_from + pairing store |
|
||||
| Per-group tool policies | ✅ | ❌ | |
|
||||
| Exec approvals | ✅ | ✅ | TUI overlay |
|
||||
| TLS 1.3 minimum | ✅ | ✅ | reqwest rustls |
|
||||
| SSRF protection | ✅ | ✅ | WASM allowlist |
|
||||
| Loopback-first | ✅ | 🚧 | HTTP binds 0.0.0.0 |
|
||||
| Docker sandbox | ✅ | ❌ | Uses WASM sandbox |
|
||||
| Docker sandbox | ✅ | ✅ | Orchestrator/worker containers |
|
||||
| WASM sandbox | ❌ | ✅ | IronClaw innovation |
|
||||
| Tool policies | ✅ | ✅ | |
|
||||
| Elevated mode | ✅ | ❌ | |
|
||||
@@ -397,6 +397,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
### P0 - Core (Already Done)
|
||||
- ✅ TUI channel with approval overlays
|
||||
- ✅ HTTP webhook channel
|
||||
- ✅ DM pairing (ironclaw pairing list/approve, host APIs)
|
||||
- ✅ WASM tool sandbox
|
||||
- ✅ Workspace/memory with hybrid search
|
||||
- ✅ Prompt injection defense
|
||||
@@ -404,23 +405,32 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
- ✅ Session management
|
||||
- ✅ Context compaction
|
||||
- ✅ Model selection
|
||||
- ✅ Gateway control plane + WebSocket
|
||||
- ✅ Web Control UI (chat, memory, jobs, logs, extensions, routines)
|
||||
- ✅ WebChat channel (web gateway)
|
||||
- ✅ Slack channel (WASM tool)
|
||||
- ✅ Telegram channel (WASM tool, MTProto)
|
||||
- ✅ Docker sandbox (orchestrator/worker)
|
||||
- ✅ Cron job scheduling (routines)
|
||||
- ✅ CLI subcommands (onboard, config, status, memory)
|
||||
- ✅ Gateway token auth
|
||||
|
||||
### P1 - High Priority
|
||||
- ❌ Slack channel (real implementation)
|
||||
- ❌ Telegram channel
|
||||
- ✅ Telegram channel (WASM, DM pairing, caption, /start)
|
||||
- ❌ WhatsApp channel
|
||||
- ❌ Multi-provider failover
|
||||
- ❌ Gateway control plane + WebSocket
|
||||
- ❌ Hooks system (beforeInbound, beforeToolCall, etc.)
|
||||
|
||||
### P2 - Medium Priority
|
||||
- ❌ Cron job scheduling
|
||||
- ❌ Web Control UI
|
||||
- ❌ WebChat channel
|
||||
- ❌ Media handling (images, PDFs)
|
||||
- 🚧 Media handling (caption support; no image/PDF processing)
|
||||
- ❌ CLI subcommands (config, status, memory, doctor)
|
||||
- ❌ Ollama/local model support
|
||||
- ❌ Configuration hot-reload
|
||||
- ❌ Webhook trigger endpoint in web gateway
|
||||
|
||||
### P3 - Lower Priority
|
||||
- ❌ Discord channel
|
||||
|
||||
@@ -43,7 +43,10 @@ IronClaw is the AI assistant you can actually trust with your personal and profe
|
||||
|
||||
### Always Available
|
||||
|
||||
- **Multi-channel** - REPL, HTTP webhooks, and extensible WASM channels (Telegram, Slack, and more)
|
||||
- **Multi-channel** - REPL, HTTP webhooks, WASM channels (Telegram, Slack), and web gateway
|
||||
- **Docker Sandbox** - Isolated container execution with per-job tokens and orchestrator/worker pattern
|
||||
- **Web Gateway** - Browser UI with real-time SSE/WebSocket streaming
|
||||
- **Routines** - Cron schedules, event triggers, webhook handlers for background automation
|
||||
- **Heartbeat System** - Proactive background execution for monitoring and maintenance tasks
|
||||
- **Parallel Jobs** - Handle multiple requests concurrently with isolated contexts
|
||||
- **Self-repair** - Automatic detection and recovery of stuck operations
|
||||
@@ -65,7 +68,7 @@ IronClaw is the AI assistant you can actually trust with your personal and profe
|
||||
### Prerequisites
|
||||
|
||||
- Rust 1.85+
|
||||
- PostgreSQL 15+ with pgvector extension
|
||||
- PostgreSQL 15+ with [pgvector](https://github.com/pgvector/pgvector) extension
|
||||
- NEAR AI account (authentication handled via setup wizard)
|
||||
|
||||
### Build
|
||||
@@ -82,6 +85,8 @@ cargo build --release
|
||||
cargo test
|
||||
```
|
||||
|
||||
For **full release** (after modifying channel sources), run `./scripts/build-all.sh` to rebuild channels first.
|
||||
|
||||
### Database Setup
|
||||
|
||||
```bash
|
||||
@@ -97,7 +102,7 @@ psql ironclaw -c "CREATE EXTENSION IF NOT EXISTS vector;"
|
||||
Run the setup wizard to configure IronClaw:
|
||||
|
||||
```bash
|
||||
ironclaw setup
|
||||
ironclaw onboard
|
||||
```
|
||||
|
||||
The wizard handles database connection, NEAR AI authentication (via browser OAuth),
|
||||
@@ -143,37 +148,42 @@ External content passes through multiple security layers:
|
||||
## Architecture
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────────────────┐
|
||||
│ Channels │
|
||||
│ ┌──────┐ ┌──────┐ ┌──────────────┐ │
|
||||
│ │ REPL │ │ HTTP │ │ WASM Channels│ │
|
||||
│ └──┬───┘ └──┬───┘ └──────┬───────┘ │
|
||||
│ └─────────┴─────────────┘ │
|
||||
│ │ │
|
||||
│ ┌────▼────┐ │
|
||||
│ │ Router │ Intent classification │
|
||||
│ └────┬────┘ │
|
||||
│ │ │
|
||||
│ ┌──────────▼──────────┐ │
|
||||
│ │ Scheduler │ Parallel job management │
|
||||
│ └──────────┬──────────┘ │
|
||||
│ │ │
|
||||
│ ┌───────────────┼───────────────┐ │
|
||||
│ ▼ ▼ ▼ │
|
||||
│ ┌─────────┐ ┌─────────┐ ┌─────────┐ │
|
||||
│ │ Worker │ │ Worker │ │ Worker │ LLM reasoning │
|
||||
│ └────┬────┘ └────┬────┘ └────┬────┘ │
|
||||
│ └───────────────┼───────────────┘ │
|
||||
│ │ │
|
||||
│ ┌──────────▼──────────┐ │
|
||||
│ │ Tool Registry │ │
|
||||
│ │ ┌───────────────┐ │ │
|
||||
│ │ │ Built-in │ │ │
|
||||
│ │ │ MCP │ │ │
|
||||
│ │ │ WASM Sandbox │ │ │
|
||||
│ │ └───────────────┘ │ │
|
||||
│ └─────────────────────┘ │
|
||||
└─────────────────────────────────────────────────────────────────┘
|
||||
┌────────────────────────────────────────────────────────────────────┐
|
||||
│ Channels │
|
||||
│ ┌──────┐ ┌──────┐ ┌─────────────┐ ┌─────────────┐ │
|
||||
│ │ REPL │ │ HTTP │ │WASM Channels│ │ Web Gateway │ │
|
||||
│ └──┬───┘ └──┬───┘ └──────┬──────┘ │ (SSE + WS) │ │
|
||||
│ │ │ │ └──────┬──────┘ │
|
||||
│ └─────────┴──────────────┴────────────────┘ │
|
||||
│ │ │
|
||||
│ ┌─────────▼─────────┐ │
|
||||
│ │ Agent Loop │ Intent routing │
|
||||
│ └────┬─────────┬────┘ │
|
||||
│ │ │ │
|
||||
│ ┌──────────▼───┐ ┌──▼──────────────┐ │
|
||||
│ │ Scheduler │ │ Routines Engine │ │
|
||||
│ │(parallel jobs)│ │(cron, event, wh) │ │
|
||||
│ └──────┬───────┘ └────────┬─────────┘ │
|
||||
│ │ │ │
|
||||
│ ┌─────────────┼───────────────────┘ │
|
||||
│ │ │ │
|
||||
│ ┌───▼────┐ ┌────▼────────────────┐ │
|
||||
│ │ Local │ │ Orchestrator │ │
|
||||
│ │Workers │ │ ┌───────────────┐ │ │
|
||||
│ │(in-proc)│ │ │ Docker Sandbox│ │ │
|
||||
│ └───┬────┘ │ │ Containers │ │ │
|
||||
│ │ │ │ ┌───────────┐ │ │ │
|
||||
│ │ │ │ │Worker / CC│ │ │ │
|
||||
│ │ │ │ └───────────┘ │ │ │
|
||||
│ │ │ └───────────────┘ │ │
|
||||
│ │ └─────────┬───────────┘ │
|
||||
│ └──────────────────┤ │
|
||||
│ │ │
|
||||
│ ┌───────────▼──────────┐ │
|
||||
│ │ Tool Registry │ │
|
||||
│ │ Built-in, MCP, WASM │ │
|
||||
│ └──────────────────────┘ │
|
||||
└────────────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### Core Components
|
||||
@@ -184,6 +194,9 @@ External content passes through multiple security layers:
|
||||
| **Router** | Classifies user intent (command, query, task) |
|
||||
| **Scheduler** | Manages parallel job execution with priorities |
|
||||
| **Worker** | Executes jobs with LLM reasoning and tool calls |
|
||||
| **Orchestrator** | Container lifecycle, LLM proxying, per-job auth |
|
||||
| **Web Gateway** | Browser UI with chat, memory, jobs, logs, extensions, routines |
|
||||
| **Routines Engine** | Scheduled (cron) and reactive (event, webhook) background tasks |
|
||||
| **Workspace** | Persistent memory with hybrid search |
|
||||
| **Safety Layer** | Prompt injection defense and content sanitization |
|
||||
|
||||
@@ -191,7 +204,7 @@ External content passes through multiple security layers:
|
||||
|
||||
```bash
|
||||
# First-time setup (configures database, auth, etc.)
|
||||
ironclaw setup
|
||||
ironclaw onboard
|
||||
|
||||
# Start interactive REPL
|
||||
cargo run
|
||||
@@ -210,12 +223,16 @@ cargo fmt
|
||||
cargo clippy --all --benches --tests --examples --all-features
|
||||
|
||||
# Run tests
|
||||
createdb ironclaw_test
|
||||
cargo test
|
||||
|
||||
# Run specific test
|
||||
cargo test test_name
|
||||
```
|
||||
|
||||
- **Telegram channel**: See [docs/TELEGRAM_SETUP.md](docs/TELEGRAM_SETUP.md) for setup and DM pairing.
|
||||
- **Changing channel sources**: Run `./channels-src/telegram/build.sh` before `cargo build` so the updated WASM is bundled.
|
||||
|
||||
## OpenClaw Heritage
|
||||
|
||||
IronClaw is a Rust reimplementation inspired by [OpenClaw](https://github.com/openclaw/openclaw). See [FEATURE_PARITY.md](FEATURE_PARITY.md) for the complete tracking matrix.
|
||||
|
||||
@@ -1,49 +0,0 @@
|
||||
[package]
|
||||
name = "ironclaw-bench"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
rust-version = "1.85"
|
||||
description = "Benchmarking harness for IronClaw agent"
|
||||
license = "MIT OR Apache-2.0"
|
||||
|
||||
[[bin]]
|
||||
name = "ironclaw-bench"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
ironclaw = { path = ".." }
|
||||
|
||||
# Async runtime
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
tokio-stream = { version = "0.1", features = ["sync"] }
|
||||
futures = "0.3"
|
||||
|
||||
# Serialization
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
toml = "0.8"
|
||||
|
||||
# CLI
|
||||
clap = { version = "4", features = ["derive"] }
|
||||
|
||||
# Core types
|
||||
uuid = { version = "1", features = ["v4", "serde"] }
|
||||
chrono = { version = "0.4", features = ["serde"] }
|
||||
rust_decimal = { version = "1", features = ["serde", "serde-with-str"] }
|
||||
|
||||
# Error handling
|
||||
thiserror = "2"
|
||||
anyhow = "1"
|
||||
|
||||
# Async traits
|
||||
async-trait = "0.1"
|
||||
|
||||
# Logging
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
||||
|
||||
# Scoring
|
||||
regex = "1"
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3"
|
||||
@@ -1,237 +0,0 @@
|
||||
use std::io::BufRead;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::error::BenchError;
|
||||
use crate::scoring;
|
||||
use crate::suite::{BenchScore, BenchSuite, BenchTask, TaskSubmission};
|
||||
|
||||
/// A single entry in the custom JSONL format.
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct CustomEntry {
|
||||
id: String,
|
||||
prompt: String,
|
||||
#[serde(default)]
|
||||
context: Option<String>,
|
||||
#[serde(default)]
|
||||
tags: Vec<String>,
|
||||
#[serde(default)]
|
||||
expected: Option<String>,
|
||||
#[serde(default)]
|
||||
expected_contains: Option<String>,
|
||||
#[serde(default)]
|
||||
expected_regex: Option<String>,
|
||||
/// "exact", "contains", "regex", or "llm" (default: "exact")
|
||||
#[serde(default = "default_scorer")]
|
||||
scorer: String,
|
||||
}
|
||||
|
||||
fn default_scorer() -> String {
|
||||
"exact".to_string()
|
||||
}
|
||||
|
||||
/// Custom JSONL benchmark suite.
|
||||
///
|
||||
/// Each line of the JSONL file is a task with `id`, `prompt`, and scoring
|
||||
/// criteria (`expected`, `expected_contains`, `expected_regex`).
|
||||
pub struct CustomSuite {
|
||||
dataset_path: PathBuf,
|
||||
}
|
||||
|
||||
impl CustomSuite {
|
||||
pub fn new(dataset_path: impl Into<PathBuf>) -> Self {
|
||||
Self {
|
||||
dataset_path: dataset_path.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BenchSuite for CustomSuite {
|
||||
fn name(&self) -> &str {
|
||||
"Custom JSONL"
|
||||
}
|
||||
|
||||
fn id(&self) -> &str {
|
||||
"custom"
|
||||
}
|
||||
|
||||
async fn load_tasks(&self) -> Result<Vec<BenchTask>, BenchError> {
|
||||
let file = std::fs::File::open(&self.dataset_path).map_err(BenchError::Io)?;
|
||||
let reader = std::io::BufReader::new(file);
|
||||
let mut tasks = Vec::new();
|
||||
|
||||
for (line_num, line) in reader.lines().enumerate() {
|
||||
let line = line?;
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let entry: CustomEntry = serde_json::from_str(trimmed)
|
||||
.map_err(|e| BenchError::Config(format!("line {}: {}", line_num + 1, e)))?;
|
||||
|
||||
let mut metadata = serde_json::json!({
|
||||
"scorer": entry.scorer,
|
||||
});
|
||||
if let Some(ref expected) = entry.expected {
|
||||
metadata["expected"] = serde_json::Value::String(expected.clone());
|
||||
}
|
||||
if let Some(ref expected_contains) = entry.expected_contains {
|
||||
metadata["expected_contains"] =
|
||||
serde_json::Value::String(expected_contains.clone());
|
||||
}
|
||||
if let Some(ref expected_regex) = entry.expected_regex {
|
||||
metadata["expected_regex"] = serde_json::Value::String(expected_regex.clone());
|
||||
}
|
||||
|
||||
tasks.push(BenchTask {
|
||||
id: entry.id,
|
||||
prompt: entry.prompt,
|
||||
context: entry.context,
|
||||
resources: vec![],
|
||||
tags: entry.tags,
|
||||
expected_turns: None,
|
||||
timeout: None,
|
||||
metadata,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(tasks)
|
||||
}
|
||||
|
||||
async fn score(
|
||||
&self,
|
||||
task: &BenchTask,
|
||||
submission: &TaskSubmission,
|
||||
) -> Result<BenchScore, BenchError> {
|
||||
let scorer = task
|
||||
.metadata
|
||||
.get("scorer")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("exact");
|
||||
|
||||
match scorer {
|
||||
"exact" => {
|
||||
if let Some(expected) = task.metadata.get("expected").and_then(|v| v.as_str()) {
|
||||
Ok(scoring::exact_match(expected, &submission.response))
|
||||
} else {
|
||||
Err(BenchError::Scoring {
|
||||
task_id: task.id.clone(),
|
||||
reason: "no 'expected' field for exact scoring".to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
"contains" => {
|
||||
if let Some(expected) = task
|
||||
.metadata
|
||||
.get("expected_contains")
|
||||
.and_then(|v| v.as_str())
|
||||
{
|
||||
Ok(scoring::contains_match(expected, &submission.response))
|
||||
} else {
|
||||
Err(BenchError::Scoring {
|
||||
task_id: task.id.clone(),
|
||||
reason: "no 'expected_contains' field for contains scoring".to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
"regex" => {
|
||||
if let Some(pattern) = task.metadata.get("expected_regex").and_then(|v| v.as_str())
|
||||
{
|
||||
Ok(scoring::regex_match(pattern, &submission.response))
|
||||
} else {
|
||||
Err(BenchError::Scoring {
|
||||
task_id: task.id.clone(),
|
||||
reason: "no 'expected_regex' field for regex scoring".to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
"llm" => {
|
||||
// TODO: LLM-as-judge scoring
|
||||
Ok(BenchScore::partial(0.5, "LLM scoring not yet implemented"))
|
||||
}
|
||||
other => Err(BenchError::Scoring {
|
||||
task_id: task.id.clone(),
|
||||
reason: format!("unknown scorer: {other}"),
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Write;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_custom_load_tasks() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("tasks.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"id": "t1", "prompt": "What is 2+2?", "expected": "4"}}"#
|
||||
)
|
||||
.unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"id": "t2", "prompt": "Say hello", "expected_contains": "hello", "scorer": "contains"}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = CustomSuite::new(&path);
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
assert_eq!(tasks.len(), 2);
|
||||
assert_eq!(tasks[0].id, "t1");
|
||||
assert_eq!(tasks[1].id, "t2");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_custom_exact_scoring() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("tasks.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"id": "t1", "prompt": "What is 2+2?", "expected": "4"}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = CustomSuite::new(&path);
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
|
||||
let submission = TaskSubmission {
|
||||
response: "4".to_string(),
|
||||
conversation: vec![],
|
||||
tool_calls: vec![],
|
||||
};
|
||||
let score = suite.score(&tasks[0], &submission).await.unwrap();
|
||||
assert_eq!(score.value, 1.0);
|
||||
assert_eq!(score.label, "pass");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_custom_contains_scoring() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("tasks.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"id": "t1", "prompt": "Greet me", "expected_contains": "hello", "scorer": "contains"}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = CustomSuite::new(&path);
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
|
||||
let submission = TaskSubmission {
|
||||
response: "Hello there!".to_string(),
|
||||
conversation: vec![],
|
||||
tool_calls: vec![],
|
||||
};
|
||||
let score = suite.score(&tasks[0], &submission).await.unwrap();
|
||||
assert_eq!(score.value, 1.0);
|
||||
}
|
||||
}
|
||||
@@ -1,183 +0,0 @@
|
||||
use std::io::BufRead;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::error::BenchError;
|
||||
use crate::scoring;
|
||||
use crate::suite::{BenchScore, BenchSuite, BenchTask, TaskResource, TaskSubmission};
|
||||
|
||||
/// GAIA dataset entry (Hugging Face JSONL format).
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct GaiaEntry {
|
||||
task_id: String,
|
||||
#[serde(alias = "Question")]
|
||||
question: String,
|
||||
#[serde(alias = "Final answer", alias = "final_answer")]
|
||||
final_answer: String,
|
||||
#[serde(alias = "Level", default)]
|
||||
level: Option<u32>,
|
||||
#[serde(alias = "file_name", default)]
|
||||
file_name: Option<String>,
|
||||
#[serde(alias = "Annotator Metadata", default)]
|
||||
annotator_metadata: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
/// GAIA benchmark suite.
|
||||
///
|
||||
/// Tasks are loaded from HuggingFace JSONL exports. Scoring uses normalized
|
||||
/// exact match against the `final_answer` field.
|
||||
pub struct GaiaSuite {
|
||||
dataset_path: PathBuf,
|
||||
attachments_dir: Option<PathBuf>,
|
||||
}
|
||||
|
||||
impl GaiaSuite {
|
||||
pub fn new(
|
||||
dataset_path: impl Into<PathBuf>,
|
||||
attachments_dir: Option<impl Into<PathBuf>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
dataset_path: dataset_path.into(),
|
||||
attachments_dir: attachments_dir.map(|d| d.into()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BenchSuite for GaiaSuite {
|
||||
fn name(&self) -> &str {
|
||||
"GAIA"
|
||||
}
|
||||
|
||||
fn id(&self) -> &str {
|
||||
"gaia"
|
||||
}
|
||||
|
||||
async fn load_tasks(&self) -> Result<Vec<BenchTask>, BenchError> {
|
||||
let file = std::fs::File::open(&self.dataset_path)?;
|
||||
let reader = std::io::BufReader::new(file);
|
||||
let mut tasks = Vec::new();
|
||||
|
||||
for (line_num, line) in reader.lines().enumerate() {
|
||||
let line = line?;
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let entry: GaiaEntry = serde_json::from_str(trimmed)
|
||||
.map_err(|e| BenchError::Config(format!("GAIA line {}: {}", line_num + 1, e)))?;
|
||||
|
||||
let mut resources = Vec::new();
|
||||
if let Some(ref file_name) = entry.file_name {
|
||||
if !file_name.is_empty() {
|
||||
if let Some(ref dir) = self.attachments_dir {
|
||||
resources.push(TaskResource {
|
||||
name: file_name.clone(),
|
||||
path: dir.join(file_name).to_string_lossy().to_string(),
|
||||
resource_type: crate::suite::ResourceType::File,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut tags = Vec::new();
|
||||
if let Some(level) = entry.level {
|
||||
tags.push(format!("level-{level}"));
|
||||
}
|
||||
|
||||
let metadata = serde_json::json!({
|
||||
"expected": entry.final_answer,
|
||||
"level": entry.level,
|
||||
});
|
||||
|
||||
tasks.push(BenchTask {
|
||||
id: entry.task_id,
|
||||
prompt: entry.question,
|
||||
context: None,
|
||||
resources,
|
||||
tags,
|
||||
expected_turns: None,
|
||||
timeout: None,
|
||||
metadata,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(tasks)
|
||||
}
|
||||
|
||||
async fn score(
|
||||
&self,
|
||||
task: &BenchTask,
|
||||
submission: &TaskSubmission,
|
||||
) -> Result<BenchScore, BenchError> {
|
||||
let expected = task
|
||||
.metadata
|
||||
.get("expected")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| BenchError::Scoring {
|
||||
task_id: task.id.clone(),
|
||||
reason: "missing expected answer in metadata".to_string(),
|
||||
})?;
|
||||
|
||||
Ok(scoring::exact_match(expected, &submission.response))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Write;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_gaia_load_tasks() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("gaia.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"task_id": "g1", "question": "What is the capital of France?", "final_answer": "Paris", "Level": 1}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = GaiaSuite::new(&path, None::<PathBuf>);
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
assert_eq!(tasks.len(), 1);
|
||||
assert_eq!(tasks[0].id, "g1");
|
||||
assert!(tasks[0].tags.contains(&"level-1".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_gaia_scoring() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("gaia.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"task_id": "g1", "question": "Capital of France?", "final_answer": "Paris"}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = GaiaSuite::new(&path, None::<PathBuf>);
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
|
||||
// Exact match (case insensitive)
|
||||
let submission = TaskSubmission {
|
||||
response: "paris".to_string(),
|
||||
conversation: vec![],
|
||||
tool_calls: vec![],
|
||||
};
|
||||
let score = suite.score(&tasks[0], &submission).await.unwrap();
|
||||
assert_eq!(score.value, 1.0);
|
||||
|
||||
// Wrong answer
|
||||
let submission = TaskSubmission {
|
||||
response: "London".to_string(),
|
||||
conversation: vec![],
|
||||
tool_calls: vec![],
|
||||
};
|
||||
let score = suite.score(&tasks[0], &submission).await.unwrap();
|
||||
assert_eq!(score.value, 0.0);
|
||||
}
|
||||
}
|
||||
@@ -1,110 +0,0 @@
|
||||
pub mod custom;
|
||||
pub mod gaia;
|
||||
pub mod swe_bench;
|
||||
pub mod tau_bench;
|
||||
|
||||
use crate::config::BenchConfig;
|
||||
use crate::error::BenchError;
|
||||
use crate::suite::BenchSuite;
|
||||
|
||||
/// List of all known suite IDs.
|
||||
pub const KNOWN_SUITES: &[(&str, &str)] = &[
|
||||
("custom", "Custom JSONL tasks"),
|
||||
("gaia", "GAIA benchmark (knowledge & reasoning)"),
|
||||
("tau_bench", "Tau-bench (multi-turn tool use)"),
|
||||
("swe_bench", "SWE-bench Pro (software engineering)"),
|
||||
];
|
||||
|
||||
/// Create a suite adapter by name.
|
||||
pub fn create_suite(name: &str, config: &BenchConfig) -> Result<Box<dyn BenchSuite>, BenchError> {
|
||||
let suite_map = config.suite_config_map();
|
||||
match name {
|
||||
"custom" => {
|
||||
let dataset_path = suite_map
|
||||
.get("dataset_path")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| {
|
||||
BenchError::Config(
|
||||
"suite_config.dataset_path is required for 'custom' suite".to_string(),
|
||||
)
|
||||
})?;
|
||||
Ok(Box::new(custom::CustomSuite::new(dataset_path)))
|
||||
}
|
||||
"gaia" => {
|
||||
let dataset_path = suite_map
|
||||
.get("dataset_path")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| {
|
||||
BenchError::Config(
|
||||
"suite_config.dataset_path is required for 'gaia' suite".to_string(),
|
||||
)
|
||||
})?;
|
||||
let attachments_dir = suite_map
|
||||
.get("attachments_dir")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
Ok(Box::new(gaia::GaiaSuite::new(
|
||||
dataset_path,
|
||||
attachments_dir,
|
||||
)))
|
||||
}
|
||||
"tau_bench" => {
|
||||
let dataset_path = suite_map
|
||||
.get("dataset_path")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| {
|
||||
BenchError::Config(
|
||||
"suite_config.dataset_path is required for 'tau_bench' suite".to_string(),
|
||||
)
|
||||
})?;
|
||||
let domain = suite_map
|
||||
.get("domain")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("retail")
|
||||
.to_string();
|
||||
Ok(Box::new(tau_bench::TauBenchSuite::new(
|
||||
dataset_path,
|
||||
domain,
|
||||
)))
|
||||
}
|
||||
"swe_bench" => {
|
||||
let dataset_path = suite_map
|
||||
.get("dataset_path")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| {
|
||||
BenchError::Config(
|
||||
"suite_config.dataset_path is required for 'swe_bench' suite".to_string(),
|
||||
)
|
||||
})?;
|
||||
let workspace_dir = suite_map
|
||||
.get("workspace_dir")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("/tmp/swe-bench")
|
||||
.to_string();
|
||||
let use_docker = suite_map
|
||||
.get("use_docker")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
Ok(Box::new(swe_bench::SweBenchSuite::new(
|
||||
dataset_path,
|
||||
workspace_dir,
|
||||
use_docker,
|
||||
)))
|
||||
}
|
||||
_ => {
|
||||
let available = KNOWN_SUITES
|
||||
.iter()
|
||||
.map(|(id, _)| *id)
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
Err(BenchError::SuiteNotFound {
|
||||
name: name.to_string(),
|
||||
available,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,291 +0,0 @@
|
||||
use std::io::BufRead;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::error::BenchError;
|
||||
use crate::suite::{BenchScore, BenchSuite, BenchTask, TaskSubmission};
|
||||
|
||||
/// SWE-bench dataset entry.
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct SweBenchEntry {
|
||||
instance_id: String,
|
||||
repo: String,
|
||||
base_commit: String,
|
||||
#[serde(default)]
|
||||
problem_statement: String,
|
||||
#[serde(default)]
|
||||
hints_text: Option<String>,
|
||||
#[serde(default)]
|
||||
test_patch: Option<String>,
|
||||
#[serde(default)]
|
||||
patch: Option<String>,
|
||||
}
|
||||
|
||||
/// SWE-bench Pro: real-world software engineering tasks.
|
||||
///
|
||||
/// Each task clones a repo at a specific commit, presents the problem statement,
|
||||
/// and expects the agent to produce a patch. Scoring runs the test suite.
|
||||
pub struct SweBenchSuite {
|
||||
dataset_path: PathBuf,
|
||||
workspace_dir: PathBuf,
|
||||
use_docker: bool,
|
||||
}
|
||||
|
||||
impl SweBenchSuite {
|
||||
pub fn new(
|
||||
dataset_path: impl Into<PathBuf>,
|
||||
workspace_dir: impl Into<PathBuf>,
|
||||
use_docker: bool,
|
||||
) -> Self {
|
||||
Self {
|
||||
dataset_path: dataset_path.into(),
|
||||
workspace_dir: workspace_dir.into(),
|
||||
use_docker,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BenchSuite for SweBenchSuite {
|
||||
fn name(&self) -> &str {
|
||||
"SWE-bench Pro"
|
||||
}
|
||||
|
||||
fn id(&self) -> &str {
|
||||
"swe_bench"
|
||||
}
|
||||
|
||||
async fn load_tasks(&self) -> Result<Vec<BenchTask>, BenchError> {
|
||||
let file = std::fs::File::open(&self.dataset_path)?;
|
||||
let reader = std::io::BufReader::new(file);
|
||||
let mut tasks = Vec::new();
|
||||
|
||||
for (line_num, line) in reader.lines().enumerate() {
|
||||
let line = line?;
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let entry: SweBenchEntry = serde_json::from_str(trimmed).map_err(|e| {
|
||||
BenchError::Config(format!("swe_bench line {}: {}", line_num + 1, e))
|
||||
})?;
|
||||
|
||||
let metadata = serde_json::json!({
|
||||
"repo": entry.repo,
|
||||
"base_commit": entry.base_commit,
|
||||
"test_patch": entry.test_patch,
|
||||
"gold_patch": entry.patch,
|
||||
"use_docker": self.use_docker,
|
||||
"workspace_dir": self.workspace_dir.to_string_lossy(),
|
||||
});
|
||||
|
||||
let prompt = if let Some(ref hints) = entry.hints_text {
|
||||
format!("{}\n\nHints:\n{}", entry.problem_statement, hints)
|
||||
} else {
|
||||
entry.problem_statement
|
||||
};
|
||||
|
||||
tasks.push(BenchTask {
|
||||
id: entry.instance_id,
|
||||
prompt,
|
||||
context: Some(format!(
|
||||
"Repository: {}, Commit: {}",
|
||||
entry.repo, entry.base_commit
|
||||
)),
|
||||
resources: vec![],
|
||||
tags: vec![format!("repo-{}", entry.repo.replace('/', "-"))],
|
||||
expected_turns: None,
|
||||
timeout: None,
|
||||
metadata,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(tasks)
|
||||
}
|
||||
|
||||
async fn setup_task(&self, task: &BenchTask) -> Result<(), BenchError> {
|
||||
let repo = task
|
||||
.metadata
|
||||
.get("repo")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| BenchError::TaskFailed {
|
||||
task_id: task.id.clone(),
|
||||
reason: "missing repo in metadata".to_string(),
|
||||
})?;
|
||||
let base_commit = task
|
||||
.metadata
|
||||
.get("base_commit")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| BenchError::TaskFailed {
|
||||
task_id: task.id.clone(),
|
||||
reason: "missing base_commit in metadata".to_string(),
|
||||
})?;
|
||||
|
||||
let task_dir = self.workspace_dir.join(&task.id);
|
||||
|
||||
// Clone repo if not already present
|
||||
if !task_dir.exists() {
|
||||
let repo_url = format!("https://github.com/{}.git", repo);
|
||||
let output = tokio::process::Command::new("git")
|
||||
.args([
|
||||
"clone",
|
||||
"--depth",
|
||||
"1",
|
||||
&repo_url,
|
||||
&task_dir.to_string_lossy(),
|
||||
])
|
||||
.output()
|
||||
.await
|
||||
.map_err(|e| BenchError::TaskFailed {
|
||||
task_id: task.id.clone(),
|
||||
reason: format!("git clone failed: {e}"),
|
||||
})?;
|
||||
|
||||
if !output.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
return Err(BenchError::TaskFailed {
|
||||
task_id: task.id.clone(),
|
||||
reason: format!("git clone failed: {stderr}"),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Checkout the base commit
|
||||
let output = tokio::process::Command::new("git")
|
||||
.args(["checkout", base_commit])
|
||||
.current_dir(&task_dir)
|
||||
.output()
|
||||
.await
|
||||
.map_err(|e| BenchError::TaskFailed {
|
||||
task_id: task.id.clone(),
|
||||
reason: format!("git checkout failed: {e}"),
|
||||
})?;
|
||||
|
||||
if !output.status.success() {
|
||||
// Shallow clone might not have the commit; fetch more history
|
||||
let _ = tokio::process::Command::new("git")
|
||||
.args(["fetch", "--unshallow"])
|
||||
.current_dir(&task_dir)
|
||||
.output()
|
||||
.await;
|
||||
|
||||
let output = tokio::process::Command::new("git")
|
||||
.args(["checkout", base_commit])
|
||||
.current_dir(&task_dir)
|
||||
.output()
|
||||
.await
|
||||
.map_err(|e| BenchError::TaskFailed {
|
||||
task_id: task.id.clone(),
|
||||
reason: format!("git checkout retry failed: {e}"),
|
||||
})?;
|
||||
|
||||
if !output.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
return Err(BenchError::TaskFailed {
|
||||
task_id: task.id.clone(),
|
||||
reason: format!("git checkout failed: {stderr}"),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn teardown_task(&self, task: &BenchTask) -> Result<(), BenchError> {
|
||||
let task_dir = self.workspace_dir.join(&task.id);
|
||||
if task_dir.exists() {
|
||||
// Reset any changes
|
||||
let _ = tokio::process::Command::new("git")
|
||||
.args(["checkout", "."])
|
||||
.current_dir(&task_dir)
|
||||
.output()
|
||||
.await;
|
||||
let _ = tokio::process::Command::new("git")
|
||||
.args(["clean", "-fdx"])
|
||||
.current_dir(&task_dir)
|
||||
.output()
|
||||
.await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn score(
|
||||
&self,
|
||||
task: &BenchTask,
|
||||
submission: &TaskSubmission,
|
||||
) -> Result<BenchScore, BenchError> {
|
||||
// For SWE-bench, scoring requires running the test patch against the agent's changes.
|
||||
// This is a simplified version that checks if the agent produced any code changes.
|
||||
|
||||
let test_patch = task.metadata.get("test_patch").and_then(|v| v.as_str());
|
||||
|
||||
if submission.response.is_empty() {
|
||||
return Ok(BenchScore::fail("no response from agent"));
|
||||
}
|
||||
|
||||
// If we have a test patch, try to verify the submission
|
||||
if let Some(_test_patch) = test_patch {
|
||||
// TODO: Apply agent's patch, then apply test patch, then run tests.
|
||||
// For now, give partial credit if the agent produced some output.
|
||||
Ok(BenchScore::partial(
|
||||
0.25,
|
||||
"test execution not yet implemented; partial credit for response",
|
||||
))
|
||||
} else {
|
||||
// No test patch available; can't automatically score
|
||||
Ok(BenchScore::partial(
|
||||
0.25,
|
||||
"no test_patch available for automated scoring",
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Write;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_swe_bench_load() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("swe.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"instance_id": "django__django-12345", "repo": "django/django", "base_commit": "abc123", "problem_statement": "Fix the ORM bug"}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = SweBenchSuite::new(&path, "/tmp/swe-test", false);
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
assert_eq!(tasks.len(), 1);
|
||||
assert_eq!(tasks[0].id, "django__django-12345");
|
||||
assert!(tasks[0].tags.contains(&"repo-django-django".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_swe_bench_scoring_no_response() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("swe.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"instance_id": "s1", "repo": "org/repo", "base_commit": "abc", "problem_statement": "Fix bug"}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = SweBenchSuite::new(&path, "/tmp/swe-test", false);
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
|
||||
let submission = TaskSubmission {
|
||||
response: String::new(),
|
||||
conversation: vec![],
|
||||
tool_calls: vec![],
|
||||
};
|
||||
let score = suite.score(&tasks[0], &submission).await.unwrap();
|
||||
assert_eq!(score.value, 0.0);
|
||||
}
|
||||
}
|
||||
@@ -1,227 +0,0 @@
|
||||
use std::io::BufRead;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::error::BenchError;
|
||||
use crate::suite::{BenchScore, BenchSuite, BenchTask, ConversationTurn, TaskSubmission};
|
||||
|
||||
/// Tau-bench task entry.
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct TauBenchEntry {
|
||||
id: String,
|
||||
#[serde(default)]
|
||||
domain: String,
|
||||
instruction: String,
|
||||
#[serde(default)]
|
||||
user_persona: Option<String>,
|
||||
#[serde(default)]
|
||||
expected_state: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
expected_actions: Vec<String>,
|
||||
#[serde(default)]
|
||||
max_turns: Option<usize>,
|
||||
}
|
||||
|
||||
/// Tau-bench: multi-turn tool-calling dialog benchmark.
|
||||
///
|
||||
/// Tests agent ability to handle customer service scenarios with simulated
|
||||
/// domain APIs (retail, airline). Scoring compares final state against expected.
|
||||
pub struct TauBenchSuite {
|
||||
dataset_path: PathBuf,
|
||||
domain: String,
|
||||
}
|
||||
|
||||
impl TauBenchSuite {
|
||||
pub fn new(dataset_path: impl Into<PathBuf>, domain: impl Into<String>) -> Self {
|
||||
Self {
|
||||
dataset_path: dataset_path.into(),
|
||||
domain: domain.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BenchSuite for TauBenchSuite {
|
||||
fn name(&self) -> &str {
|
||||
"Tau-bench"
|
||||
}
|
||||
|
||||
fn id(&self) -> &str {
|
||||
"tau_bench"
|
||||
}
|
||||
|
||||
async fn load_tasks(&self) -> Result<Vec<BenchTask>, BenchError> {
|
||||
let file = std::fs::File::open(&self.dataset_path)?;
|
||||
let reader = std::io::BufReader::new(file);
|
||||
let mut tasks = Vec::new();
|
||||
|
||||
for (line_num, line) in reader.lines().enumerate() {
|
||||
let line = line?;
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let entry: TauBenchEntry = serde_json::from_str(trimmed).map_err(|e| {
|
||||
BenchError::Config(format!("tau_bench line {}: {}", line_num + 1, e))
|
||||
})?;
|
||||
|
||||
let domain = if entry.domain.is_empty() {
|
||||
self.domain.clone()
|
||||
} else {
|
||||
entry.domain.clone()
|
||||
};
|
||||
|
||||
let metadata = serde_json::json!({
|
||||
"domain": domain,
|
||||
"user_persona": entry.user_persona,
|
||||
"expected_state": entry.expected_state,
|
||||
"expected_actions": entry.expected_actions,
|
||||
});
|
||||
|
||||
tasks.push(BenchTask {
|
||||
id: entry.id,
|
||||
prompt: entry.instruction,
|
||||
context: entry.user_persona.clone(),
|
||||
resources: vec![],
|
||||
tags: vec![format!("domain-{domain}")],
|
||||
expected_turns: entry.max_turns,
|
||||
timeout: None,
|
||||
metadata,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(tasks)
|
||||
}
|
||||
|
||||
async fn score(
|
||||
&self,
|
||||
task: &BenchTask,
|
||||
submission: &TaskSubmission,
|
||||
) -> Result<BenchScore, BenchError> {
|
||||
// Score based on expected actions completion
|
||||
let expected_actions: Vec<String> = task
|
||||
.metadata
|
||||
.get("expected_actions")
|
||||
.and_then(|v| serde_json::from_value(v.clone()).ok())
|
||||
.unwrap_or_default();
|
||||
|
||||
if expected_actions.is_empty() {
|
||||
// No expected actions defined; score based on whether agent responded
|
||||
if submission.response.is_empty() {
|
||||
return Ok(BenchScore::fail("no response"));
|
||||
}
|
||||
return Ok(BenchScore::partial(
|
||||
0.5,
|
||||
"no expected_actions to evaluate against",
|
||||
));
|
||||
}
|
||||
|
||||
// Check which expected actions were actually called
|
||||
let called: std::collections::HashSet<&str> =
|
||||
submission.tool_calls.iter().map(|s| s.as_str()).collect();
|
||||
let matched = expected_actions
|
||||
.iter()
|
||||
.filter(|a| called.contains(a.as_str()))
|
||||
.count();
|
||||
|
||||
let ratio = matched as f64 / expected_actions.len() as f64;
|
||||
if ratio >= 1.0 {
|
||||
Ok(BenchScore::pass())
|
||||
} else if ratio > 0.0 {
|
||||
Ok(BenchScore::partial(
|
||||
ratio,
|
||||
format!(
|
||||
"{}/{} expected actions completed",
|
||||
matched,
|
||||
expected_actions.len()
|
||||
),
|
||||
))
|
||||
} else {
|
||||
Ok(BenchScore::fail(format!(
|
||||
"0/{} expected actions completed",
|
||||
expected_actions.len()
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
async fn next_user_message(
|
||||
&self,
|
||||
task: &BenchTask,
|
||||
conversation: &[ConversationTurn],
|
||||
) -> Result<Option<String>, BenchError> {
|
||||
// Check if we've exceeded max turns
|
||||
if let Some(max) = task.expected_turns {
|
||||
let user_turns = conversation
|
||||
.iter()
|
||||
.filter(|t| matches!(t.role, crate::suite::TurnRole::User))
|
||||
.count();
|
||||
if user_turns >= max {
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
|
||||
// For now, multi-turn simulation requires an LLM (not implemented yet).
|
||||
// Return None to end after the first turn.
|
||||
// TODO: Use LLM to simulate customer based on user_persona.
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Write;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_tau_bench_load() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("tau.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"id": "t1", "instruction": "Return my order", "expected_actions": ["lookup_order", "process_return"], "max_turns": 3}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = TauBenchSuite::new(&path, "retail");
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
assert_eq!(tasks.len(), 1);
|
||||
assert_eq!(tasks[0].expected_turns, Some(3));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_tau_bench_scoring() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("tau.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"id": "t1", "instruction": "Return order", "expected_actions": ["lookup_order", "process_return"]}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = TauBenchSuite::new(&path, "retail");
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
|
||||
// Partial completion
|
||||
let submission = TaskSubmission {
|
||||
response: "I found your order.".to_string(),
|
||||
conversation: vec![],
|
||||
tool_calls: vec!["lookup_order".to_string()],
|
||||
};
|
||||
let score = suite.score(&tasks[0], &submission).await.unwrap();
|
||||
assert_eq!(score.value, 0.5);
|
||||
assert_eq!(score.label, "partial");
|
||||
|
||||
// Full completion
|
||||
let submission = TaskSubmission {
|
||||
response: "Return processed.".to_string(),
|
||||
conversation: vec![],
|
||||
tool_calls: vec!["lookup_order".to_string(), "process_return".to_string()],
|
||||
};
|
||||
let score = suite.score(&tasks[0], &submission).await.unwrap();
|
||||
assert_eq!(score.value, 1.0);
|
||||
}
|
||||
}
|
||||
@@ -1,229 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use tokio::sync::{Mutex, mpsc};
|
||||
use tokio_stream::wrappers::ReceiverStream;
|
||||
|
||||
use ironclaw::channels::{Channel, IncomingMessage, MessageStream, OutgoingResponse, StatusUpdate};
|
||||
use ironclaw::error::ChannelError;
|
||||
|
||||
use crate::results::TraceToolCall;
|
||||
use crate::suite::ConversationTurn;
|
||||
|
||||
/// Captured state from a benchmark channel run.
|
||||
#[derive(Debug, Default)]
|
||||
pub struct ChannelCapture {
|
||||
/// All responses the agent sent back.
|
||||
pub responses: Vec<String>,
|
||||
/// Tool calls observed (name, success, duration_ms).
|
||||
pub tool_calls: Vec<TraceToolCall>,
|
||||
/// Full conversation turns for multi-turn scoring.
|
||||
pub conversation: Vec<ConversationTurn>,
|
||||
/// Status messages (for debugging).
|
||||
pub status_log: Vec<String>,
|
||||
}
|
||||
|
||||
/// A headless Channel implementation for benchmarking.
|
||||
///
|
||||
/// Modeled after `ReplChannel`: uses mpsc to inject messages and captures
|
||||
/// all responses and tool status events. Auto-approves tool execution
|
||||
/// so benchmarks run without user interaction.
|
||||
pub struct BenchChannel {
|
||||
/// Sender to inject messages into the agent loop.
|
||||
msg_tx: mpsc::Sender<IncomingMessage>,
|
||||
/// Receiver the agent loop reads from (taken once by `start()`).
|
||||
msg_rx: Mutex<Option<mpsc::Receiver<IncomingMessage>>>,
|
||||
/// Accumulated capture data.
|
||||
capture: Arc<Mutex<ChannelCapture>>,
|
||||
}
|
||||
|
||||
impl BenchChannel {
|
||||
pub fn new() -> (Self, mpsc::Sender<IncomingMessage>) {
|
||||
let (tx, rx) = mpsc::channel(64);
|
||||
let channel = Self {
|
||||
msg_tx: tx.clone(),
|
||||
msg_rx: Mutex::new(Some(rx)),
|
||||
capture: Arc::new(Mutex::new(ChannelCapture::default())),
|
||||
};
|
||||
(channel, tx)
|
||||
}
|
||||
|
||||
/// Get a handle to the capture data.
|
||||
pub fn capture(&self) -> Arc<Mutex<ChannelCapture>> {
|
||||
Arc::clone(&self.capture)
|
||||
}
|
||||
|
||||
/// Get a clone of the message sender for injecting follow-up messages.
|
||||
pub fn sender(&self) -> mpsc::Sender<IncomingMessage> {
|
||||
self.msg_tx.clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Channel for BenchChannel {
|
||||
fn name(&self) -> &str {
|
||||
"bench"
|
||||
}
|
||||
|
||||
async fn start(&self) -> Result<MessageStream, ChannelError> {
|
||||
let rx = self
|
||||
.msg_rx
|
||||
.lock()
|
||||
.await
|
||||
.take()
|
||||
.ok_or_else(|| ChannelError::StartupFailed {
|
||||
name: "bench".to_string(),
|
||||
reason: "start() already called".to_string(),
|
||||
})?;
|
||||
Ok(Box::pin(ReceiverStream::new(rx)))
|
||||
}
|
||||
|
||||
async fn respond(
|
||||
&self,
|
||||
_msg: &IncomingMessage,
|
||||
response: OutgoingResponse,
|
||||
) -> Result<(), ChannelError> {
|
||||
let mut cap = self.capture.lock().await;
|
||||
cap.responses.push(response.content.clone());
|
||||
cap.conversation.push(ConversationTurn {
|
||||
role: crate::suite::TurnRole::Assistant,
|
||||
content: response.content,
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_status(
|
||||
&self,
|
||||
status: StatusUpdate,
|
||||
_metadata: &serde_json::Value,
|
||||
) -> Result<(), ChannelError> {
|
||||
let mut cap = self.capture.lock().await;
|
||||
|
||||
match status {
|
||||
StatusUpdate::ToolCompleted { ref name, success } => {
|
||||
cap.tool_calls.push(TraceToolCall {
|
||||
name: name.clone(),
|
||||
duration_ms: 0, // We don't have precise per-tool timing here
|
||||
success,
|
||||
});
|
||||
cap.status_log
|
||||
.push(format!("tool_completed: {name} success={success}"));
|
||||
}
|
||||
StatusUpdate::ApprovalNeeded { ref request_id, .. } => {
|
||||
// Auto-approve all tools during benchmarks
|
||||
cap.status_log.push(format!("auto_approved: {request_id}"));
|
||||
drop(cap); // Release lock before sending
|
||||
let approval = IncomingMessage::new("bench", "bench-user", "always");
|
||||
let _ = self.msg_tx.send(approval).await;
|
||||
return Ok(());
|
||||
}
|
||||
StatusUpdate::Thinking(ref msg) => {
|
||||
cap.status_log.push(format!("thinking: {msg}"));
|
||||
}
|
||||
StatusUpdate::ToolStarted { ref name } => {
|
||||
cap.status_log.push(format!("tool_started: {name}"));
|
||||
}
|
||||
StatusUpdate::ToolResult {
|
||||
ref name,
|
||||
ref preview,
|
||||
} => {
|
||||
cap.status_log.push(format!(
|
||||
"tool_result: {name} -> {}",
|
||||
&preview[..preview.len().min(100)]
|
||||
));
|
||||
}
|
||||
StatusUpdate::StreamChunk(_) => {}
|
||||
StatusUpdate::Status(ref msg) => {
|
||||
cap.status_log.push(format!("status: {msg}"));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn broadcast(
|
||||
&self,
|
||||
_user_id: &str,
|
||||
response: OutgoingResponse,
|
||||
) -> Result<(), ChannelError> {
|
||||
let mut cap = self.capture.lock().await;
|
||||
cap.status_log.push(format!(
|
||||
"broadcast: {}",
|
||||
&response.content[..response.content.len().min(100)]
|
||||
));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> Result<(), ChannelError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn shutdown(&self) -> Result<(), ChannelError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bench_channel_captures_responses() {
|
||||
let (channel, _tx) = BenchChannel::new();
|
||||
let capture = channel.capture();
|
||||
|
||||
let msg = IncomingMessage::new("bench", "user", "hello");
|
||||
let response = OutgoingResponse::text("world");
|
||||
channel.respond(&msg, response).await.unwrap();
|
||||
|
||||
let cap = capture.lock().await;
|
||||
assert_eq!(cap.responses.len(), 1);
|
||||
assert_eq!(cap.responses[0], "world");
|
||||
assert_eq!(cap.conversation.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bench_channel_auto_approves() {
|
||||
let (channel, _tx) = BenchChannel::new();
|
||||
// start() to consume the receiver
|
||||
let _stream = channel.start().await.unwrap();
|
||||
|
||||
let status = StatusUpdate::ApprovalNeeded {
|
||||
request_id: "req-1".to_string(),
|
||||
tool_name: "shell".to_string(),
|
||||
description: "run ls".to_string(),
|
||||
parameters: serde_json::json!({}),
|
||||
};
|
||||
channel
|
||||
.send_status(status, &serde_json::Value::Null)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// The approval message was sent through msg_tx,
|
||||
// which means the stream would receive it.
|
||||
// We can't easily read from the stream in this test without
|
||||
// consuming it, but we can verify the status log.
|
||||
let capture_arc = channel.capture();
|
||||
let cap = capture_arc.lock().await;
|
||||
assert!(cap.status_log.iter().any(|s| s.contains("auto_approved")));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bench_channel_captures_tool_events() {
|
||||
let (channel, _tx) = BenchChannel::new();
|
||||
|
||||
let status = StatusUpdate::ToolCompleted {
|
||||
name: "echo".to_string(),
|
||||
success: true,
|
||||
};
|
||||
channel
|
||||
.send_status(status, &serde_json::Value::Null)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let capture_arc = channel.capture();
|
||||
let cap = capture_arc.lock().await;
|
||||
assert_eq!(cap.tool_calls.len(), 1);
|
||||
assert_eq!(cap.tool_calls[0].name, "echo");
|
||||
assert!(cap.tool_calls[0].success);
|
||||
}
|
||||
}
|
||||
@@ -1,203 +0,0 @@
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::Duration;
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::error::BenchError;
|
||||
|
||||
/// Top-level bench configuration, loaded from TOML.
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct BenchConfig {
|
||||
/// Where to write results. Default: "./bench-results".
|
||||
#[serde(default = "default_results_dir")]
|
||||
pub results_dir: PathBuf,
|
||||
|
||||
/// Per-task timeout. Default: "300s".
|
||||
#[serde(
|
||||
default = "default_task_timeout",
|
||||
deserialize_with = "deserialize_duration"
|
||||
)]
|
||||
pub task_timeout: Duration,
|
||||
|
||||
/// Max agent iterations per task. Default: 15.
|
||||
#[serde(default = "default_max_iterations")]
|
||||
pub max_iterations: usize,
|
||||
|
||||
/// How many tasks to run in parallel. Default: 1.
|
||||
#[serde(default = "default_parallelism")]
|
||||
pub parallelism: usize,
|
||||
|
||||
/// Model/config matrix entries. At least one required.
|
||||
#[serde(default)]
|
||||
pub matrix: Vec<MatrixEntry>,
|
||||
|
||||
/// Suite-specific configuration (passed through to adapter).
|
||||
#[serde(default = "default_suite_config")]
|
||||
pub suite_config: toml::Value,
|
||||
}
|
||||
|
||||
/// A single model/config combination to benchmark.
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct MatrixEntry {
|
||||
/// Label for this configuration (used in results).
|
||||
pub label: String,
|
||||
|
||||
/// Model identifier.
|
||||
#[serde(default)]
|
||||
pub model: Option<String>,
|
||||
|
||||
/// Optional tool allowlist. If set, only these tools are available.
|
||||
#[serde(default)]
|
||||
pub tools: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
impl BenchConfig {
|
||||
/// Load from a TOML file.
|
||||
pub fn from_file(path: &Path) -> Result<Self, BenchError> {
|
||||
if !path.exists() {
|
||||
return Err(BenchError::ConfigNotFound {
|
||||
path: path.to_path_buf(),
|
||||
});
|
||||
}
|
||||
let content = std::fs::read_to_string(path)?;
|
||||
let config: BenchConfig = toml::from_str(&content)?;
|
||||
Ok(config)
|
||||
}
|
||||
|
||||
/// Create a minimal config for when no config file is provided.
|
||||
/// Uses defaults and optional CLI overrides.
|
||||
pub fn minimal(model: Option<String>) -> Self {
|
||||
let label = model.as_deref().unwrap_or("default").to_string();
|
||||
Self {
|
||||
results_dir: default_results_dir(),
|
||||
task_timeout: default_task_timeout(),
|
||||
max_iterations: default_max_iterations(),
|
||||
parallelism: default_parallelism(),
|
||||
matrix: vec![MatrixEntry {
|
||||
label,
|
||||
model,
|
||||
tools: None,
|
||||
}],
|
||||
suite_config: toml::Value::Table(toml::map::Map::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the suite_config as a generic map for adapter use.
|
||||
pub fn suite_config_map(&self) -> toml::map::Map<String, toml::Value> {
|
||||
match &self.suite_config {
|
||||
toml::Value::Table(map) => map.clone(),
|
||||
_ => toml::map::Map::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get a string value from suite_config.
|
||||
pub fn suite_config_str(&self, key: &str) -> Option<String> {
|
||||
self.suite_config_map()
|
||||
.get(key)
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
fn default_suite_config() -> toml::Value {
|
||||
toml::Value::Table(toml::map::Map::new())
|
||||
}
|
||||
|
||||
fn default_results_dir() -> PathBuf {
|
||||
PathBuf::from("./bench-results")
|
||||
}
|
||||
|
||||
fn default_task_timeout() -> Duration {
|
||||
Duration::from_secs(300)
|
||||
}
|
||||
|
||||
fn default_max_iterations() -> usize {
|
||||
15
|
||||
}
|
||||
|
||||
fn default_parallelism() -> usize {
|
||||
1
|
||||
}
|
||||
|
||||
/// Deserialize a duration from a string like "300s", "5m", etc.
|
||||
fn deserialize_duration<'de, D>(deserializer: D) -> Result<Duration, D::Error>
|
||||
where
|
||||
D: serde::Deserializer<'de>,
|
||||
{
|
||||
let s = String::deserialize(deserializer)?;
|
||||
parse_duration(&s).map_err(serde::de::Error::custom)
|
||||
}
|
||||
|
||||
fn parse_duration(s: &str) -> Result<Duration, String> {
|
||||
let s = s.trim();
|
||||
if let Some(secs) = s.strip_suffix('s') {
|
||||
secs.trim()
|
||||
.parse::<u64>()
|
||||
.map(Duration::from_secs)
|
||||
.map_err(|e| format!("invalid seconds: {e}"))
|
||||
} else if let Some(mins) = s.strip_suffix('m') {
|
||||
mins.trim()
|
||||
.parse::<u64>()
|
||||
.map(|m| Duration::from_secs(m * 60))
|
||||
.map_err(|e| format!("invalid minutes: {e}"))
|
||||
} else {
|
||||
// Assume seconds if no suffix
|
||||
s.parse::<u64>()
|
||||
.map(Duration::from_secs)
|
||||
.map_err(|e| format!("invalid duration '{s}': {e}"))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_parse_duration() {
|
||||
assert_eq!(parse_duration("300s").unwrap(), Duration::from_secs(300));
|
||||
assert_eq!(parse_duration("5m").unwrap(), Duration::from_secs(300));
|
||||
assert_eq!(parse_duration("60").unwrap(), Duration::from_secs(60));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_minimal_config() {
|
||||
let config = BenchConfig::minimal(Some("test-model".to_string()));
|
||||
assert_eq!(config.matrix.len(), 1);
|
||||
assert_eq!(config.matrix[0].label, "test-model");
|
||||
assert_eq!(config.max_iterations, 15);
|
||||
assert_eq!(config.parallelism, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_config_from_toml() {
|
||||
let toml_str = r#"
|
||||
results_dir = "./my-results"
|
||||
task_timeout = "60s"
|
||||
max_iterations = 10
|
||||
parallelism = 2
|
||||
|
||||
[[matrix]]
|
||||
label = "fast"
|
||||
model = "gpt-4o-mini"
|
||||
|
||||
[[matrix]]
|
||||
label = "full"
|
||||
model = "claude-3-5-sonnet"
|
||||
tools = ["echo", "time"]
|
||||
|
||||
[suite_config]
|
||||
dataset_path = "./data/test.jsonl"
|
||||
"#;
|
||||
let config: BenchConfig = toml::from_str(toml_str).unwrap();
|
||||
assert_eq!(config.results_dir, PathBuf::from("./my-results"));
|
||||
assert_eq!(config.task_timeout, Duration::from_secs(60));
|
||||
assert_eq!(config.max_iterations, 10);
|
||||
assert_eq!(config.parallelism, 2);
|
||||
assert_eq!(config.matrix.len(), 2);
|
||||
assert_eq!(config.matrix[1].tools.as_ref().unwrap().len(), 2);
|
||||
assert_eq!(
|
||||
config.suite_config_str("dataset_path").unwrap(),
|
||||
"./data/test.jsonl"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,40 +0,0 @@
|
||||
use std::path::PathBuf;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum BenchError {
|
||||
#[error("Config error: {0}")]
|
||||
Config(String),
|
||||
|
||||
#[error("Config file not found: {path}")]
|
||||
ConfigNotFound { path: PathBuf },
|
||||
|
||||
#[error("Suite {name} not found. Available: {available}")]
|
||||
SuiteNotFound { name: String, available: String },
|
||||
|
||||
#[error("Task {task_id} failed: {reason}")]
|
||||
TaskFailed { task_id: String, reason: String },
|
||||
|
||||
#[error("Timeout after {seconds}s for task {task_id}")]
|
||||
Timeout { task_id: String, seconds: u64 },
|
||||
|
||||
#[error("Scoring error for task {task_id}: {reason}")]
|
||||
Scoring { task_id: String, reason: String },
|
||||
|
||||
#[error("IO error: {0}")]
|
||||
Io(#[from] std::io::Error),
|
||||
|
||||
#[error("JSON error: {0}")]
|
||||
Json(#[from] serde_json::Error),
|
||||
|
||||
#[error("TOML parse error: {0}")]
|
||||
Toml(#[from] toml::de::Error),
|
||||
|
||||
#[error("Agent error: {0}")]
|
||||
Agent(#[from] ironclaw::Error),
|
||||
|
||||
#[error("Results directory error: {0}")]
|
||||
ResultsDir(String),
|
||||
|
||||
#[error("Resume failed: no completed tasks found in {path}")]
|
||||
ResumeEmpty { path: PathBuf },
|
||||
}
|
||||
@@ -1,248 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU32, Ordering};
|
||||
use std::time::Instant;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use rust_decimal::Decimal;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use ironclaw::error::LlmError;
|
||||
use ironclaw::llm::{
|
||||
CompletionRequest, CompletionResponse, LlmProvider, ToolCompletionRequest,
|
||||
ToolCompletionResponse,
|
||||
};
|
||||
|
||||
/// Recorded metrics from a single LLM call.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct LlmCallRecord {
|
||||
pub input_tokens: u32,
|
||||
pub output_tokens: u32,
|
||||
pub duration_ms: u64,
|
||||
pub had_tool_calls: bool,
|
||||
}
|
||||
|
||||
/// Wraps an `LlmProvider` to record per-call metrics.
|
||||
///
|
||||
/// The wrapper is transparent to the agent: it delegates every call
|
||||
/// to the inner provider and captures token counts and timings.
|
||||
pub struct InstrumentedLlm {
|
||||
inner: Arc<dyn LlmProvider>,
|
||||
records: Mutex<Vec<LlmCallRecord>>,
|
||||
total_input_tokens: AtomicU32,
|
||||
total_output_tokens: AtomicU32,
|
||||
call_count: AtomicU32,
|
||||
}
|
||||
|
||||
impl InstrumentedLlm {
|
||||
pub fn new(inner: Arc<dyn LlmProvider>) -> Self {
|
||||
Self {
|
||||
inner,
|
||||
records: Mutex::new(Vec::new()),
|
||||
total_input_tokens: AtomicU32::new(0),
|
||||
total_output_tokens: AtomicU32::new(0),
|
||||
call_count: AtomicU32::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
/// Take all recorded call metrics, clearing the internal buffer.
|
||||
pub async fn take_records(&self) -> Vec<LlmCallRecord> {
|
||||
let mut records = self.records.lock().await;
|
||||
std::mem::take(&mut *records)
|
||||
}
|
||||
|
||||
/// Snapshot of total tokens without clearing.
|
||||
pub fn total_input_tokens(&self) -> u32 {
|
||||
self.total_input_tokens.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
pub fn total_output_tokens(&self) -> u32 {
|
||||
self.total_output_tokens.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
pub fn call_count(&self) -> u32 {
|
||||
self.call_count.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Estimated cost using the inner provider's cost-per-token rates.
|
||||
pub fn estimated_cost(&self) -> f64 {
|
||||
let (input_rate, output_rate) = self.inner.cost_per_token();
|
||||
let input_cost =
|
||||
input_rate * Decimal::from(self.total_input_tokens.load(Ordering::Relaxed));
|
||||
let output_cost =
|
||||
output_rate * Decimal::from(self.total_output_tokens.load(Ordering::Relaxed));
|
||||
let total = input_cost + output_cost;
|
||||
// Convert Decimal to f64 for the trace (benchmarks don't need exact precision)
|
||||
total.to_string().parse::<f64>().unwrap_or(0.0)
|
||||
}
|
||||
|
||||
/// Reset all counters and records.
|
||||
pub async fn reset(&self) {
|
||||
self.records.lock().await.clear();
|
||||
self.total_input_tokens.store(0, Ordering::Relaxed);
|
||||
self.total_output_tokens.store(0, Ordering::Relaxed);
|
||||
self.call_count.store(0, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
async fn record(
|
||||
&self,
|
||||
input_tokens: u32,
|
||||
output_tokens: u32,
|
||||
duration_ms: u64,
|
||||
had_tool_calls: bool,
|
||||
) {
|
||||
self.total_input_tokens
|
||||
.fetch_add(input_tokens, Ordering::Relaxed);
|
||||
self.total_output_tokens
|
||||
.fetch_add(output_tokens, Ordering::Relaxed);
|
||||
self.call_count.fetch_add(1, Ordering::Relaxed);
|
||||
self.records.lock().await.push(LlmCallRecord {
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
duration_ms,
|
||||
had_tool_calls,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LlmProvider for InstrumentedLlm {
|
||||
fn model_name(&self) -> &str {
|
||||
self.inner.model_name()
|
||||
}
|
||||
|
||||
fn cost_per_token(&self) -> (Decimal, Decimal) {
|
||||
self.inner.cost_per_token()
|
||||
}
|
||||
|
||||
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||
let start = Instant::now();
|
||||
let response = self.inner.complete(request).await?;
|
||||
let elapsed = start.elapsed().as_millis() as u64;
|
||||
self.record(
|
||||
response.input_tokens,
|
||||
response.output_tokens,
|
||||
elapsed,
|
||||
false,
|
||||
)
|
||||
.await;
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
async fn complete_with_tools(
|
||||
&self,
|
||||
request: ToolCompletionRequest,
|
||||
) -> Result<ToolCompletionResponse, LlmError> {
|
||||
let start = Instant::now();
|
||||
let response = self.inner.complete_with_tools(request).await?;
|
||||
let elapsed = start.elapsed().as_millis() as u64;
|
||||
let had_tool_calls = !response.tool_calls.is_empty();
|
||||
self.record(
|
||||
response.input_tokens,
|
||||
response.output_tokens,
|
||||
elapsed,
|
||||
had_tool_calls,
|
||||
)
|
||||
.await;
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
async fn list_models(&self) -> Result<Vec<String>, LlmError> {
|
||||
self.inner.list_models().await
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use ironclaw::llm::{ChatMessage, CompletionRequest, CompletionResponse, FinishReason};
|
||||
|
||||
/// Fake LLM that returns a canned response with known token counts.
|
||||
struct FakeLlm;
|
||||
|
||||
#[async_trait]
|
||||
impl LlmProvider for FakeLlm {
|
||||
fn model_name(&self) -> &str {
|
||||
"fake-model"
|
||||
}
|
||||
|
||||
fn cost_per_token(&self) -> (Decimal, Decimal) {
|
||||
(
|
||||
Decimal::new(3, 6), // $0.000003 per input token
|
||||
Decimal::new(15, 6), // $0.000015 per output token
|
||||
)
|
||||
}
|
||||
|
||||
async fn complete(
|
||||
&self,
|
||||
_request: CompletionRequest,
|
||||
) -> Result<CompletionResponse, LlmError> {
|
||||
Ok(CompletionResponse {
|
||||
content: "test response".to_string(),
|
||||
input_tokens: 100,
|
||||
output_tokens: 50,
|
||||
finish_reason: FinishReason::Stop,
|
||||
})
|
||||
}
|
||||
|
||||
async fn complete_with_tools(
|
||||
&self,
|
||||
_request: ToolCompletionRequest,
|
||||
) -> Result<ToolCompletionResponse, LlmError> {
|
||||
Ok(ToolCompletionResponse {
|
||||
content: Some("tool response".to_string()),
|
||||
tool_calls: vec![],
|
||||
input_tokens: 200,
|
||||
output_tokens: 100,
|
||||
finish_reason: FinishReason::Stop,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_instrumented_records_metrics() {
|
||||
let inner = Arc::new(FakeLlm);
|
||||
let instrumented = InstrumentedLlm::new(inner);
|
||||
|
||||
let request = CompletionRequest::new(vec![ChatMessage::user("hello")]);
|
||||
let _ = instrumented.complete(request).await.unwrap();
|
||||
|
||||
assert_eq!(instrumented.call_count(), 1);
|
||||
assert_eq!(instrumented.total_input_tokens(), 100);
|
||||
assert_eq!(instrumented.total_output_tokens(), 50);
|
||||
|
||||
let records = instrumented.take_records().await;
|
||||
assert_eq!(records.len(), 1);
|
||||
assert_eq!(records[0].input_tokens, 100);
|
||||
assert!(!records[0].had_tool_calls);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_instrumented_cost_calculation() {
|
||||
let inner = Arc::new(FakeLlm);
|
||||
let instrumented = InstrumentedLlm::new(inner);
|
||||
|
||||
let request = CompletionRequest::new(vec![ChatMessage::user("hello")]);
|
||||
let _ = instrumented.complete(request).await.unwrap();
|
||||
|
||||
// 100 * 0.000003 + 50 * 0.000015 = 0.0003 + 0.00075 = 0.00105
|
||||
let cost = instrumented.estimated_cost();
|
||||
assert!((cost - 0.00105).abs() < 0.0001);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_instrumented_reset() {
|
||||
let inner = Arc::new(FakeLlm);
|
||||
let instrumented = InstrumentedLlm::new(inner);
|
||||
|
||||
let request = CompletionRequest::new(vec![ChatMessage::user("hello")]);
|
||||
let _ = instrumented.complete(request).await.unwrap();
|
||||
assert_eq!(instrumented.call_count(), 1);
|
||||
|
||||
instrumented.reset().await;
|
||||
assert_eq!(instrumented.call_count(), 0);
|
||||
assert_eq!(instrumented.total_input_tokens(), 0);
|
||||
|
||||
let records = instrumented.take_records().await;
|
||||
assert!(records.is_empty());
|
||||
}
|
||||
}
|
||||
@@ -1,314 +0,0 @@
|
||||
mod adapters;
|
||||
mod channel;
|
||||
mod config;
|
||||
mod error;
|
||||
mod instrumented_llm;
|
||||
mod results;
|
||||
mod runner;
|
||||
mod scoring;
|
||||
mod suite;
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use clap::{Parser, Subcommand};
|
||||
use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::config::BenchConfig;
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(name = "ironclaw-bench", about = "IronClaw benchmarking harness")]
|
||||
struct Cli {
|
||||
#[command(subcommand)]
|
||||
command: Commands,
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
enum Commands {
|
||||
/// Run a benchmark suite.
|
||||
Run {
|
||||
/// Suite to run (custom, gaia, tau_bench, swe_bench).
|
||||
#[arg(long)]
|
||||
suite: String,
|
||||
|
||||
/// Path to bench config TOML.
|
||||
#[arg(long)]
|
||||
config: Option<PathBuf>,
|
||||
|
||||
/// Override model for all matrix entries.
|
||||
#[arg(long)]
|
||||
model: Option<String>,
|
||||
|
||||
/// Max tasks to run in parallel.
|
||||
#[arg(long)]
|
||||
parallelism: Option<usize>,
|
||||
|
||||
/// Sample N tasks from the suite (for quick testing).
|
||||
#[arg(long)]
|
||||
sample: Option<usize>,
|
||||
|
||||
/// Only run these task IDs (comma-separated).
|
||||
#[arg(long, value_delimiter = ',')]
|
||||
task_ids: Option<Vec<String>>,
|
||||
|
||||
/// Only run tasks with these tags (comma-separated).
|
||||
#[arg(long, value_delimiter = ',')]
|
||||
tags: Option<Vec<String>>,
|
||||
|
||||
/// Per-task timeout in seconds.
|
||||
#[arg(long)]
|
||||
timeout_secs: Option<u64>,
|
||||
|
||||
/// Override results directory.
|
||||
#[arg(long)]
|
||||
results_dir: Option<PathBuf>,
|
||||
|
||||
/// Resume a previous run by ID.
|
||||
#[arg(long)]
|
||||
resume: Option<Uuid>,
|
||||
},
|
||||
|
||||
/// Show results for a run.
|
||||
Results {
|
||||
/// Run ID or "latest".
|
||||
#[arg(default_value = "latest")]
|
||||
run_id: String,
|
||||
|
||||
/// Output format.
|
||||
#[arg(long, default_value = "table")]
|
||||
format: ResultsFormat,
|
||||
|
||||
/// Override results directory.
|
||||
#[arg(long)]
|
||||
results_dir: Option<PathBuf>,
|
||||
},
|
||||
|
||||
/// Compare two runs.
|
||||
Compare {
|
||||
/// Baseline run ID.
|
||||
baseline: Uuid,
|
||||
|
||||
/// Comparison run ID.
|
||||
comparison: Uuid,
|
||||
|
||||
/// Override results directory.
|
||||
#[arg(long)]
|
||||
results_dir: Option<PathBuf>,
|
||||
},
|
||||
|
||||
/// List available benchmark suites.
|
||||
List,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, clap::ValueEnum)]
|
||||
enum ResultsFormat {
|
||||
Table,
|
||||
Json,
|
||||
Csv,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
let cli = Cli::parse();
|
||||
|
||||
tracing_subscriber::registry()
|
||||
.with(
|
||||
EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| EnvFilter::new("ironclaw_bench=info,ironclaw=warn")),
|
||||
)
|
||||
.with(tracing_subscriber::fmt::layer().with_target(false))
|
||||
.init();
|
||||
|
||||
match cli.command {
|
||||
Commands::List => {
|
||||
println!("Available benchmark suites:\n");
|
||||
for (id, desc) in adapters::KNOWN_SUITES {
|
||||
println!(" {:<15} {}", id, desc);
|
||||
}
|
||||
println!();
|
||||
}
|
||||
Commands::Run {
|
||||
suite,
|
||||
config: config_path,
|
||||
model,
|
||||
parallelism,
|
||||
sample,
|
||||
task_ids,
|
||||
tags,
|
||||
timeout_secs,
|
||||
results_dir,
|
||||
resume,
|
||||
} => {
|
||||
// Load or create config
|
||||
let mut bench_config = if let Some(ref path) = config_path {
|
||||
BenchConfig::from_file(path)?
|
||||
} else {
|
||||
BenchConfig::minimal(model.clone())
|
||||
};
|
||||
|
||||
// Apply CLI overrides
|
||||
if let Some(p) = parallelism {
|
||||
bench_config.parallelism = p;
|
||||
}
|
||||
if let Some(t) = timeout_secs {
|
||||
bench_config.task_timeout = std::time::Duration::from_secs(t);
|
||||
}
|
||||
if let Some(ref dir) = results_dir {
|
||||
bench_config.results_dir = dir.clone();
|
||||
}
|
||||
|
||||
// If model override specified and we have matrix entries, update them
|
||||
if let Some(ref m) = model {
|
||||
for entry in &mut bench_config.matrix {
|
||||
entry.model = Some(m.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// Create suite
|
||||
let bench_suite = adapters::create_suite(&suite, &bench_config)?;
|
||||
|
||||
// Initialize ironclaw LLM provider
|
||||
let ironclaw_config = ironclaw::Config::from_env().map_err(|e| {
|
||||
anyhow::anyhow!(
|
||||
"Failed to load ironclaw config: {}. Make sure .env is configured.",
|
||||
e
|
||||
)
|
||||
})?;
|
||||
|
||||
let session = ironclaw::llm::create_session_manager(ironclaw::llm::SessionConfig {
|
||||
auth_base_url: ironclaw_config.llm.nearai.auth_base_url.clone(),
|
||||
session_path: ironclaw_config.llm.nearai.session_path.clone(),
|
||||
..Default::default()
|
||||
})
|
||||
.await;
|
||||
session.ensure_authenticated().await?;
|
||||
|
||||
let llm = ironclaw::llm::create_llm_provider(&ironclaw_config.llm, session)?;
|
||||
let safety = Arc::new(ironclaw::safety::SafetyLayer::new(&ironclaw_config.safety));
|
||||
|
||||
let runner = runner::BenchRunner::new(bench_suite, bench_config.clone(), llm, safety);
|
||||
|
||||
// Run for each matrix entry
|
||||
for matrix_entry in &bench_config.matrix {
|
||||
let run_id = runner
|
||||
.run(
|
||||
matrix_entry,
|
||||
sample,
|
||||
task_ids.as_deref(),
|
||||
tags.as_deref(),
|
||||
resume,
|
||||
)
|
||||
.await?;
|
||||
println!("Run complete: {}", run_id);
|
||||
}
|
||||
}
|
||||
Commands::Results {
|
||||
run_id,
|
||||
format,
|
||||
results_dir,
|
||||
} => {
|
||||
let base = results_dir.unwrap_or_else(|| PathBuf::from("./bench-results"));
|
||||
let uuid = if run_id == "latest" {
|
||||
results::find_latest_run(&base)?
|
||||
.ok_or_else(|| anyhow::anyhow!("No runs found in {}", base.display()))?
|
||||
} else {
|
||||
Uuid::parse_str(&run_id)?
|
||||
};
|
||||
|
||||
let json_path = results::run_json_path(&base, uuid);
|
||||
let jsonl_path = results::tasks_jsonl_path(&base, uuid);
|
||||
|
||||
let run = results::read_run_result(&json_path)?;
|
||||
let tasks = results::read_task_results(&jsonl_path)?;
|
||||
|
||||
match format {
|
||||
ResultsFormat::Table => {
|
||||
results::print_results_table(&tasks, &run);
|
||||
}
|
||||
ResultsFormat::Json => {
|
||||
let output = serde_json::json!({
|
||||
"run": run,
|
||||
"tasks": tasks,
|
||||
});
|
||||
println!("{}", serde_json::to_string_pretty(&output)?);
|
||||
}
|
||||
ResultsFormat::Csv => {
|
||||
println!("task_id,score,label,tokens,cost,turns,time_s");
|
||||
for task in &tasks {
|
||||
println!(
|
||||
"{},{:.3},{},{},{:.4},{},{:.1}",
|
||||
task.task_id,
|
||||
task.score.value,
|
||||
task.score.label,
|
||||
task.trace.input_tokens + task.trace.output_tokens,
|
||||
task.trace.estimated_cost_usd,
|
||||
task.trace.turns,
|
||||
task.trace.wall_time_ms as f64 / 1000.0,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Commands::Compare {
|
||||
baseline,
|
||||
comparison,
|
||||
results_dir,
|
||||
} => {
|
||||
let base = results_dir.unwrap_or_else(|| PathBuf::from("./bench-results"));
|
||||
|
||||
let baseline_run = results::read_run_result(&results::run_json_path(&base, baseline))?;
|
||||
let comparison_run =
|
||||
results::read_run_result(&results::run_json_path(&base, comparison))?;
|
||||
|
||||
println!("\nComparison: {} vs {}\n", baseline, comparison);
|
||||
println!(
|
||||
"{:<20} {:>12} {:>12} {:>10}",
|
||||
"Metric", "Baseline", "Comparison", "Delta"
|
||||
);
|
||||
println!("{}", "-".repeat(58));
|
||||
|
||||
let pass_delta = comparison_run.pass_rate - baseline_run.pass_rate;
|
||||
println!(
|
||||
"{:<20} {:>11.1}% {:>11.1}% {:>+9.1}%",
|
||||
"Pass rate",
|
||||
baseline_run.pass_rate * 100.0,
|
||||
comparison_run.pass_rate * 100.0,
|
||||
pass_delta * 100.0,
|
||||
);
|
||||
|
||||
let score_delta = comparison_run.avg_score - baseline_run.avg_score;
|
||||
println!(
|
||||
"{:<20} {:>12.3} {:>12.3} {:>+10.3}",
|
||||
"Avg score", baseline_run.avg_score, comparison_run.avg_score, score_delta,
|
||||
);
|
||||
|
||||
let cost_delta = comparison_run.total_cost_usd - baseline_run.total_cost_usd;
|
||||
println!(
|
||||
"{:<20} {:>11.4}$ {:>11.4}$ {:>+9.4}$",
|
||||
"Total cost",
|
||||
baseline_run.total_cost_usd,
|
||||
comparison_run.total_cost_usd,
|
||||
cost_delta,
|
||||
);
|
||||
|
||||
let time_b = baseline_run.total_wall_time_ms as f64 / 1000.0;
|
||||
let time_c = comparison_run.total_wall_time_ms as f64 / 1000.0;
|
||||
println!(
|
||||
"{:<20} {:>11.1}s {:>11.1}s {:>+9.1}s",
|
||||
"Total time",
|
||||
time_b,
|
||||
time_c,
|
||||
time_c - time_b,
|
||||
);
|
||||
|
||||
println!(
|
||||
"{:<20} {:>12} {:>12}",
|
||||
"Model", baseline_run.model, comparison_run.model,
|
||||
);
|
||||
println!();
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,392 +0,0 @@
|
||||
use std::collections::HashSet;
|
||||
use std::io::{BufRead, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::Duration;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::error::BenchError;
|
||||
use crate::suite::BenchScore;
|
||||
|
||||
/// Metrics from a single task run: LLM usage, timing, tool calls.
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct Trace {
|
||||
pub wall_time_ms: u64,
|
||||
pub llm_calls: u32,
|
||||
pub input_tokens: u32,
|
||||
pub output_tokens: u32,
|
||||
pub estimated_cost_usd: f64,
|
||||
pub tool_calls: Vec<TraceToolCall>,
|
||||
pub turns: u32,
|
||||
pub hit_iteration_limit: bool,
|
||||
pub hit_timeout: bool,
|
||||
}
|
||||
|
||||
impl Trace {
|
||||
pub fn wall_time(&self) -> Duration {
|
||||
Duration::from_millis(self.wall_time_ms)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct TraceToolCall {
|
||||
pub name: String,
|
||||
pub duration_ms: u64,
|
||||
pub success: bool,
|
||||
}
|
||||
|
||||
/// Result of running a single benchmark task.
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct TaskResult {
|
||||
pub task_id: String,
|
||||
pub suite_id: String,
|
||||
pub score: BenchScore,
|
||||
pub trace: Trace,
|
||||
pub response: String,
|
||||
pub started_at: DateTime<Utc>,
|
||||
pub finished_at: DateTime<Utc>,
|
||||
pub config_label: String,
|
||||
#[serde(default)]
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
/// Aggregate results for a full benchmark run.
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct RunResult {
|
||||
pub run_id: Uuid,
|
||||
pub suite_id: String,
|
||||
pub config_label: String,
|
||||
pub model: String,
|
||||
pub pass_rate: f64,
|
||||
pub avg_score: f64,
|
||||
pub total_tasks: usize,
|
||||
pub completed_tasks: usize,
|
||||
pub total_cost_usd: f64,
|
||||
pub total_wall_time_ms: u64,
|
||||
pub started_at: DateTime<Utc>,
|
||||
pub finished_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
impl RunResult {
|
||||
/// Build aggregate from individual task results.
|
||||
pub fn from_tasks(
|
||||
run_id: Uuid,
|
||||
suite_id: &str,
|
||||
config_label: &str,
|
||||
model: &str,
|
||||
total_tasks: usize,
|
||||
tasks: &[TaskResult],
|
||||
started_at: DateTime<Utc>,
|
||||
) -> Self {
|
||||
let pass_count = tasks.iter().filter(|t| t.score.value >= 1.0).count();
|
||||
let pass_rate = if tasks.is_empty() {
|
||||
0.0
|
||||
} else {
|
||||
pass_count as f64 / tasks.len() as f64
|
||||
};
|
||||
let avg_score = if tasks.is_empty() {
|
||||
0.0
|
||||
} else {
|
||||
tasks.iter().map(|t| t.score.value).sum::<f64>() / tasks.len() as f64
|
||||
};
|
||||
let total_cost: f64 = tasks.iter().map(|t| t.trace.estimated_cost_usd).sum();
|
||||
let total_wall: u64 = tasks.iter().map(|t| t.trace.wall_time_ms).sum();
|
||||
|
||||
Self {
|
||||
run_id,
|
||||
suite_id: suite_id.to_string(),
|
||||
config_label: config_label.to_string(),
|
||||
model: model.to_string(),
|
||||
pass_rate,
|
||||
avg_score,
|
||||
total_tasks,
|
||||
completed_tasks: tasks.len(),
|
||||
total_cost_usd: total_cost,
|
||||
total_wall_time_ms: total_wall,
|
||||
started_at,
|
||||
finished_at: Utc::now(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn total_wall_time(&self) -> Duration {
|
||||
Duration::from_millis(self.total_wall_time_ms)
|
||||
}
|
||||
}
|
||||
|
||||
/// Append a single task result as one JSON line to the JSONL file.
|
||||
pub fn append_task_result(path: &Path, result: &TaskResult) -> Result<(), BenchError> {
|
||||
let mut file = std::fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(path)?;
|
||||
let line = serde_json::to_string(result)?;
|
||||
writeln!(file, "{line}")?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Read all task results from a JSONL file.
|
||||
pub fn read_task_results(path: &Path) -> Result<Vec<TaskResult>, BenchError> {
|
||||
if !path.exists() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let file = std::fs::File::open(path)?;
|
||||
let reader = std::io::BufReader::new(file);
|
||||
let mut results = Vec::new();
|
||||
for line in reader.lines() {
|
||||
let line = line?;
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let result: TaskResult = serde_json::from_str(trimmed)?;
|
||||
results.push(result);
|
||||
}
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
/// Write the aggregate run result as JSON.
|
||||
pub fn write_run_result(path: &Path, result: &RunResult) -> Result<(), BenchError> {
|
||||
let json = serde_json::to_string_pretty(result)?;
|
||||
std::fs::write(path, json)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Read the aggregate run result from JSON.
|
||||
pub fn read_run_result(path: &Path) -> Result<RunResult, BenchError> {
|
||||
let json = std::fs::read_to_string(path)?;
|
||||
let result: RunResult = serde_json::from_str(&json)?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Get the set of already-completed task IDs from a JSONL file (for resume).
|
||||
pub fn completed_task_ids(path: &Path) -> Result<HashSet<String>, BenchError> {
|
||||
let results = read_task_results(path)?;
|
||||
Ok(results.into_iter().map(|r| r.task_id).collect())
|
||||
}
|
||||
|
||||
/// Get the results directory for a specific run.
|
||||
pub fn run_dir(base: &Path, run_id: Uuid) -> PathBuf {
|
||||
base.join(run_id.to_string())
|
||||
}
|
||||
|
||||
/// Get the tasks JSONL path for a run.
|
||||
pub fn tasks_jsonl_path(base: &Path, run_id: Uuid) -> PathBuf {
|
||||
run_dir(base, run_id).join("tasks.jsonl")
|
||||
}
|
||||
|
||||
/// Get the run JSON path for a run.
|
||||
pub fn run_json_path(base: &Path, run_id: Uuid) -> PathBuf {
|
||||
run_dir(base, run_id).join("run.json")
|
||||
}
|
||||
|
||||
/// Find the latest run directory (by modification time).
|
||||
pub fn find_latest_run(base: &Path) -> Result<Option<Uuid>, BenchError> {
|
||||
if !base.exists() {
|
||||
return Ok(None);
|
||||
}
|
||||
let mut entries: Vec<_> = std::fs::read_dir(base)?
|
||||
.filter_map(|e| e.ok())
|
||||
.filter(|e| e.file_type().map(|ft| ft.is_dir()).unwrap_or(false))
|
||||
.filter_map(|e| {
|
||||
let name = e.file_name().to_string_lossy().to_string();
|
||||
let uuid = Uuid::parse_str(&name).ok()?;
|
||||
let modified = e.metadata().ok()?.modified().ok()?;
|
||||
Some((uuid, modified))
|
||||
})
|
||||
.collect();
|
||||
entries.sort_by(|a, b| b.1.cmp(&a.1));
|
||||
Ok(entries.first().map(|(uuid, _)| *uuid))
|
||||
}
|
||||
|
||||
/// Print a summary table of task results.
|
||||
pub fn print_results_table(tasks: &[TaskResult], run: &RunResult) {
|
||||
println!();
|
||||
println!(
|
||||
"Run: {} | Suite: {} | Config: {} | Model: {}",
|
||||
run.run_id, run.suite_id, run.config_label, run.model
|
||||
);
|
||||
println!(
|
||||
"Pass rate: {:.1}% | Avg score: {:.3} | Tasks: {}/{} | Cost: ${:.4} | Time: {:.1}s",
|
||||
run.pass_rate * 100.0,
|
||||
run.avg_score,
|
||||
run.completed_tasks,
|
||||
run.total_tasks,
|
||||
run.total_cost_usd,
|
||||
run.total_wall_time_ms as f64 / 1000.0,
|
||||
);
|
||||
println!();
|
||||
|
||||
// Header
|
||||
println!(
|
||||
"{:<30} {:>6} {:>7} {:>8} {:>10} {:>6} {:>8}",
|
||||
"Task ID", "Score", "Label", "Tokens", "Cost", "Turns", "Time"
|
||||
);
|
||||
println!("{}", "-".repeat(80));
|
||||
|
||||
for task in tasks {
|
||||
let total_tokens = task.trace.input_tokens + task.trace.output_tokens;
|
||||
let task_id_display = if task.task_id.len() > 28 {
|
||||
format!("{}...", &task.task_id[..25])
|
||||
} else {
|
||||
task.task_id.clone()
|
||||
};
|
||||
println!(
|
||||
"{:<30} {:>6.3} {:>7} {:>8} {:>10.4} {:>6} {:>7.1}s",
|
||||
task_id_display,
|
||||
task.score.value,
|
||||
task.score.label,
|
||||
total_tokens,
|
||||
task.trace.estimated_cost_usd,
|
||||
task.trace.turns,
|
||||
task.trace.wall_time_ms as f64 / 1000.0,
|
||||
);
|
||||
}
|
||||
println!();
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_run_result_from_tasks() {
|
||||
let tasks = vec![
|
||||
TaskResult {
|
||||
task_id: "t1".to_string(),
|
||||
suite_id: "custom".to_string(),
|
||||
score: BenchScore {
|
||||
value: 1.0,
|
||||
label: "pass".to_string(),
|
||||
details: None,
|
||||
},
|
||||
trace: Trace {
|
||||
wall_time_ms: 1000,
|
||||
llm_calls: 2,
|
||||
input_tokens: 100,
|
||||
output_tokens: 50,
|
||||
estimated_cost_usd: 0.01,
|
||||
tool_calls: vec![],
|
||||
turns: 1,
|
||||
hit_iteration_limit: false,
|
||||
hit_timeout: false,
|
||||
},
|
||||
response: "answer".to_string(),
|
||||
started_at: Utc::now(),
|
||||
finished_at: Utc::now(),
|
||||
config_label: "default".to_string(),
|
||||
error: None,
|
||||
},
|
||||
TaskResult {
|
||||
task_id: "t2".to_string(),
|
||||
suite_id: "custom".to_string(),
|
||||
score: BenchScore {
|
||||
value: 0.0,
|
||||
label: "fail".to_string(),
|
||||
details: Some("wrong".to_string()),
|
||||
},
|
||||
trace: Trace {
|
||||
wall_time_ms: 2000,
|
||||
llm_calls: 3,
|
||||
input_tokens: 200,
|
||||
output_tokens: 100,
|
||||
estimated_cost_usd: 0.02,
|
||||
tool_calls: vec![],
|
||||
turns: 2,
|
||||
hit_iteration_limit: false,
|
||||
hit_timeout: false,
|
||||
},
|
||||
response: "wrong answer".to_string(),
|
||||
started_at: Utc::now(),
|
||||
finished_at: Utc::now(),
|
||||
config_label: "default".to_string(),
|
||||
error: None,
|
||||
},
|
||||
];
|
||||
|
||||
let run = RunResult::from_tasks(
|
||||
Uuid::new_v4(),
|
||||
"custom",
|
||||
"default",
|
||||
"test-model",
|
||||
2,
|
||||
&tasks,
|
||||
Utc::now(),
|
||||
);
|
||||
|
||||
assert_eq!(run.pass_rate, 0.5);
|
||||
assert_eq!(run.avg_score, 0.5);
|
||||
assert_eq!(run.total_tasks, 2);
|
||||
assert_eq!(run.completed_tasks, 2);
|
||||
assert!((run.total_cost_usd - 0.03).abs() < f64::EPSILON);
|
||||
assert_eq!(run.total_wall_time_ms, 3000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_jsonl_roundtrip() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let path = dir.path().join("tasks.jsonl");
|
||||
|
||||
let result = TaskResult {
|
||||
task_id: "round-trip-test".to_string(),
|
||||
suite_id: "custom".to_string(),
|
||||
score: BenchScore::pass(),
|
||||
trace: Trace {
|
||||
wall_time_ms: 500,
|
||||
llm_calls: 1,
|
||||
input_tokens: 10,
|
||||
output_tokens: 5,
|
||||
estimated_cost_usd: 0.001,
|
||||
tool_calls: vec![],
|
||||
turns: 1,
|
||||
hit_iteration_limit: false,
|
||||
hit_timeout: false,
|
||||
},
|
||||
response: "hello".to_string(),
|
||||
started_at: Utc::now(),
|
||||
finished_at: Utc::now(),
|
||||
config_label: "test".to_string(),
|
||||
error: None,
|
||||
};
|
||||
|
||||
append_task_result(&path, &result).expect("append");
|
||||
append_task_result(&path, &result).expect("append");
|
||||
|
||||
let loaded = read_task_results(&path).expect("read");
|
||||
assert_eq!(loaded.len(), 2);
|
||||
assert_eq!(loaded[0].task_id, "round-trip-test");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_completed_task_ids() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let path = dir.path().join("tasks.jsonl");
|
||||
|
||||
let result = TaskResult {
|
||||
task_id: "unique-id-1".to_string(),
|
||||
suite_id: "custom".to_string(),
|
||||
score: BenchScore::pass(),
|
||||
trace: Trace {
|
||||
wall_time_ms: 100,
|
||||
llm_calls: 1,
|
||||
input_tokens: 10,
|
||||
output_tokens: 5,
|
||||
estimated_cost_usd: 0.0,
|
||||
tool_calls: vec![],
|
||||
turns: 1,
|
||||
hit_iteration_limit: false,
|
||||
hit_timeout: false,
|
||||
},
|
||||
response: "x".to_string(),
|
||||
started_at: Utc::now(),
|
||||
finished_at: Utc::now(),
|
||||
config_label: "test".to_string(),
|
||||
error: None,
|
||||
};
|
||||
append_task_result(&path, &result).expect("append");
|
||||
|
||||
let ids = completed_task_ids(&path).expect("ids");
|
||||
assert!(ids.contains("unique-id-1"));
|
||||
assert!(!ids.contains("unique-id-2"));
|
||||
}
|
||||
}
|
||||
@@ -1,460 +0,0 @@
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
use chrono::Utc;
|
||||
use tokio::sync::Mutex;
|
||||
use uuid::Uuid;
|
||||
|
||||
use ironclaw::agent::{Agent, AgentDeps};
|
||||
use ironclaw::channels::{ChannelManager, IncomingMessage};
|
||||
use ironclaw::config::AgentConfig;
|
||||
use ironclaw::llm::LlmProvider;
|
||||
use ironclaw::safety::SafetyLayer;
|
||||
use ironclaw::tools::ToolRegistry;
|
||||
|
||||
use crate::channel::BenchChannel;
|
||||
use crate::config::{BenchConfig, MatrixEntry};
|
||||
use crate::error::BenchError;
|
||||
use crate::instrumented_llm::InstrumentedLlm;
|
||||
use crate::results::{
|
||||
RunResult, TaskResult, Trace, append_task_result, completed_task_ids, run_dir, run_json_path,
|
||||
tasks_jsonl_path, write_run_result,
|
||||
};
|
||||
use crate::suite::{BenchSuite, BenchTask, ConversationTurn, TaskSubmission, TurnRole};
|
||||
|
||||
/// Parameters for running a single task in isolation.
|
||||
struct TaskRunParams<'a> {
|
||||
task: &'a BenchTask,
|
||||
suite_id: &'a str,
|
||||
config_label: &'a str,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
timeout: std::time::Duration,
|
||||
additional_tools: &'a [Arc<dyn ironclaw::tools::Tool>],
|
||||
}
|
||||
|
||||
/// Orchestrates benchmark execution: loads tasks, runs agent per task,
|
||||
/// scores results, writes JSONL output.
|
||||
pub struct BenchRunner {
|
||||
suite: Box<dyn BenchSuite>,
|
||||
config: BenchConfig,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
}
|
||||
|
||||
impl BenchRunner {
|
||||
pub fn new(
|
||||
suite: Box<dyn BenchSuite>,
|
||||
config: BenchConfig,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
) -> Self {
|
||||
Self {
|
||||
suite,
|
||||
config,
|
||||
llm,
|
||||
safety,
|
||||
}
|
||||
}
|
||||
|
||||
/// Run the benchmark for one matrix entry.
|
||||
///
|
||||
/// Returns the run_id for result retrieval.
|
||||
pub async fn run(
|
||||
&self,
|
||||
matrix: &MatrixEntry,
|
||||
sample: Option<usize>,
|
||||
task_filter: Option<&[String]>,
|
||||
tag_filter: Option<&[String]>,
|
||||
resume_run_id: Option<Uuid>,
|
||||
) -> Result<Uuid, BenchError> {
|
||||
let run_id = resume_run_id.unwrap_or_else(Uuid::new_v4);
|
||||
let results_base = &self.config.results_dir;
|
||||
let dir = run_dir(results_base, run_id);
|
||||
std::fs::create_dir_all(&dir)?;
|
||||
|
||||
let jsonl_path = tasks_jsonl_path(results_base, run_id);
|
||||
let json_path = run_json_path(results_base, run_id);
|
||||
|
||||
// Load completed task IDs for resume support
|
||||
let completed: HashSet<String> = if resume_run_id.is_some() {
|
||||
completed_task_ids(&jsonl_path)?
|
||||
} else {
|
||||
HashSet::new()
|
||||
};
|
||||
|
||||
if !completed.is_empty() {
|
||||
tracing::info!(
|
||||
"Resuming run {}: {} tasks already completed",
|
||||
run_id,
|
||||
completed.len()
|
||||
);
|
||||
}
|
||||
|
||||
// Load and filter tasks
|
||||
let mut tasks = self.suite.load_tasks().await?;
|
||||
|
||||
if let Some(ids) = task_filter {
|
||||
let id_set: HashSet<&str> = ids.iter().map(|s| s.as_str()).collect();
|
||||
tasks.retain(|t| id_set.contains(t.id.as_str()));
|
||||
}
|
||||
|
||||
if let Some(tags) = tag_filter {
|
||||
let tag_set: HashSet<&str> = tags.iter().map(|s| s.as_str()).collect();
|
||||
tasks.retain(|t| t.tags.iter().any(|tag| tag_set.contains(tag.as_str())));
|
||||
}
|
||||
|
||||
// Filter out already-completed tasks
|
||||
tasks.retain(|t| !completed.contains(&t.id));
|
||||
|
||||
// Sample if requested
|
||||
if let Some(n) = sample {
|
||||
tasks.truncate(n);
|
||||
}
|
||||
|
||||
let total_tasks = tasks.len() + completed.len();
|
||||
tracing::info!(
|
||||
"Running {} tasks for suite '{}' (config: {}, run: {})",
|
||||
tasks.len(),
|
||||
self.suite.id(),
|
||||
matrix.label,
|
||||
run_id
|
||||
);
|
||||
|
||||
let started_at = Utc::now();
|
||||
let all_results: Arc<Mutex<Vec<TaskResult>>> =
|
||||
Arc::new(Mutex::new(Vec::with_capacity(tasks.len())));
|
||||
|
||||
if self.config.parallelism <= 1 {
|
||||
// Sequential execution
|
||||
let additional_tools = self.suite.additional_tools();
|
||||
for (i, task) in tasks.iter().enumerate() {
|
||||
tracing::info!(
|
||||
"[{}/{}] Running task: {}",
|
||||
i + 1 + completed.len(),
|
||||
total_tasks,
|
||||
task.id
|
||||
);
|
||||
let params = TaskRunParams {
|
||||
task,
|
||||
suite_id: self.suite.id(),
|
||||
config_label: &matrix.label,
|
||||
llm: Arc::clone(&self.llm),
|
||||
safety: Arc::clone(&self.safety),
|
||||
timeout: task.timeout.unwrap_or(self.config.task_timeout),
|
||||
additional_tools: &additional_tools,
|
||||
};
|
||||
let result = run_task_isolated(params).await;
|
||||
append_task_result(&jsonl_path, &result)?;
|
||||
all_results.lock().await.push(result);
|
||||
}
|
||||
} else {
|
||||
// Parallel execution with bounded concurrency
|
||||
let semaphore = Arc::new(tokio::sync::Semaphore::new(self.config.parallelism));
|
||||
|
||||
let mut handles = Vec::new();
|
||||
for (i, task) in tasks.into_iter().enumerate() {
|
||||
let sem = Arc::clone(&semaphore);
|
||||
let suite_id = self.suite.id().to_string();
|
||||
let config_label = matrix.label.clone();
|
||||
let llm = Arc::clone(&self.llm);
|
||||
let safety = Arc::clone(&self.safety);
|
||||
let timeout = task.timeout.unwrap_or(self.config.task_timeout);
|
||||
let results_ref = Arc::clone(&all_results);
|
||||
let jsonl = jsonl_path.clone();
|
||||
let completed_count = completed.len();
|
||||
let total = total_tasks;
|
||||
let additional_tools = self.suite.additional_tools();
|
||||
|
||||
handles.push(tokio::spawn(async move {
|
||||
let _permit = sem.acquire().await.expect("semaphore closed");
|
||||
tracing::info!(
|
||||
"[{}/{}] Running task: {}",
|
||||
i + 1 + completed_count,
|
||||
total,
|
||||
task.id
|
||||
);
|
||||
let params = TaskRunParams {
|
||||
task: &task,
|
||||
suite_id: &suite_id,
|
||||
config_label: &config_label,
|
||||
llm,
|
||||
safety,
|
||||
timeout,
|
||||
additional_tools: &additional_tools,
|
||||
};
|
||||
let result = run_task_isolated(params).await;
|
||||
if let Err(e) = append_task_result(&jsonl, &result) {
|
||||
tracing::error!("Failed to write result for {}: {}", task.id, e);
|
||||
}
|
||||
results_ref.lock().await.push(result);
|
||||
}));
|
||||
}
|
||||
|
||||
for handle in handles {
|
||||
if let Err(e) = handle.await {
|
||||
tracing::error!("Task panicked: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Score all results
|
||||
let results = all_results.lock().await;
|
||||
let mut scored: Vec<TaskResult> = Vec::with_capacity(results.len());
|
||||
for result in results.iter() {
|
||||
let task_opt = self
|
||||
.suite
|
||||
.load_tasks()
|
||||
.await?
|
||||
.into_iter()
|
||||
.find(|t| t.id == result.task_id);
|
||||
|
||||
if let Some(task) = task_opt {
|
||||
let submission = TaskSubmission {
|
||||
response: result.response.clone(),
|
||||
conversation: vec![],
|
||||
tool_calls: result
|
||||
.trace
|
||||
.tool_calls
|
||||
.iter()
|
||||
.map(|tc| tc.name.clone())
|
||||
.collect(),
|
||||
};
|
||||
match self.suite.score(&task, &submission).await {
|
||||
Ok(score) => {
|
||||
let mut scored_result = result.clone();
|
||||
scored_result.score = score;
|
||||
scored.push(scored_result);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Scoring failed for {}: {}", result.task_id, e);
|
||||
scored.push(result.clone());
|
||||
}
|
||||
}
|
||||
} else {
|
||||
scored.push(result.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// Combine with any previously completed results for the aggregate
|
||||
let mut all_for_aggregate = crate::results::read_task_results(&jsonl_path)?;
|
||||
// De-duplicate (prefer the newer scored versions)
|
||||
let scored_ids: HashSet<String> = scored.iter().map(|r| r.task_id.clone()).collect();
|
||||
all_for_aggregate.retain(|r| !scored_ids.contains(&r.task_id));
|
||||
all_for_aggregate.extend(scored);
|
||||
|
||||
let model_name = matrix.model.as_deref().unwrap_or(self.llm.model_name());
|
||||
|
||||
let run_result = RunResult::from_tasks(
|
||||
run_id,
|
||||
self.suite.id(),
|
||||
&matrix.label,
|
||||
model_name,
|
||||
total_tasks,
|
||||
&all_for_aggregate,
|
||||
started_at,
|
||||
);
|
||||
|
||||
write_run_result(&json_path, &run_result)?;
|
||||
|
||||
tracing::info!(
|
||||
"Run {} complete: {:.1}% pass rate, {:.3} avg score, ${:.4} cost",
|
||||
run_id,
|
||||
run_result.pass_rate * 100.0,
|
||||
run_result.avg_score,
|
||||
run_result.total_cost_usd,
|
||||
);
|
||||
|
||||
Ok(run_id)
|
||||
}
|
||||
}
|
||||
|
||||
/// Run a single benchmark task in complete isolation.
|
||||
///
|
||||
/// Creates a fresh Agent + BenchChannel + InstrumentedLlm for the task,
|
||||
/// injects the prompt, waits for the response, and returns the result.
|
||||
async fn run_task_isolated(params: TaskRunParams<'_>) -> TaskResult {
|
||||
let TaskRunParams {
|
||||
task,
|
||||
suite_id,
|
||||
config_label,
|
||||
llm,
|
||||
safety,
|
||||
timeout,
|
||||
additional_tools,
|
||||
} = params;
|
||||
|
||||
let started_at = Utc::now();
|
||||
let start = Instant::now();
|
||||
|
||||
// Wrap LLM with instrumentation
|
||||
let instrumented = Arc::new(InstrumentedLlm::new(llm));
|
||||
|
||||
// Create bench channel
|
||||
let (bench_channel, msg_tx) = BenchChannel::new();
|
||||
let capture = bench_channel.capture();
|
||||
|
||||
// Build tool registry
|
||||
let tools = Arc::new(ToolRegistry::new());
|
||||
tools.register_builtin_tools();
|
||||
|
||||
// Register additional suite-specific tools
|
||||
for tool in additional_tools {
|
||||
tools.register(Arc::clone(tool)).await;
|
||||
}
|
||||
|
||||
// Build agent config (minimal, headless)
|
||||
let agent_config = AgentConfig {
|
||||
name: format!("bench-{}", task.id),
|
||||
max_parallel_jobs: 1,
|
||||
job_timeout: timeout,
|
||||
stuck_threshold: timeout,
|
||||
repair_check_interval: timeout + std::time::Duration::from_secs(999),
|
||||
max_repair_attempts: 0,
|
||||
use_planning: false,
|
||||
session_idle_timeout: timeout,
|
||||
};
|
||||
|
||||
let deps = AgentDeps {
|
||||
store: None,
|
||||
llm: instrumented.clone() as Arc<dyn LlmProvider>,
|
||||
safety,
|
||||
tools,
|
||||
workspace: None,
|
||||
extension_manager: None,
|
||||
};
|
||||
|
||||
let mut channels = ChannelManager::new();
|
||||
channels.add(Box::new(bench_channel));
|
||||
|
||||
let agent = Agent::new(agent_config, deps, channels, None, None, None);
|
||||
|
||||
// Build the full prompt with context
|
||||
let full_prompt = if let Some(ref ctx) = task.context {
|
||||
format!("{}\n\nContext:\n{}", task.prompt, ctx)
|
||||
} else {
|
||||
task.prompt.clone()
|
||||
};
|
||||
|
||||
// Inject the task prompt
|
||||
let incoming = IncomingMessage::new("bench", "bench-user", &full_prompt);
|
||||
if msg_tx.send(incoming).await.is_err() {
|
||||
return make_error_result(
|
||||
task,
|
||||
suite_id,
|
||||
config_label,
|
||||
started_at,
|
||||
"failed to send prompt",
|
||||
);
|
||||
}
|
||||
|
||||
// Record prompt in conversation
|
||||
{
|
||||
let mut cap = capture.lock().await;
|
||||
cap.conversation.push(ConversationTurn {
|
||||
role: TurnRole::User,
|
||||
content: full_prompt,
|
||||
});
|
||||
}
|
||||
|
||||
// Run agent with timeout.
|
||||
// After the first response, send /quit to end the session.
|
||||
let quit_tx = msg_tx.clone();
|
||||
let capture_for_quit = Arc::clone(&capture);
|
||||
let quit_handle = tokio::spawn(async move {
|
||||
// Poll for first response
|
||||
loop {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
let cap = capture_for_quit.lock().await;
|
||||
if !cap.responses.is_empty() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
// Give a small grace period for any final status events
|
||||
tokio::time::sleep(std::time::Duration::from_millis(200)).await;
|
||||
let quit = IncomingMessage::new("bench", "bench-user", "/quit");
|
||||
let _ = quit_tx.send(quit).await;
|
||||
});
|
||||
|
||||
let agent_result = tokio::time::timeout(timeout, agent.run()).await;
|
||||
|
||||
quit_handle.abort();
|
||||
|
||||
let wall_time = start.elapsed();
|
||||
let hit_timeout = agent_result.is_err();
|
||||
|
||||
if let Ok(Err(e)) = &agent_result {
|
||||
tracing::warn!("Agent error for task {}: {}", task.id, e);
|
||||
}
|
||||
|
||||
// Extract results from capture
|
||||
let cap = capture.lock().await;
|
||||
let response = cap.responses.last().cloned().unwrap_or_default();
|
||||
|
||||
let trace = Trace {
|
||||
wall_time_ms: wall_time.as_millis() as u64,
|
||||
llm_calls: instrumented.call_count(),
|
||||
input_tokens: instrumented.total_input_tokens(),
|
||||
output_tokens: instrumented.total_output_tokens(),
|
||||
estimated_cost_usd: instrumented.estimated_cost(),
|
||||
tool_calls: cap.tool_calls.clone(),
|
||||
turns: cap.responses.len() as u32,
|
||||
hit_iteration_limit: false,
|
||||
hit_timeout,
|
||||
};
|
||||
|
||||
let error = if hit_timeout {
|
||||
Some(format!("timeout after {}s", timeout.as_secs()))
|
||||
} else if let Ok(Err(e)) = &agent_result {
|
||||
Some(e.to_string())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
TaskResult {
|
||||
task_id: task.id.clone(),
|
||||
suite_id: suite_id.to_string(),
|
||||
score: crate::suite::BenchScore {
|
||||
value: 0.0,
|
||||
label: "pending".to_string(),
|
||||
details: None,
|
||||
},
|
||||
trace,
|
||||
response,
|
||||
started_at,
|
||||
finished_at: Utc::now(),
|
||||
config_label: config_label.to_string(),
|
||||
error,
|
||||
}
|
||||
}
|
||||
|
||||
fn make_error_result(
|
||||
task: &BenchTask,
|
||||
suite_id: &str,
|
||||
config_label: &str,
|
||||
started_at: chrono::DateTime<Utc>,
|
||||
reason: &str,
|
||||
) -> TaskResult {
|
||||
TaskResult {
|
||||
task_id: task.id.clone(),
|
||||
suite_id: suite_id.to_string(),
|
||||
score: crate::suite::BenchScore::fail(reason),
|
||||
trace: Trace {
|
||||
wall_time_ms: 0,
|
||||
llm_calls: 0,
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
estimated_cost_usd: 0.0,
|
||||
tool_calls: vec![],
|
||||
turns: 0,
|
||||
hit_iteration_limit: false,
|
||||
hit_timeout: false,
|
||||
},
|
||||
response: String::new(),
|
||||
started_at,
|
||||
finished_at: Utc::now(),
|
||||
config_label: config_label.to_string(),
|
||||
error: Some(reason.to_string()),
|
||||
}
|
||||
}
|
||||
@@ -1,113 +0,0 @@
|
||||
use regex::Regex;
|
||||
|
||||
use crate::suite::BenchScore;
|
||||
|
||||
/// Normalize an answer string for comparison: lowercase, trim whitespace,
|
||||
/// strip trailing punctuation, collapse internal whitespace.
|
||||
pub fn normalize_answer(s: &str) -> String {
|
||||
let trimmed = s.trim().to_lowercase();
|
||||
let collapsed: String = trimmed.split_whitespace().collect::<Vec<_>>().join(" ");
|
||||
collapsed.trim_end_matches(['.', ',', ';', '!']).to_string()
|
||||
}
|
||||
|
||||
/// Exact match after normalization.
|
||||
pub fn exact_match(expected: &str, actual: &str) -> BenchScore {
|
||||
let norm_expected = normalize_answer(expected);
|
||||
let norm_actual = normalize_answer(actual);
|
||||
if norm_expected == norm_actual {
|
||||
BenchScore::pass()
|
||||
} else {
|
||||
BenchScore::fail(format!(
|
||||
"expected \"{norm_expected}\", got \"{norm_actual}\""
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if the actual answer contains the expected substring (normalized).
|
||||
pub fn contains_match(expected_substring: &str, actual: &str) -> BenchScore {
|
||||
let norm_expected = normalize_answer(expected_substring);
|
||||
let norm_actual = normalize_answer(actual);
|
||||
if norm_actual.contains(&norm_expected) {
|
||||
BenchScore::pass()
|
||||
} else {
|
||||
BenchScore::fail(format!("response does not contain \"{norm_expected}\""))
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if the actual answer matches a regex pattern.
|
||||
pub fn regex_match(pattern: &str, actual: &str) -> BenchScore {
|
||||
match Regex::new(pattern) {
|
||||
Ok(re) => {
|
||||
if re.is_match(actual) {
|
||||
BenchScore::pass()
|
||||
} else {
|
||||
BenchScore::fail(format!("response does not match pattern /{pattern}/"))
|
||||
}
|
||||
}
|
||||
Err(e) => BenchScore::fail(format!("invalid regex pattern: {e}")),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_normalize_answer() {
|
||||
assert_eq!(normalize_answer(" Hello World. "), "hello world");
|
||||
assert_eq!(normalize_answer("Yes!"), "yes");
|
||||
assert_eq!(normalize_answer("42"), "42");
|
||||
assert_eq!(normalize_answer(" "), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exact_match_pass() {
|
||||
let score = exact_match("Hello World", " hello world. ");
|
||||
assert_eq!(score.value, 1.0);
|
||||
assert_eq!(score.label, "pass");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exact_match_fail() {
|
||||
let score = exact_match("hello", "world");
|
||||
assert_eq!(score.value, 0.0);
|
||||
assert_eq!(score.label, "fail");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_contains_match_pass() {
|
||||
let score = contains_match("world", "Hello World!");
|
||||
assert_eq!(score.value, 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_contains_match_fail() {
|
||||
let score = contains_match("xyz", "Hello World!");
|
||||
assert_eq!(score.value, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_regex_match_pass() {
|
||||
let score = regex_match(r"\d{4}", "The year is 2024.");
|
||||
assert_eq!(score.value, 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_regex_match_fail() {
|
||||
let score = regex_match(r"\d{4}", "No numbers here.");
|
||||
assert_eq!(score.value, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_regex_match_invalid_pattern() {
|
||||
let score = regex_match(r"[invalid", "anything");
|
||||
assert_eq!(score.value, 0.0);
|
||||
assert!(
|
||||
score
|
||||
.details
|
||||
.as_deref()
|
||||
.unwrap_or("")
|
||||
.contains("invalid regex")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,156 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::error::BenchError;
|
||||
|
||||
/// A single task in a benchmark suite.
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct BenchTask {
|
||||
pub id: String,
|
||||
pub prompt: String,
|
||||
#[serde(default)]
|
||||
pub context: Option<String>,
|
||||
#[serde(default)]
|
||||
pub resources: Vec<TaskResource>,
|
||||
#[serde(default)]
|
||||
pub tags: Vec<String>,
|
||||
#[serde(default)]
|
||||
pub expected_turns: Option<usize>,
|
||||
#[serde(default)]
|
||||
pub timeout: Option<Duration>,
|
||||
#[serde(default)]
|
||||
pub metadata: serde_json::Value,
|
||||
}
|
||||
|
||||
/// A resource attached to a benchmark task (file, URL, etc.).
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct TaskResource {
|
||||
pub name: String,
|
||||
pub path: String,
|
||||
#[serde(default)]
|
||||
pub resource_type: ResourceType,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ResourceType {
|
||||
#[default]
|
||||
File,
|
||||
Url,
|
||||
Directory,
|
||||
}
|
||||
|
||||
/// What the agent produced for scoring.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TaskSubmission {
|
||||
pub response: String,
|
||||
pub conversation: Vec<ConversationTurn>,
|
||||
pub tool_calls: Vec<String>,
|
||||
}
|
||||
|
||||
/// A single turn in a multi-turn conversation.
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct ConversationTurn {
|
||||
pub role: TurnRole,
|
||||
pub content: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum TurnRole {
|
||||
User,
|
||||
Assistant,
|
||||
System,
|
||||
}
|
||||
|
||||
/// Score for a single task.
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct BenchScore {
|
||||
/// 0.0 to 1.0 (1.0 = perfect).
|
||||
pub value: f64,
|
||||
/// "pass" / "fail" / "partial".
|
||||
pub label: String,
|
||||
#[serde(default)]
|
||||
pub details: Option<String>,
|
||||
}
|
||||
|
||||
impl BenchScore {
|
||||
pub fn pass() -> Self {
|
||||
Self {
|
||||
value: 1.0,
|
||||
label: "pass".to_string(),
|
||||
details: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn fail(details: impl Into<String>) -> Self {
|
||||
Self {
|
||||
value: 0.0,
|
||||
label: "fail".to_string(),
|
||||
details: Some(details.into()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn partial(value: f64, details: impl Into<String>) -> Self {
|
||||
Self {
|
||||
value: value.clamp(0.0, 1.0),
|
||||
label: "partial".to_string(),
|
||||
details: Some(details.into()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait for benchmark suite adapters.
|
||||
///
|
||||
/// Each suite (GAIA, Tau-bench, custom, etc.) implements this trait
|
||||
/// to provide task loading, scoring, and optional lifecycle hooks.
|
||||
#[async_trait]
|
||||
pub trait BenchSuite: Send + Sync {
|
||||
/// Human-readable name (e.g., "GAIA Validation").
|
||||
fn name(&self) -> &str;
|
||||
|
||||
/// Machine ID (e.g., "gaia").
|
||||
fn id(&self) -> &str;
|
||||
|
||||
/// Load all tasks from the suite's data source.
|
||||
async fn load_tasks(&self) -> Result<Vec<BenchTask>, BenchError>;
|
||||
|
||||
/// Score the agent's submission against the expected answer.
|
||||
async fn score(
|
||||
&self,
|
||||
task: &BenchTask,
|
||||
submission: &TaskSubmission,
|
||||
) -> Result<BenchScore, BenchError>;
|
||||
|
||||
/// Optional: set up environment before running a task (clone repo, init DB, etc.).
|
||||
async fn setup_task(&self, _task: &BenchTask) -> Result<(), BenchError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Optional: tear down environment after a task completes.
|
||||
async fn teardown_task(&self, _task: &BenchTask) -> Result<(), BenchError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Optional: additional tools to register for this suite's tasks.
|
||||
fn additional_tools(&self) -> Vec<Arc<dyn ironclaw::tools::Tool>> {
|
||||
vec![]
|
||||
}
|
||||
|
||||
/// Optional: restrict which tools the agent can use (allowlist).
|
||||
fn tool_whitelist(&self) -> Option<Vec<String>> {
|
||||
None
|
||||
}
|
||||
|
||||
/// Multi-turn: generate next simulated user message based on conversation so far.
|
||||
/// Return `None` to end the conversation.
|
||||
async fn next_user_message(
|
||||
&self,
|
||||
_task: &BenchTask,
|
||||
_conversation: &[ConversationTurn],
|
||||
) -> Result<Option<String>, BenchError> {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
//! Build script: compile Telegram channel WASM from source.
|
||||
//!
|
||||
//! Do not commit compiled WASM binaries — they are a supply chain risk.
|
||||
//! This script builds telegram.wasm from channels-src/telegram before the main crate compiles.
|
||||
//!
|
||||
//! Reproducible build:
|
||||
//! cargo build --release
|
||||
//! (build.rs invokes the channel build automatically)
|
||||
//!
|
||||
//! Prerequisites: rustup target add wasm32-wasip2, cargo install wasm-tools
|
||||
|
||||
use std::env;
|
||||
use std::path::PathBuf;
|
||||
use std::process::Command;
|
||||
|
||||
fn main() {
|
||||
let manifest_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
|
||||
let root = PathBuf::from(&manifest_dir);
|
||||
let channel_dir = root.join("channels-src/telegram");
|
||||
let wasm_out = channel_dir.join("telegram.wasm");
|
||||
|
||||
// Rerun when channel source or build script changes
|
||||
println!("cargo:rerun-if-changed=channels-src/telegram/src");
|
||||
println!("cargo:rerun-if-changed=channels-src/telegram/Cargo.toml");
|
||||
println!("cargo:rerun-if-changed=wit/channel.wit");
|
||||
|
||||
if !channel_dir.is_dir() {
|
||||
return;
|
||||
}
|
||||
|
||||
// Build WASM module
|
||||
let status = match Command::new("cargo")
|
||||
.args([
|
||||
"build",
|
||||
"--release",
|
||||
"--target",
|
||||
"wasm32-wasip2",
|
||||
"--manifest-path",
|
||||
channel_dir.join("Cargo.toml").to_str().unwrap(),
|
||||
])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
{
|
||||
Ok(s) => s,
|
||||
Err(_) => {
|
||||
eprintln!(
|
||||
"cargo:warning=Telegram channel build failed. Run: ./channels-src/telegram/build.sh"
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
if !status.success() {
|
||||
eprintln!(
|
||||
"cargo:warning=Telegram channel build failed. Run: ./channels-src/telegram/build.sh"
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
let raw_wasm = channel_dir.join("target/wasm32-wasip2/release/telegram_channel.wasm");
|
||||
if !raw_wasm.exists() {
|
||||
eprintln!(
|
||||
"cargo:warning=Telegram WASM output not found at {:?}",
|
||||
raw_wasm
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
// Convert to component and strip (wasm-tools)
|
||||
let component_ok = Command::new("wasm-tools")
|
||||
.args([
|
||||
"component",
|
||||
"new",
|
||||
raw_wasm.to_str().unwrap(),
|
||||
"-o",
|
||||
wasm_out.to_str().unwrap(),
|
||||
])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.map(|s| s.success())
|
||||
.unwrap_or(false);
|
||||
|
||||
if !component_ok {
|
||||
// Fallback: copy raw module if wasm-tools unavailable
|
||||
if std::fs::copy(&raw_wasm, &wasm_out).is_err() {
|
||||
eprintln!("cargo:warning=wasm-tools not found. Run: cargo install wasm-tools");
|
||||
}
|
||||
} else {
|
||||
// Strip debug info (use temp file to avoid clobbering)
|
||||
let stripped = wasm_out.with_extension("wasm.stripped");
|
||||
let strip_ok = Command::new("wasm-tools")
|
||||
.args([
|
||||
"strip",
|
||||
wasm_out.to_str().unwrap(),
|
||||
"-o",
|
||||
stripped.to_str().unwrap(),
|
||||
])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.map(|s| s.success())
|
||||
.unwrap_or(false);
|
||||
if strip_ok {
|
||||
let _ = std::fs::rename(&stripped, &wasm_out);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -30,7 +30,7 @@ if [ -f "$WASM_PATH" ]; then
|
||||
wasm-tools strip slack.wasm -o slack.wasm
|
||||
|
||||
echo "Built: slack.wasm ($(du -h slack.wasm | cut -f1))"
|
||||
echo "Copy slack.wasm and slack.capabilities.json to ~/.near-agent/channels/"
|
||||
echo "Copy slack.wasm and slack.capabilities.json to ~/.ironclaw/channels/"
|
||||
else
|
||||
echo "Error: WASM output not found at $WASM_PATH"
|
||||
exit 1
|
||||
|
||||
@@ -108,7 +108,10 @@ struct SlackPostMessageResponse {
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct SlackConfig {
|
||||
/// Name of secret containing signing secret (for verification by host).
|
||||
/// Parsed from config for forward compatibility; not yet used in WASM
|
||||
/// (host handles signature verification).
|
||||
#[serde(default = "default_signing_secret_name")]
|
||||
#[allow(dead_code)]
|
||||
signing_secret_name: String,
|
||||
}
|
||||
|
||||
@@ -175,11 +178,7 @@ impl Guest for SlackChannel {
|
||||
// Actual event callback
|
||||
"event_callback" => {
|
||||
if let Some(event) = event_wrapper.event {
|
||||
handle_slack_event(
|
||||
event,
|
||||
event_wrapper.team_id,
|
||||
event_wrapper.event_id,
|
||||
);
|
||||
handle_slack_event(event, event_wrapper.team_id, event_wrapper.event_id);
|
||||
}
|
||||
// Always respond 200 quickly to Slack (they have a 3s timeout)
|
||||
json_response(200, serde_json::json!({"ok": true}))
|
||||
@@ -230,6 +229,7 @@ impl Guest for SlackChannel {
|
||||
"https://slack.com/api/chat.postMessage",
|
||||
&headers.to_string(),
|
||||
Some(&payload_bytes),
|
||||
None,
|
||||
);
|
||||
|
||||
match result {
|
||||
@@ -243,14 +243,15 @@ impl Guest for SlackChannel {
|
||||
|
||||
// Parse Slack response
|
||||
let slack_response: SlackPostMessageResponse =
|
||||
serde_json::from_slice(&http_response.body).map_err(|e| {
|
||||
format!("Failed to parse Slack response: {}", e)
|
||||
})?;
|
||||
serde_json::from_slice(&http_response.body)
|
||||
.map_err(|e| format!("Failed to parse Slack response: {}", e))?;
|
||||
|
||||
if !slack_response.ok {
|
||||
return Err(format!(
|
||||
"Slack API error: {}",
|
||||
slack_response.error.unwrap_or_else(|| "unknown".to_string())
|
||||
slack_response
|
||||
.error
|
||||
.unwrap_or_else(|| "unknown".to_string())
|
||||
));
|
||||
}
|
||||
|
||||
@@ -277,17 +278,16 @@ impl Guest for SlackChannel {
|
||||
}
|
||||
|
||||
/// Handle a Slack event and emit message if applicable.
|
||||
fn handle_slack_event(
|
||||
event: SlackEvent,
|
||||
team_id: Option<String>,
|
||||
_event_id: Option<String>,
|
||||
) {
|
||||
fn handle_slack_event(event: SlackEvent, team_id: Option<String>, _event_id: Option<String>) {
|
||||
match event.event_type.as_str() {
|
||||
// Direct mention of the bot
|
||||
"app_mention" => {
|
||||
if let (Some(user), Some(channel), Some(text), Some(ts)) =
|
||||
(event.user, event.channel.clone(), event.text, event.ts.clone())
|
||||
{
|
||||
if let (Some(user), Some(channel), Some(text), Some(ts)) = (
|
||||
event.user,
|
||||
event.channel.clone(),
|
||||
event.text,
|
||||
event.ts.clone(),
|
||||
) {
|
||||
emit_message(user, text, channel, event.thread_ts.or(Some(ts)), team_id);
|
||||
}
|
||||
}
|
||||
@@ -299,9 +299,12 @@ fn handle_slack_event(
|
||||
return;
|
||||
}
|
||||
|
||||
if let (Some(user), Some(channel), Some(text), Some(ts)) =
|
||||
(event.user, event.channel.clone(), event.text, event.ts.clone())
|
||||
{
|
||||
if let (Some(user), Some(channel), Some(text), Some(ts)) = (
|
||||
event.user,
|
||||
event.channel.clone(),
|
||||
event.text,
|
||||
event.ts.clone(),
|
||||
) {
|
||||
// Only process DMs (channel IDs starting with D)
|
||||
if channel.starts_with('D') {
|
||||
emit_message(user, text, channel, event.thread_ts.or(Some(ts)), team_id);
|
||||
@@ -335,8 +338,7 @@ fn emit_message(
|
||||
team_id,
|
||||
};
|
||||
|
||||
let metadata_json =
|
||||
serde_json::to_string(&metadata).unwrap_or_else(|_| "{}".to_string());
|
||||
let metadata_json = serde_json::to_string(&metadata).unwrap_or_else(|_| "{}".to_string());
|
||||
|
||||
// Strip @ mentions of the bot from the text for cleaner messages
|
||||
let cleaned_text = strip_bot_mention(&text);
|
||||
|
||||
@@ -32,8 +32,8 @@ if [ -f "$WASM_PATH" ]; then
|
||||
echo "Built: telegram.wasm ($(du -h telegram.wasm | cut -f1))"
|
||||
echo ""
|
||||
echo "To install:"
|
||||
echo " mkdir -p ~/.near-agent/channels"
|
||||
echo " cp telegram.wasm telegram.capabilities.json ~/.near-agent/channels/"
|
||||
echo " mkdir -p ~/.ironclaw/channels"
|
||||
echo " cp telegram.wasm telegram.capabilities.json ~/.ironclaw/channels/"
|
||||
echo ""
|
||||
echo "Then add your bot token to secrets:"
|
||||
echo " # Set TELEGRAM_BOT_TOKEN in your environment or secrets store"
|
||||
|
||||
+402
-101
@@ -72,6 +72,10 @@ struct TelegramMessage {
|
||||
/// Message text.
|
||||
text: Option<String>,
|
||||
|
||||
/// Caption for media (photo, video, document, etc.).
|
||||
#[serde(default)]
|
||||
caption: Option<String>,
|
||||
|
||||
/// Original message if this is a reply.
|
||||
reply_to_message: Option<Box<TelegramMessage>>,
|
||||
|
||||
@@ -160,6 +164,21 @@ const POLLING_STATE_PATH: &str = "state/last_update_id";
|
||||
/// Workspace path for persisting owner_id across WASM callbacks.
|
||||
const OWNER_ID_PATH: &str = "state/owner_id";
|
||||
|
||||
/// Workspace path for persisting dm_policy across WASM callbacks.
|
||||
const DM_POLICY_PATH: &str = "state/dm_policy";
|
||||
|
||||
/// Workspace path for persisting allow_from (JSON array) across WASM callbacks.
|
||||
const ALLOW_FROM_PATH: &str = "state/allow_from";
|
||||
|
||||
/// Channel name for pairing store (used by pairing host APIs).
|
||||
const CHANNEL_NAME: &str = "telegram";
|
||||
|
||||
/// Workspace path for persisting bot_username for mention detection in groups.
|
||||
const BOT_USERNAME_PATH: &str = "state/bot_username";
|
||||
|
||||
/// Workspace path for persisting respond_to_all_group_messages flag.
|
||||
const RESPOND_TO_ALL_GROUP_PATH: &str = "state/respond_to_all_group_messages";
|
||||
|
||||
// ============================================================================
|
||||
// Channel Metadata
|
||||
// ============================================================================
|
||||
@@ -196,6 +215,14 @@ struct TelegramConfig {
|
||||
#[serde(default)]
|
||||
owner_id: Option<i64>,
|
||||
|
||||
/// DM policy: "pairing" (default), "allowlist", or "open".
|
||||
#[serde(default)]
|
||||
dm_policy: Option<String>,
|
||||
|
||||
/// Allowed sender IDs/usernames from config (merged with pairing-approved store).
|
||||
#[serde(default)]
|
||||
allow_from: Option<Vec<String>>,
|
||||
|
||||
/// Whether to respond to all group messages (not just mentions).
|
||||
#[serde(default)]
|
||||
respond_to_all_group_messages: bool,
|
||||
@@ -257,6 +284,28 @@ impl Guest for TelegramChannel {
|
||||
);
|
||||
}
|
||||
|
||||
// Persist dm_policy and allow_from for DM pairing in handle_message
|
||||
let dm_policy = config
|
||||
.dm_policy
|
||||
.as_deref()
|
||||
.unwrap_or("pairing")
|
||||
.to_string();
|
||||
let _ = channel_host::workspace_write(DM_POLICY_PATH, &dm_policy);
|
||||
|
||||
let allow_from_json = serde_json::to_string(&config.allow_from.unwrap_or_default())
|
||||
.unwrap_or_else(|_| "[]".to_string());
|
||||
let _ = channel_host::workspace_write(ALLOW_FROM_PATH, &allow_from_json);
|
||||
|
||||
// Persist bot_username and respond_to_all_group_messages for group handling
|
||||
let _ = channel_host::workspace_write(
|
||||
BOT_USERNAME_PATH,
|
||||
&config.bot_username.unwrap_or_default(),
|
||||
);
|
||||
let _ = channel_host::workspace_write(
|
||||
RESPOND_TO_ALL_GROUP_PATH,
|
||||
&config.respond_to_all_group_messages.to_string(),
|
||||
);
|
||||
|
||||
// Mode is determined by whether the host injected a tunnel_url
|
||||
// If tunnel is configured, use webhooks. Otherwise, use polling.
|
||||
let webhook_mode = config.tunnel_url.is_some();
|
||||
@@ -388,7 +437,9 @@ impl Guest for TelegramChannel {
|
||||
|
||||
let headers = serde_json::json!({});
|
||||
|
||||
let result = channel_host::http_request("GET", &url, &headers.to_string(), None);
|
||||
// 35s HTTP timeout outlives Telegram's 30s server-side long-poll
|
||||
let result =
|
||||
channel_host::http_request("GET", &url, &headers.to_string(), None, Some(35_000));
|
||||
|
||||
match result {
|
||||
Ok(response) => {
|
||||
@@ -461,72 +512,52 @@ impl Guest for TelegramChannel {
|
||||
}
|
||||
|
||||
fn on_respond(response: AgentResponse) -> Result<(), String> {
|
||||
// Parse metadata to get chat info
|
||||
let metadata: TelegramMessageMetadata = serde_json::from_str(&response.metadata_json)
|
||||
.map_err(|e| format!("Failed to parse metadata: {}", e))?;
|
||||
|
||||
// Build sendMessage payload
|
||||
let mut payload = serde_json::json!({
|
||||
"chat_id": metadata.chat_id,
|
||||
"text": response.content,
|
||||
"parse_mode": "Markdown",
|
||||
});
|
||||
|
||||
// Reply to the original message for context
|
||||
payload["reply_to_message_id"] = serde_json::Value::Number(metadata.message_id.into());
|
||||
|
||||
let payload_bytes = serde_json::to_vec(&payload)
|
||||
.map_err(|e| format!("Failed to serialize payload: {}", e))?;
|
||||
|
||||
// Make HTTP request to Telegram API
|
||||
// The bot token is injected into the URL by the host
|
||||
let headers = serde_json::json!({
|
||||
"Content-Type": "application/json"
|
||||
});
|
||||
|
||||
let result = channel_host::http_request(
|
||||
"POST",
|
||||
"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/sendMessage",
|
||||
&headers.to_string(),
|
||||
Some(&payload_bytes),
|
||||
// Try sending with Markdown first; fall back to plain text if Telegram
|
||||
// can't parse the entities (e.g. model leaked <tool_call> with underscores).
|
||||
let result = send_message(
|
||||
metadata.chat_id,
|
||||
&response.content,
|
||||
metadata.message_id,
|
||||
Some("Markdown"),
|
||||
);
|
||||
|
||||
match result {
|
||||
Ok(http_response) => {
|
||||
if http_response.status != 200 {
|
||||
let body_str = String::from_utf8_lossy(&http_response.body);
|
||||
return Err(format!(
|
||||
"Telegram API returned status {}: {}",
|
||||
http_response.status, body_str
|
||||
));
|
||||
}
|
||||
|
||||
// Parse Telegram response
|
||||
let api_response: TelegramApiResponse<SentMessage> =
|
||||
serde_json::from_slice(&http_response.body)
|
||||
.map_err(|e| format!("Failed to parse Telegram response: {}", e))?;
|
||||
|
||||
if !api_response.ok {
|
||||
return Err(format!(
|
||||
"Telegram API error: {}",
|
||||
api_response
|
||||
.description
|
||||
.unwrap_or_else(|| "unknown".to_string())
|
||||
));
|
||||
}
|
||||
|
||||
Ok(msg_id) => {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Debug,
|
||||
&format!(
|
||||
"Sent message to chat {}: message_id={}",
|
||||
metadata.chat_id,
|
||||
api_response.result.map(|r| r.message_id).unwrap_or(0)
|
||||
metadata.chat_id, msg_id
|
||||
),
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => Err(format!("HTTP request failed: {}", e)),
|
||||
Err(SendError::ParseEntities(detail)) => {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Warn,
|
||||
&format!("Markdown parse failed ({}), retrying as plain text", detail),
|
||||
);
|
||||
let msg_id = send_message(
|
||||
metadata.chat_id,
|
||||
&response.content,
|
||||
metadata.message_id,
|
||||
None,
|
||||
)
|
||||
.map_err(|e| format!("Plain-text retry also failed: {}", e))?;
|
||||
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Debug,
|
||||
&format!(
|
||||
"Sent plain-text message to chat {}: message_id={}",
|
||||
metadata.chat_id, msg_id
|
||||
),
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => Err(e.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -568,6 +599,7 @@ impl Guest for TelegramChannel {
|
||||
"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/sendChatAction",
|
||||
&headers.to_string(),
|
||||
Some(&payload_bytes),
|
||||
None,
|
||||
);
|
||||
|
||||
if let Err(e) = result {
|
||||
@@ -586,6 +618,101 @@ impl Guest for TelegramChannel {
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Send Message Helper
|
||||
// ============================================================================
|
||||
|
||||
/// Errors from send_message, split so callers can match on parse-entity failures.
|
||||
enum SendError {
|
||||
/// Telegram returned 400 with "can't parse entities" (Markdown issue).
|
||||
ParseEntities(String),
|
||||
/// Any other failure.
|
||||
Other(String),
|
||||
}
|
||||
|
||||
impl std::fmt::Display for SendError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
SendError::ParseEntities(detail) => write!(f, "parse entities error: {}", detail),
|
||||
SendError::Other(msg) => write!(f, "{}", msg),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a message via the Telegram Bot API.
|
||||
///
|
||||
/// Returns the sent message_id on success. When `parse_mode` is set and
|
||||
/// Telegram returns a 400 "can't parse entities" error, returns
|
||||
/// `SendError::ParseEntities` so the caller can retry without formatting.
|
||||
fn send_message(
|
||||
chat_id: i64,
|
||||
text: &str,
|
||||
reply_to_message_id: i64,
|
||||
parse_mode: Option<&str>,
|
||||
) -> Result<i64, SendError> {
|
||||
let mut payload = serde_json::json!({
|
||||
"chat_id": chat_id,
|
||||
"text": text,
|
||||
"reply_to_message_id": reply_to_message_id,
|
||||
});
|
||||
|
||||
if let Some(mode) = parse_mode {
|
||||
payload["parse_mode"] = serde_json::Value::String(mode.to_string());
|
||||
}
|
||||
|
||||
let payload_bytes = serde_json::to_vec(&payload)
|
||||
.map_err(|e| SendError::Other(format!("Failed to serialize payload: {}", e)))?;
|
||||
|
||||
let headers = serde_json::json!({ "Content-Type": "application/json" });
|
||||
|
||||
let result = channel_host::http_request(
|
||||
"POST",
|
||||
"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/sendMessage",
|
||||
&headers.to_string(),
|
||||
Some(&payload_bytes),
|
||||
None,
|
||||
);
|
||||
|
||||
match result {
|
||||
Ok(http_response) => {
|
||||
if http_response.status == 400 {
|
||||
let body_str = String::from_utf8_lossy(&http_response.body);
|
||||
if body_str.contains("can't parse entities") {
|
||||
return Err(SendError::ParseEntities(body_str.to_string()));
|
||||
}
|
||||
return Err(SendError::Other(format!(
|
||||
"Telegram API returned 400: {}",
|
||||
body_str
|
||||
)));
|
||||
}
|
||||
|
||||
if http_response.status != 200 {
|
||||
let body_str = String::from_utf8_lossy(&http_response.body);
|
||||
return Err(SendError::Other(format!(
|
||||
"Telegram API returned status {}: {}",
|
||||
http_response.status, body_str
|
||||
)));
|
||||
}
|
||||
|
||||
let api_response: TelegramApiResponse<SentMessage> =
|
||||
serde_json::from_slice(&http_response.body)
|
||||
.map_err(|e| SendError::Other(format!("Failed to parse response: {}", e)))?;
|
||||
|
||||
if !api_response.ok {
|
||||
return Err(SendError::Other(format!(
|
||||
"Telegram API error: {}",
|
||||
api_response
|
||||
.description
|
||||
.unwrap_or_else(|| "unknown".to_string())
|
||||
)));
|
||||
}
|
||||
|
||||
Ok(api_response.result.map(|r| r.message_id).unwrap_or(0))
|
||||
}
|
||||
Err(e) => Err(SendError::Other(format!("HTTP request failed: {}", e))),
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Webhook Management
|
||||
// ============================================================================
|
||||
@@ -604,6 +731,7 @@ fn delete_webhook() -> Result<(), String> {
|
||||
"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/deleteWebhook",
|
||||
&headers.to_string(),
|
||||
None,
|
||||
None,
|
||||
);
|
||||
|
||||
match result {
|
||||
@@ -666,6 +794,7 @@ fn register_webhook(tunnel_url: &str, webhook_secret: Option<&str>) -> Result<()
|
||||
"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/setWebhook",
|
||||
&headers.to_string(),
|
||||
Some(&body_bytes),
|
||||
None,
|
||||
);
|
||||
|
||||
match result {
|
||||
@@ -700,6 +829,47 @@ fn register_webhook(tunnel_url: &str, webhook_secret: Option<&str>) -> Result<()
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Pairing Reply
|
||||
// ============================================================================
|
||||
|
||||
/// Send a pairing code message to a chat. Used when an unknown user DMs the bot.
|
||||
fn send_pairing_reply(chat_id: i64, code: &str) -> Result<(), String> {
|
||||
let payload = serde_json::json!({
|
||||
"chat_id": chat_id,
|
||||
"text": format!(
|
||||
"To pair with this bot, run: `ironclaw pairing approve telegram {}`",
|
||||
code
|
||||
),
|
||||
"parse_mode": "Markdown",
|
||||
});
|
||||
|
||||
let payload_bytes = serde_json::to_vec(&payload)
|
||||
.map_err(|e| format!("Failed to serialize payload: {}", e))?;
|
||||
|
||||
let headers = serde_json::json!({
|
||||
"Content-Type": "application/json"
|
||||
});
|
||||
|
||||
let result = channel_host::http_request(
|
||||
"POST",
|
||||
"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/sendMessage",
|
||||
&headers.to_string(),
|
||||
Some(&payload_bytes),
|
||||
);
|
||||
|
||||
match result {
|
||||
Ok(response) => {
|
||||
if response.status != 200 {
|
||||
let body_str = String::from_utf8_lossy(&response.body);
|
||||
return Err(format!("HTTP {}: {}", response.status, body_str));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => Err(format!("HTTP request failed: {}", e)),
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Update Handling
|
||||
// ============================================================================
|
||||
@@ -719,11 +889,16 @@ fn handle_update(update: TelegramUpdate) {
|
||||
|
||||
/// Process a single message.
|
||||
fn handle_message(message: TelegramMessage) {
|
||||
// Skip messages without text
|
||||
let text = match message.text {
|
||||
Some(t) if !t.is_empty() => t,
|
||||
_ => return,
|
||||
};
|
||||
// Use text or caption (for media messages)
|
||||
let content = message
|
||||
.text
|
||||
.filter(|t| !t.is_empty())
|
||||
.or_else(|| message.caption.filter(|c| !c.is_empty()))
|
||||
.unwrap_or_default();
|
||||
|
||||
if content.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
// Skip messages without a sender (channel posts)
|
||||
let from = match message.from {
|
||||
@@ -736,41 +911,111 @@ fn handle_message(message: TelegramMessage) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Owner validation: silently drop messages from non-owner users
|
||||
if let Some(owner_id_str) = channel_host::workspace_read(OWNER_ID_PATH) {
|
||||
if !owner_id_str.is_empty() {
|
||||
if let Ok(owner_id) = owner_id_str.parse::<i64>() {
|
||||
if from.id != owner_id {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Debug,
|
||||
&format!(
|
||||
"Dropping message from non-owner user {} (owner: {})",
|
||||
from.id, owner_id
|
||||
),
|
||||
);
|
||||
return;
|
||||
let is_private = message.chat.chat_type == "private";
|
||||
|
||||
// Owner validation: when owner_id is set, only that user can message
|
||||
let owner_configured = channel_host::workspace_read(OWNER_ID_PATH)
|
||||
.map(|s| !s.is_empty())
|
||||
.unwrap_or(false);
|
||||
|
||||
if owner_configured {
|
||||
if let Ok(owner_id) = channel_host::workspace_read(OWNER_ID_PATH)
|
||||
.unwrap()
|
||||
.parse::<i64>()
|
||||
{
|
||||
if from.id != owner_id {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Debug,
|
||||
&format!(
|
||||
"Dropping message from non-owner user {} (owner: {})",
|
||||
from.id, owner_id
|
||||
),
|
||||
);
|
||||
return;
|
||||
}
|
||||
}
|
||||
} else if is_private {
|
||||
// No owner_id: apply dm_policy for private chats
|
||||
let dm_policy = channel_host::workspace_read(DM_POLICY_PATH)
|
||||
.unwrap_or_else(|| "pairing".to_string());
|
||||
|
||||
if dm_policy != "open" {
|
||||
// Build effective allow list: config allow_from + pairing store
|
||||
let mut allowed: Vec<String> = channel_host::workspace_read(ALLOW_FROM_PATH)
|
||||
.and_then(|s| serde_json::from_str(&s).ok())
|
||||
.unwrap_or_default();
|
||||
|
||||
if let Ok(store_allowed) = channel_host::pairing_read_allow_from(CHANNEL_NAME) {
|
||||
allowed.extend(store_allowed);
|
||||
}
|
||||
|
||||
let id_str = from.id.to_string();
|
||||
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()));
|
||||
|
||||
if !is_allowed {
|
||||
if dm_policy == "pairing" {
|
||||
// Upsert pairing request and send reply
|
||||
let meta = serde_json::json!({
|
||||
"chat_id": message.chat.id,
|
||||
"user_id": from.id,
|
||||
"username": username_opt,
|
||||
})
|
||||
.to_string();
|
||||
|
||||
match channel_host::pairing_upsert_request(CHANNEL_NAME, &id_str, &meta) {
|
||||
Ok(result) => {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Info,
|
||||
&format!(
|
||||
"Pairing request for user {} (chat {}): code {}",
|
||||
from.id, message.chat.id, result.code
|
||||
),
|
||||
);
|
||||
if result.created {
|
||||
let _ = send_pairing_reply(message.chat.id, &result.code);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Error,
|
||||
&format!("Pairing upsert failed: {}", e),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let is_private = message.chat.chat_type == "private";
|
||||
|
||||
// For group chats, check if the bot was mentioned
|
||||
// TODO: Read bot_username from config and check mentions
|
||||
// For now, process all messages in private chats and groups
|
||||
// For group chats, only respond if bot was mentioned or respond_to_all is enabled
|
||||
if !is_private {
|
||||
// In groups, only respond if there's a bot mention or command
|
||||
// This is a simplified check - proper implementation would use entities
|
||||
let has_command = text.starts_with('/');
|
||||
let has_mention = text.contains('@');
|
||||
let respond_to_all = channel_host::workspace_read(RESPOND_TO_ALL_GROUP_PATH)
|
||||
.as_deref()
|
||||
.unwrap_or("false")
|
||||
== "true";
|
||||
|
||||
if !has_command && !has_mention {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Debug,
|
||||
&format!("Ignoring group message without mention: {}", text),
|
||||
);
|
||||
return;
|
||||
if !respond_to_all {
|
||||
let has_command = content.starts_with('/');
|
||||
let bot_username = channel_host::workspace_read(BOT_USERNAME_PATH)
|
||||
.unwrap_or_default();
|
||||
let has_bot_mention = if bot_username.is_empty() {
|
||||
content.contains('@')
|
||||
} else {
|
||||
let mention = format!("@{}", bot_username);
|
||||
content.to_lowercase().contains(&mention.to_lowercase())
|
||||
};
|
||||
|
||||
if !has_command && !has_bot_mention {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Debug,
|
||||
&format!("Ignoring group message without mention: {}", content),
|
||||
);
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -792,17 +1037,30 @@ fn handle_message(message: TelegramMessage) {
|
||||
let metadata_json = serde_json::to_string(&metadata).unwrap_or_else(|_| "{}".to_string());
|
||||
|
||||
// Clean the message text (strip bot mentions and commands)
|
||||
let cleaned_text = clean_message_text(&text);
|
||||
let bot_username = channel_host::workspace_read(BOT_USERNAME_PATH).unwrap_or_default();
|
||||
let cleaned_text = clean_message_text(
|
||||
&content,
|
||||
if bot_username.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(bot_username.as_str())
|
||||
},
|
||||
);
|
||||
|
||||
if cleaned_text.is_empty() {
|
||||
// For /start with no args, emit placeholder so agent can respond with welcome
|
||||
let content_to_emit = if cleaned_text.is_empty() && content.trim().starts_with('/') {
|
||||
"[User started the bot]".to_string()
|
||||
} else if cleaned_text.is_empty() {
|
||||
return;
|
||||
}
|
||||
} else {
|
||||
cleaned_text
|
||||
};
|
||||
|
||||
// Emit the message to the agent
|
||||
channel_host::emit_message(&EmittedMessage {
|
||||
user_id: from.id.to_string(),
|
||||
user_name: Some(user_name),
|
||||
content: cleaned_text,
|
||||
content: content_to_emit,
|
||||
thread_id: None, // Telegram doesn't have threads in the same way
|
||||
metadata_json,
|
||||
});
|
||||
@@ -817,7 +1075,8 @@ fn handle_message(message: TelegramMessage) {
|
||||
}
|
||||
|
||||
/// Clean message text by removing bot commands and @mentions at the start.
|
||||
fn clean_message_text(text: &str) -> String {
|
||||
/// When bot_username is set, only strips that specific mention; otherwise strips any leading @mention.
|
||||
fn clean_message_text(text: &str, bot_username: Option<&str>) -> String {
|
||||
let mut result = text.trim().to_string();
|
||||
|
||||
// Remove leading /command
|
||||
@@ -832,11 +1091,30 @@ fn clean_message_text(text: &str) -> String {
|
||||
|
||||
// Remove leading @mention
|
||||
if result.starts_with('@') {
|
||||
if let Some(space_idx) = result.find(' ') {
|
||||
result = result[space_idx..].trim_start().to_string();
|
||||
if let Some(bot) = bot_username {
|
||||
let mention = format!("@{}", bot);
|
||||
let mention_lower = mention.to_lowercase();
|
||||
let result_lower = result.to_lowercase();
|
||||
if result_lower.starts_with(&mention_lower) {
|
||||
let rest = result[mention.len()..].trim_start();
|
||||
if rest.is_empty() {
|
||||
return String::new();
|
||||
}
|
||||
result = rest.to_string();
|
||||
} else if let Some(space_idx) = result.find(' ') {
|
||||
// Different leading @mention - only strip if it's the bot
|
||||
let first_word = &result[..space_idx];
|
||||
if first_word.eq_ignore_ascii_case(&mention) {
|
||||
result = result[space_idx..].trim_start().to_string();
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Just a mention with no text
|
||||
return String::new();
|
||||
// No bot_username: strip any leading @mention
|
||||
if let Some(space_idx) = result.find(' ') {
|
||||
result = result[space_idx..].trim_start().to_string();
|
||||
} else {
|
||||
return String::new();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -872,12 +1150,22 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_clean_message_text() {
|
||||
assert_eq!(clean_message_text("/start hello"), "hello");
|
||||
assert_eq!(clean_message_text("@bot hello world"), "hello world");
|
||||
assert_eq!(clean_message_text("/start"), "");
|
||||
assert_eq!(clean_message_text("@botname"), "");
|
||||
assert_eq!(clean_message_text("just text"), "just text");
|
||||
assert_eq!(clean_message_text(" spaced "), "spaced");
|
||||
// Without bot_username: strips any leading @mention
|
||||
assert_eq!(clean_message_text("/start hello", None), "hello");
|
||||
assert_eq!(clean_message_text("@bot hello world", None), "hello world");
|
||||
assert_eq!(clean_message_text("/start", None), "");
|
||||
assert_eq!(clean_message_text("@botname", None), "");
|
||||
assert_eq!(clean_message_text("just text", None), "just text");
|
||||
assert_eq!(clean_message_text(" spaced ", None), "spaced");
|
||||
|
||||
// With bot_username: only strips @MyBot, not @alice
|
||||
assert_eq!(clean_message_text("@MyBot hello", Some("MyBot")), "hello");
|
||||
assert_eq!(clean_message_text("@mybot hi", Some("MyBot")), "hi");
|
||||
assert_eq!(
|
||||
clean_message_text("@alice hello", Some("MyBot")),
|
||||
"@alice hello"
|
||||
);
|
||||
assert_eq!(clean_message_text("@MyBot", Some("MyBot")), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -945,4 +1233,17 @@ mod tests {
|
||||
assert_eq!(from.id, 789);
|
||||
assert_eq!(from.first_name, "John");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_message_with_caption() {
|
||||
let json = r#"{
|
||||
"message_id": 1,
|
||||
"from": {"id": 1, "is_bot": false, "first_name": "A"},
|
||||
"chat": {"id": 1, "type": "private"},
|
||||
"caption": "What's in this image?"
|
||||
}"#;
|
||||
let msg: TelegramMessage = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(msg.text, None);
|
||||
assert_eq!(msg.caption.as_deref(), Some("What's in this image?"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,43 +1 @@
|
||||
{
|
||||
"type": "channel",
|
||||
"name": "telegram",
|
||||
"description": "Telegram Bot API channel for receiving and responding to Telegram messages",
|
||||
"capabilities": {
|
||||
"http": {
|
||||
"allowlist": [
|
||||
{ "host": "api.telegram.org", "path_prefix": "/bot" }
|
||||
],
|
||||
"credentials": {
|
||||
"telegram_bot": {
|
||||
"secret_name": "telegram_bot_token",
|
||||
"location": { "type": "url_path", "placeholder": "{TELEGRAM_BOT_TOKEN}" },
|
||||
"host_patterns": ["api.telegram.org"]
|
||||
}
|
||||
},
|
||||
"rate_limit": {
|
||||
"requests_per_minute": 30,
|
||||
"requests_per_hour": 1000
|
||||
}
|
||||
},
|
||||
"secrets": {
|
||||
"allowed_names": ["telegram_*"]
|
||||
},
|
||||
"channel": {
|
||||
"allowed_paths": ["/webhook/telegram"],
|
||||
"allow_polling": true,
|
||||
"min_poll_interval_ms": 30000,
|
||||
"workspace_prefix": "channels/telegram/",
|
||||
"emit_rate_limit": {
|
||||
"messages_per_minute": 100,
|
||||
"messages_per_hour": 5000
|
||||
}
|
||||
}
|
||||
},
|
||||
"config": {
|
||||
"bot_username": null,
|
||||
"owner_id": null,
|
||||
"respond_to_all_group_messages": false,
|
||||
"polling_enabled": false,
|
||||
"poll_interval_ms": 30000
|
||||
}
|
||||
}
|
||||
{"type":"channel","name":"telegram","description":"Telegram Bot API channel for receiving and responding to Telegram messages","capabilities":{"http":{"allowlist":[{"host":"api.telegram.org","path_prefix":"/bot"}],"credentials":{"telegram_bot":{"secret_name":"telegram_bot_token","location":{"type":"url_path","placeholder":"{TELEGRAM_BOT_TOKEN}"},"host_patterns":["api.telegram.org"]}},"rate_limit":{"requests_per_minute":30,"requests_per_hour":1000}},"secrets":{"allowed_names":["telegram_*"]},"channel":{"allowed_paths":["/webhook/telegram"],"allow_polling":true,"min_poll_interval_ms":30000,"workspace_prefix":"channels/telegram/","emit_rate_limit":{"messages_per_minute":100,"messages_per_hour":5000}}},"config":{"bot_username":null,"owner_id":null,"respond_to_all_group_messages":false,"polling_enabled":false,"poll_interval_ms":30000,"dm_policy":"pairing","allow_from":[]}}
|
||||
|
||||
Binary file not shown.
@@ -361,6 +361,7 @@ impl Guest for WhatsAppChannel {
|
||||
&api_url,
|
||||
&headers.to_string(),
|
||||
Some(&payload_bytes),
|
||||
None,
|
||||
);
|
||||
|
||||
match result {
|
||||
|
||||
@@ -246,13 +246,46 @@ Create `my-channel.capabilities.json`:
|
||||
|
||||
## Building and Deploying
|
||||
|
||||
### Supply Chain Security: No Committed Binaries
|
||||
|
||||
**Do not commit compiled WASM binaries.** They are a supply chain risk — the binary in a PR may not match the source. IronClaw builds channels from source:
|
||||
|
||||
- `cargo build` automatically builds `telegram.wasm` via `build.rs`
|
||||
- The built binary is in `.gitignore` and is not committed
|
||||
- CI should run `cargo build` (or `./scripts/build-all.sh`) to produce releases
|
||||
|
||||
**Reproducible build:**
|
||||
```bash
|
||||
cargo build --release
|
||||
```
|
||||
|
||||
Prerequisites: `rustup target add wasm32-wasip2`, `cargo install wasm-tools` (optional; fallback copies raw WASM if unavailable).
|
||||
|
||||
### Telegram Channel (Manual Build)
|
||||
|
||||
```bash
|
||||
# Add WASM target if needed
|
||||
rustup target add wasm32-wasip2
|
||||
|
||||
# Build Telegram channel
|
||||
./channels-src/telegram/build.sh
|
||||
|
||||
# Install (or use ironclaw onboard to install bundled channel)
|
||||
mkdir -p ~/.ironclaw/channels
|
||||
cp channels-src/telegram/telegram.wasm channels-src/telegram/telegram.capabilities.json ~/.ironclaw/channels/
|
||||
```
|
||||
|
||||
**Note**: The main IronClaw binary bundles `telegram.wasm` via `include_bytes!`. When modifying the Telegram channel source, run `./channels-src/telegram/build.sh` **before** building the main crate, so the updated WASM is included.
|
||||
|
||||
### Other Channels
|
||||
|
||||
```bash
|
||||
# Build the WASM component
|
||||
cd channels/my-channel
|
||||
cargo component build --release
|
||||
cd channels-src/my-channel
|
||||
cargo build --release --target wasm32-wasip2
|
||||
|
||||
# Deploy to ~/.ironclaw/channels/
|
||||
cp target/wasm32-wasip1/release/my_channel.wasm ~/.ironclaw/channels/my-channel.wasm
|
||||
cp target/wasm32-wasip2/release/my_channel.wasm ~/.ironclaw/channels/my-channel.wasm
|
||||
cp my-channel.capabilities.json ~/.ironclaw/channels/
|
||||
```
|
||||
|
||||
|
||||
@@ -0,0 +1,135 @@
|
||||
# Telegram Channel Setup
|
||||
|
||||
This guide covers configuring the Telegram channel for IronClaw, including DM pairing for access control.
|
||||
|
||||
## Overview
|
||||
|
||||
The Telegram channel lets you interact with IronClaw via Telegram DMs and groups. It supports:
|
||||
|
||||
- **Webhook mode** (recommended): Instant delivery via tunnel
|
||||
- **Polling mode**: No tunnel required; ~30s delay
|
||||
- **DM pairing**: Approve unknown users before they can message the agent
|
||||
- **Group mentions**: `@YourBot` or `/command` to trigger in groups
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- IronClaw installed and configured (`ironclaw onboard`)
|
||||
- A Telegram bot token from [@BotFather](https://t.me/BotFather)
|
||||
|
||||
## Quick Start
|
||||
|
||||
### 1. Create a Bot
|
||||
|
||||
1. Message [@BotFather](https://t.me/BotFather) on Telegram
|
||||
2. Send `/newbot` and follow the prompts
|
||||
3. Copy the bot token (e.g., `123456789:ABCdefGHIjklMNOpqrsTUVwxyz`)
|
||||
|
||||
### 2. Configure via Setup Wizard
|
||||
|
||||
```bash
|
||||
ironclaw onboard
|
||||
```
|
||||
|
||||
When prompted, enable the Telegram channel and paste your bot token. The wizard will:
|
||||
|
||||
- Validate the token
|
||||
- Optionally configure a webhook secret
|
||||
- Set up tunnel (if you want webhook mode)
|
||||
|
||||
### 3. (Optional) Configure Tunnel for Webhooks
|
||||
|
||||
For instant message delivery, expose your agent via a tunnel:
|
||||
|
||||
```bash
|
||||
# ngrok
|
||||
ngrok http 8080
|
||||
|
||||
# Cloudflare
|
||||
cloudflared tunnel --url http://localhost:8080
|
||||
```
|
||||
|
||||
Set the tunnel URL in settings or via `TUNNEL_URL` env var. Without a tunnel, the channel uses polling (~30s delay).
|
||||
|
||||
## DM Pairing
|
||||
|
||||
When an unknown user DMs your bot, they receive a pairing code. You must approve them before they can message the agent.
|
||||
|
||||
### Flow
|
||||
|
||||
1. Unknown user sends a message to your bot
|
||||
2. Bot replies: `To pair with this bot, run: ironclaw pairing approve telegram ABC12345`
|
||||
3. You run: `ironclaw pairing approve telegram ABC12345`
|
||||
4. User is added to the allow list; future messages are delivered
|
||||
|
||||
### Commands
|
||||
|
||||
```bash
|
||||
# List pending pairing requests
|
||||
ironclaw pairing list telegram
|
||||
|
||||
# List as JSON
|
||||
ironclaw pairing list telegram --json
|
||||
|
||||
# Approve a user by code
|
||||
ironclaw pairing approve telegram ABC12345
|
||||
```
|
||||
|
||||
### Configuration
|
||||
|
||||
Edit `~/.ironclaw/channels/telegram.capabilities.json` (or the config injected by the host):
|
||||
|
||||
| Option | Values | Default | Description |
|
||||
|--------|--------|---------|-------------|
|
||||
| `dm_policy` | `open`, `allowlist`, `pairing` | `pairing` | `open` = allow all; `allowlist` = config + approved only; `pairing` = allowlist + send pairing reply to unknown |
|
||||
| `allow_from` | `["user_id", "username", "*"]` | `[]` | Pre-approved IDs/usernames. `*` allows everyone. |
|
||||
| `owner_id` | Telegram user ID | `null` | When set, only this user can message (overrides dm_policy) |
|
||||
| `bot_username` | Bot username (no @) | `null` | Used for mention detection in groups; when set, only strips this mention from messages |
|
||||
| `respond_to_all_group_messages` | `true`/`false` | `false` | When true, respond to all group messages; when false, only @mentions and /commands |
|
||||
|
||||
## Manual Installation
|
||||
|
||||
If the channel isn't installed via the wizard:
|
||||
|
||||
```bash
|
||||
# Build the Telegram channel (requires wasm32-wasip2 target)
|
||||
rustup target add wasm32-wasip2
|
||||
./channels-src/telegram/build.sh
|
||||
|
||||
# Install
|
||||
mkdir -p ~/.ironclaw/channels
|
||||
cp channels-src/telegram/telegram.wasm channels-src/telegram/telegram.capabilities.json ~/.ironclaw/channels/
|
||||
```
|
||||
|
||||
## Secrets
|
||||
|
||||
The channel expects a secret named `telegram_bot_token`. Configure via:
|
||||
|
||||
- **Setup wizard**: Saves to encrypted secrets store
|
||||
- **Environment**: `TELEGRAM_BOT_TOKEN=your_token`
|
||||
- **Secrets store**: `ironclaw` CLI (if available)
|
||||
|
||||
## Webhook Secret (Optional)
|
||||
|
||||
For webhook validation, set `telegram_webhook_secret` in secrets. Telegram will send `X-Telegram-Bot-Api-Secret-Token` with each request; the host validates it before forwarding.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Messages not delivered
|
||||
|
||||
- **Polling mode**: Check logs for `getUpdates` errors. Ensure the bot token is valid.
|
||||
- **Webhook mode**: Verify tunnel is running and `TUNNEL_URL` is correct. Telegram requires HTTPS.
|
||||
|
||||
### Pairing code not received
|
||||
|
||||
- Verify the channel can send messages (HTTP allowlist includes `api.telegram.org`)
|
||||
- Check `dm_policy` is `pairing` (not `allowlist` which blocks without reply)
|
||||
|
||||
### Group mentions not working
|
||||
|
||||
- Set `bot_username` in config to your bot's username (e.g., `MyIronClawBot`)
|
||||
- Ensure the message contains `@YourBot` or starts with `/`
|
||||
|
||||
### "Connection refused" when starting
|
||||
|
||||
- For webhook mode: Start your tunnel before `ironclaw run`
|
||||
- For polling only: No tunnel needed; ignore tunnel-related warnings
|
||||
@@ -0,0 +1,122 @@
|
||||
//! Standalone heartbeat test.
|
||||
//!
|
||||
//! Exercises the heartbeat system in isolation: connects to the real
|
||||
//! database, reads the real HEARTBEAT.md, calls the real LLM, and prints
|
||||
//! every step so you can see exactly where it breaks.
|
||||
//!
|
||||
//! Usage:
|
||||
//! cargo run --example test_heartbeat
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use ironclaw::{
|
||||
agent::HeartbeatRunner,
|
||||
config::Config,
|
||||
history::Store,
|
||||
llm::{SessionConfig, create_llm_provider, create_session_manager},
|
||||
workspace::Workspace,
|
||||
};
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
// Load .env and set up logging
|
||||
let _ = dotenvy::dotenv();
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter("ironclaw=debug")
|
||||
.init();
|
||||
|
||||
println!("=== Heartbeat Integration Test ===\n");
|
||||
|
||||
// 1. Load config
|
||||
let config = Config::from_env()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Config: {}", e))?;
|
||||
println!("[1/6] Config loaded");
|
||||
println!(" heartbeat.enabled = {}", config.heartbeat.enabled);
|
||||
println!(
|
||||
" heartbeat.interval_secs = {}",
|
||||
config.heartbeat.interval_secs
|
||||
);
|
||||
println!(
|
||||
" heartbeat.notify_channel = {:?}",
|
||||
config.heartbeat.notify_channel
|
||||
);
|
||||
println!(
|
||||
" heartbeat.notify_user = {:?}",
|
||||
config.heartbeat.notify_user
|
||||
);
|
||||
|
||||
// 2. Connect to database
|
||||
let store = Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
println!("[2/6] Database connected");
|
||||
|
||||
// 3. Create workspace
|
||||
let workspace = Arc::new(Workspace::new("default", store.pool()));
|
||||
println!("[3/6] Workspace created");
|
||||
|
||||
// 4. Read HEARTBEAT.md
|
||||
let checklist = workspace.heartbeat_checklist().await;
|
||||
match &checklist {
|
||||
Ok(Some(content)) => {
|
||||
let preview: String = content.chars().take(200).collect();
|
||||
println!("[4/6] HEARTBEAT.md found ({} chars)", content.len());
|
||||
println!(" Preview: {}...", preview);
|
||||
}
|
||||
Ok(None) => {
|
||||
println!("[4/6] HEARTBEAT.md is None (no file, no seed fallback)");
|
||||
println!(" Heartbeat will return Skipped.");
|
||||
}
|
||||
Err(e) => {
|
||||
println!("[4/6] HEARTBEAT.md read error: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
// Check if the checklist would be considered "effectively empty"
|
||||
if let Ok(Some(_)) = checklist {
|
||||
println!(" (Will verify via runner below)");
|
||||
}
|
||||
|
||||
// 5. Create LLM provider
|
||||
let session = create_session_manager(SessionConfig {
|
||||
auth_base_url: config.llm.nearai.auth_base_url.clone(),
|
||||
session_path: config.llm.nearai.session_path.clone(),
|
||||
..Default::default()
|
||||
})
|
||||
.await;
|
||||
let llm = create_llm_provider(&config.llm, session)?;
|
||||
println!("[5/6] LLM provider created (model: {})", llm.model_name());
|
||||
|
||||
// 6. Run heartbeat check
|
||||
println!("[6/6] Running check_heartbeat()...\n");
|
||||
|
||||
let hb_config = ironclaw::agent::HeartbeatConfig::default();
|
||||
let runner = HeartbeatRunner::new(hb_config, workspace, llm);
|
||||
|
||||
let result = runner.check_heartbeat().await;
|
||||
|
||||
println!("=== Result ===\n");
|
||||
match &result {
|
||||
ironclaw::agent::HeartbeatResult::Ok => {
|
||||
println!("HeartbeatResult::Ok");
|
||||
println!(" LLM responded HEARTBEAT_OK, nothing needs attention.");
|
||||
}
|
||||
ironclaw::agent::HeartbeatResult::NeedsAttention(msg) => {
|
||||
println!("HeartbeatResult::NeedsAttention");
|
||||
println!(" Message:\n{}", msg);
|
||||
}
|
||||
ironclaw::agent::HeartbeatResult::Skipped => {
|
||||
println!("HeartbeatResult::Skipped");
|
||||
println!(" No checklist found, or checklist was effectively empty.");
|
||||
println!(" This means the HEARTBEAT.md either:");
|
||||
println!(" - Does not exist in the workspace database");
|
||||
println!(" - Contains only headers, comments, and empty checkboxes");
|
||||
}
|
||||
ironclaw::agent::HeartbeatResult::Failed(err) => {
|
||||
println!("HeartbeatResult::Failed");
|
||||
println!(" Error: {}", err);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
-- Add project_dir and user_id columns for sandbox job tracking.
|
||||
-- user_id was previously hardcoded to "default" in the Rust layer;
|
||||
-- now it's persisted so we can filter per-user.
|
||||
|
||||
ALTER TABLE agent_jobs ADD COLUMN IF NOT EXISTS project_dir TEXT;
|
||||
ALTER TABLE agent_jobs ADD COLUMN IF NOT EXISTS user_id TEXT NOT NULL DEFAULT 'default';
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_jobs_source ON agent_jobs(source);
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_jobs_user ON agent_jobs(user_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_agent_jobs_created ON agent_jobs(created_at DESC);
|
||||
@@ -0,0 +1,14 @@
|
||||
-- Track which mode a sandbox job uses (worker vs claude_code).
|
||||
ALTER TABLE agent_jobs ADD COLUMN IF NOT EXISTS job_mode TEXT NOT NULL DEFAULT 'worker';
|
||||
|
||||
-- Persist Claude Code streaming events so they survive restarts and can be
|
||||
-- loaded when the frontend opens a job detail view after the fact.
|
||||
CREATE TABLE IF NOT EXISTS claude_code_events (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
job_id UUID NOT NULL REFERENCES agent_jobs(id),
|
||||
event_type TEXT NOT NULL,
|
||||
data JSONB NOT NULL,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_cc_events_job ON claude_code_events(job_id, id);
|
||||
@@ -0,0 +1,73 @@
|
||||
-- Routines: scheduled and reactive job system.
|
||||
--
|
||||
-- A routine is a named, persistent, user-owned task with a trigger and an action.
|
||||
-- Triggers fire independently (cron, event, webhook, manual) so only the
|
||||
-- relevant routine's prompt hits the LLM, not the whole checklist.
|
||||
|
||||
CREATE TABLE routines (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
name TEXT NOT NULL,
|
||||
description TEXT NOT NULL DEFAULT '',
|
||||
user_id TEXT NOT NULL,
|
||||
enabled BOOLEAN NOT NULL DEFAULT true,
|
||||
|
||||
-- Trigger definition
|
||||
trigger_type TEXT NOT NULL, -- 'cron', 'event', 'webhook', 'manual'
|
||||
trigger_config JSONB NOT NULL, -- type-specific config (schedule, pattern, etc.)
|
||||
|
||||
-- Action definition
|
||||
action_type TEXT NOT NULL, -- 'lightweight', 'full_job'
|
||||
action_config JSONB NOT NULL, -- prompt, context_paths, max_tokens / title, max_iterations
|
||||
|
||||
-- Guardrails
|
||||
cooldown_secs INTEGER NOT NULL DEFAULT 300,
|
||||
max_concurrent INTEGER NOT NULL DEFAULT 1,
|
||||
dedup_window_secs INTEGER, -- NULL = no dedup
|
||||
|
||||
-- Notification preferences
|
||||
notify_channel TEXT, -- NULL = use default
|
||||
notify_user TEXT NOT NULL DEFAULT 'default',
|
||||
notify_on_success BOOLEAN NOT NULL DEFAULT false,
|
||||
notify_on_failure BOOLEAN NOT NULL DEFAULT true,
|
||||
notify_on_attention BOOLEAN NOT NULL DEFAULT true,
|
||||
|
||||
-- Runtime state (updated by engine)
|
||||
state JSONB NOT NULL DEFAULT '{}',
|
||||
last_run_at TIMESTAMPTZ,
|
||||
next_fire_at TIMESTAMPTZ, -- pre-computed for cron triggers
|
||||
run_count BIGINT NOT NULL DEFAULT 0,
|
||||
consecutive_failures INTEGER NOT NULL DEFAULT 0,
|
||||
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
|
||||
UNIQUE (user_id, name)
|
||||
);
|
||||
|
||||
-- Fast lookup: "which cron routines need to fire right now?"
|
||||
CREATE INDEX idx_routines_next_fire
|
||||
ON routines (next_fire_at)
|
||||
WHERE enabled AND next_fire_at IS NOT NULL;
|
||||
|
||||
-- Fast lookup: event triggers for a user
|
||||
CREATE INDEX idx_routines_event_triggers
|
||||
ON routines (user_id)
|
||||
WHERE enabled AND trigger_type = 'event';
|
||||
|
||||
-- Audit log of individual routine executions.
|
||||
CREATE TABLE routine_runs (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
routine_id UUID NOT NULL REFERENCES routines(id) ON DELETE CASCADE,
|
||||
trigger_type TEXT NOT NULL,
|
||||
trigger_detail TEXT, -- e.g. matched message preview, cron expression
|
||||
started_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
completed_at TIMESTAMPTZ,
|
||||
status TEXT NOT NULL DEFAULT 'running', -- running, ok, attention, failed
|
||||
result_summary TEXT,
|
||||
tokens_used INTEGER,
|
||||
job_id UUID REFERENCES agent_jobs(id), -- non-NULL for full_job runs
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
|
||||
CREATE INDEX idx_routine_runs_routine ON routine_runs (routine_id);
|
||||
CREATE INDEX idx_routine_runs_status ON routine_runs (status) WHERE status = 'running';
|
||||
@@ -0,0 +1,3 @@
|
||||
-- Rename claude_code_events to job_events (generic for all sandbox job types).
|
||||
ALTER TABLE claude_code_events RENAME TO job_events;
|
||||
ALTER INDEX idx_cc_events_job RENAME TO idx_job_events_job;
|
||||
@@ -0,0 +1,16 @@
|
||||
-- Settings table: key-value store for all user configuration.
|
||||
--
|
||||
-- Replaces ~/.ironclaw/settings.json, session.json, and mcp-servers.json.
|
||||
-- Keys use dotted paths matching the existing Settings.get()/set() convention
|
||||
-- (e.g., "agent.name", "sandbox.enabled", "mcp_servers").
|
||||
-- One row per setting so individual values can be updated atomically.
|
||||
|
||||
CREATE TABLE IF NOT EXISTS settings (
|
||||
user_id TEXT NOT NULL,
|
||||
key TEXT NOT NULL,
|
||||
value JSONB NOT NULL,
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
PRIMARY KEY (user_id, key)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_settings_user ON settings (user_id);
|
||||
Executable
+21
@@ -0,0 +1,21 @@
|
||||
#!/usr/bin/env bash
|
||||
# Build IronClaw and all bundled channels.
|
||||
#
|
||||
# Run this before release or when channel sources have changed.
|
||||
# The main binary bundles telegram.wasm via include_bytes!; it must exist.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
cd "$(dirname "$0")/.."
|
||||
|
||||
echo "Building bundled channels..."
|
||||
if [ -d "channels-src/telegram" ]; then
|
||||
./channels-src/telegram/build.sh
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "Building IronClaw..."
|
||||
cargo build --release
|
||||
|
||||
echo ""
|
||||
echo "Done. Binary: target/release/ironclaw"
|
||||
+641
-117
@@ -9,13 +9,14 @@ use uuid::Uuid;
|
||||
use crate::agent::compaction::ContextCompactor;
|
||||
use crate::agent::context_monitor::ContextMonitor;
|
||||
use crate::agent::heartbeat::spawn_heartbeat;
|
||||
use crate::agent::routine_engine::{RoutineEngine, spawn_cron_ticker};
|
||||
use crate::agent::self_repair::{DefaultSelfRepair, RepairResult, SelfRepair};
|
||||
use crate::agent::session::{PendingApproval, Session, ThreadState};
|
||||
use crate::agent::session_manager::SessionManager;
|
||||
use crate::agent::submission::{Submission, SubmissionParser, SubmissionResult};
|
||||
use crate::agent::{HeartbeatConfig as AgentHeartbeatConfig, MessageIntent, Router, Scheduler};
|
||||
use crate::channels::{ChannelManager, IncomingMessage, OutgoingResponse, StatusUpdate};
|
||||
use crate::config::{AgentConfig, HeartbeatConfig};
|
||||
use crate::config::{AgentConfig, HeartbeatConfig, RoutineConfig};
|
||||
use crate::context::ContextManager;
|
||||
use crate::context::JobContext;
|
||||
use crate::error::Error;
|
||||
@@ -77,6 +78,7 @@ pub struct Agent {
|
||||
session_manager: Arc<SessionManager>,
|
||||
context_monitor: ContextMonitor,
|
||||
heartbeat_config: Option<HeartbeatConfig>,
|
||||
routine_config: Option<RoutineConfig>,
|
||||
}
|
||||
|
||||
impl Agent {
|
||||
@@ -89,6 +91,7 @@ impl Agent {
|
||||
deps: AgentDeps,
|
||||
channels: ChannelManager,
|
||||
heartbeat_config: Option<HeartbeatConfig>,
|
||||
routine_config: Option<RoutineConfig>,
|
||||
context_manager: Option<Arc<ContextManager>>,
|
||||
session_manager: Option<Arc<SessionManager>>,
|
||||
) -> Self {
|
||||
@@ -116,6 +119,7 @@ impl Agent {
|
||||
session_manager,
|
||||
context_monitor: ContextMonitor::new(),
|
||||
heartbeat_config,
|
||||
routine_config,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -256,53 +260,28 @@ impl Agent {
|
||||
let channels = self.channels.clone();
|
||||
tokio::spawn(async move {
|
||||
while let Some(response) = notify_rx.recv().await {
|
||||
// Route notification to configured channel/user, or broadcast to all
|
||||
match (¬ify_channel, ¬ify_user) {
|
||||
(Some(channel), Some(user)) => {
|
||||
// Send to specific channel and user
|
||||
if let Err(e) =
|
||||
channels.broadcast(channel, user, response.clone()).await
|
||||
{
|
||||
let user = notify_user.as_deref().unwrap_or("default");
|
||||
|
||||
// Try the configured channel first, fall back to
|
||||
// broadcasting on all channels.
|
||||
let targeted_ok = if let Some(ref channel) = notify_channel {
|
||||
channels
|
||||
.broadcast(channel, user, response.clone())
|
||||
.await
|
||||
.is_ok()
|
||||
} else {
|
||||
false
|
||||
};
|
||||
|
||||
if !targeted_ok {
|
||||
let results = channels.broadcast_all(user, response).await;
|
||||
for (ch, result) in results {
|
||||
if let Err(e) = result {
|
||||
tracing::warn!(
|
||||
"Failed to send heartbeat to {}/{}: {}",
|
||||
channel,
|
||||
user,
|
||||
"Failed to broadcast heartbeat to {}: {}",
|
||||
ch,
|
||||
e
|
||||
);
|
||||
} else {
|
||||
tracing::debug!(
|
||||
"Heartbeat notification sent to {}/{}",
|
||||
channel,
|
||||
user
|
||||
);
|
||||
}
|
||||
}
|
||||
(None, Some(user)) => {
|
||||
// Broadcast to all channels for this user
|
||||
let results = channels.broadcast_all(user, response).await;
|
||||
for (ch, result) in results {
|
||||
if let Err(e) = result {
|
||||
tracing::warn!(
|
||||
"Failed to broadcast heartbeat to {}: {}",
|
||||
ch,
|
||||
e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
// No explicit target, broadcast to all channels
|
||||
// for the default user so notifications actually
|
||||
// reach someone instead of vanishing into logs.
|
||||
let results = channels.broadcast_all("default", response).await;
|
||||
for (ch, result) in results {
|
||||
if let Err(e) = result {
|
||||
tracing::warn!(
|
||||
"Failed to broadcast heartbeat to {}: {}",
|
||||
ch,
|
||||
e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -330,6 +309,85 @@ impl Agent {
|
||||
None
|
||||
};
|
||||
|
||||
// Spawn routine engine if enabled
|
||||
let routine_handle = if let Some(ref rt_config) = self.routine_config {
|
||||
if rt_config.enabled {
|
||||
if let (Some(store), Some(workspace)) = (self.store(), self.workspace()) {
|
||||
// Set up notification channel (same pattern as heartbeat)
|
||||
let (notify_tx, mut notify_rx) =
|
||||
tokio::sync::mpsc::channel::<OutgoingResponse>(32);
|
||||
|
||||
let engine = Arc::new(RoutineEngine::new(
|
||||
rt_config.clone(),
|
||||
Arc::clone(store),
|
||||
self.llm().clone(),
|
||||
Arc::clone(workspace),
|
||||
notify_tx,
|
||||
));
|
||||
|
||||
// Register routine tools
|
||||
self.deps
|
||||
.tools
|
||||
.register_routine_tools(Arc::clone(store), Arc::clone(&engine));
|
||||
|
||||
// Load initial event cache
|
||||
engine.refresh_event_cache().await;
|
||||
|
||||
// Spawn notification forwarder
|
||||
let channels = self.channels.clone();
|
||||
tokio::spawn(async move {
|
||||
while let Some(response) = notify_rx.recv().await {
|
||||
let user = response
|
||||
.metadata
|
||||
.get("notify_user")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("default")
|
||||
.to_string();
|
||||
let results = channels.broadcast_all(&user, response).await;
|
||||
for (ch, result) in results {
|
||||
if let Err(e) = result {
|
||||
tracing::warn!(
|
||||
"Failed to broadcast routine notification to {}: {}",
|
||||
ch,
|
||||
e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Spawn cron ticker
|
||||
let cron_interval =
|
||||
std::time::Duration::from_secs(rt_config.cron_check_interval_secs);
|
||||
let cron_handle = spawn_cron_ticker(Arc::clone(&engine), cron_interval);
|
||||
|
||||
// Store engine reference for event trigger checking
|
||||
// Safety: we're in run() which takes self, no other reference exists
|
||||
let engine_ref = Arc::clone(&engine);
|
||||
// SAFETY: self is consumed by run(), we can smuggle the engine in
|
||||
// via a local to use in the message loop below.
|
||||
|
||||
tracing::info!(
|
||||
"Routines enabled: cron ticker every {}s, max {} concurrent",
|
||||
rt_config.cron_check_interval_secs,
|
||||
rt_config.max_concurrent_routines
|
||||
);
|
||||
|
||||
Some((cron_handle, engine_ref))
|
||||
} else {
|
||||
tracing::warn!("Routines enabled but store/workspace not available");
|
||||
None
|
||||
}
|
||||
} else {
|
||||
None
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Extract engine ref for use in message loop
|
||||
let routine_engine_for_loop = routine_handle.as_ref().map(|(_, e)| Arc::clone(e));
|
||||
|
||||
// Main message loop
|
||||
tracing::info!("Agent {} ready and listening", self.config.name);
|
||||
|
||||
@@ -374,6 +432,14 @@ impl Agent {
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
// Check event triggers (cheap in-memory regex, fires async if matched)
|
||||
if let Some(ref engine) = routine_engine_for_loop {
|
||||
let fired = engine.check_event_triggers(&message).await;
|
||||
if fired > 0 {
|
||||
tracing::debug!("Fired {} event-triggered routines", fired);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Cleanup
|
||||
@@ -383,6 +449,9 @@ impl Agent {
|
||||
if let Some(handle) = heartbeat_handle {
|
||||
handle.abort();
|
||||
}
|
||||
if let Some((cron_handle, _)) = routine_handle {
|
||||
cron_handle.abort();
|
||||
}
|
||||
self.scheduler.stop_all().await;
|
||||
self.channels.shutdown_all().await?;
|
||||
|
||||
@@ -393,6 +462,11 @@ impl Agent {
|
||||
// Parse submission type first
|
||||
let submission = SubmissionParser::parse(&message.content);
|
||||
|
||||
// Hydrate thread from DB if it's a historical thread not in memory
|
||||
if let Some(ref external_thread_id) = message.thread_id {
|
||||
self.maybe_hydrate_thread(message, external_thread_id).await;
|
||||
}
|
||||
|
||||
// Resolve session and thread
|
||||
let (session, thread_id) = self
|
||||
.session_manager
|
||||
@@ -444,6 +518,9 @@ impl Agent {
|
||||
self.process_user_input(message, session, thread_id, &content)
|
||||
.await
|
||||
}
|
||||
Submission::SystemCommand { command, args } => {
|
||||
self.handle_system_command(&command, &args).await
|
||||
}
|
||||
Submission::Undo => self.process_undo(session, thread_id).await,
|
||||
Submission::Redo => self.process_redo(session, thread_id).await,
|
||||
Submission::Interrupt => self.process_interrupt(session, thread_id).await,
|
||||
@@ -515,6 +592,107 @@ impl Agent {
|
||||
}
|
||||
}
|
||||
|
||||
/// Hydrate a historical thread from DB into memory if not already present.
|
||||
///
|
||||
/// Called before `resolve_thread` so that the session manager finds the
|
||||
/// thread on lookup instead of creating a new one.
|
||||
///
|
||||
/// Creates an in-memory thread with the exact UUID the frontend sent,
|
||||
/// even when the conversation has zero messages (e.g. a brand-new
|
||||
/// assistant thread). Without this, `resolve_thread` would mint a
|
||||
/// fresh UUID and all messages would land in the wrong conversation.
|
||||
async fn maybe_hydrate_thread(&self, message: &IncomingMessage, external_thread_id: &str) {
|
||||
// Only hydrate UUID-shaped thread IDs (web gateway uses UUIDs)
|
||||
let thread_uuid = match Uuid::parse_str(external_thread_id) {
|
||||
Ok(id) => id,
|
||||
Err(_) => return,
|
||||
};
|
||||
|
||||
// Check if already in memory
|
||||
let session = self
|
||||
.session_manager
|
||||
.get_or_create_session(&message.user_id)
|
||||
.await;
|
||||
{
|
||||
let sess = session.lock().await;
|
||||
if sess.threads.contains_key(&thread_uuid) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// Load history from DB (may be empty for a newly created thread).
|
||||
let mut chat_messages: Vec<ChatMessage> = Vec::new();
|
||||
let msg_count;
|
||||
|
||||
if let Some(store) = self.store() {
|
||||
let db_messages = store
|
||||
.list_conversation_messages(thread_uuid)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
msg_count = db_messages.len();
|
||||
chat_messages = db_messages
|
||||
.iter()
|
||||
.filter_map(|m| match m.role.as_str() {
|
||||
"user" => Some(ChatMessage::user(&m.content)),
|
||||
"assistant" => Some(ChatMessage::assistant(&m.content)),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
} else {
|
||||
msg_count = 0;
|
||||
}
|
||||
|
||||
// Create thread with the historical ID and restore messages
|
||||
let session_id = {
|
||||
let sess = session.lock().await;
|
||||
sess.id
|
||||
};
|
||||
|
||||
let mut thread = crate::agent::session::Thread::with_id(thread_uuid, session_id);
|
||||
if !chat_messages.is_empty() {
|
||||
thread.restore_from_messages(chat_messages);
|
||||
}
|
||||
|
||||
// Restore response chain from conversation metadata
|
||||
if let Some(store) = self.store() {
|
||||
if let Ok(Some(metadata)) = store.get_conversation_metadata(thread_uuid).await {
|
||||
if let Some(rid) = metadata
|
||||
.get("last_response_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from)
|
||||
{
|
||||
thread.last_response_id = Some(rid.clone());
|
||||
self.llm()
|
||||
.seed_response_chain(&thread_uuid.to_string(), rid);
|
||||
tracing::debug!("Restored response chain for thread {}", thread_uuid);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Insert into session and register with session manager
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
sess.threads.insert(thread_uuid, thread);
|
||||
sess.active_thread = Some(thread_uuid);
|
||||
sess.last_active_at = chrono::Utc::now();
|
||||
}
|
||||
|
||||
self.session_manager
|
||||
.register_thread(
|
||||
&message.user_id,
|
||||
&message.channel,
|
||||
thread_uuid,
|
||||
Arc::clone(&session),
|
||||
)
|
||||
.await;
|
||||
|
||||
tracing::debug!(
|
||||
"Hydrated thread {} from DB ({} messages)",
|
||||
thread_uuid,
|
||||
msg_count
|
||||
);
|
||||
}
|
||||
|
||||
async fn process_user_input(
|
||||
&self,
|
||||
message: &IncomingMessage,
|
||||
@@ -694,6 +872,7 @@ impl Agent {
|
||||
match result {
|
||||
Ok(AgenticLoopResult::Response(response)) => {
|
||||
thread.complete_turn(&response);
|
||||
self.persist_response_chain(thread);
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
@@ -702,6 +881,10 @@ impl Agent {
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
|
||||
// Fire-and-forget: persist turn to DB
|
||||
self.persist_turn(thread_id, &message.user_id, content, Some(&response));
|
||||
|
||||
Ok(SubmissionResult::response(response))
|
||||
}
|
||||
Ok(AgenticLoopResult::NeedApproval { pending }) => {
|
||||
@@ -728,11 +911,95 @@ impl Agent {
|
||||
}
|
||||
Err(e) => {
|
||||
thread.fail_turn(e.to_string());
|
||||
|
||||
// Persist the user message even on failure
|
||||
self.persist_turn(thread_id, &message.user_id, content, None);
|
||||
|
||||
Ok(SubmissionResult::error(e.to_string()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Fire-and-forget: persist a turn (user message + optional assistant response) to the DB.
|
||||
fn persist_turn(
|
||||
&self,
|
||||
thread_id: Uuid,
|
||||
user_id: &str,
|
||||
user_input: &str,
|
||||
response: Option<&str>,
|
||||
) {
|
||||
let store = match self.store() {
|
||||
Some(s) => Arc::clone(s),
|
||||
None => return,
|
||||
};
|
||||
|
||||
let user_id = user_id.to_string();
|
||||
let user_input = user_input.to_string();
|
||||
let response = response.map(String::from);
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = store
|
||||
.ensure_conversation(thread_id, "gateway", &user_id, None)
|
||||
.await
|
||||
{
|
||||
tracing::warn!("Failed to ensure conversation {}: {}", thread_id, e);
|
||||
return;
|
||||
}
|
||||
|
||||
if let Err(e) = store
|
||||
.add_conversation_message(thread_id, "user", &user_input)
|
||||
.await
|
||||
{
|
||||
tracing::warn!("Failed to persist user message: {}", e);
|
||||
return;
|
||||
}
|
||||
|
||||
if let Some(ref resp) = response {
|
||||
if let Err(e) = store
|
||||
.add_conversation_message(thread_id, "assistant", resp)
|
||||
.await
|
||||
{
|
||||
tracing::warn!("Failed to persist assistant message: {}", e);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Sync the provider's response chain ID to the thread and DB metadata.
|
||||
///
|
||||
/// Call after a successful agentic loop to persist the latest
|
||||
/// `previous_response_id` so chaining survives restarts.
|
||||
fn persist_response_chain(&self, thread: &mut crate::agent::session::Thread) {
|
||||
let tid = thread.id.to_string();
|
||||
let response_id = match self.llm().get_response_chain_id(&tid) {
|
||||
Some(rid) => rid,
|
||||
None => return,
|
||||
};
|
||||
|
||||
// Update in-memory thread
|
||||
thread.last_response_id = Some(response_id.clone());
|
||||
|
||||
// Fire-and-forget DB write
|
||||
let store = match self.store() {
|
||||
Some(s) => Arc::clone(s),
|
||||
None => return,
|
||||
};
|
||||
let thread_id = thread.id;
|
||||
tokio::spawn(async move {
|
||||
let val = serde_json::json!(response_id);
|
||||
if let Err(e) = store
|
||||
.update_conversation_metadata_field(thread_id, "last_response_id", &val)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
"Failed to persist response chain for thread {}: {}",
|
||||
thread_id,
|
||||
e
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Run the agentic loop: call LLM, execute tools, repeat until text response.
|
||||
///
|
||||
/// Returns `AgenticLoopResult::Response` on completion, or
|
||||
@@ -782,7 +1049,7 @@ impl Agent {
|
||||
iteration += 1;
|
||||
if iteration > MAX_TOOL_ITERATIONS {
|
||||
return Err(crate::error::LlmError::InvalidResponse {
|
||||
provider: "nearai".to_string(),
|
||||
provider: "agent".to_string(),
|
||||
reason: format!("Exceeded maximum tool iterations ({})", MAX_TOOL_ITERATIONS),
|
||||
}
|
||||
.into());
|
||||
@@ -808,7 +1075,12 @@ impl Agent {
|
||||
// Call LLM with current context
|
||||
let context = ReasoningContext::new()
|
||||
.with_messages(context_messages.clone())
|
||||
.with_tools(tool_defs);
|
||||
.with_tools(tool_defs)
|
||||
.with_metadata({
|
||||
let mut m = std::collections::HashMap::new();
|
||||
m.insert("thread_id".to_string(), thread_id.to_string());
|
||||
m
|
||||
});
|
||||
|
||||
let result = reasoning.respond_with_tools(&context).await?;
|
||||
|
||||
@@ -832,13 +1104,16 @@ impl Agent {
|
||||
// Tools have been executed or we've tried multiple times, return response
|
||||
return Ok(AgenticLoopResult::Response(text));
|
||||
}
|
||||
RespondResult::ToolCalls(tool_calls) => {
|
||||
RespondResult::ToolCalls {
|
||||
tool_calls,
|
||||
content,
|
||||
} => {
|
||||
tools_executed = true;
|
||||
|
||||
// Add the assistant message with tool_calls to context.
|
||||
// OpenAI-compatible APIs require this before tool-result messages.
|
||||
// OpenAI protocol requires this before tool-result messages.
|
||||
context_messages.push(ChatMessage::assistant_with_tool_calls(
|
||||
"",
|
||||
content,
|
||||
tool_calls.clone(),
|
||||
));
|
||||
|
||||
@@ -960,10 +1235,26 @@ impl Agent {
|
||||
if let Some((ext_name, instructions)) =
|
||||
detect_auth_awaiting(&tc.name, &tool_result)
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
thread.enter_auth_mode(ext_name);
|
||||
let auth_data = parse_auth_result(&tool_result);
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
thread.enter_auth_mode(ext_name.clone());
|
||||
}
|
||||
}
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::AuthRequired {
|
||||
extension_name: ext_name,
|
||||
instructions: Some(instructions.clone()),
|
||||
auth_url: auth_data.auth_url,
|
||||
setup_url: auth_data.setup_url,
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
return Ok(AgenticLoopResult::Response(instructions));
|
||||
}
|
||||
|
||||
@@ -1024,19 +1315,59 @@ impl Agent {
|
||||
.into());
|
||||
}
|
||||
|
||||
// Execute with timeout
|
||||
let result = tokio::time::timeout(std::time::Duration::from_secs(60), async {
|
||||
tracing::debug!(
|
||||
tool = %tool_name,
|
||||
params = %params,
|
||||
"Tool call started"
|
||||
);
|
||||
|
||||
// Execute with per-tool timeout
|
||||
let timeout = tool.execution_timeout();
|
||||
let start = std::time::Instant::now();
|
||||
let result = tokio::time::timeout(timeout, async {
|
||||
tool.execute(params.clone(), job_ctx).await
|
||||
})
|
||||
.await
|
||||
.map_err(|_| crate::error::ToolError::Timeout {
|
||||
name: tool_name.to_string(),
|
||||
timeout: std::time::Duration::from_secs(60),
|
||||
})?
|
||||
.map_err(|e| crate::error::ToolError::ExecutionFailed {
|
||||
name: tool_name.to_string(),
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
.await;
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
match &result {
|
||||
Ok(Ok(output)) => {
|
||||
let result_str = serde_json::to_string(&output.result)
|
||||
.unwrap_or_else(|_| "<serialize error>".to_string());
|
||||
tracing::debug!(
|
||||
tool = %tool_name,
|
||||
elapsed_ms = elapsed.as_millis() as u64,
|
||||
result = %result_str,
|
||||
"Tool call succeeded"
|
||||
);
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
tracing::debug!(
|
||||
tool = %tool_name,
|
||||
elapsed_ms = elapsed.as_millis() as u64,
|
||||
error = %e,
|
||||
"Tool call failed"
|
||||
);
|
||||
}
|
||||
Err(_) => {
|
||||
tracing::debug!(
|
||||
tool = %tool_name,
|
||||
elapsed_ms = elapsed.as_millis() as u64,
|
||||
timeout_secs = timeout.as_secs(),
|
||||
"Tool call timed out"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let result = result
|
||||
.map_err(|_| crate::error::ToolError::Timeout {
|
||||
name: tool_name.to_string(),
|
||||
timeout,
|
||||
})?
|
||||
.map_err(|e| crate::error::ToolError::ExecutionFailed {
|
||||
name: tool_name.to_string(),
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
// Convert result to string
|
||||
serde_json::to_string_pretty(&result.result).map_err(|e| {
|
||||
@@ -1380,10 +1711,11 @@ impl Agent {
|
||||
if let Some((ext_name, instructions)) =
|
||||
detect_auth_awaiting(&pending.tool_name, &tool_result)
|
||||
{
|
||||
let auth_data = parse_auth_result(&tool_result);
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
thread.enter_auth_mode(ext_name);
|
||||
thread.enter_auth_mode(ext_name.clone());
|
||||
thread.complete_turn(&instructions);
|
||||
}
|
||||
}
|
||||
@@ -1391,7 +1723,12 @@ impl Agent {
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::Status("Awaiting token".into()),
|
||||
StatusUpdate::AuthRequired {
|
||||
extension_name: ext_name,
|
||||
instructions: Some(instructions.clone()),
|
||||
auth_url: auth_data.auth_url,
|
||||
setup_url: auth_data.setup_url,
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
@@ -1434,6 +1771,7 @@ impl Agent {
|
||||
match result {
|
||||
Ok(AgenticLoopResult::Response(response)) => {
|
||||
thread.complete_turn(&response);
|
||||
self.persist_response_chain(thread);
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
@@ -1532,16 +1870,6 @@ impl Agent {
|
||||
pending.extension_name
|
||||
);
|
||||
|
||||
// Notify via channel status so the response doesn't echo the token
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::Status("Authenticated, loading tools...".into()),
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
|
||||
// Auto-activate so tools are available immediately after auth
|
||||
match ext_mgr.activate(&pending.extension_name).await {
|
||||
Ok(activate_result) => {
|
||||
@@ -1551,10 +1879,23 @@ impl Agent {
|
||||
} else {
|
||||
format!("\n\nTools: {}", activate_result.tools_loaded.join(", "))
|
||||
};
|
||||
Ok(Some(format!(
|
||||
let msg = format!(
|
||||
"{} authenticated and activated ({} tools loaded).{}",
|
||||
pending.extension_name, tool_count, tool_list
|
||||
)))
|
||||
);
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::AuthCompleted {
|
||||
extension_name: pending.extension_name.clone(),
|
||||
success: true,
|
||||
message: msg.clone(),
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
Ok(Some(msg))
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
@@ -1562,16 +1903,29 @@ impl Agent {
|
||||
pending.extension_name,
|
||||
e
|
||||
);
|
||||
Ok(Some(format!(
|
||||
let msg = format!(
|
||||
"{} authenticated successfully, but activation failed: {}. \
|
||||
Try activating manually.",
|
||||
pending.extension_name, e
|
||||
)))
|
||||
);
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::AuthCompleted {
|
||||
extension_name: pending.extension_name.clone(),
|
||||
success: true,
|
||||
message: msg.clone(),
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
Ok(Some(msg))
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(result) => {
|
||||
// Unexpected state, re-enter auth mode
|
||||
// Invalid token, re-enter auth mode
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
@@ -1580,13 +1934,43 @@ impl Agent {
|
||||
}
|
||||
let msg = result
|
||||
.instructions
|
||||
.clone()
|
||||
.unwrap_or_else(|| "Invalid token. Please try again.".to_string());
|
||||
// Re-emit AuthRequired so web UI re-shows the card
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::AuthRequired {
|
||||
extension_name: pending.extension_name.clone(),
|
||||
instructions: Some(msg.clone()),
|
||||
auth_url: result.auth_url,
|
||||
setup_url: result.setup_url,
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
Ok(Some(msg))
|
||||
}
|
||||
Err(e) => {
|
||||
let msg = format!(
|
||||
"Authentication failed for {}: {}",
|
||||
pending.extension_name, e
|
||||
);
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::AuthCompleted {
|
||||
extension_name: pending.extension_name.clone(),
|
||||
success: false,
|
||||
message: msg.clone(),
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
Ok(Some(msg))
|
||||
}
|
||||
Err(e) => Ok(Some(format!(
|
||||
"Authentication failed for {}: {}",
|
||||
pending.extension_name, e
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1937,40 +2321,49 @@ impl Agent {
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_command(
|
||||
/// Handle system commands that bypass thread-state checks entirely.
|
||||
async fn handle_system_command(
|
||||
&self,
|
||||
command: &str,
|
||||
_args: &[String],
|
||||
) -> Result<Option<String>, Error> {
|
||||
args: &[String],
|
||||
) -> Result<SubmissionResult, Error> {
|
||||
match command {
|
||||
"help" => Ok(Some(
|
||||
r#"Commands:
|
||||
/job <desc> - Create a job
|
||||
/status [id] - Check job status
|
||||
/cancel <id> - Cancel a job
|
||||
/list - List all jobs
|
||||
/help <job_id> - Help a stuck job
|
||||
"help" => Ok(SubmissionResult::response(concat!(
|
||||
"System:\n",
|
||||
" /help Show this help\n",
|
||||
" /model [name] Show or switch the active model\n",
|
||||
" /version Show version info\n",
|
||||
" /tools List available tools\n",
|
||||
" /debug Toggle debug mode\n",
|
||||
" /ping Connectivity check\n",
|
||||
"\n",
|
||||
"Jobs:\n",
|
||||
" /job <desc> Create a new job\n",
|
||||
" /status [id] Check job status\n",
|
||||
" /cancel <id> Cancel a job\n",
|
||||
" /list List all jobs\n",
|
||||
"\n",
|
||||
"Session:\n",
|
||||
" /undo Undo last turn\n",
|
||||
" /redo Redo undone turn\n",
|
||||
" /compact Compress context window\n",
|
||||
" /clear Clear current thread\n",
|
||||
" /interrupt Stop current operation\n",
|
||||
" /new New conversation thread\n",
|
||||
" /thread <id> Switch to thread\n",
|
||||
" /resume <id> Resume from checkpoint\n",
|
||||
"\n",
|
||||
"Agent:\n",
|
||||
" /heartbeat Run heartbeat check\n",
|
||||
" /summarize Summarize current thread\n",
|
||||
" /suggest Suggest next steps\n",
|
||||
"\n",
|
||||
" /quit Exit",
|
||||
))),
|
||||
|
||||
/undo - Undo last turn
|
||||
/redo - Redo undone turn
|
||||
/compact - Compress context
|
||||
/clear - Clear thread
|
||||
/interrupt - Stop current turn
|
||||
/thread new - New thread
|
||||
/thread <id> - Switch thread
|
||||
/resume <id> - Resume checkpoint
|
||||
"ping" => Ok(SubmissionResult::response("pong!")),
|
||||
|
||||
/heartbeat - Run heartbeat check now
|
||||
/summarize - Summarize current thread
|
||||
/suggest - Suggest next steps
|
||||
|
||||
/quit - Exit"#
|
||||
.to_string(),
|
||||
)),
|
||||
|
||||
"ping" => Ok(Some("pong!".to_string())),
|
||||
|
||||
"version" => Ok(Some(format!(
|
||||
"version" => Ok(SubmissionResult::response(format!(
|
||||
"{} v{}",
|
||||
env!("CARGO_PKG_NAME"),
|
||||
env!("CARGO_PKG_VERSION")
|
||||
@@ -1978,12 +2371,113 @@ impl Agent {
|
||||
|
||||
"tools" => {
|
||||
let tools = self.tools().list().await;
|
||||
Ok(Some(format!("Available tools: {}", tools.join(", "))))
|
||||
Ok(SubmissionResult::response(format!(
|
||||
"Available tools: {}",
|
||||
tools.join(", ")
|
||||
)))
|
||||
}
|
||||
|
||||
_ => Ok(Some(format!("Unknown command: {}. Try /help", command))),
|
||||
"debug" => {
|
||||
// Debug toggle is handled client-side in the REPL.
|
||||
// For non-REPL channels, just acknowledge.
|
||||
Ok(SubmissionResult::ok_with_message(
|
||||
"Debug toggle is handled by your client.",
|
||||
))
|
||||
}
|
||||
|
||||
"model" => {
|
||||
if args.is_empty() {
|
||||
// Show current model
|
||||
let name = self.llm().active_model_name();
|
||||
Ok(SubmissionResult::response(format!(
|
||||
"Active model: {}",
|
||||
name
|
||||
)))
|
||||
} else {
|
||||
let requested = &args[0];
|
||||
|
||||
// Validate the model exists
|
||||
match self.llm().list_models().await {
|
||||
Ok(models) if !models.is_empty() => {
|
||||
if !models.iter().any(|m| m == requested) {
|
||||
return Ok(SubmissionResult::error(format!(
|
||||
"Unknown model: {}. Available models:\n {}",
|
||||
requested,
|
||||
models.join("\n ")
|
||||
)));
|
||||
}
|
||||
}
|
||||
Ok(_) => {
|
||||
// Empty model list, can't validate but try anyway
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Could not fetch model list for validation: {}", e);
|
||||
// Proceed anyway, the provider will error on the next call if invalid
|
||||
}
|
||||
}
|
||||
|
||||
match self.llm().set_model(requested) {
|
||||
Ok(()) => Ok(SubmissionResult::response(format!(
|
||||
"Switched model to: {}",
|
||||
requested
|
||||
))),
|
||||
Err(e) => Ok(SubmissionResult::error(format!(
|
||||
"Failed to switch model: {}",
|
||||
e
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
_ => Ok(SubmissionResult::error(format!(
|
||||
"Unknown command: {}. Try /help",
|
||||
command
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle legacy command routing from the Router (job commands that go through
|
||||
/// process_user_input -> router -> handle_job_or_command -> here).
|
||||
async fn handle_command(
|
||||
&self,
|
||||
command: &str,
|
||||
args: &[String],
|
||||
) -> Result<Option<String>, Error> {
|
||||
// System commands are now handled directly via Submission::SystemCommand,
|
||||
// but the router may still send us unknown /commands.
|
||||
match self.handle_system_command(command, args).await? {
|
||||
SubmissionResult::Response { content } => Ok(Some(content)),
|
||||
SubmissionResult::Ok { message } => Ok(message),
|
||||
SubmissionResult::Error { message } => Ok(Some(format!("Error: {}", message))),
|
||||
_ => Ok(None),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Parsed auth result fields for emitting StatusUpdate::AuthRequired.
|
||||
struct ParsedAuthData {
|
||||
auth_url: Option<String>,
|
||||
setup_url: Option<String>,
|
||||
}
|
||||
|
||||
/// Extract auth_url and setup_url from a tool_auth result JSON string.
|
||||
fn parse_auth_result(result: &Result<String, Error>) -> ParsedAuthData {
|
||||
let parsed = result
|
||||
.as_ref()
|
||||
.ok()
|
||||
.and_then(|s| serde_json::from_str::<serde_json::Value>(s).ok());
|
||||
ParsedAuthData {
|
||||
auth_url: parsed
|
||||
.as_ref()
|
||||
.and_then(|v| v.get("auth_url"))
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string()),
|
||||
setup_url: parsed
|
||||
.as_ref()
|
||||
.and_then(|v| v.get("setup_url"))
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if a tool_auth result indicates the extension is awaiting a token.
|
||||
@@ -1994,7 +2488,7 @@ fn detect_auth_awaiting(
|
||||
tool_name: &str,
|
||||
result: &Result<String, Error>,
|
||||
) -> Option<(String, String)> {
|
||||
if tool_name != "tool_auth" {
|
||||
if tool_name != "tool_auth" && tool_name != "tool_activate" {
|
||||
return None;
|
||||
}
|
||||
let output = result.as_ref().ok()?;
|
||||
@@ -2078,4 +2572,34 @@ mod tests {
|
||||
let (_, instructions) = detect_auth_awaiting("tool_auth", &result).unwrap();
|
||||
assert_eq!(instructions, "Please provide your API token/key.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_auth_awaiting_tool_activate() {
|
||||
let result: Result<String, Error> = Ok(serde_json::json!({
|
||||
"name": "slack",
|
||||
"kind": "McpServer",
|
||||
"awaiting_token": true,
|
||||
"status": "awaiting_token",
|
||||
"instructions": "Provide your Slack Bot token."
|
||||
})
|
||||
.to_string());
|
||||
|
||||
let detected = detect_auth_awaiting("tool_activate", &result);
|
||||
assert!(detected.is_some());
|
||||
let (name, instructions) = detected.unwrap();
|
||||
assert_eq!(name, "slack");
|
||||
assert!(instructions.contains("Slack Bot"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_auth_awaiting_tool_activate_not_awaiting() {
|
||||
let result: Result<String, Error> = Ok(serde_json::json!({
|
||||
"name": "slack",
|
||||
"tools_loaded": ["slack_post_message"],
|
||||
"message": "Activated"
|
||||
})
|
||||
.to_string());
|
||||
|
||||
assert!(detect_auth_awaiting("tool_activate", &result).is_none());
|
||||
}
|
||||
}
|
||||
|
||||
+34
-3
@@ -29,7 +29,7 @@ use std::time::Duration;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::channels::OutgoingResponse;
|
||||
use crate::llm::{ChatMessage, CompletionRequest, LlmProvider};
|
||||
use crate::llm::{ChatMessage, CompletionRequest, FinishReason, LlmProvider};
|
||||
use crate::workspace::Workspace;
|
||||
|
||||
/// Configuration for the heartbeat runner.
|
||||
@@ -217,9 +217,26 @@ impl HeartbeatRunner {
|
||||
]
|
||||
};
|
||||
|
||||
// Use the model's context_length to set max_tokens. The API returns
|
||||
// the total context window; we cap output at half of that (the rest is
|
||||
// the prompt) with a floor of 4096.
|
||||
let max_tokens = match self.llm.model_metadata().await {
|
||||
Ok(meta) => {
|
||||
let from_api = meta.context_length.map(|ctx| ctx / 2).unwrap_or(4096);
|
||||
from_api.max(4096)
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Could not fetch model metadata, using default max_tokens: {}",
|
||||
e
|
||||
);
|
||||
4096
|
||||
}
|
||||
};
|
||||
|
||||
let request = CompletionRequest::new(messages)
|
||||
.with_max_tokens(1024)
|
||||
.with_temperature(0.3); // Lower temperature for more focused responses
|
||||
.with_max_tokens(max_tokens)
|
||||
.with_temperature(0.3);
|
||||
|
||||
let response = match self.llm.complete(request).await {
|
||||
Ok(r) => r,
|
||||
@@ -228,6 +245,20 @@ impl HeartbeatRunner {
|
||||
|
||||
let content = response.content.trim();
|
||||
|
||||
// Guard against empty content. Reasoning models (e.g. GLM-4.7) may
|
||||
// burn all output tokens on chain-of-thought and return content: null.
|
||||
if content.is_empty() {
|
||||
return if response.finish_reason == FinishReason::Length {
|
||||
HeartbeatResult::Failed(
|
||||
"LLM response was truncated (finish_reason=length) with no content. \
|
||||
The model may have exhausted its token budget on reasoning."
|
||||
.to_string(),
|
||||
)
|
||||
} else {
|
||||
HeartbeatResult::Failed("LLM returned empty content.".to_string())
|
||||
};
|
||||
}
|
||||
|
||||
// Check if nothing needs attention
|
||||
if content == "HEARTBEAT_OK" || content.contains("HEARTBEAT_OK") {
|
||||
return HeartbeatResult::Ok;
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
//! - Tool invocation with safety
|
||||
//! - Self-repair for stuck jobs
|
||||
//! - Proactive heartbeat execution
|
||||
//! - Routine-based scheduled and reactive jobs
|
||||
//! - Turn-based session management with undo
|
||||
//! - Context compaction for long conversations
|
||||
|
||||
@@ -14,6 +15,8 @@ pub mod compaction;
|
||||
pub mod context_monitor;
|
||||
mod heartbeat;
|
||||
mod router;
|
||||
pub mod routine;
|
||||
pub mod routine_engine;
|
||||
mod scheduler;
|
||||
mod self_repair;
|
||||
pub mod session;
|
||||
@@ -28,6 +31,8 @@ pub use compaction::{CompactionResult, ContextCompactor};
|
||||
pub use context_monitor::{CompactionStrategy, ContextBreakdown, ContextMonitor};
|
||||
pub use heartbeat::{HeartbeatConfig, HeartbeatResult, HeartbeatRunner, spawn_heartbeat};
|
||||
pub use router::{MessageIntent, Router};
|
||||
pub use routine::{Routine, RoutineAction, RoutineRun, Trigger};
|
||||
pub use routine_engine::RoutineEngine;
|
||||
pub use scheduler::Scheduler;
|
||||
pub use self_repair::{BrokenTool, RepairResult, RepairTask, SelfRepair, StuckJob};
|
||||
pub use session::{PendingApproval, PendingAuth, Session, Thread, ThreadState, Turn, TurnState};
|
||||
|
||||
@@ -0,0 +1,509 @@
|
||||
//! Core types for the routines system.
|
||||
//!
|
||||
//! A routine is a named, persistent, user-owned task with a trigger and an action.
|
||||
//! Each routine fires independently when its trigger condition is met, with only
|
||||
//! that routine's prompt and context sent to the LLM.
|
||||
//!
|
||||
//! ```text
|
||||
//! ┌──────────┐ ┌─────────┐ ┌──────────────────┐
|
||||
//! │ Trigger │────▶│ Engine │────▶│ Execution Mode │
|
||||
//! │ cron/event│ │guardrail│ │lightweight│full_job│
|
||||
//! │ webhook │ │ check │ └──────────────────┘
|
||||
//! │ manual │ └─────────┘ │
|
||||
//! └──────────┘ ▼
|
||||
//! ┌──────────────┐
|
||||
//! │ Notify user │
|
||||
//! │ if needed │
|
||||
//! └──────────────┘
|
||||
//! ```
|
||||
|
||||
use std::collections::hash_map::DefaultHasher;
|
||||
use std::hash::{Hash, Hasher};
|
||||
use std::str::FromStr;
|
||||
use std::time::Duration;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use uuid::Uuid;
|
||||
|
||||
/// A routine is a named, persistent, user-owned task with a trigger and an action.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Routine {
|
||||
pub id: Uuid,
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
pub user_id: String,
|
||||
pub enabled: bool,
|
||||
pub trigger: Trigger,
|
||||
pub action: RoutineAction,
|
||||
pub guardrails: RoutineGuardrails,
|
||||
pub notify: NotifyConfig,
|
||||
|
||||
// Runtime state (DB-managed)
|
||||
pub last_run_at: Option<DateTime<Utc>>,
|
||||
pub next_fire_at: Option<DateTime<Utc>>,
|
||||
pub run_count: u64,
|
||||
pub consecutive_failures: u32,
|
||||
pub state: serde_json::Value,
|
||||
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// When a routine should fire.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum Trigger {
|
||||
/// Fire on a cron schedule (e.g. "0 9 * * MON-FRI" or "every 2h").
|
||||
Cron { schedule: String },
|
||||
/// Fire when a channel message matches a pattern.
|
||||
Event {
|
||||
/// Optional channel filter (e.g. "telegram", "slack").
|
||||
channel: Option<String>,
|
||||
/// Regex pattern to match against message content.
|
||||
pattern: String,
|
||||
},
|
||||
/// Fire on incoming webhook POST to /hooks/routine/{id}.
|
||||
Webhook {
|
||||
/// Optional webhook path suffix (defaults to routine id).
|
||||
path: Option<String>,
|
||||
/// Optional shared secret for HMAC validation.
|
||||
secret: Option<String>,
|
||||
},
|
||||
/// Only fires via tool call or CLI.
|
||||
Manual,
|
||||
}
|
||||
|
||||
impl Trigger {
|
||||
/// The string tag stored in the DB trigger_type column.
|
||||
pub fn type_tag(&self) -> &'static str {
|
||||
match self {
|
||||
Trigger::Cron { .. } => "cron",
|
||||
Trigger::Event { .. } => "event",
|
||||
Trigger::Webhook { .. } => "webhook",
|
||||
Trigger::Manual => "manual",
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse a trigger from its DB representation.
|
||||
pub fn from_db(trigger_type: &str, config: serde_json::Value) -> Result<Self, String> {
|
||||
match trigger_type {
|
||||
"cron" => {
|
||||
let schedule = config
|
||||
.get("schedule")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or("cron trigger missing 'schedule'")?
|
||||
.to_string();
|
||||
Ok(Trigger::Cron { schedule })
|
||||
}
|
||||
"event" => {
|
||||
let pattern = config
|
||||
.get("pattern")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or("event trigger missing 'pattern'")?
|
||||
.to_string();
|
||||
let channel = config
|
||||
.get("channel")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from);
|
||||
Ok(Trigger::Event { channel, pattern })
|
||||
}
|
||||
"webhook" => {
|
||||
let path = config
|
||||
.get("path")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from);
|
||||
let secret = config
|
||||
.get("secret")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from);
|
||||
Ok(Trigger::Webhook { path, secret })
|
||||
}
|
||||
"manual" => Ok(Trigger::Manual),
|
||||
other => Err(format!("unknown trigger type: {other}")),
|
||||
}
|
||||
}
|
||||
|
||||
/// Serialize trigger-specific config to JSON for DB storage.
|
||||
pub fn to_config_json(&self) -> serde_json::Value {
|
||||
match self {
|
||||
Trigger::Cron { schedule } => serde_json::json!({ "schedule": schedule }),
|
||||
Trigger::Event { channel, pattern } => serde_json::json!({
|
||||
"pattern": pattern,
|
||||
"channel": channel,
|
||||
}),
|
||||
Trigger::Webhook { path, secret } => serde_json::json!({
|
||||
"path": path,
|
||||
"secret": secret,
|
||||
}),
|
||||
Trigger::Manual => serde_json::json!({}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What happens when a routine fires.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum RoutineAction {
|
||||
/// Single LLM call, no tools. Cheap and fast.
|
||||
Lightweight {
|
||||
/// The prompt sent to the LLM.
|
||||
prompt: String,
|
||||
/// Workspace paths to load as context (e.g. ["context/priorities.md"]).
|
||||
#[serde(default)]
|
||||
context_paths: Vec<String>,
|
||||
/// Max output tokens (default: 4096).
|
||||
#[serde(default = "default_max_tokens")]
|
||||
max_tokens: u32,
|
||||
},
|
||||
/// Full multi-turn worker job with tool access.
|
||||
FullJob {
|
||||
/// Job title for the scheduler.
|
||||
title: String,
|
||||
/// Job description / initial prompt.
|
||||
description: String,
|
||||
/// Max reasoning iterations (default: 10).
|
||||
#[serde(default = "default_max_iterations")]
|
||||
max_iterations: u32,
|
||||
},
|
||||
}
|
||||
|
||||
fn default_max_tokens() -> u32 {
|
||||
4096
|
||||
}
|
||||
|
||||
fn default_max_iterations() -> u32 {
|
||||
10
|
||||
}
|
||||
|
||||
impl RoutineAction {
|
||||
/// The string tag stored in the DB action_type column.
|
||||
pub fn type_tag(&self) -> &'static str {
|
||||
match self {
|
||||
RoutineAction::Lightweight { .. } => "lightweight",
|
||||
RoutineAction::FullJob { .. } => "full_job",
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse an action from its DB representation.
|
||||
pub fn from_db(action_type: &str, config: serde_json::Value) -> Result<Self, String> {
|
||||
match action_type {
|
||||
"lightweight" => {
|
||||
let prompt = config
|
||||
.get("prompt")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or("lightweight action missing 'prompt'")?
|
||||
.to_string();
|
||||
let context_paths = config
|
||||
.get("context_paths")
|
||||
.and_then(|v| v.as_array())
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.filter_map(|v| v.as_str().map(String::from))
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
let max_tokens = config
|
||||
.get("max_tokens")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(default_max_tokens() as u64) as u32;
|
||||
Ok(RoutineAction::Lightweight {
|
||||
prompt,
|
||||
context_paths,
|
||||
max_tokens,
|
||||
})
|
||||
}
|
||||
"full_job" => {
|
||||
let title = config
|
||||
.get("title")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or("full_job action missing 'title'")?
|
||||
.to_string();
|
||||
let description = config
|
||||
.get("description")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or("full_job action missing 'description'")?
|
||||
.to_string();
|
||||
let max_iterations = config
|
||||
.get("max_iterations")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(default_max_iterations() as u64)
|
||||
as u32;
|
||||
Ok(RoutineAction::FullJob {
|
||||
title,
|
||||
description,
|
||||
max_iterations,
|
||||
})
|
||||
}
|
||||
other => Err(format!("unknown action type: {other}")),
|
||||
}
|
||||
}
|
||||
|
||||
/// Serialize action config to JSON for DB storage.
|
||||
pub fn to_config_json(&self) -> serde_json::Value {
|
||||
match self {
|
||||
RoutineAction::Lightweight {
|
||||
prompt,
|
||||
context_paths,
|
||||
max_tokens,
|
||||
} => serde_json::json!({
|
||||
"prompt": prompt,
|
||||
"context_paths": context_paths,
|
||||
"max_tokens": max_tokens,
|
||||
}),
|
||||
RoutineAction::FullJob {
|
||||
title,
|
||||
description,
|
||||
max_iterations,
|
||||
} => serde_json::json!({
|
||||
"title": title,
|
||||
"description": description,
|
||||
"max_iterations": max_iterations,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Guardrails to prevent runaway execution.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct RoutineGuardrails {
|
||||
/// Minimum time between fires.
|
||||
pub cooldown: Duration,
|
||||
/// Max simultaneous runs of this routine.
|
||||
pub max_concurrent: u32,
|
||||
/// Window for content-hash dedup (event triggers). None = no dedup.
|
||||
pub dedup_window: Option<Duration>,
|
||||
}
|
||||
|
||||
impl Default for RoutineGuardrails {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
cooldown: Duration::from_secs(300),
|
||||
max_concurrent: 1,
|
||||
dedup_window: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Notification preferences for a routine.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct NotifyConfig {
|
||||
/// Channel to notify on (None = default/broadcast all).
|
||||
pub channel: Option<String>,
|
||||
/// User to notify.
|
||||
pub user: String,
|
||||
/// Notify when routine produces actionable output.
|
||||
pub on_attention: bool,
|
||||
/// Notify when routine errors.
|
||||
pub on_failure: bool,
|
||||
/// Notify when routine runs with no findings.
|
||||
pub on_success: bool,
|
||||
}
|
||||
|
||||
impl Default for NotifyConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
channel: None,
|
||||
user: "default".to_string(),
|
||||
on_attention: true,
|
||||
on_failure: true,
|
||||
on_success: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Status of a routine run.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum RunStatus {
|
||||
Running,
|
||||
Ok,
|
||||
Attention,
|
||||
Failed,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for RunStatus {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
RunStatus::Running => write!(f, "running"),
|
||||
RunStatus::Ok => write!(f, "ok"),
|
||||
RunStatus::Attention => write!(f, "attention"),
|
||||
RunStatus::Failed => write!(f, "failed"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl FromStr for RunStatus {
|
||||
type Err = String;
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
match s {
|
||||
"running" => Ok(RunStatus::Running),
|
||||
"ok" => Ok(RunStatus::Ok),
|
||||
"attention" => Ok(RunStatus::Attention),
|
||||
"failed" => Ok(RunStatus::Failed),
|
||||
other => Err(format!("unknown run status: {other}")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A single execution of a routine.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct RoutineRun {
|
||||
pub id: Uuid,
|
||||
pub routine_id: Uuid,
|
||||
pub trigger_type: String,
|
||||
pub trigger_detail: Option<String>,
|
||||
pub started_at: DateTime<Utc>,
|
||||
pub completed_at: Option<DateTime<Utc>>,
|
||||
pub status: RunStatus,
|
||||
pub result_summary: Option<String>,
|
||||
pub tokens_used: Option<i32>,
|
||||
pub job_id: Option<Uuid>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// Compute a content hash for event dedup.
|
||||
pub fn content_hash(content: &str) -> u64 {
|
||||
let mut hasher = DefaultHasher::new();
|
||||
content.hash(&mut hasher);
|
||||
hasher.finish()
|
||||
}
|
||||
|
||||
/// Parse a cron expression and compute the next fire time from now.
|
||||
pub fn next_cron_fire(schedule: &str) -> Result<Option<DateTime<Utc>>, String> {
|
||||
let cron_schedule =
|
||||
cron::Schedule::from_str(schedule).map_err(|e| format!("invalid cron: {e}"))?;
|
||||
Ok(cron_schedule.upcoming(Utc).next())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::agent::routine::{
|
||||
RoutineAction, RoutineGuardrails, RunStatus, Trigger, content_hash, next_cron_fire,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn test_trigger_roundtrip() {
|
||||
let trigger = Trigger::Cron {
|
||||
schedule: "0 9 * * MON-FRI".to_string(),
|
||||
};
|
||||
let json = trigger.to_config_json();
|
||||
let parsed = Trigger::from_db("cron", json).expect("parse cron");
|
||||
assert!(matches!(parsed, Trigger::Cron { schedule } if schedule == "0 9 * * MON-FRI"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_event_trigger_roundtrip() {
|
||||
let trigger = Trigger::Event {
|
||||
channel: Some("telegram".to_string()),
|
||||
pattern: r"deploy\s+\w+".to_string(),
|
||||
};
|
||||
let json = trigger.to_config_json();
|
||||
let parsed = Trigger::from_db("event", json).expect("parse event");
|
||||
assert!(matches!(parsed, Trigger::Event { channel, pattern }
|
||||
if channel == Some("telegram".to_string()) && pattern == r"deploy\s+\w+"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_action_lightweight_roundtrip() {
|
||||
let action = RoutineAction::Lightweight {
|
||||
prompt: "Check PRs".to_string(),
|
||||
context_paths: vec!["context/priorities.md".to_string()],
|
||||
max_tokens: 2048,
|
||||
};
|
||||
let json = action.to_config_json();
|
||||
let parsed = RoutineAction::from_db("lightweight", json).expect("parse lightweight");
|
||||
assert!(
|
||||
matches!(parsed, RoutineAction::Lightweight { prompt, context_paths, max_tokens }
|
||||
if prompt == "Check PRs" && context_paths.len() == 1 && max_tokens == 2048)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_action_full_job_roundtrip() {
|
||||
let action = RoutineAction::FullJob {
|
||||
title: "Deploy review".to_string(),
|
||||
description: "Review and deploy pending changes".to_string(),
|
||||
max_iterations: 5,
|
||||
};
|
||||
let json = action.to_config_json();
|
||||
let parsed = RoutineAction::from_db("full_job", json).expect("parse full_job");
|
||||
assert!(
|
||||
matches!(parsed, RoutineAction::FullJob { title, max_iterations, .. }
|
||||
if title == "Deploy review" && max_iterations == 5)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_run_status_display_parse() {
|
||||
for status in [
|
||||
RunStatus::Running,
|
||||
RunStatus::Ok,
|
||||
RunStatus::Attention,
|
||||
RunStatus::Failed,
|
||||
] {
|
||||
let s = status.to_string();
|
||||
let parsed: RunStatus = s.parse().expect("parse status");
|
||||
assert_eq!(parsed, status);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_content_hash_deterministic() {
|
||||
let h1 = content_hash("deploy production");
|
||||
let h2 = content_hash("deploy production");
|
||||
assert_eq!(h1, h2);
|
||||
|
||||
let h3 = content_hash("deploy staging");
|
||||
assert_ne!(h1, h3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_next_cron_fire_valid() {
|
||||
// Every minute should always have a next fire
|
||||
let next = next_cron_fire("* * * * * *").expect("valid cron");
|
||||
assert!(next.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_next_cron_fire_invalid() {
|
||||
let result = next_cron_fire("not a cron");
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_guardrails_default() {
|
||||
let g = RoutineGuardrails::default();
|
||||
assert_eq!(g.cooldown.as_secs(), 300);
|
||||
assert_eq!(g.max_concurrent, 1);
|
||||
assert!(g.dedup_window.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_trigger_type_tag() {
|
||||
assert_eq!(
|
||||
Trigger::Cron {
|
||||
schedule: String::new()
|
||||
}
|
||||
.type_tag(),
|
||||
"cron"
|
||||
);
|
||||
assert_eq!(
|
||||
Trigger::Event {
|
||||
channel: None,
|
||||
pattern: String::new()
|
||||
}
|
||||
.type_tag(),
|
||||
"event"
|
||||
);
|
||||
assert_eq!(
|
||||
Trigger::Webhook {
|
||||
path: None,
|
||||
secret: None
|
||||
}
|
||||
.type_tag(),
|
||||
"webhook"
|
||||
);
|
||||
assert_eq!(Trigger::Manual.type_tag(), "manual");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,606 @@
|
||||
//! Routine execution engine.
|
||||
//!
|
||||
//! Handles loading routines, checking triggers, enforcing guardrails,
|
||||
//! and executing both lightweight (single LLM call) and full-job routines.
|
||||
//!
|
||||
//! The engine runs two independent loops:
|
||||
//! - A **cron ticker** that polls the DB every N seconds for due cron routines
|
||||
//! - An **event matcher** called synchronously from the agent main loop
|
||||
//!
|
||||
//! Lightweight routines execute inline (single LLM call, no scheduler slot).
|
||||
//! Full-job routines are delegated to the existing `Scheduler`.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use chrono::Utc;
|
||||
use regex::Regex;
|
||||
use tokio::sync::{RwLock, mpsc};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::agent::routine::{
|
||||
NotifyConfig, Routine, RoutineAction, RoutineRun, RunStatus, Trigger, next_cron_fire,
|
||||
};
|
||||
use crate::channels::{IncomingMessage, OutgoingResponse};
|
||||
use crate::config::RoutineConfig;
|
||||
use crate::history::Store;
|
||||
use crate::llm::{ChatMessage, CompletionRequest, FinishReason, LlmProvider};
|
||||
use crate::workspace::Workspace;
|
||||
|
||||
/// The routine execution engine.
|
||||
pub struct RoutineEngine {
|
||||
config: RoutineConfig,
|
||||
store: Arc<Store>,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
workspace: Arc<Workspace>,
|
||||
/// Sender for notifications (routed to channel manager).
|
||||
notify_tx: mpsc::Sender<OutgoingResponse>,
|
||||
/// Currently running routine count (across all routines).
|
||||
running_count: Arc<RwLock<usize>>,
|
||||
/// Compiled event regex cache: routine_id -> compiled regex.
|
||||
event_cache: Arc<RwLock<Vec<(Uuid, Routine, Regex)>>>,
|
||||
}
|
||||
|
||||
impl RoutineEngine {
|
||||
pub fn new(
|
||||
config: RoutineConfig,
|
||||
store: Arc<Store>,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
workspace: Arc<Workspace>,
|
||||
notify_tx: mpsc::Sender<OutgoingResponse>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config,
|
||||
store,
|
||||
llm,
|
||||
workspace,
|
||||
notify_tx,
|
||||
running_count: Arc::new(RwLock::new(0)),
|
||||
event_cache: Arc::new(RwLock::new(Vec::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Refresh the in-memory event trigger cache from DB.
|
||||
pub async fn refresh_event_cache(&self) {
|
||||
match self.store.list_event_routines().await {
|
||||
Ok(routines) => {
|
||||
let mut cache = Vec::new();
|
||||
for routine in routines {
|
||||
if let Trigger::Event { ref pattern, .. } = routine.trigger {
|
||||
match Regex::new(pattern) {
|
||||
Ok(re) => cache.push((routine.id, routine.clone(), re)),
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
routine = %routine.name,
|
||||
"Invalid event regex '{}': {}",
|
||||
pattern, e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let count = cache.len();
|
||||
*self.event_cache.write().await = cache;
|
||||
tracing::debug!("Refreshed event cache: {} routines", count);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("Failed to refresh event cache: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Check incoming message against event triggers. Returns number of routines fired.
|
||||
///
|
||||
/// Called synchronously from the main loop after handle_message(). The actual
|
||||
/// execution is spawned async so this returns quickly.
|
||||
pub async fn check_event_triggers(&self, message: &IncomingMessage) -> usize {
|
||||
let cache = self.event_cache.read().await;
|
||||
let mut fired = 0;
|
||||
|
||||
for (_, routine, re) in cache.iter() {
|
||||
// Channel filter
|
||||
if let Trigger::Event {
|
||||
channel: Some(ch), ..
|
||||
} = &routine.trigger
|
||||
{
|
||||
if ch != &message.channel {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
// Regex match
|
||||
if !re.is_match(&message.content) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Cooldown check
|
||||
if !self.check_cooldown(routine) {
|
||||
tracing::debug!(routine = %routine.name, "Skipped: cooldown active");
|
||||
continue;
|
||||
}
|
||||
|
||||
// Concurrent run check
|
||||
if !self.check_concurrent(routine).await {
|
||||
tracing::debug!(routine = %routine.name, "Skipped: max concurrent reached");
|
||||
continue;
|
||||
}
|
||||
|
||||
// Global capacity check
|
||||
if *self.running_count.read().await >= self.config.max_concurrent_routines {
|
||||
tracing::warn!(routine = %routine.name, "Skipped: global max concurrent reached");
|
||||
continue;
|
||||
}
|
||||
|
||||
let detail = truncate(&message.content, 200);
|
||||
self.spawn_fire(routine.clone(), "event", Some(detail));
|
||||
fired += 1;
|
||||
}
|
||||
|
||||
fired
|
||||
}
|
||||
|
||||
/// Check all due cron routines and fire them. Called by the cron ticker.
|
||||
pub async fn check_cron_triggers(&self) {
|
||||
let routines = match self.store.list_due_cron_routines().await {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
tracing::error!("Failed to load due cron routines: {}", e);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
for routine in routines {
|
||||
if *self.running_count.read().await >= self.config.max_concurrent_routines {
|
||||
tracing::warn!("Global max concurrent routines reached, skipping remaining");
|
||||
break;
|
||||
}
|
||||
|
||||
if !self.check_cooldown(&routine) {
|
||||
continue;
|
||||
}
|
||||
|
||||
if !self.check_concurrent(&routine).await {
|
||||
continue;
|
||||
}
|
||||
|
||||
let detail = if let Trigger::Cron { ref schedule } = routine.trigger {
|
||||
Some(schedule.clone())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
self.spawn_fire(routine, "cron", detail);
|
||||
}
|
||||
}
|
||||
|
||||
/// Fire a routine manually (from tool call or CLI).
|
||||
pub async fn fire_manual(&self, routine_id: Uuid) -> Result<Uuid, String> {
|
||||
let routine = self
|
||||
.store
|
||||
.get_routine(routine_id)
|
||||
.await
|
||||
.map_err(|e| format!("DB error: {e}"))?
|
||||
.ok_or_else(|| format!("routine {routine_id} not found"))?;
|
||||
|
||||
if !routine.enabled {
|
||||
return Err(format!("routine '{}' is disabled", routine.name));
|
||||
}
|
||||
|
||||
if !self.check_concurrent(&routine).await {
|
||||
return Err(format!(
|
||||
"routine '{}' already at max concurrent runs",
|
||||
routine.name
|
||||
));
|
||||
}
|
||||
|
||||
let run_id = Uuid::new_v4();
|
||||
let run = RoutineRun {
|
||||
id: run_id,
|
||||
routine_id: routine.id,
|
||||
trigger_type: "manual".to_string(),
|
||||
trigger_detail: None,
|
||||
started_at: Utc::now(),
|
||||
completed_at: None,
|
||||
status: RunStatus::Running,
|
||||
result_summary: None,
|
||||
tokens_used: None,
|
||||
job_id: None,
|
||||
created_at: Utc::now(),
|
||||
};
|
||||
|
||||
if let Err(e) = self.store.create_routine_run(&run).await {
|
||||
return Err(format!("failed to create run record: {e}"));
|
||||
}
|
||||
|
||||
// Execute inline for manual triggers (caller wants to wait)
|
||||
let engine = EngineContext {
|
||||
store: self.store.clone(),
|
||||
llm: self.llm.clone(),
|
||||
workspace: self.workspace.clone(),
|
||||
notify_tx: self.notify_tx.clone(),
|
||||
running_count: self.running_count.clone(),
|
||||
max_lightweight_tokens: self.config.max_lightweight_tokens,
|
||||
};
|
||||
|
||||
tokio::spawn(async move {
|
||||
execute_routine(engine, routine, run).await;
|
||||
});
|
||||
|
||||
Ok(run_id)
|
||||
}
|
||||
|
||||
/// Spawn a fire in a background task.
|
||||
fn spawn_fire(&self, routine: Routine, trigger_type: &str, trigger_detail: Option<String>) {
|
||||
let run = RoutineRun {
|
||||
id: Uuid::new_v4(),
|
||||
routine_id: routine.id,
|
||||
trigger_type: trigger_type.to_string(),
|
||||
trigger_detail,
|
||||
started_at: Utc::now(),
|
||||
completed_at: None,
|
||||
status: RunStatus::Running,
|
||||
result_summary: None,
|
||||
tokens_used: None,
|
||||
job_id: None,
|
||||
created_at: Utc::now(),
|
||||
};
|
||||
|
||||
let engine = EngineContext {
|
||||
store: self.store.clone(),
|
||||
llm: self.llm.clone(),
|
||||
workspace: self.workspace.clone(),
|
||||
notify_tx: self.notify_tx.clone(),
|
||||
running_count: self.running_count.clone(),
|
||||
max_lightweight_tokens: self.config.max_lightweight_tokens,
|
||||
};
|
||||
|
||||
// Record the run in DB, then spawn execution
|
||||
let store = self.store.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = store.create_routine_run(&run).await {
|
||||
tracing::error!(routine = %routine.name, "Failed to record run: {}", e);
|
||||
return;
|
||||
}
|
||||
execute_routine(engine, routine, run).await;
|
||||
});
|
||||
}
|
||||
|
||||
fn check_cooldown(&self, routine: &Routine) -> bool {
|
||||
if let Some(last_run) = routine.last_run_at {
|
||||
let elapsed = Utc::now().signed_duration_since(last_run);
|
||||
let cooldown = chrono::Duration::from_std(routine.guardrails.cooldown)
|
||||
.unwrap_or(chrono::Duration::seconds(300));
|
||||
if elapsed < cooldown {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
async fn check_concurrent(&self, routine: &Routine) -> bool {
|
||||
match self.store.count_running_routine_runs(routine.id).await {
|
||||
Ok(count) => count < routine.guardrails.max_concurrent as i64,
|
||||
Err(e) => {
|
||||
tracing::error!(
|
||||
routine = %routine.name,
|
||||
"Failed to check concurrent runs: {}", e
|
||||
);
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Shared context passed to the execution function.
|
||||
struct EngineContext {
|
||||
store: Arc<Store>,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
workspace: Arc<Workspace>,
|
||||
notify_tx: mpsc::Sender<OutgoingResponse>,
|
||||
running_count: Arc<RwLock<usize>>,
|
||||
max_lightweight_tokens: u32,
|
||||
}
|
||||
|
||||
/// Execute a routine run. Handles both lightweight and full_job modes.
|
||||
async fn execute_routine(ctx: EngineContext, routine: Routine, run: RoutineRun) {
|
||||
// Increment running count
|
||||
{
|
||||
let mut count = ctx.running_count.write().await;
|
||||
*count += 1;
|
||||
}
|
||||
|
||||
let result = match &routine.action {
|
||||
RoutineAction::Lightweight {
|
||||
prompt,
|
||||
context_paths,
|
||||
max_tokens,
|
||||
} => execute_lightweight(&ctx, &routine, prompt, context_paths, *max_tokens).await,
|
||||
RoutineAction::FullJob { description, .. } => {
|
||||
// Full job mode: for now, execute as lightweight with the description
|
||||
// as prompt. Full scheduler integration will come as a follow-up.
|
||||
tracing::info!(
|
||||
routine = %routine.name,
|
||||
"FullJob mode executing as lightweight (scheduler integration pending)"
|
||||
);
|
||||
execute_lightweight(&ctx, &routine, description, &[], ctx.max_lightweight_tokens).await
|
||||
}
|
||||
};
|
||||
|
||||
// Decrement running count
|
||||
{
|
||||
let mut count = ctx.running_count.write().await;
|
||||
*count = count.saturating_sub(1);
|
||||
}
|
||||
|
||||
// Process result
|
||||
let (status, summary, tokens) = match result {
|
||||
Ok(execution) => execution,
|
||||
Err(e) => {
|
||||
tracing::error!(routine = %routine.name, "Execution failed: {}", e);
|
||||
(RunStatus::Failed, Some(e), None)
|
||||
}
|
||||
};
|
||||
|
||||
// Complete the run record
|
||||
if let Err(e) = ctx
|
||||
.store
|
||||
.complete_routine_run(run.id, status, summary.as_deref(), tokens)
|
||||
.await
|
||||
{
|
||||
tracing::error!(routine = %routine.name, "Failed to complete run record: {}", e);
|
||||
}
|
||||
|
||||
// Update routine runtime state
|
||||
let now = Utc::now();
|
||||
let next_fire = if let Trigger::Cron { ref schedule } = routine.trigger {
|
||||
next_cron_fire(schedule).unwrap_or(None)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let new_failures = if status == RunStatus::Failed {
|
||||
routine.consecutive_failures + 1
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
if let Err(e) = ctx
|
||||
.store
|
||||
.update_routine_runtime(
|
||||
routine.id,
|
||||
now,
|
||||
next_fire,
|
||||
routine.run_count + 1,
|
||||
new_failures,
|
||||
&routine.state,
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::error!(routine = %routine.name, "Failed to update runtime state: {}", e);
|
||||
}
|
||||
|
||||
// Send notifications based on config
|
||||
send_notification(
|
||||
&ctx.notify_tx,
|
||||
&routine.notify,
|
||||
&routine.name,
|
||||
status,
|
||||
summary.as_deref(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
/// Execute a lightweight routine (single LLM call).
|
||||
async fn execute_lightweight(
|
||||
ctx: &EngineContext,
|
||||
routine: &Routine,
|
||||
prompt: &str,
|
||||
context_paths: &[String],
|
||||
max_tokens: u32,
|
||||
) -> Result<(RunStatus, Option<String>, Option<i32>), String> {
|
||||
// Load context from workspace
|
||||
let mut context_parts = Vec::new();
|
||||
for path in context_paths {
|
||||
match ctx.workspace.read(path).await {
|
||||
Ok(doc) => {
|
||||
context_parts.push(format!("## {}\n\n{}", path, doc.content));
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::debug!(
|
||||
routine = %routine.name,
|
||||
"Failed to read context path {}: {}", path, e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Load routine state from workspace
|
||||
let state_path = format!("routines/{}/state.md", routine.name);
|
||||
let state_content = match ctx.workspace.read(&state_path).await {
|
||||
Ok(doc) => Some(doc.content),
|
||||
Err(_) => None,
|
||||
};
|
||||
|
||||
// Build the prompt
|
||||
let mut full_prompt = String::new();
|
||||
full_prompt.push_str(prompt);
|
||||
|
||||
if !context_parts.is_empty() {
|
||||
full_prompt.push_str("\n\n---\n\n# Context\n\n");
|
||||
full_prompt.push_str(&context_parts.join("\n\n"));
|
||||
}
|
||||
|
||||
if let Some(state) = &state_content {
|
||||
full_prompt.push_str("\n\n---\n\n# Previous State\n\n");
|
||||
full_prompt.push_str(state);
|
||||
}
|
||||
|
||||
full_prompt.push_str(
|
||||
"\n\n---\n\nIf nothing needs attention, reply EXACTLY with: ROUTINE_OK\n\
|
||||
If something needs attention, provide a concise summary.",
|
||||
);
|
||||
|
||||
// Get system prompt
|
||||
let system_prompt = match ctx.workspace.system_prompt().await {
|
||||
Ok(p) => p,
|
||||
Err(e) => {
|
||||
tracing::warn!(routine = %routine.name, "Failed to get system prompt: {}", e);
|
||||
String::new()
|
||||
}
|
||||
};
|
||||
|
||||
let messages = if system_prompt.is_empty() {
|
||||
vec![ChatMessage::user(&full_prompt)]
|
||||
} else {
|
||||
vec![
|
||||
ChatMessage::system(&system_prompt),
|
||||
ChatMessage::user(&full_prompt),
|
||||
]
|
||||
};
|
||||
|
||||
// Determine max_tokens from model metadata with fallback
|
||||
let effective_max_tokens = match ctx.llm.model_metadata().await {
|
||||
Ok(meta) => {
|
||||
let from_api = meta.context_length.map(|ctx| ctx / 2).unwrap_or(max_tokens);
|
||||
from_api.max(max_tokens)
|
||||
}
|
||||
Err(_) => max_tokens,
|
||||
};
|
||||
|
||||
let request = CompletionRequest::new(messages)
|
||||
.with_max_tokens(effective_max_tokens)
|
||||
.with_temperature(0.3);
|
||||
|
||||
let response = ctx
|
||||
.llm
|
||||
.complete(request)
|
||||
.await
|
||||
.map_err(|e| format!("LLM call failed: {e}"))?;
|
||||
|
||||
let content = response.content.trim();
|
||||
let tokens_used = Some((response.input_tokens + response.output_tokens) as i32);
|
||||
|
||||
// Empty content guard (same as heartbeat)
|
||||
if content.is_empty() {
|
||||
return if response.finish_reason == FinishReason::Length {
|
||||
Err(
|
||||
"LLM response truncated (finish_reason=length) with no content. \
|
||||
Model may have exhausted token budget on reasoning."
|
||||
.to_string(),
|
||||
)
|
||||
} else {
|
||||
Err("LLM returned empty content.".to_string())
|
||||
};
|
||||
}
|
||||
|
||||
// Check for the "nothing to do" sentinel
|
||||
if content == "ROUTINE_OK" || content.contains("ROUTINE_OK") {
|
||||
return Ok((RunStatus::Ok, None, tokens_used));
|
||||
}
|
||||
|
||||
Ok((RunStatus::Attention, Some(content.to_string()), tokens_used))
|
||||
}
|
||||
|
||||
/// Send a notification based on the routine's notify config and run status.
|
||||
async fn send_notification(
|
||||
tx: &mpsc::Sender<OutgoingResponse>,
|
||||
notify: &NotifyConfig,
|
||||
routine_name: &str,
|
||||
status: RunStatus,
|
||||
summary: Option<&str>,
|
||||
) {
|
||||
let should_notify = match status {
|
||||
RunStatus::Ok => notify.on_success,
|
||||
RunStatus::Attention => notify.on_attention,
|
||||
RunStatus::Failed => notify.on_failure,
|
||||
RunStatus::Running => false,
|
||||
};
|
||||
|
||||
if !should_notify {
|
||||
return;
|
||||
}
|
||||
|
||||
let icon = match status {
|
||||
RunStatus::Ok => "✅",
|
||||
RunStatus::Attention => "🔔",
|
||||
RunStatus::Failed => "❌",
|
||||
RunStatus::Running => "⏳",
|
||||
};
|
||||
|
||||
let message = match summary {
|
||||
Some(s) => format!("{} *Routine '{}'*: {}\n\n{}", icon, routine_name, status, s),
|
||||
None => format!("{} *Routine '{}'*: {}", icon, routine_name, status),
|
||||
};
|
||||
|
||||
let response = OutgoingResponse {
|
||||
content: message,
|
||||
thread_id: None,
|
||||
metadata: serde_json::json!({
|
||||
"source": "routine",
|
||||
"routine_name": routine_name,
|
||||
"status": status.to_string(),
|
||||
}),
|
||||
};
|
||||
|
||||
if let Err(e) = tx.send(response).await {
|
||||
tracing::error!(routine = %routine_name, "Failed to send notification: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
/// Spawn the cron ticker background task.
|
||||
pub fn spawn_cron_ticker(
|
||||
engine: Arc<RoutineEngine>,
|
||||
interval: Duration,
|
||||
) -> tokio::task::JoinHandle<()> {
|
||||
tokio::spawn(async move {
|
||||
let mut ticker = tokio::time::interval(interval);
|
||||
// Skip immediate first tick
|
||||
ticker.tick().await;
|
||||
|
||||
loop {
|
||||
ticker.tick().await;
|
||||
engine.check_cron_triggers().await;
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn truncate(s: &str, max: usize) -> String {
|
||||
if s.len() <= max {
|
||||
s.to_string()
|
||||
} else {
|
||||
format!("{}...", &s[..max])
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::agent::routine::{NotifyConfig, RunStatus};
|
||||
|
||||
#[test]
|
||||
fn test_notification_gating() {
|
||||
let config = NotifyConfig {
|
||||
on_success: false,
|
||||
on_failure: true,
|
||||
on_attention: true,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
// on_success = false means Ok status should not notify
|
||||
assert!(!config.on_success);
|
||||
assert!(config.on_failure);
|
||||
assert!(config.on_attention);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_run_status_icons() {
|
||||
// Just verify the mapping doesn't panic
|
||||
for status in [
|
||||
RunStatus::Ok,
|
||||
RunStatus::Attention,
|
||||
RunStatus::Failed,
|
||||
RunStatus::Running,
|
||||
] {
|
||||
let _ = status.to_string();
|
||||
}
|
||||
}
|
||||
}
|
||||
+17
-17
@@ -373,23 +373,23 @@ impl Scheduler {
|
||||
.into());
|
||||
}
|
||||
|
||||
// Execute with timeout
|
||||
let result = tokio::time::timeout(Duration::from_secs(60), async {
|
||||
tool.execute(params, &job_ctx).await
|
||||
})
|
||||
.await
|
||||
.map_err(|_| {
|
||||
Error::Tool(crate::error::ToolError::Timeout {
|
||||
name: tool_name.to_string(),
|
||||
timeout: Duration::from_secs(60),
|
||||
})
|
||||
})?
|
||||
.map_err(|e| {
|
||||
Error::Tool(crate::error::ToolError::ExecutionFailed {
|
||||
name: tool_name.to_string(),
|
||||
reason: e.to_string(),
|
||||
})
|
||||
})?;
|
||||
// Execute with per-tool timeout
|
||||
let tool_timeout = tool.execution_timeout();
|
||||
let result =
|
||||
tokio::time::timeout(tool_timeout, async { tool.execute(params, &job_ctx).await })
|
||||
.await
|
||||
.map_err(|_| {
|
||||
Error::Tool(crate::error::ToolError::Timeout {
|
||||
name: tool_name.to_string(),
|
||||
timeout: tool_timeout,
|
||||
})
|
||||
})?
|
||||
.map_err(|e| {
|
||||
Error::Tool(crate::error::ToolError::ExecutionFailed {
|
||||
name: tool_name.to_string(),
|
||||
reason: e.to_string(),
|
||||
})
|
||||
})?;
|
||||
|
||||
Ok(TaskOutput::new(result.result, start.elapsed()))
|
||||
}
|
||||
|
||||
@@ -173,6 +173,10 @@ pub struct Thread {
|
||||
/// Pending auth token request (thread is in auth mode).
|
||||
#[serde(default)]
|
||||
pub pending_auth: Option<PendingAuth>,
|
||||
/// Last NEAR AI response ID for response chaining. Persisted to DB
|
||||
/// metadata so we can resume chaining across restarts.
|
||||
#[serde(default)]
|
||||
pub last_response_id: Option<String>,
|
||||
}
|
||||
|
||||
impl Thread {
|
||||
@@ -189,6 +193,24 @@ impl Thread {
|
||||
metadata: serde_json::Value::Null,
|
||||
pending_approval: None,
|
||||
pending_auth: None,
|
||||
last_response_id: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a thread with a specific ID (for DB hydration).
|
||||
pub fn with_id(id: Uuid, session_id: Uuid) -> Self {
|
||||
let now = Utc::now();
|
||||
Self {
|
||||
id,
|
||||
session_id,
|
||||
state: ThreadState::Idle,
|
||||
turns: Vec::new(),
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
metadata: serde_json::Value::Null,
|
||||
pending_approval: None,
|
||||
pending_auth: None,
|
||||
last_response_id: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -593,4 +615,386 @@ mod tests {
|
||||
let restored: Thread = serde_json::from_str(&json).expect("should deserialize");
|
||||
assert!(restored.pending_auth.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_thread_with_id() {
|
||||
let specific_id = Uuid::new_v4();
|
||||
let session_id = Uuid::new_v4();
|
||||
let thread = Thread::with_id(specific_id, session_id);
|
||||
|
||||
assert_eq!(thread.id, specific_id);
|
||||
assert_eq!(thread.session_id, session_id);
|
||||
assert_eq!(thread.state, ThreadState::Idle);
|
||||
assert!(thread.turns.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_thread_with_id_restore_messages() {
|
||||
let thread_id = Uuid::new_v4();
|
||||
let session_id = Uuid::new_v4();
|
||||
let mut thread = Thread::with_id(thread_id, session_id);
|
||||
|
||||
let messages = vec![
|
||||
ChatMessage::user("Hello from DB"),
|
||||
ChatMessage::assistant("Restored response"),
|
||||
];
|
||||
thread.restore_from_messages(messages);
|
||||
|
||||
assert_eq!(thread.id, thread_id);
|
||||
assert_eq!(thread.turns.len(), 1);
|
||||
assert_eq!(thread.turns[0].user_input, "Hello from DB");
|
||||
assert_eq!(
|
||||
thread.turns[0].response,
|
||||
Some("Restored response".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_restore_from_messages_empty() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
// Add a turn first, then restore with empty vec
|
||||
thread.start_turn("hello");
|
||||
thread.complete_turn("hi");
|
||||
assert_eq!(thread.turns.len(), 1);
|
||||
|
||||
thread.restore_from_messages(Vec::new());
|
||||
|
||||
// Should clear all turns and stay idle
|
||||
assert!(thread.turns.is_empty());
|
||||
assert_eq!(thread.state, ThreadState::Idle);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_restore_from_messages_only_assistant_messages() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
// Only assistant messages (no user messages to anchor turns)
|
||||
let messages = vec![
|
||||
ChatMessage::assistant("I'm here"),
|
||||
ChatMessage::assistant("Still here"),
|
||||
];
|
||||
|
||||
thread.restore_from_messages(messages);
|
||||
|
||||
// Assistant-only messages have no user turn to attach to, so
|
||||
// they should be skipped entirely.
|
||||
assert!(thread.turns.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_restore_from_messages_multiple_user_messages_in_a_row() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
// Two user messages with no assistant response between them
|
||||
let messages = vec![
|
||||
ChatMessage::user("first"),
|
||||
ChatMessage::user("second"),
|
||||
ChatMessage::assistant("reply to second"),
|
||||
];
|
||||
|
||||
thread.restore_from_messages(messages);
|
||||
|
||||
// First user message becomes a turn with no response,
|
||||
// second user message pairs with the assistant response.
|
||||
assert_eq!(thread.turns.len(), 2);
|
||||
assert_eq!(thread.turns[0].user_input, "first");
|
||||
assert!(thread.turns[0].response.is_none());
|
||||
assert_eq!(thread.turns[1].user_input, "second");
|
||||
assert_eq!(
|
||||
thread.turns[1].response,
|
||||
Some("reply to second".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_thread_switch() {
|
||||
let mut session = Session::new("user-1");
|
||||
|
||||
let t1_id = session.create_thread().id;
|
||||
let t2_id = session.create_thread().id;
|
||||
|
||||
// After creating two threads, active should be the last one
|
||||
assert_eq!(session.active_thread, Some(t2_id));
|
||||
|
||||
// Switch back to the first
|
||||
assert!(session.switch_thread(t1_id));
|
||||
assert_eq!(session.active_thread, Some(t1_id));
|
||||
|
||||
// Switching to a nonexistent thread should fail
|
||||
let fake_id = Uuid::new_v4();
|
||||
assert!(!session.switch_thread(fake_id));
|
||||
// Active thread should remain unchanged
|
||||
assert_eq!(session.active_thread, Some(t1_id));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_or_create_thread_idempotent() {
|
||||
let mut session = Session::new("user-1");
|
||||
|
||||
let tid1 = session.get_or_create_thread().id;
|
||||
let tid2 = session.get_or_create_thread().id;
|
||||
|
||||
// Should return the same thread (not create a new one each time)
|
||||
assert_eq!(tid1, tid2);
|
||||
assert_eq!(session.threads.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_turns() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
for i in 0..5 {
|
||||
thread.start_turn(format!("msg-{}", i));
|
||||
thread.complete_turn(format!("resp-{}", i));
|
||||
}
|
||||
assert_eq!(thread.turns.len(), 5);
|
||||
|
||||
thread.truncate_turns(3);
|
||||
assert_eq!(thread.turns.len(), 3);
|
||||
|
||||
// Should keep the most recent turns
|
||||
assert_eq!(thread.turns[0].user_input, "msg-2");
|
||||
assert_eq!(thread.turns[1].user_input, "msg-3");
|
||||
assert_eq!(thread.turns[2].user_input, "msg-4");
|
||||
|
||||
// Turn numbers should be re-indexed
|
||||
assert_eq!(thread.turns[0].turn_number, 0);
|
||||
assert_eq!(thread.turns[1].turn_number, 1);
|
||||
assert_eq!(thread.turns[2].turn_number, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_turns_noop_when_fewer() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
thread.start_turn("only one");
|
||||
thread.complete_turn("response");
|
||||
|
||||
thread.truncate_turns(10);
|
||||
assert_eq!(thread.turns.len(), 1);
|
||||
assert_eq!(thread.turns[0].user_input, "only one");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_thread_interrupt_and_resume() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
thread.start_turn("do something");
|
||||
assert_eq!(thread.state, ThreadState::Processing);
|
||||
|
||||
thread.interrupt();
|
||||
assert_eq!(thread.state, ThreadState::Interrupted);
|
||||
|
||||
let last_turn = thread.last_turn().unwrap();
|
||||
assert_eq!(last_turn.state, TurnState::Interrupted);
|
||||
assert!(last_turn.completed_at.is_some());
|
||||
|
||||
thread.resume();
|
||||
assert_eq!(thread.state, ThreadState::Idle);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resume_only_from_interrupted() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
// Idle thread: resume should be a no-op
|
||||
assert_eq!(thread.state, ThreadState::Idle);
|
||||
thread.resume();
|
||||
assert_eq!(thread.state, ThreadState::Idle);
|
||||
|
||||
// Processing thread: resume should not change state
|
||||
thread.start_turn("work");
|
||||
assert_eq!(thread.state, ThreadState::Processing);
|
||||
thread.resume();
|
||||
assert_eq!(thread.state, ThreadState::Processing);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_turn_fail() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
thread.start_turn("risky operation");
|
||||
thread.fail_turn("connection timed out");
|
||||
|
||||
assert_eq!(thread.state, ThreadState::Idle);
|
||||
|
||||
let turn = thread.last_turn().unwrap();
|
||||
assert_eq!(turn.state, TurnState::Failed);
|
||||
assert_eq!(turn.error, Some("connection timed out".to_string()));
|
||||
assert!(turn.response.is_none());
|
||||
assert!(turn.completed_at.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_messages_with_incomplete_last_turn() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
thread.start_turn("first");
|
||||
thread.complete_turn("first reply");
|
||||
thread.start_turn("second (in progress)");
|
||||
|
||||
let messages = thread.messages();
|
||||
// Should have 3 messages: user, assistant, user (no assistant for in-progress)
|
||||
assert_eq!(messages.len(), 3);
|
||||
assert_eq!(messages[0].content, "first");
|
||||
assert_eq!(messages[1].content, "first reply");
|
||||
assert_eq!(messages[2].content, "second (in progress)");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_thread_serialization_round_trip() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
thread.start_turn("hello");
|
||||
thread.complete_turn("world");
|
||||
thread.last_response_id = Some("resp_abc123".to_string());
|
||||
|
||||
let json = serde_json::to_string(&thread).unwrap();
|
||||
let restored: Thread = serde_json::from_str(&json).unwrap();
|
||||
|
||||
assert_eq!(restored.id, thread.id);
|
||||
assert_eq!(restored.session_id, thread.session_id);
|
||||
assert_eq!(restored.turns.len(), 1);
|
||||
assert_eq!(restored.turns[0].user_input, "hello");
|
||||
assert_eq!(restored.turns[0].response, Some("world".to_string()));
|
||||
assert_eq!(restored.last_response_id, Some("resp_abc123".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_session_serialization_round_trip() {
|
||||
let mut session = Session::new("user-ser");
|
||||
session.create_thread();
|
||||
session.auto_approve_tool("echo");
|
||||
|
||||
let json = serde_json::to_string(&session).unwrap();
|
||||
let restored: Session = serde_json::from_str(&json).unwrap();
|
||||
|
||||
assert_eq!(restored.user_id, "user-ser");
|
||||
assert_eq!(restored.threads.len(), 1);
|
||||
assert!(restored.is_tool_auto_approved("echo"));
|
||||
assert!(!restored.is_tool_auto_approved("shell"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_auto_approved_tools() {
|
||||
let mut session = Session::new("user-1");
|
||||
|
||||
assert!(!session.is_tool_auto_approved("shell"));
|
||||
session.auto_approve_tool("shell");
|
||||
assert!(session.is_tool_auto_approved("shell"));
|
||||
|
||||
// Idempotent
|
||||
session.auto_approve_tool("shell");
|
||||
assert_eq!(session.auto_approved_tools.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_turn_tool_call_error() {
|
||||
let mut turn = Turn::new(0, "test");
|
||||
turn.record_tool_call("http", serde_json::json!({"url": "example.com"}));
|
||||
turn.record_tool_error("timeout");
|
||||
|
||||
assert_eq!(turn.tool_calls.len(), 1);
|
||||
assert_eq!(turn.tool_calls[0].error, Some("timeout".to_string()));
|
||||
assert!(turn.tool_calls[0].result.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_turn_number_increments() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
// Before any turns, turn_number() is 1 (1-indexed for display)
|
||||
assert_eq!(thread.turn_number(), 1);
|
||||
|
||||
thread.start_turn("first");
|
||||
thread.complete_turn("done");
|
||||
assert_eq!(thread.turn_number(), 2);
|
||||
|
||||
thread.start_turn("second");
|
||||
assert_eq!(thread.turn_number(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_complete_turn_on_empty_thread() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
// Completing a turn when there are no turns should be a safe no-op
|
||||
thread.complete_turn("phantom response");
|
||||
assert_eq!(thread.state, ThreadState::Idle);
|
||||
assert!(thread.turns.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fail_turn_on_empty_thread() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
// Failing a turn when there are no turns should be a safe no-op
|
||||
thread.fail_turn("phantom error");
|
||||
assert_eq!(thread.state, ThreadState::Idle);
|
||||
assert!(thread.turns.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pending_approval_flow() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
let approval = PendingApproval {
|
||||
request_id: Uuid::new_v4(),
|
||||
tool_name: "shell".to_string(),
|
||||
parameters: serde_json::json!({"command": "rm -rf /"}),
|
||||
description: "dangerous command".to_string(),
|
||||
tool_call_id: "call_123".to_string(),
|
||||
context_messages: vec![ChatMessage::user("do it")],
|
||||
};
|
||||
|
||||
thread.await_approval(approval);
|
||||
assert_eq!(thread.state, ThreadState::AwaitingApproval);
|
||||
assert!(thread.pending_approval.is_some());
|
||||
|
||||
let taken = thread.take_pending_approval();
|
||||
assert!(taken.is_some());
|
||||
assert_eq!(taken.unwrap().tool_name, "shell");
|
||||
assert!(thread.pending_approval.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_clear_pending_approval() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
let approval = PendingApproval {
|
||||
request_id: Uuid::new_v4(),
|
||||
tool_name: "http".to_string(),
|
||||
parameters: serde_json::json!({}),
|
||||
description: "test".to_string(),
|
||||
tool_call_id: "call_456".to_string(),
|
||||
context_messages: vec![],
|
||||
};
|
||||
|
||||
thread.await_approval(approval);
|
||||
thread.clear_pending_approval();
|
||||
|
||||
assert_eq!(thread.state, ThreadState::Idle);
|
||||
assert!(thread.pending_approval.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_active_thread_accessors() {
|
||||
let mut session = Session::new("user-1");
|
||||
|
||||
assert!(session.active_thread().is_none());
|
||||
assert!(session.active_thread_mut().is_none());
|
||||
|
||||
let tid = session.create_thread().id;
|
||||
|
||||
assert!(session.active_thread().is_some());
|
||||
assert_eq!(session.active_thread().unwrap().id, tid);
|
||||
|
||||
// Mutably modify through accessor
|
||||
session.active_thread_mut().unwrap().start_turn("test");
|
||||
assert_eq!(
|
||||
session.active_thread().unwrap().state,
|
||||
ThreadState::Processing
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -110,6 +110,41 @@ impl SessionManager {
|
||||
(session, thread_id)
|
||||
}
|
||||
|
||||
/// Register a hydrated thread so subsequent `resolve_thread` calls find it.
|
||||
///
|
||||
/// Inserts into the thread_map and creates an undo manager for the thread.
|
||||
pub async fn register_thread(
|
||||
&self,
|
||||
user_id: &str,
|
||||
channel: &str,
|
||||
thread_id: Uuid,
|
||||
session: Arc<Mutex<Session>>,
|
||||
) {
|
||||
let key = ThreadKey {
|
||||
user_id: user_id.to_string(),
|
||||
channel: channel.to_string(),
|
||||
external_thread_id: Some(thread_id.to_string()),
|
||||
};
|
||||
|
||||
{
|
||||
let mut thread_map = self.thread_map.write().await;
|
||||
thread_map.insert(key, thread_id);
|
||||
}
|
||||
|
||||
{
|
||||
let mut undo_managers = self.undo_managers.write().await;
|
||||
undo_managers
|
||||
.entry(thread_id)
|
||||
.or_insert_with(|| Arc::new(Mutex::new(UndoManager::new())));
|
||||
}
|
||||
|
||||
// Ensure the session is tracked
|
||||
{
|
||||
let mut sessions = self.sessions.write().await;
|
||||
sessions.entry(user_id.to_string()).or_insert(session);
|
||||
}
|
||||
}
|
||||
|
||||
/// Get undo manager for a thread.
|
||||
pub async fn get_undo_manager(&self, thread_id: Uuid) -> Arc<Mutex<UndoManager>> {
|
||||
// Fast path
|
||||
@@ -296,4 +331,344 @@ mod tests {
|
||||
.await;
|
||||
assert_eq!(pruned, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_register_thread() {
|
||||
use crate::agent::session::{Session, Thread};
|
||||
|
||||
let manager = SessionManager::new();
|
||||
let thread_id = Uuid::new_v4();
|
||||
|
||||
// Create a session with a hydrated thread
|
||||
let session = Arc::new(Mutex::new(Session::new("user-hydrate")));
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(thread_id, sess.id);
|
||||
sess.threads.insert(thread_id, thread);
|
||||
sess.active_thread = Some(thread_id);
|
||||
}
|
||||
|
||||
// Register the thread
|
||||
manager
|
||||
.register_thread("user-hydrate", "gateway", thread_id, Arc::clone(&session))
|
||||
.await;
|
||||
|
||||
// resolve_thread should find it (using the UUID as external_thread_id)
|
||||
let (resolved_session, resolved_tid) = manager
|
||||
.resolve_thread("user-hydrate", "gateway", Some(&thread_id.to_string()))
|
||||
.await;
|
||||
assert_eq!(resolved_tid, thread_id);
|
||||
|
||||
// Should be the same session object
|
||||
let sess = resolved_session.lock().await;
|
||||
assert!(sess.threads.contains_key(&thread_id));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_thread_with_explicit_external_id() {
|
||||
let manager = SessionManager::new();
|
||||
|
||||
// Two calls with the same explicit external thread ID should resolve
|
||||
// to the same internal thread.
|
||||
let (_, t1) = manager
|
||||
.resolve_thread("user-1", "gateway", Some("ext-abc"))
|
||||
.await;
|
||||
let (_, t2) = manager
|
||||
.resolve_thread("user-1", "gateway", Some("ext-abc"))
|
||||
.await;
|
||||
assert_eq!(t1, t2);
|
||||
|
||||
// A different external ID on the same channel/user gets a new thread.
|
||||
let (_, t3) = manager
|
||||
.resolve_thread("user-1", "gateway", Some("ext-xyz"))
|
||||
.await;
|
||||
assert_ne!(t1, t3);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_thread_none_vs_some_external_id() {
|
||||
let manager = SessionManager::new();
|
||||
|
||||
// None external_thread_id is a distinct key from Some("ext-1").
|
||||
let (_, t_none) = manager.resolve_thread("user-1", "cli", None).await;
|
||||
let (_, t_some) = manager.resolve_thread("user-1", "cli", Some("ext-1")).await;
|
||||
assert_ne!(t_none, t_some);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_thread_different_users_isolated() {
|
||||
let manager = SessionManager::new();
|
||||
|
||||
let (_, t1) = manager
|
||||
.resolve_thread("user-a", "gateway", Some("same-ext"))
|
||||
.await;
|
||||
let (_, t2) = manager
|
||||
.resolve_thread("user-b", "gateway", Some("same-ext"))
|
||||
.await;
|
||||
|
||||
// Same channel + same external ID but different users = different threads
|
||||
assert_ne!(t1, t2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_thread_different_channels_isolated() {
|
||||
let manager = SessionManager::new();
|
||||
|
||||
let (_, t1) = manager
|
||||
.resolve_thread("user-1", "gateway", Some("thread-x"))
|
||||
.await;
|
||||
let (_, t2) = manager
|
||||
.resolve_thread("user-1", "telegram", Some("thread-x"))
|
||||
.await;
|
||||
|
||||
// Same user + same external ID but different channels = different threads
|
||||
assert_ne!(t1, t2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_thread_stale_mapping_creates_new_thread() {
|
||||
let manager = SessionManager::new();
|
||||
|
||||
// Create a thread normally
|
||||
let (session, original_tid) = manager
|
||||
.resolve_thread("user-1", "gateway", Some("ext-1"))
|
||||
.await;
|
||||
|
||||
// Simulate the thread being removed from the session (e.g. pruned)
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
sess.threads.remove(&original_tid);
|
||||
}
|
||||
|
||||
// Next resolve should detect the stale mapping and create a fresh thread
|
||||
let (_, new_tid) = manager
|
||||
.resolve_thread("user-1", "gateway", Some("ext-1"))
|
||||
.await;
|
||||
assert_ne!(original_tid, new_tid);
|
||||
|
||||
// The new thread should actually exist in the session
|
||||
let sess = session.lock().await;
|
||||
assert!(sess.threads.contains_key(&new_tid));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_register_thread_preserves_uuid_on_resolve() {
|
||||
use crate::agent::session::{Session, Thread};
|
||||
|
||||
let manager = SessionManager::new();
|
||||
let known_uuid = Uuid::new_v4();
|
||||
|
||||
let session = Arc::new(Mutex::new(Session::new("user-web")));
|
||||
let session_id = {
|
||||
let sess = session.lock().await;
|
||||
sess.id
|
||||
};
|
||||
|
||||
// Simulate hydration: create thread with a known UUID
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(known_uuid, session_id);
|
||||
sess.threads.insert(known_uuid, thread);
|
||||
}
|
||||
|
||||
// Register it
|
||||
manager
|
||||
.register_thread("user-web", "gateway", known_uuid, Arc::clone(&session))
|
||||
.await;
|
||||
|
||||
// resolve_thread with UUID as external_thread_id MUST return the same UUID,
|
||||
// not mint a new one (this was the root cause of the "wrong conversation" bug)
|
||||
let (_, resolved) = manager
|
||||
.resolve_thread("user-web", "gateway", Some(&known_uuid.to_string()))
|
||||
.await;
|
||||
assert_eq!(resolved, known_uuid);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_register_thread_idempotent() {
|
||||
use crate::agent::session::{Session, Thread};
|
||||
|
||||
let manager = SessionManager::new();
|
||||
let tid = Uuid::new_v4();
|
||||
|
||||
let session = Arc::new(Mutex::new(Session::new("user-idem")));
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(tid, sess.id);
|
||||
sess.threads.insert(tid, thread);
|
||||
}
|
||||
|
||||
// Register twice
|
||||
manager
|
||||
.register_thread("user-idem", "gateway", tid, Arc::clone(&session))
|
||||
.await;
|
||||
manager
|
||||
.register_thread("user-idem", "gateway", tid, Arc::clone(&session))
|
||||
.await;
|
||||
|
||||
// Should still resolve to the same thread
|
||||
let (_, resolved) = manager
|
||||
.resolve_thread("user-idem", "gateway", Some(&tid.to_string()))
|
||||
.await;
|
||||
assert_eq!(resolved, tid);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_register_thread_creates_undo_manager() {
|
||||
use crate::agent::session::{Session, Thread};
|
||||
|
||||
let manager = SessionManager::new();
|
||||
let tid = Uuid::new_v4();
|
||||
|
||||
let session = Arc::new(Mutex::new(Session::new("user-undo")));
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(tid, sess.id);
|
||||
sess.threads.insert(tid, thread);
|
||||
}
|
||||
|
||||
manager
|
||||
.register_thread("user-undo", "gateway", tid, Arc::clone(&session))
|
||||
.await;
|
||||
|
||||
// Undo manager should exist for the registered thread
|
||||
let undo = manager.get_undo_manager(tid).await;
|
||||
let undo2 = manager.get_undo_manager(tid).await;
|
||||
assert!(Arc::ptr_eq(&undo, &undo2));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_register_thread_stores_session() {
|
||||
use crate::agent::session::{Session, Thread};
|
||||
|
||||
let manager = SessionManager::new();
|
||||
let tid = Uuid::new_v4();
|
||||
|
||||
let session = Arc::new(Mutex::new(Session::new("user-new")));
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(tid, sess.id);
|
||||
sess.threads.insert(tid, thread);
|
||||
}
|
||||
|
||||
// The user has no session yet in the manager
|
||||
{
|
||||
let sessions = manager.sessions.read().await;
|
||||
assert!(!sessions.contains_key("user-new"));
|
||||
}
|
||||
|
||||
manager
|
||||
.register_thread("user-new", "gateway", tid, Arc::clone(&session))
|
||||
.await;
|
||||
|
||||
// Now the session should be tracked
|
||||
{
|
||||
let sessions = manager.sessions.read().await;
|
||||
assert!(sessions.contains_key("user-new"));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multiple_threads_per_user() {
|
||||
let manager = SessionManager::new();
|
||||
|
||||
let (_, t1) = manager
|
||||
.resolve_thread("user-1", "gateway", Some("thread-a"))
|
||||
.await;
|
||||
let (_, t2) = manager
|
||||
.resolve_thread("user-1", "gateway", Some("thread-b"))
|
||||
.await;
|
||||
let (session, t3) = manager
|
||||
.resolve_thread("user-1", "gateway", Some("thread-c"))
|
||||
.await;
|
||||
|
||||
// All three should be distinct
|
||||
assert_ne!(t1, t2);
|
||||
assert_ne!(t2, t3);
|
||||
assert_ne!(t1, t3);
|
||||
|
||||
// All three should exist in the same session
|
||||
let sess = session.lock().await;
|
||||
assert!(sess.threads.contains_key(&t1));
|
||||
assert!(sess.threads.contains_key(&t2));
|
||||
assert!(sess.threads.contains_key(&t3));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_prune_cleans_thread_map_and_undo_managers() {
|
||||
let manager = SessionManager::new();
|
||||
|
||||
let (stale_session, stale_tid) = manager.resolve_thread("user-stale", "cli", None).await;
|
||||
|
||||
// Backdate the session
|
||||
{
|
||||
let mut sess = stale_session.lock().await;
|
||||
sess.last_active_at = chrono::Utc::now() - chrono::TimeDelta::seconds(86400 * 30);
|
||||
}
|
||||
|
||||
// Verify thread_map and undo_managers have entries
|
||||
{
|
||||
let tm = manager.thread_map.read().await;
|
||||
assert!(!tm.is_empty());
|
||||
}
|
||||
{
|
||||
let um = manager.undo_managers.read().await;
|
||||
assert!(um.contains_key(&stale_tid));
|
||||
}
|
||||
|
||||
let pruned = manager
|
||||
.prune_stale_sessions(std::time::Duration::from_secs(86400 * 7))
|
||||
.await;
|
||||
assert_eq!(pruned, 1);
|
||||
|
||||
// Thread map and undo managers should be cleaned up
|
||||
{
|
||||
let tm = manager.thread_map.read().await;
|
||||
assert!(tm.is_empty());
|
||||
}
|
||||
{
|
||||
let um = manager.undo_managers.read().await;
|
||||
assert!(!um.contains_key(&stale_tid));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_thread_active_thread_set() {
|
||||
let manager = SessionManager::new();
|
||||
|
||||
let (session, thread_id) = manager
|
||||
.resolve_thread("user-1", "gateway", Some("ext-1"))
|
||||
.await;
|
||||
|
||||
// The resolved thread should be set as the active thread
|
||||
let sess = session.lock().await;
|
||||
assert_eq!(sess.active_thread, Some(thread_id));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_register_then_resolve_different_channel_creates_new() {
|
||||
use crate::agent::session::{Session, Thread};
|
||||
|
||||
let manager = SessionManager::new();
|
||||
let tid = Uuid::new_v4();
|
||||
|
||||
let session = Arc::new(Mutex::new(Session::new("user-cross")));
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
let thread = Thread::with_id(tid, sess.id);
|
||||
sess.threads.insert(tid, thread);
|
||||
}
|
||||
|
||||
// Register on "gateway" channel
|
||||
manager
|
||||
.register_thread("user-cross", "gateway", tid, Arc::clone(&session))
|
||||
.await;
|
||||
|
||||
// Resolve on a different channel with the same UUID string should NOT
|
||||
// find the registered thread (channel is part of the key)
|
||||
let (_, resolved) = manager
|
||||
.resolve_thread("user-cross", "telegram", Some(&tid.to_string()))
|
||||
.await;
|
||||
assert_ne!(resolved, tid);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -43,6 +43,49 @@ impl SubmissionParser {
|
||||
if lower == "/thread new" || lower == "/new" {
|
||||
return Submission::NewThread;
|
||||
}
|
||||
// System commands (bypass thread-state checks)
|
||||
if lower == "/help" || lower == "/?" {
|
||||
return Submission::SystemCommand {
|
||||
command: "help".to_string(),
|
||||
args: vec![],
|
||||
};
|
||||
}
|
||||
if lower == "/version" {
|
||||
return Submission::SystemCommand {
|
||||
command: "version".to_string(),
|
||||
args: vec![],
|
||||
};
|
||||
}
|
||||
if lower == "/tools" {
|
||||
return Submission::SystemCommand {
|
||||
command: "tools".to_string(),
|
||||
args: vec![],
|
||||
};
|
||||
}
|
||||
if lower == "/ping" {
|
||||
return Submission::SystemCommand {
|
||||
command: "ping".to_string(),
|
||||
args: vec![],
|
||||
};
|
||||
}
|
||||
if lower == "/debug" {
|
||||
return Submission::SystemCommand {
|
||||
command: "debug".to_string(),
|
||||
args: vec![],
|
||||
};
|
||||
}
|
||||
if lower.starts_with("/model") {
|
||||
let args: Vec<String> = trimmed
|
||||
.split_whitespace()
|
||||
.skip(1)
|
||||
.map(|s| s.to_string())
|
||||
.collect();
|
||||
return Submission::SystemCommand {
|
||||
command: "model".to_string(),
|
||||
args,
|
||||
};
|
||||
}
|
||||
|
||||
if lower == "/quit" || lower == "/exit" || lower == "/shutdown" {
|
||||
return Submission::Quit;
|
||||
}
|
||||
@@ -172,6 +215,15 @@ pub enum Submission {
|
||||
|
||||
/// Quit the agent. Bypasses thread-state checks.
|
||||
Quit,
|
||||
|
||||
/// System command (help, model, version, tools, ping, debug).
|
||||
/// Bypasses thread-state checks and safety validation.
|
||||
SystemCommand {
|
||||
/// The command name (e.g. "help", "model", "version").
|
||||
command: String,
|
||||
/// Arguments to the command.
|
||||
args: Vec<String>,
|
||||
},
|
||||
}
|
||||
|
||||
impl Submission {
|
||||
@@ -238,6 +290,7 @@ impl Submission {
|
||||
| Self::Heartbeat
|
||||
| Self::Summarize
|
||||
| Self::Suggest
|
||||
| Self::SystemCommand { .. }
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -504,6 +557,84 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parser_system_command_help() {
|
||||
let submission = SubmissionParser::parse("/help");
|
||||
assert!(
|
||||
matches!(submission, Submission::SystemCommand { command, args } if command == "help" && args.is_empty())
|
||||
);
|
||||
|
||||
let submission = SubmissionParser::parse("/?");
|
||||
assert!(
|
||||
matches!(submission, Submission::SystemCommand { command, .. } if command == "help")
|
||||
);
|
||||
|
||||
let submission = SubmissionParser::parse("/HELP");
|
||||
assert!(
|
||||
matches!(submission, Submission::SystemCommand { command, .. } if command == "help")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parser_system_command_model() {
|
||||
// No args: show current model
|
||||
let submission = SubmissionParser::parse("/model");
|
||||
assert!(
|
||||
matches!(submission, Submission::SystemCommand { command, args } if command == "model" && args.is_empty())
|
||||
);
|
||||
|
||||
// With args: switch model
|
||||
let submission = SubmissionParser::parse("/model gpt-4o");
|
||||
assert!(
|
||||
matches!(submission, Submission::SystemCommand { command, args } if command == "model" && args == vec!["gpt-4o"])
|
||||
);
|
||||
|
||||
// Case insensitive command, preserves arg case
|
||||
let submission = SubmissionParser::parse("/MODEL Claude-3.5");
|
||||
assert!(
|
||||
matches!(submission, Submission::SystemCommand { command, args } if command == "model" && args == vec!["Claude-3.5"])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parser_system_command_version() {
|
||||
let submission = SubmissionParser::parse("/version");
|
||||
assert!(
|
||||
matches!(submission, Submission::SystemCommand { command, args } if command == "version" && args.is_empty())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parser_system_command_tools() {
|
||||
let submission = SubmissionParser::parse("/tools");
|
||||
assert!(
|
||||
matches!(submission, Submission::SystemCommand { command, args } if command == "tools" && args.is_empty())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parser_system_command_ping() {
|
||||
let submission = SubmissionParser::parse("/ping");
|
||||
assert!(
|
||||
matches!(submission, Submission::SystemCommand { command, args } if command == "ping" && args.is_empty())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parser_system_command_debug() {
|
||||
let submission = SubmissionParser::parse("/debug");
|
||||
assert!(
|
||||
matches!(submission, Submission::SystemCommand { command, args } if command == "debug" && args.is_empty())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parser_system_command_is_control() {
|
||||
let submission = SubmissionParser::parse("/help");
|
||||
assert!(submission.is_control());
|
||||
assert!(!submission.starts_turn());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parser_quit() {
|
||||
assert!(matches!(SubmissionParser::parse("/quit"), Submission::Quit));
|
||||
|
||||
+52
-4
@@ -272,7 +272,10 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
));
|
||||
}
|
||||
}
|
||||
RespondResult::ToolCalls(tool_calls) => {
|
||||
RespondResult::ToolCalls {
|
||||
tool_calls,
|
||||
content,
|
||||
} => {
|
||||
// Model returned tool calls - execute them
|
||||
tracing::debug!(
|
||||
"Job {} respond_with_tools returned {} tool calls",
|
||||
@@ -280,6 +283,14 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
tool_calls.len()
|
||||
);
|
||||
|
||||
// Add assistant message with tool_calls (OpenAI protocol)
|
||||
reason_ctx
|
||||
.messages
|
||||
.push(ChatMessage::assistant_with_tool_calls(
|
||||
content,
|
||||
tool_calls.clone(),
|
||||
));
|
||||
|
||||
for tc in tool_calls {
|
||||
let result = self.execute_tool(&tc.name, &tc.arguments).await;
|
||||
|
||||
@@ -417,14 +428,51 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
.into());
|
||||
}
|
||||
|
||||
// Execute with timeout and timing
|
||||
tracing::debug!(
|
||||
tool = %tool_name,
|
||||
params = %params,
|
||||
job = %job_id,
|
||||
"Tool call started"
|
||||
);
|
||||
|
||||
// Execute with per-tool timeout and timing
|
||||
let tool_timeout = tool.execution_timeout();
|
||||
let start = std::time::Instant::now();
|
||||
let result = tokio::time::timeout(Duration::from_secs(60), async {
|
||||
let result = tokio::time::timeout(tool_timeout, async {
|
||||
tool.execute(params.clone(), &job_ctx).await
|
||||
})
|
||||
.await;
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
match &result {
|
||||
Ok(Ok(output)) => {
|
||||
let result_str = serde_json::to_string(&output.result)
|
||||
.unwrap_or_else(|_| "<serialize error>".to_string());
|
||||
tracing::debug!(
|
||||
tool = %tool_name,
|
||||
elapsed_ms = elapsed.as_millis() as u64,
|
||||
result = %result_str,
|
||||
"Tool call succeeded"
|
||||
);
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
tracing::debug!(
|
||||
tool = %tool_name,
|
||||
elapsed_ms = elapsed.as_millis() as u64,
|
||||
error = %e,
|
||||
"Tool call failed"
|
||||
);
|
||||
}
|
||||
Err(_) => {
|
||||
tracing::debug!(
|
||||
tool = %tool_name,
|
||||
elapsed_ms = elapsed.as_millis() as u64,
|
||||
timeout_secs = tool_timeout.as_secs(),
|
||||
"Tool call timed out"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Record action in memory and get the ActionRecord for persistence
|
||||
let action = match &result {
|
||||
Ok(Ok(output)) => {
|
||||
@@ -479,7 +527,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
let output = result
|
||||
.map_err(|_| crate::error::ToolError::Timeout {
|
||||
name: tool_name.to_string(),
|
||||
timeout: Duration::from_secs(60),
|
||||
timeout: tool_timeout,
|
||||
})?
|
||||
.map_err(|e| crate::error::ToolError::ExecutionFailed {
|
||||
name: tool_name.to_string(),
|
||||
|
||||
@@ -0,0 +1,325 @@
|
||||
//! Bootstrap configuration for IronClaw.
|
||||
//!
|
||||
//! These are the only settings that MUST live on disk because they're needed
|
||||
//! before the database connection is established. Everything else lives in the
|
||||
//! `settings` table in PostgreSQL.
|
||||
//!
|
||||
//! File: `~/.ironclaw/bootstrap.json`
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::settings::KeySource;
|
||||
|
||||
/// Minimal config needed to connect to the database and decrypt secrets.
|
||||
///
|
||||
/// This is the only JSON file IronClaw reads from disk at startup.
|
||||
/// All other configuration lives in the `settings` table in PostgreSQL.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct BootstrapConfig {
|
||||
/// Database connection URL (postgres://...).
|
||||
#[serde(default)]
|
||||
pub database_url: Option<String>,
|
||||
|
||||
/// Database connection pool size.
|
||||
#[serde(default)]
|
||||
pub database_pool_size: Option<usize>,
|
||||
|
||||
/// Source for the secrets master key.
|
||||
#[serde(default)]
|
||||
pub secrets_master_key_source: KeySource,
|
||||
|
||||
/// Whether onboarding wizard has been completed.
|
||||
#[serde(default)]
|
||||
pub onboard_completed: bool,
|
||||
}
|
||||
|
||||
impl Default for BootstrapConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
database_url: None,
|
||||
database_pool_size: None,
|
||||
secrets_master_key_source: KeySource::None,
|
||||
onboard_completed: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl BootstrapConfig {
|
||||
/// Default bootstrap file path: `~/.ironclaw/bootstrap.json`.
|
||||
pub fn default_path() -> PathBuf {
|
||||
dirs::home_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join(".ironclaw")
|
||||
.join("bootstrap.json")
|
||||
}
|
||||
|
||||
/// Legacy settings.json path (for migration detection).
|
||||
pub fn legacy_settings_path() -> PathBuf {
|
||||
dirs::home_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join(".ironclaw")
|
||||
.join("settings.json")
|
||||
}
|
||||
|
||||
/// Load from the default path, falling back to legacy settings.json,
|
||||
/// then to defaults if neither exists.
|
||||
pub fn load() -> Self {
|
||||
let bootstrap_path = Self::default_path();
|
||||
if bootstrap_path.exists() {
|
||||
return Self::load_from(&bootstrap_path);
|
||||
}
|
||||
|
||||
// Fall back to legacy settings.json (extract just the 4 bootstrap fields)
|
||||
let legacy_path = Self::legacy_settings_path();
|
||||
if legacy_path.exists() {
|
||||
return Self::load_from_legacy(&legacy_path);
|
||||
}
|
||||
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Load from a specific path.
|
||||
pub fn load_from(path: &PathBuf) -> Self {
|
||||
match std::fs::read_to_string(path) {
|
||||
Ok(data) => serde_json::from_str(&data).unwrap_or_default(),
|
||||
Err(_) => Self::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract bootstrap fields from a legacy settings.json.
|
||||
fn load_from_legacy(path: &PathBuf) -> Self {
|
||||
match std::fs::read_to_string(path) {
|
||||
Ok(data) => {
|
||||
// The legacy Settings struct is a superset; serde will ignore extra fields.
|
||||
serde_json::from_str(&data).unwrap_or_default()
|
||||
}
|
||||
Err(_) => Self::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Save to the default path.
|
||||
pub fn save(&self) -> std::io::Result<()> {
|
||||
self.save_to(&Self::default_path())
|
||||
}
|
||||
|
||||
/// Save to a specific path.
|
||||
pub fn save_to(&self, path: &PathBuf) -> std::io::Result<()> {
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
let json = serde_json::to_string_pretty(self)
|
||||
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string()))?;
|
||||
std::fs::write(path, json)
|
||||
}
|
||||
}
|
||||
|
||||
/// One-time migration from disk config files to the database settings table.
|
||||
///
|
||||
/// On first boot after upgrade, checks if:
|
||||
/// 1. `~/.ironclaw/settings.json` exists
|
||||
/// 2. The DB settings table is empty for this user
|
||||
///
|
||||
/// If both conditions hold, migrates settings, MCP servers, and session data
|
||||
/// to the database, writes `bootstrap.json`, and renames old files to `.migrated`.
|
||||
pub async fn migrate_disk_to_db(
|
||||
store: &crate::history::Store,
|
||||
user_id: &str,
|
||||
) -> Result<(), MigrationError> {
|
||||
let legacy_settings_path = BootstrapConfig::legacy_settings_path();
|
||||
if !legacy_settings_path.exists() {
|
||||
tracing::debug!("No legacy settings.json found, skipping disk-to-DB migration");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Only migrate if DB is empty for this user
|
||||
let has_settings = store.has_settings(user_id).await.map_err(|e| {
|
||||
MigrationError::Database(format!("Failed to check existing settings: {}", e))
|
||||
})?;
|
||||
if has_settings {
|
||||
tracing::debug!(
|
||||
"DB already has settings for user '{}', skipping migration",
|
||||
user_id
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
tracing::info!("Migrating disk settings to database...");
|
||||
|
||||
// 1. Load and migrate settings.json
|
||||
let settings = crate::settings::Settings::load_from(&legacy_settings_path);
|
||||
let db_map = settings.to_db_map();
|
||||
if !db_map.is_empty() {
|
||||
store
|
||||
.set_all_settings(user_id, &db_map)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
MigrationError::Database(format!("Failed to write settings to DB: {}", e))
|
||||
})?;
|
||||
tracing::info!("Migrated {} settings to database", db_map.len());
|
||||
}
|
||||
|
||||
// 2. Write bootstrap.json with the 4 essential fields
|
||||
let bootstrap = BootstrapConfig {
|
||||
database_url: settings.database_url.clone(),
|
||||
database_pool_size: settings.database_pool_size,
|
||||
secrets_master_key_source: settings.secrets_master_key_source,
|
||||
onboard_completed: settings.onboard_completed,
|
||||
};
|
||||
bootstrap
|
||||
.save()
|
||||
.map_err(|e| MigrationError::Io(format!("Failed to write bootstrap.json: {}", e)))?;
|
||||
tracing::info!("Wrote bootstrap.json");
|
||||
|
||||
// 3. Migrate mcp-servers.json if it exists
|
||||
let ironclaw_dir = dirs::home_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join(".ironclaw");
|
||||
let mcp_path = ironclaw_dir.join("mcp-servers.json");
|
||||
if mcp_path.exists() {
|
||||
match std::fs::read_to_string(&mcp_path) {
|
||||
Ok(content) => match serde_json::from_str::<serde_json::Value>(&content) {
|
||||
Ok(value) => {
|
||||
store
|
||||
.set_setting(user_id, "mcp_servers", &value)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
MigrationError::Database(format!(
|
||||
"Failed to write MCP servers to DB: {}",
|
||||
e
|
||||
))
|
||||
})?;
|
||||
tracing::info!("Migrated mcp-servers.json to database");
|
||||
|
||||
rename_to_migrated(&mcp_path);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to parse mcp-servers.json: {}", e);
|
||||
}
|
||||
},
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to read mcp-servers.json: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 4. Migrate session.json if it exists
|
||||
let session_path = ironclaw_dir.join("session.json");
|
||||
if session_path.exists() {
|
||||
match std::fs::read_to_string(&session_path) {
|
||||
Ok(content) => match serde_json::from_str::<serde_json::Value>(&content) {
|
||||
Ok(value) => {
|
||||
store
|
||||
.set_setting(user_id, "nearai.session", &value)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
MigrationError::Database(format!(
|
||||
"Failed to write session to DB: {}",
|
||||
e
|
||||
))
|
||||
})?;
|
||||
tracing::info!("Migrated session.json to database");
|
||||
|
||||
rename_to_migrated(&session_path);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to parse session.json: {}", e);
|
||||
}
|
||||
},
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to read session.json: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 5. Rename settings.json to .migrated (don't delete, safety net)
|
||||
rename_to_migrated(&legacy_settings_path);
|
||||
|
||||
tracing::info!("Disk-to-DB migration complete");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Rename a file to `<name>.migrated` as a safety net.
|
||||
fn rename_to_migrated(path: &PathBuf) {
|
||||
let mut migrated = path.as_os_str().to_owned();
|
||||
migrated.push(".migrated");
|
||||
if let Err(e) = std::fs::rename(path, &migrated) {
|
||||
tracing::warn!("Failed to rename {} to .migrated: {}", path.display(), e);
|
||||
}
|
||||
}
|
||||
|
||||
/// Errors that can occur during disk-to-DB migration.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum MigrationError {
|
||||
#[error("Database error: {0}")]
|
||||
Database(String),
|
||||
#[error("IO error: {0}")]
|
||||
Io(String),
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::tempdir;
|
||||
|
||||
#[test]
|
||||
fn test_bootstrap_save_load() {
|
||||
let dir = tempdir().unwrap();
|
||||
let path = dir.path().join("bootstrap.json");
|
||||
|
||||
let config = BootstrapConfig {
|
||||
database_url: Some("postgres://localhost/test".to_string()),
|
||||
database_pool_size: Some(5),
|
||||
secrets_master_key_source: KeySource::Keychain,
|
||||
onboard_completed: true,
|
||||
};
|
||||
|
||||
config.save_to(&path).unwrap();
|
||||
|
||||
let loaded = BootstrapConfig::load_from(&path);
|
||||
assert_eq!(
|
||||
loaded.database_url,
|
||||
Some("postgres://localhost/test".to_string())
|
||||
);
|
||||
assert_eq!(loaded.database_pool_size, Some(5));
|
||||
assert_eq!(loaded.secrets_master_key_source, KeySource::Keychain);
|
||||
assert!(loaded.onboard_completed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bootstrap_from_legacy_settings() {
|
||||
let dir = tempdir().unwrap();
|
||||
let path = dir.path().join("settings.json");
|
||||
|
||||
// Write a legacy settings.json with many extra fields
|
||||
let legacy = serde_json::json!({
|
||||
"database_url": "postgres://localhost/ironclaw",
|
||||
"database_pool_size": 10,
|
||||
"secrets_master_key_source": "keychain",
|
||||
"onboard_completed": true,
|
||||
"selected_model": "claude-3-5-sonnet",
|
||||
"agent": { "name": "testbot", "max_parallel_jobs": 3 },
|
||||
"heartbeat": { "enabled": true }
|
||||
});
|
||||
std::fs::write(&path, serde_json::to_string_pretty(&legacy).unwrap()).unwrap();
|
||||
|
||||
let config = BootstrapConfig::load_from_legacy(&path);
|
||||
assert_eq!(
|
||||
config.database_url,
|
||||
Some("postgres://localhost/ironclaw".to_string())
|
||||
);
|
||||
assert_eq!(config.database_pool_size, Some(10));
|
||||
assert_eq!(config.secrets_master_key_source, KeySource::Keychain);
|
||||
assert!(config.onboard_completed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bootstrap_defaults() {
|
||||
let config = BootstrapConfig::default();
|
||||
assert!(config.database_url.is_none());
|
||||
assert!(config.database_pool_size.is_none());
|
||||
assert_eq!(config.secrets_master_key_source, KeySource::None);
|
||||
assert!(!config.onboard_completed);
|
||||
}
|
||||
}
|
||||
@@ -114,6 +114,12 @@ pub enum StatusUpdate {
|
||||
StreamChunk(String),
|
||||
/// General status message.
|
||||
Status(String),
|
||||
/// A sandbox job has started (shown as a clickable card in the UI).
|
||||
JobStarted {
|
||||
job_id: String,
|
||||
title: String,
|
||||
browse_url: String,
|
||||
},
|
||||
/// Tool requires user approval before execution.
|
||||
ApprovalNeeded {
|
||||
request_id: String,
|
||||
@@ -121,6 +127,19 @@ pub enum StatusUpdate {
|
||||
description: String,
|
||||
parameters: serde_json::Value,
|
||||
},
|
||||
/// Extension needs user authentication (token or OAuth).
|
||||
AuthRequired {
|
||||
extension_name: String,
|
||||
instructions: Option<String>,
|
||||
auth_url: Option<String>,
|
||||
setup_url: Option<String>,
|
||||
},
|
||||
/// Extension authentication completed.
|
||||
AuthCompleted {
|
||||
extension_name: String,
|
||||
success: bool,
|
||||
message: String,
|
||||
},
|
||||
}
|
||||
|
||||
/// Trait for message channels.
|
||||
|
||||
+52
-2
@@ -42,12 +42,25 @@ const SLASH_COMMANDS: &[&str] = &[
|
||||
"/quit",
|
||||
"/exit",
|
||||
"/debug",
|
||||
"/model",
|
||||
"/undo",
|
||||
"/redo",
|
||||
"/clear",
|
||||
"/compact",
|
||||
"/new",
|
||||
"/interrupt",
|
||||
"/version",
|
||||
"/tools",
|
||||
"/ping",
|
||||
"/job",
|
||||
"/status",
|
||||
"/cancel",
|
||||
"/list",
|
||||
"/heartbeat",
|
||||
"/summarize",
|
||||
"/suggest",
|
||||
"/thread",
|
||||
"/resume",
|
||||
];
|
||||
|
||||
/// Rustyline helper for slash-command tab completion.
|
||||
@@ -295,10 +308,11 @@ impl Channel for ReplChannel {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Handle local REPL commands
|
||||
// Handle local REPL commands (only commands that need
|
||||
// immediate local handling stay here)
|
||||
match line.to_lowercase().as_str() {
|
||||
"/quit" | "/exit" => break,
|
||||
"/help" | "/?" => {
|
||||
"/help" => {
|
||||
print_help();
|
||||
continue;
|
||||
}
|
||||
@@ -413,6 +427,15 @@ impl Channel for ReplChannel {
|
||||
print!("{chunk}");
|
||||
let _ = io::stdout().flush();
|
||||
}
|
||||
StatusUpdate::JobStarted {
|
||||
job_id,
|
||||
title,
|
||||
browse_url,
|
||||
} => {
|
||||
eprintln!(
|
||||
" \x1b[36m[job]\x1b[0m {title} \x1b[90m({job_id})\x1b[0m \x1b[4m{browse_url}\x1b[0m"
|
||||
);
|
||||
}
|
||||
StatusUpdate::Status(msg) => {
|
||||
if debug || msg.contains("approval") || msg.contains("Approval") {
|
||||
eprintln!(" \x1b[90m{msg}\x1b[0m");
|
||||
@@ -472,6 +495,33 @@ impl Channel for ReplChannel {
|
||||
eprintln!(" {bot_border}");
|
||||
eprintln!();
|
||||
}
|
||||
StatusUpdate::AuthRequired {
|
||||
extension_name,
|
||||
instructions,
|
||||
setup_url,
|
||||
..
|
||||
} => {
|
||||
eprintln!();
|
||||
eprintln!("\x1b[33m Authentication required for {extension_name}\x1b[0m");
|
||||
if let Some(ref instr) = instructions {
|
||||
eprintln!(" {instr}");
|
||||
}
|
||||
if let Some(ref url) = setup_url {
|
||||
eprintln!(" \x1b[4m{url}\x1b[0m");
|
||||
}
|
||||
eprintln!();
|
||||
}
|
||||
StatusUpdate::AuthCompleted {
|
||||
extension_name,
|
||||
success,
|
||||
message,
|
||||
} => {
|
||||
if success {
|
||||
eprintln!("\x1b[32m {extension_name}: {message}\x1b[0m");
|
||||
} else {
|
||||
eprintln!("\x1b[31m {extension_name}: {message}\x1b[0m");
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
+109
-48
@@ -1,68 +1,125 @@
|
||||
//! Bundled WASM channels that can be installed locally.
|
||||
//! Known WASM channels that can be installed from build artifacts.
|
||||
//!
|
||||
//! Instead of embedding WASM binaries in the host binary via include_bytes!,
|
||||
//! channels are compiled separately and installed from their build output
|
||||
//! directories during onboarding.
|
||||
//!
|
||||
//! Channel source layout:
|
||||
//! channels-src/<name>/
|
||||
//! target/wasm32-wasip2/release/<name>_channel.wasm
|
||||
//! <name>.capabilities.json
|
||||
|
||||
use std::path::Path;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use tokio::fs;
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
struct BundledChannel {
|
||||
name: &'static str,
|
||||
wasm: &'static [u8],
|
||||
capabilities: &'static [u8],
|
||||
/// Compile-time project root, used to locate channels-src/ in dev builds.
|
||||
const CARGO_MANIFEST_DIR: &str = env!("CARGO_MANIFEST_DIR");
|
||||
|
||||
/// Known channel names and their crate names (for locating build artifacts).
|
||||
const KNOWN_CHANNELS: &[(&str, &str)] = &[
|
||||
("telegram", "telegram_channel"),
|
||||
("slack", "slack_channel"),
|
||||
("whatsapp", "whatsapp_channel"),
|
||||
];
|
||||
|
||||
/// Names of known channels that can be installed.
|
||||
pub fn bundled_channel_names() -> Vec<&'static str> {
|
||||
KNOWN_CHANNELS.iter().map(|(name, _)| *name).collect()
|
||||
}
|
||||
|
||||
/// Names of bundled channels shipped with IronClaw.
|
||||
pub fn bundled_channel_names() -> &'static [&'static str] {
|
||||
&["telegram"]
|
||||
/// Resolve the channels source directory.
|
||||
///
|
||||
/// Checks (in order):
|
||||
/// 1. `IRONCLAW_CHANNELS_SRC` env var
|
||||
/// 2. `<CARGO_MANIFEST_DIR>/channels-src/` (dev builds)
|
||||
fn channels_src_dir() -> PathBuf {
|
||||
if let Ok(dir) = std::env::var("IRONCLAW_CHANNELS_SRC") {
|
||||
return PathBuf::from(dir);
|
||||
}
|
||||
PathBuf::from(CARGO_MANIFEST_DIR).join("channels-src")
|
||||
}
|
||||
|
||||
/// Install a bundled channel into a channels directory.
|
||||
/// Locate the build artifacts for a channel.
|
||||
///
|
||||
/// Returns (wasm_path, capabilities_path) or an error if files are missing.
|
||||
fn locate_channel_artifacts(name: &str) -> Result<(PathBuf, PathBuf), String> {
|
||||
let (_, crate_name) = KNOWN_CHANNELS
|
||||
.iter()
|
||||
.find(|(n, _)| *n == name)
|
||||
.ok_or_else(|| format!("Unknown channel '{}'", name))?;
|
||||
|
||||
let src_dir = channels_src_dir();
|
||||
let channel_dir = src_dir.join(name);
|
||||
|
||||
let wasm_path = channel_dir
|
||||
.join("target/wasm32-wasip2/release")
|
||||
.join(format!("{}.wasm", crate_name));
|
||||
|
||||
let caps_path = channel_dir.join(format!("{}.capabilities.json", name));
|
||||
|
||||
if !wasm_path.exists() {
|
||||
return Err(format!(
|
||||
"Channel '{}' WASM not found at {}. Build it first:\n \
|
||||
cd {} && cargo build --target wasm32-wasip2 --release",
|
||||
name,
|
||||
wasm_path.display(),
|
||||
channel_dir.display()
|
||||
));
|
||||
}
|
||||
|
||||
if !caps_path.exists() {
|
||||
return Err(format!(
|
||||
"Channel '{}' capabilities not found at {}",
|
||||
name,
|
||||
caps_path.display()
|
||||
));
|
||||
}
|
||||
|
||||
Ok((wasm_path, caps_path))
|
||||
}
|
||||
|
||||
/// Install a channel from build artifacts into the channels directory.
|
||||
pub async fn install_bundled_channel(
|
||||
name: &str,
|
||||
target_dir: &Path,
|
||||
force: bool,
|
||||
) -> Result<(), String> {
|
||||
let channel = bundled_channel(name)
|
||||
.ok_or_else(|| format!("Unknown bundled channel '{}'", name.to_lowercase()))?;
|
||||
let (wasm_src, caps_src) = locate_channel_artifacts(name)?;
|
||||
|
||||
fs::create_dir_all(target_dir)
|
||||
.await
|
||||
.map_err(|e| format!("Failed to create channels directory: {}", e))?;
|
||||
|
||||
let wasm_path = target_dir.join(format!("{}.wasm", channel.name));
|
||||
let caps_path = target_dir.join(format!("{}.capabilities.json", channel.name));
|
||||
let wasm_dst = target_dir.join(format!("{}.wasm", name));
|
||||
let caps_dst = target_dir.join(format!("{}.capabilities.json", name));
|
||||
|
||||
let has_existing = wasm_path.exists() || caps_path.exists();
|
||||
let has_existing = wasm_dst.exists() || caps_dst.exists();
|
||||
if has_existing && !force {
|
||||
return Err(format!(
|
||||
"Channel '{}' already exists at {}",
|
||||
channel.name,
|
||||
name,
|
||||
target_dir.display()
|
||||
));
|
||||
}
|
||||
|
||||
fs::write(&wasm_path, channel.wasm)
|
||||
fs::copy(&wasm_src, &wasm_dst)
|
||||
.await
|
||||
.map_err(|e| format!("Failed to write {}: {}", wasm_path.display(), e))?;
|
||||
fs::write(&caps_path, channel.capabilities)
|
||||
.map_err(|e| format!("Failed to copy {}: {}", wasm_src.display(), e))?;
|
||||
fs::copy(&caps_src, &caps_dst)
|
||||
.await
|
||||
.map_err(|e| format!("Failed to write {}: {}", caps_path.display(), e))?;
|
||||
.map_err(|e| format!("Failed to copy {}: {}", caps_src.display(), e))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn bundled_channel(name: &str) -> Option<BundledChannel> {
|
||||
if name.eq_ignore_ascii_case("telegram") {
|
||||
Some(BundledChannel {
|
||||
name: "telegram",
|
||||
wasm: include_bytes!("../../../channels-src/telegram/telegram.wasm"),
|
||||
capabilities: include_bytes!(
|
||||
"../../../channels-src/telegram/telegram.capabilities.json"
|
||||
),
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
/// Check which known channels have build artifacts available.
|
||||
pub fn available_channel_names() -> Vec<&'static str> {
|
||||
KNOWN_CHANNELS
|
||||
.iter()
|
||||
.filter(|(name, _)| locate_channel_artifacts(name).is_ok())
|
||||
.map(|(name, _)| *name)
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -73,31 +130,35 @@ mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_bundled_channel_names_contains_telegram() {
|
||||
assert!(bundled_channel_names().contains(&"telegram"));
|
||||
fn test_known_channels_includes_all_three() {
|
||||
let names = bundled_channel_names();
|
||||
assert!(names.contains(&"telegram"));
|
||||
assert!(names.contains(&"slack"));
|
||||
assert!(names.contains(&"whatsapp"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_channels_src_dir_default() {
|
||||
let dir = channels_src_dir();
|
||||
assert!(dir.ends_with("channels-src"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_locate_unknown_channel_errors() {
|
||||
assert!(locate_channel_artifacts("nonexistent").is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_install_bundled_channel_writes_files() {
|
||||
let dir = tempdir().unwrap();
|
||||
|
||||
install_bundled_channel("telegram", dir.path(), false)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(dir.path().join("telegram.wasm").exists());
|
||||
assert!(dir.path().join("telegram.capabilities.json").exists());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_install_bundled_channel_refuses_overwrite_without_force() {
|
||||
async fn test_install_refuses_overwrite_without_force() {
|
||||
let dir = tempdir().unwrap();
|
||||
let wasm_path = dir.path().join("telegram.wasm");
|
||||
fs::write(&wasm_path, b"custom").await.unwrap();
|
||||
|
||||
let result = install_bundled_channel("telegram", dir.path(), false).await;
|
||||
// Either fails because artifacts missing OR because file exists
|
||||
assert!(result.is_err());
|
||||
|
||||
// Original file should be untouched
|
||||
let existing = fs::read(&wasm_path).await.unwrap();
|
||||
assert_eq!(existing, b"custom");
|
||||
}
|
||||
|
||||
@@ -16,16 +16,21 @@ use crate::channels::wasm::error::WasmChannelError;
|
||||
use crate::channels::wasm::runtime::WasmChannelRuntime;
|
||||
use crate::channels::wasm::schema::ChannelCapabilitiesFile;
|
||||
use crate::channels::wasm::wrapper::WasmChannel;
|
||||
use crate::pairing::PairingStore;
|
||||
|
||||
/// Loads WASM channels from the filesystem.
|
||||
pub struct WasmChannelLoader {
|
||||
runtime: Arc<WasmChannelRuntime>,
|
||||
pairing_store: Arc<PairingStore>,
|
||||
}
|
||||
|
||||
impl WasmChannelLoader {
|
||||
/// Create a new loader with the given runtime.
|
||||
pub fn new(runtime: Arc<WasmChannelRuntime>) -> Self {
|
||||
Self { runtime }
|
||||
/// Create a new loader with the given runtime and pairing store.
|
||||
pub fn new(runtime: Arc<WasmChannelRuntime>, pairing_store: Arc<PairingStore>) -> Self {
|
||||
Self {
|
||||
runtime,
|
||||
pairing_store,
|
||||
}
|
||||
}
|
||||
|
||||
/// Load a single WASM channel from a file pair.
|
||||
@@ -114,7 +119,13 @@ impl WasmChannelLoader {
|
||||
.await?;
|
||||
|
||||
// Create the channel
|
||||
let channel = WasmChannel::new(self.runtime.clone(), prepared, capabilities, config_json);
|
||||
let channel = WasmChannel::new(
|
||||
self.runtime.clone(),
|
||||
prepared,
|
||||
capabilities,
|
||||
config_json,
|
||||
self.pairing_store.clone(),
|
||||
);
|
||||
|
||||
tracing::info!(
|
||||
name = name,
|
||||
@@ -352,6 +363,7 @@ mod tests {
|
||||
|
||||
use crate::channels::wasm::loader::{WasmChannelLoader, discover_channels};
|
||||
use crate::channels::wasm::runtime::{WasmChannelRuntime, WasmChannelRuntimeConfig};
|
||||
use crate::pairing::PairingStore;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[tokio::test]
|
||||
@@ -408,7 +420,7 @@ mod tests {
|
||||
async fn test_loader_invalid_name() {
|
||||
let config = WasmChannelRuntimeConfig::for_testing();
|
||||
let runtime = Arc::new(WasmChannelRuntime::new(config).unwrap());
|
||||
let loader = WasmChannelLoader::new(runtime);
|
||||
let loader = WasmChannelLoader::new(runtime, Arc::new(PairingStore::new()));
|
||||
|
||||
let dir = TempDir::new().unwrap();
|
||||
let wasm_path = dir.path().join("test.wasm");
|
||||
|
||||
@@ -89,7 +89,7 @@ mod schema;
|
||||
mod wrapper;
|
||||
|
||||
// Core types
|
||||
pub use bundled::{bundled_channel_names, install_bundled_channel};
|
||||
pub use bundled::{available_channel_names, bundled_channel_names, install_bundled_channel};
|
||||
pub use capabilities::{ChannelCapabilities, EmitRateLimitConfig, HttpEndpointConfig, PollConfig};
|
||||
pub use error::WasmChannelError;
|
||||
pub use host::{ChannelEmitRateLimiter, ChannelHostState, EmittedMessage};
|
||||
|
||||
@@ -478,6 +478,7 @@ mod tests {
|
||||
PreparedChannelModule, WasmChannelRuntime, WasmChannelRuntimeConfig,
|
||||
};
|
||||
use crate::channels::wasm::wrapper::WasmChannel;
|
||||
use crate::pairing::PairingStore;
|
||||
use crate::tools::wasm::ResourceLimits;
|
||||
|
||||
fn create_test_channel(name: &str) -> Arc<WasmChannel> {
|
||||
@@ -499,6 +500,7 @@ mod tests {
|
||||
prepared,
|
||||
capabilities,
|
||||
"{}".to_string(),
|
||||
Arc::new(PairingStore::new()),
|
||||
))
|
||||
}
|
||||
|
||||
|
||||
+271
-22
@@ -48,6 +48,7 @@ use crate::channels::wasm::runtime::{PreparedChannelModule, WasmChannelRuntime};
|
||||
use crate::channels::wasm::schema::ChannelConfig;
|
||||
use crate::channels::{Channel, IncomingMessage, MessageStream, OutgoingResponse, StatusUpdate};
|
||||
use crate::error::ChannelError;
|
||||
use crate::pairing::PairingStore;
|
||||
use crate::safety::LeakDetector;
|
||||
use crate::tools::wasm::LogLevel;
|
||||
use crate::tools::wasm::WasmResourceLimiter;
|
||||
@@ -73,6 +74,8 @@ struct ChannelStoreData {
|
||||
/// Injected credentials for URL substitution (e.g., bot tokens).
|
||||
/// Keys are placeholder names like "TELEGRAM_BOT_TOKEN".
|
||||
credentials: HashMap<String, String>,
|
||||
/// Pairing store for DM pairing (guest access control).
|
||||
pairing_store: Arc<PairingStore>,
|
||||
}
|
||||
|
||||
impl ChannelStoreData {
|
||||
@@ -81,6 +84,7 @@ impl ChannelStoreData {
|
||||
channel_name: &str,
|
||||
capabilities: ChannelCapabilities,
|
||||
credentials: HashMap<String, String>,
|
||||
pairing_store: Arc<PairingStore>,
|
||||
) -> Self {
|
||||
// Create a minimal WASI context (no filesystem, no env vars for security)
|
||||
let wasi = WasiCtxBuilder::new().build();
|
||||
@@ -91,6 +95,7 @@ impl ChannelStoreData {
|
||||
wasi,
|
||||
table: ResourceTable::new(),
|
||||
credentials,
|
||||
pairing_store,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -141,6 +146,22 @@ impl ChannelStoreData {
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
/// Replace injected credential values with `[REDACTED]` in text.
|
||||
///
|
||||
/// Prevents credentials from leaking through error messages, logs, or
|
||||
/// return values to WASM. reqwest::Error includes the full URL in its
|
||||
/// Display output, so any error from an injected-URL request will
|
||||
/// contain the raw credential unless we scrub it.
|
||||
fn redact_credentials(&self, text: &str) -> String {
|
||||
let mut result = text.to_string();
|
||||
for (name, value) in &self.credentials {
|
||||
if !value.is_empty() {
|
||||
result = result.replace(value, &format!("[REDACTED:{}]", name));
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
// Implement WasiView to provide WASI context and resource table
|
||||
@@ -187,6 +208,7 @@ impl near::agent::channel_host::Host for ChannelStoreData {
|
||||
url: String,
|
||||
headers_json: String,
|
||||
body: Option<Vec<u8>>,
|
||||
timeout_ms: Option<u32>,
|
||||
) -> Result<near::agent::channel_host::HttpResponse, String> {
|
||||
tracing::info!(
|
||||
method = %method,
|
||||
@@ -276,12 +298,21 @@ impl near::agent::channel_host::Host for ChannelStoreData {
|
||||
request = request.body(body_bytes);
|
||||
}
|
||||
|
||||
// Send request with timeout
|
||||
let response = request
|
||||
.timeout(std::time::Duration::from_secs(30))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("HTTP request failed: {}", e))?;
|
||||
// Send request with caller-specified timeout (default 30s).
|
||||
// Cap at callback_timeout to prevent outliving the host wrapper.
|
||||
let timeout = std::time::Duration::from_millis(timeout_ms.unwrap_or(30_000) as u64);
|
||||
let response = request.timeout(timeout).send().await.map_err(|e| {
|
||||
// Walk the full error chain so we get the actual root cause
|
||||
// (DNS, TLS, connection refused, etc.) instead of just
|
||||
// "error sending request for url (...)".
|
||||
let mut chain = format!("HTTP request failed: {}", e);
|
||||
let mut source = std::error::Error::source(&e);
|
||||
while let Some(cause) = source {
|
||||
chain.push_str(&format!(" -> {}", cause));
|
||||
source = cause.source();
|
||||
}
|
||||
chain
|
||||
})?;
|
||||
|
||||
let status = response.status().as_u16();
|
||||
let response_headers: std::collections::HashMap<String, String> = response
|
||||
@@ -330,6 +361,11 @@ impl near::agent::channel_host::Host for ChannelStoreData {
|
||||
})
|
||||
});
|
||||
|
||||
// Scrub credential values from error messages before logging or returning
|
||||
// to WASM. reqwest::Error includes the full URL (with injected credentials)
|
||||
// in its Display output.
|
||||
let result = result.map_err(|e| self.redact_credentials(&e));
|
||||
|
||||
match &result {
|
||||
Ok(resp) => {
|
||||
tracing::info!(status = resp.status, "http_request completed successfully");
|
||||
@@ -372,6 +408,43 @@ impl near::agent::channel_host::Host for ChannelStoreData {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn pairing_upsert_request(
|
||||
&mut self,
|
||||
channel: String,
|
||||
id: String,
|
||||
meta_json: String,
|
||||
) -> Result<near::agent::channel_host::PairingUpsertResult, String> {
|
||||
let meta = if meta_json.is_empty() {
|
||||
None
|
||||
} else {
|
||||
serde_json::from_str(&meta_json).ok()
|
||||
};
|
||||
match self.pairing_store.upsert_request(&channel, &id, meta) {
|
||||
Ok(r) => Ok(near::agent::channel_host::PairingUpsertResult {
|
||||
code: r.code,
|
||||
created: r.created,
|
||||
}),
|
||||
Err(e) => Err(e.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn pairing_is_allowed(
|
||||
&mut self,
|
||||
channel: String,
|
||||
id: String,
|
||||
username: Option<String>,
|
||||
) -> Result<bool, String> {
|
||||
self.pairing_store
|
||||
.is_sender_allowed(&channel, &id, username.as_deref())
|
||||
.map_err(|e| e.to_string())
|
||||
}
|
||||
|
||||
fn pairing_read_allow_from(&mut self, channel: String) -> Result<Vec<String>, String> {
|
||||
self.pairing_store
|
||||
.read_allow_from(&channel)
|
||||
.map_err(|e| e.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
/// A WASM-based channel implementing the Channel trait.
|
||||
@@ -424,6 +497,9 @@ pub struct WasmChannel {
|
||||
/// Background task that repeats typing indicators every 4 seconds.
|
||||
/// Telegram's "typing..." indicator expires after ~5s, so we refresh it.
|
||||
typing_task: RwLock<Option<tokio::task::JoinHandle<()>>>,
|
||||
|
||||
/// Pairing store for DM pairing (guest access control).
|
||||
pairing_store: Arc<PairingStore>,
|
||||
}
|
||||
|
||||
impl WasmChannel {
|
||||
@@ -433,6 +509,7 @@ impl WasmChannel {
|
||||
prepared: Arc<PreparedChannelModule>,
|
||||
capabilities: ChannelCapabilities,
|
||||
config_json: String,
|
||||
pairing_store: Arc<PairingStore>,
|
||||
) -> Self {
|
||||
let name = prepared.name.clone();
|
||||
let rate_limiter = ChannelEmitRateLimiter::new(capabilities.emit_rate_limit.clone());
|
||||
@@ -452,6 +529,7 @@ impl WasmChannel {
|
||||
endpoints: RwLock::new(Vec::new()),
|
||||
credentials: Arc::new(RwLock::new(HashMap::new())),
|
||||
typing_task: RwLock::new(None),
|
||||
pairing_store,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -533,6 +611,7 @@ impl WasmChannel {
|
||||
prepared: &PreparedChannelModule,
|
||||
capabilities: &ChannelCapabilities,
|
||||
credentials: HashMap<String, String>,
|
||||
pairing_store: Arc<PairingStore>,
|
||||
) -> Result<Store<ChannelStoreData>, WasmChannelError> {
|
||||
let engine = runtime.engine();
|
||||
let limits = &prepared.limits;
|
||||
@@ -543,6 +622,7 @@ impl WasmChannel {
|
||||
&prepared.name,
|
||||
capabilities.clone(),
|
||||
credentials,
|
||||
pairing_store,
|
||||
);
|
||||
let mut store = Store::new(engine, store_data);
|
||||
|
||||
@@ -643,12 +723,18 @@ impl WasmChannel {
|
||||
let timeout = self.runtime.config().callback_timeout;
|
||||
let channel_name = self.name.clone();
|
||||
let credentials = self.get_credentials().await;
|
||||
let pairing_store = self.pairing_store.clone();
|
||||
|
||||
// Execute in blocking task with timeout
|
||||
let result = tokio::time::timeout(timeout, async move {
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let mut store =
|
||||
Self::create_store(&runtime, &prepared, &capabilities, credentials)?;
|
||||
let mut store = Self::create_store(
|
||||
&runtime,
|
||||
&prepared,
|
||||
&capabilities,
|
||||
credentials,
|
||||
pairing_store,
|
||||
)?;
|
||||
let instance = Self::instantiate_component(&runtime, &prepared, &mut store)?;
|
||||
|
||||
// Call on_start using the generated typed interface
|
||||
@@ -753,6 +839,7 @@ impl WasmChannel {
|
||||
let capabilities = self.capabilities.clone();
|
||||
let timeout = self.runtime.config().callback_timeout;
|
||||
let credentials = self.get_credentials().await;
|
||||
let pairing_store = self.pairing_store.clone();
|
||||
|
||||
// Prepare request data
|
||||
let method = method.to_string();
|
||||
@@ -766,8 +853,13 @@ impl WasmChannel {
|
||||
// Execute in blocking task with timeout
|
||||
let result = tokio::time::timeout(timeout, async move {
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let mut store =
|
||||
Self::create_store(&runtime, &prepared, &capabilities, credentials)?;
|
||||
let mut store = Self::create_store(
|
||||
&runtime,
|
||||
&prepared,
|
||||
&capabilities,
|
||||
credentials,
|
||||
pairing_store,
|
||||
)?;
|
||||
let instance = Self::instantiate_component(&runtime, &prepared, &mut store)?;
|
||||
|
||||
// Build the WIT request type
|
||||
@@ -840,12 +932,18 @@ impl WasmChannel {
|
||||
let timeout = self.runtime.config().callback_timeout;
|
||||
let channel_name = self.name.clone();
|
||||
let credentials = self.get_credentials().await;
|
||||
let pairing_store = self.pairing_store.clone();
|
||||
|
||||
// Execute in blocking task with timeout
|
||||
let result = tokio::time::timeout(timeout, async move {
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let mut store =
|
||||
Self::create_store(&runtime, &prepared, &capabilities, credentials)?;
|
||||
let mut store = Self::create_store(
|
||||
&runtime,
|
||||
&prepared,
|
||||
&capabilities,
|
||||
credentials,
|
||||
pairing_store,
|
||||
)?;
|
||||
let instance = Self::instantiate_component(&runtime, &prepared, &mut store)?;
|
||||
|
||||
// Call on_poll using the generated typed interface
|
||||
@@ -929,6 +1027,7 @@ impl WasmChannel {
|
||||
let timeout = self.runtime.config().callback_timeout;
|
||||
let channel_name = self.name.clone();
|
||||
let credentials = self.get_credentials().await;
|
||||
let pairing_store = self.pairing_store.clone();
|
||||
|
||||
// Prepare response data
|
||||
let message_id_str = message_id.to_string();
|
||||
@@ -942,8 +1041,13 @@ impl WasmChannel {
|
||||
let result = tokio::time::timeout(timeout, async move {
|
||||
tokio::task::spawn_blocking(move || {
|
||||
tracing::info!("Creating WASM store for on_respond");
|
||||
let mut store =
|
||||
Self::create_store(&runtime, &prepared, &capabilities, credentials)?;
|
||||
let mut store = Self::create_store(
|
||||
&runtime,
|
||||
&prepared,
|
||||
&capabilities,
|
||||
credentials,
|
||||
pairing_store,
|
||||
)?;
|
||||
|
||||
tracing::info!("Instantiating WASM component for on_respond");
|
||||
let instance = Self::instantiate_component(&runtime, &prepared, &mut store)?;
|
||||
@@ -1036,13 +1140,19 @@ impl WasmChannel {
|
||||
let timeout = self.runtime.config().callback_timeout;
|
||||
let channel_name = self.name.clone();
|
||||
let credentials = self.get_credentials().await;
|
||||
let pairing_store = self.pairing_store.clone();
|
||||
|
||||
let wit_update = status_to_wit(status, metadata);
|
||||
|
||||
let result = tokio::time::timeout(timeout, async move {
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let mut store =
|
||||
Self::create_store(&runtime, &prepared, &capabilities, credentials)?;
|
||||
let mut store = Self::create_store(
|
||||
&runtime,
|
||||
&prepared,
|
||||
&capabilities,
|
||||
credentials,
|
||||
pairing_store,
|
||||
)?;
|
||||
let instance = Self::instantiate_component(&runtime, &prepared, &mut store)?;
|
||||
|
||||
let channel_iface = instance.near_agent_channel();
|
||||
@@ -1080,12 +1190,14 @@ impl WasmChannel {
|
||||
///
|
||||
/// Static method for use by the background typing repeat task (which
|
||||
/// doesn't have access to `&self`).
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn execute_status(
|
||||
channel_name: &str,
|
||||
runtime: &Arc<WasmChannelRuntime>,
|
||||
prepared: &Arc<PreparedChannelModule>,
|
||||
capabilities: &ChannelCapabilities,
|
||||
credentials: &RwLock<HashMap<String, String>>,
|
||||
pairing_store: Arc<PairingStore>,
|
||||
timeout: Duration,
|
||||
wit_update: wit_channel::StatusUpdate,
|
||||
) -> Result<(), WasmChannelError> {
|
||||
@@ -1101,8 +1213,13 @@ impl WasmChannel {
|
||||
|
||||
let result = tokio::time::timeout(timeout, async move {
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let mut store =
|
||||
Self::create_store(&runtime, &prepared, &capabilities, credentials_snapshot)?;
|
||||
let mut store = Self::create_store(
|
||||
&runtime,
|
||||
&prepared,
|
||||
&capabilities,
|
||||
credentials_snapshot,
|
||||
pairing_store,
|
||||
)?;
|
||||
let instance = Self::instantiate_component(&runtime, &prepared, &mut store)?;
|
||||
|
||||
let channel_iface = instance.near_agent_channel();
|
||||
@@ -1170,6 +1287,7 @@ impl WasmChannel {
|
||||
let prepared = Arc::clone(&self.prepared);
|
||||
let capabilities = self.capabilities.clone();
|
||||
let credentials = self.credentials.clone();
|
||||
let pairing_store = self.pairing_store.clone();
|
||||
let callback_timeout = self.runtime.config().callback_timeout;
|
||||
let wit_update = status_to_wit(&status, metadata);
|
||||
|
||||
@@ -1189,6 +1307,7 @@ impl WasmChannel {
|
||||
&prepared,
|
||||
&capabilities,
|
||||
&credentials,
|
||||
pairing_store.clone(),
|
||||
callback_timeout,
|
||||
wit_update_clone,
|
||||
)
|
||||
@@ -1319,6 +1438,7 @@ impl WasmChannel {
|
||||
let message_tx = self.message_tx.clone();
|
||||
let rate_limiter = self.rate_limiter.clone();
|
||||
let credentials = self.credentials.clone();
|
||||
let pairing_store = self.pairing_store.clone();
|
||||
let callback_timeout = self.runtime.config().callback_timeout;
|
||||
|
||||
tokio::spawn(async move {
|
||||
@@ -1340,6 +1460,7 @@ impl WasmChannel {
|
||||
&prepared,
|
||||
&capabilities,
|
||||
&credentials,
|
||||
pairing_store.clone(),
|
||||
callback_timeout,
|
||||
).await;
|
||||
|
||||
@@ -1391,6 +1512,7 @@ impl WasmChannel {
|
||||
prepared: &Arc<PreparedChannelModule>,
|
||||
capabilities: &ChannelCapabilities,
|
||||
credentials: &RwLock<HashMap<String, String>>,
|
||||
pairing_store: Arc<PairingStore>,
|
||||
timeout: Duration,
|
||||
) -> Result<Vec<EmittedMessage>, WasmChannelError> {
|
||||
// Skip if no WASM bytes (testing mode)
|
||||
@@ -1411,8 +1533,13 @@ impl WasmChannel {
|
||||
// Execute in blocking task with timeout
|
||||
let result = tokio::time::timeout(timeout, async move {
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let mut store =
|
||||
Self::create_store(&runtime, &prepared, &capabilities, credentials_snapshot)?;
|
||||
let mut store = Self::create_store(
|
||||
&runtime,
|
||||
&prepared,
|
||||
&capabilities,
|
||||
credentials_snapshot,
|
||||
pairing_store,
|
||||
)?;
|
||||
let instance = Self::instantiate_component(&runtime, &prepared, &mut store)?;
|
||||
|
||||
// Call on_poll using the generated typed interface
|
||||
@@ -1858,6 +1985,29 @@ fn status_to_wit(status: &StatusUpdate, metadata: &serde_json::Value) -> wit_cha
|
||||
message: format!("Approval needed: {} - {}", tool_name, description),
|
||||
metadata_json,
|
||||
},
|
||||
StatusUpdate::JobStarted { job_id, title, .. } => wit_channel::StatusUpdate {
|
||||
status: wit_channel::StatusType::Thinking,
|
||||
message: format!("Job started: {} ({})", title, job_id),
|
||||
metadata_json,
|
||||
},
|
||||
StatusUpdate::AuthRequired { extension_name, .. } => wit_channel::StatusUpdate {
|
||||
status: wit_channel::StatusType::Thinking,
|
||||
message: format!("Auth required: {}", extension_name),
|
||||
metadata_json,
|
||||
},
|
||||
StatusUpdate::AuthCompleted {
|
||||
extension_name,
|
||||
success,
|
||||
..
|
||||
} => wit_channel::StatusUpdate {
|
||||
status: wit_channel::StatusType::Thinking,
|
||||
message: format!(
|
||||
"Auth {}: {}",
|
||||
if *success { "completed" } else { "failed" },
|
||||
extension_name
|
||||
),
|
||||
metadata_json,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1929,6 +2079,7 @@ mod tests {
|
||||
PreparedChannelModule, WasmChannelRuntime, WasmChannelRuntimeConfig,
|
||||
};
|
||||
use crate::channels::wasm::wrapper::{HttpResponse, WasmChannel};
|
||||
use crate::pairing::PairingStore;
|
||||
use crate::tools::wasm::ResourceLimits;
|
||||
|
||||
fn create_test_channel() -> WasmChannel {
|
||||
@@ -1944,7 +2095,13 @@ mod tests {
|
||||
|
||||
let capabilities = ChannelCapabilities::for_channel("test").with_path("/webhook/test");
|
||||
|
||||
WasmChannel::new(runtime, prepared, capabilities, "{}".to_string())
|
||||
WasmChannel::new(
|
||||
runtime,
|
||||
prepared,
|
||||
capabilities,
|
||||
"{}".to_string(),
|
||||
Arc::new(PairingStore::new()),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -2019,6 +2176,7 @@ mod tests {
|
||||
&prepared,
|
||||
&capabilities,
|
||||
&credentials,
|
||||
Arc::new(PairingStore::new()),
|
||||
timeout,
|
||||
)
|
||||
.await;
|
||||
@@ -2112,7 +2270,13 @@ mod tests {
|
||||
.with_path("/webhook/poll")
|
||||
.with_polling(1000);
|
||||
|
||||
let channel = WasmChannel::new(runtime, prepared, capabilities, "{}".to_string());
|
||||
let channel = WasmChannel::new(
|
||||
runtime,
|
||||
prepared,
|
||||
capabilities,
|
||||
"{}".to_string(),
|
||||
Arc::new(PairingStore::new()),
|
||||
);
|
||||
|
||||
// Start the channel
|
||||
let _stream = channel.start().await.expect("Channel should start");
|
||||
@@ -2350,4 +2514,89 @@ mod tests {
|
||||
assert_eq!(cloned.message, "hello");
|
||||
assert_eq!(cloned.metadata_json, "{\"a\":1}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_redact_credentials_replaces_values() {
|
||||
use super::ChannelStoreData;
|
||||
|
||||
let mut creds = std::collections::HashMap::new();
|
||||
creds.insert(
|
||||
"TELEGRAM_BOT_TOKEN".to_string(),
|
||||
"8218490433:AAEZeUxwqZ5OO3mOCXv7fKvpdhDgsmBBNis".to_string(),
|
||||
);
|
||||
creds.insert("OTHER_SECRET".to_string(), "s3cret".to_string());
|
||||
|
||||
let store = ChannelStoreData::new(
|
||||
1024 * 1024,
|
||||
"test",
|
||||
ChannelCapabilities::default(),
|
||||
creds,
|
||||
Arc::new(PairingStore::new()),
|
||||
);
|
||||
|
||||
let error = "HTTP request failed: error sending request for url \
|
||||
(https://api.telegram.org/bot8218490433:AAEZeUxwqZ5OO3mOCXv7fKvpdhDgsmBBNis/getUpdates)";
|
||||
|
||||
let redacted = store.redact_credentials(error);
|
||||
|
||||
assert!(
|
||||
!redacted.contains("8218490433:AAEZeUxwqZ5OO3mOCXv7fKvpdhDgsmBBNis"),
|
||||
"credential value should be redacted"
|
||||
);
|
||||
assert!(
|
||||
redacted.contains("[REDACTED:TELEGRAM_BOT_TOKEN]"),
|
||||
"redacted text should contain placeholder name"
|
||||
);
|
||||
assert!(
|
||||
!redacted.contains("s3cret"),
|
||||
"other credentials should also be redacted"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_redact_credentials_no_op_without_credentials() {
|
||||
use super::ChannelStoreData;
|
||||
|
||||
let store = ChannelStoreData::new(
|
||||
1024 * 1024,
|
||||
"test",
|
||||
ChannelCapabilities::default(),
|
||||
std::collections::HashMap::new(),
|
||||
Arc::new(PairingStore::new()),
|
||||
);
|
||||
|
||||
let input = "some error message";
|
||||
assert_eq!(store.redact_credentials(input), input);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_redact_credentials_skips_empty_values() {
|
||||
use super::ChannelStoreData;
|
||||
|
||||
let mut creds = std::collections::HashMap::new();
|
||||
creds.insert("EMPTY_TOKEN".to_string(), String::new());
|
||||
|
||||
let store = ChannelStoreData::new(
|
||||
1024 * 1024,
|
||||
"test",
|
||||
ChannelCapabilities::default(),
|
||||
creds,
|
||||
Arc::new(PairingStore::new()),
|
||||
);
|
||||
|
||||
let input = "should not match anything";
|
||||
assert_eq!(store.redact_credentials(input), input);
|
||||
}
|
||||
|
||||
/// Verify that the block_on-inside-spawn_blocking pattern used by the WASM
|
||||
/// channel HTTP host function doesn't deadlock or panic.
|
||||
#[tokio::test]
|
||||
async fn test_block_on_inside_spawn_blocking_does_not_deadlock() {
|
||||
let result = tokio::task::spawn_blocking(|| {
|
||||
tokio::runtime::Handle::current().block_on(async { 42 })
|
||||
})
|
||||
.await
|
||||
.expect("spawn_blocking panicked");
|
||||
assert_eq!(result, 42);
|
||||
}
|
||||
}
|
||||
|
||||
+98
-24
@@ -10,7 +10,7 @@
|
||||
//! ◄── GET /api/chat/events ── SSE stream
|
||||
//! ─── GET /api/chat/ws ─────► WebSocket (bidirectional)
|
||||
//! ─── GET /api/memory/* ────► Workspace
|
||||
//! ─── GET /api/jobs/* ──────► ContextManager
|
||||
//! ─── GET /api/jobs/* ──────► Database
|
||||
//! ◄── GET / ───────────────── Static HTML/CSS/JS
|
||||
//! ```
|
||||
|
||||
@@ -31,9 +31,10 @@ use tokio_stream::wrappers::ReceiverStream;
|
||||
use crate::agent::SessionManager;
|
||||
use crate::channels::{Channel, IncomingMessage, MessageStream, OutgoingResponse, StatusUpdate};
|
||||
use crate::config::GatewayConfig;
|
||||
use crate::context::ContextManager;
|
||||
use crate::error::ChannelError;
|
||||
use crate::extensions::ExtensionManager;
|
||||
use crate::history::Store;
|
||||
use crate::orchestrator::job_manager::ContainerJobManager;
|
||||
use crate::tools::ToolRegistry;
|
||||
use crate::workspace::Workspace;
|
||||
|
||||
@@ -70,11 +71,13 @@ impl GatewayChannel {
|
||||
msg_tx: tokio::sync::RwLock::new(None),
|
||||
sse: SseManager::new(),
|
||||
workspace: None,
|
||||
context_manager: None,
|
||||
session_manager: None,
|
||||
log_broadcaster: None,
|
||||
extension_manager: None,
|
||||
tool_registry: None,
|
||||
store: None,
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
user_id: config.user_id.clone(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(ws::WsConnectionTracker::new())),
|
||||
@@ -93,11 +96,13 @@ impl GatewayChannel {
|
||||
msg_tx: tokio::sync::RwLock::new(None),
|
||||
sse: SseManager::new(),
|
||||
workspace: self.state.workspace.clone(),
|
||||
context_manager: self.state.context_manager.clone(),
|
||||
session_manager: self.state.session_manager.clone(),
|
||||
log_broadcaster: self.state.log_broadcaster.clone(),
|
||||
extension_manager: self.state.extension_manager.clone(),
|
||||
tool_registry: self.state.tool_registry.clone(),
|
||||
store: self.state.store.clone(),
|
||||
job_manager: self.state.job_manager.clone(),
|
||||
prompt_queue: self.state.prompt_queue.clone(),
|
||||
user_id: self.state.user_id.clone(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: self.state.ws_tracker.clone(),
|
||||
@@ -112,12 +117,6 @@ impl GatewayChannel {
|
||||
self
|
||||
}
|
||||
|
||||
/// Inject the context manager for the jobs API.
|
||||
pub fn with_context_manager(mut self, cm: Arc<ContextManager>) -> Self {
|
||||
self.rebuild_state(|s| s.context_manager = Some(cm));
|
||||
self
|
||||
}
|
||||
|
||||
/// Inject the session manager for thread/session info.
|
||||
pub fn with_session_manager(mut self, sm: Arc<SessionManager>) -> Self {
|
||||
self.rebuild_state(|s| s.session_manager = Some(sm));
|
||||
@@ -142,6 +141,34 @@ impl GatewayChannel {
|
||||
self
|
||||
}
|
||||
|
||||
/// Inject the database store for sandbox job persistence.
|
||||
pub fn with_store(mut self, store: Arc<Store>) -> Self {
|
||||
self.rebuild_state(|s| s.store = Some(store));
|
||||
self
|
||||
}
|
||||
|
||||
/// Inject the container job manager for sandbox operations.
|
||||
pub fn with_job_manager(mut self, jm: Arc<ContainerJobManager>) -> Self {
|
||||
self.rebuild_state(|s| s.job_manager = Some(jm));
|
||||
self
|
||||
}
|
||||
|
||||
/// Inject the prompt queue for Claude Code follow-up prompts.
|
||||
pub fn with_prompt_queue(
|
||||
mut self,
|
||||
pq: Arc<
|
||||
tokio::sync::Mutex<
|
||||
std::collections::HashMap<
|
||||
uuid::Uuid,
|
||||
std::collections::VecDeque<crate::orchestrator::api::PendingPrompt>,
|
||||
>,
|
||||
>,
|
||||
>,
|
||||
) -> Self {
|
||||
self.rebuild_state(|s| s.prompt_queue = Some(pq));
|
||||
self
|
||||
}
|
||||
|
||||
/// Get the auth token (for printing to console on startup).
|
||||
pub fn auth_token(&self) -> &str {
|
||||
&self.auth_token
|
||||
@@ -173,11 +200,7 @@ impl Channel for GatewayChannel {
|
||||
),
|
||||
})?;
|
||||
|
||||
let bound_addr =
|
||||
server::start_server(addr, self.state.clone(), self.auth_token.clone()).await?;
|
||||
|
||||
tracing::info!("Web gateway listening on http://{}", bound_addr);
|
||||
tracing::info!("Auth token: {}", self.auth_token);
|
||||
server::start_server(addr, self.state.clone(), self.auth_token.clone()).await?;
|
||||
|
||||
Ok(Box::pin(ReceiverStream::new(rx)))
|
||||
}
|
||||
@@ -200,17 +223,48 @@ impl Channel for GatewayChannel {
|
||||
async fn send_status(
|
||||
&self,
|
||||
status: StatusUpdate,
|
||||
_metadata: &serde_json::Value,
|
||||
metadata: &serde_json::Value,
|
||||
) -> Result<(), ChannelError> {
|
||||
let thread_id = metadata
|
||||
.get("thread_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from);
|
||||
let event = match status {
|
||||
StatusUpdate::Thinking(msg) => SseEvent::Thinking { message: msg },
|
||||
StatusUpdate::ToolStarted { name } => SseEvent::ToolStarted { name },
|
||||
StatusUpdate::ToolCompleted { name, success } => {
|
||||
SseEvent::ToolCompleted { name, success }
|
||||
}
|
||||
StatusUpdate::ToolResult { name, preview } => SseEvent::ToolResult { name, preview },
|
||||
StatusUpdate::StreamChunk(content) => SseEvent::StreamChunk { content },
|
||||
StatusUpdate::Status(msg) => SseEvent::Status { message: msg },
|
||||
StatusUpdate::Thinking(msg) => SseEvent::Thinking {
|
||||
message: msg,
|
||||
thread_id: thread_id.clone(),
|
||||
},
|
||||
StatusUpdate::ToolStarted { name } => SseEvent::ToolStarted {
|
||||
name,
|
||||
thread_id: thread_id.clone(),
|
||||
},
|
||||
StatusUpdate::ToolCompleted { name, success } => SseEvent::ToolCompleted {
|
||||
name,
|
||||
success,
|
||||
thread_id: thread_id.clone(),
|
||||
},
|
||||
StatusUpdate::ToolResult { name, preview } => SseEvent::ToolResult {
|
||||
name,
|
||||
preview,
|
||||
thread_id: thread_id.clone(),
|
||||
},
|
||||
StatusUpdate::StreamChunk(content) => SseEvent::StreamChunk {
|
||||
content,
|
||||
thread_id: thread_id.clone(),
|
||||
},
|
||||
StatusUpdate::Status(msg) => SseEvent::Status {
|
||||
message: msg,
|
||||
thread_id: thread_id.clone(),
|
||||
},
|
||||
StatusUpdate::JobStarted {
|
||||
job_id,
|
||||
title,
|
||||
browse_url,
|
||||
} => SseEvent::JobStarted {
|
||||
job_id,
|
||||
title,
|
||||
browse_url,
|
||||
},
|
||||
StatusUpdate::ApprovalNeeded {
|
||||
request_id,
|
||||
tool_name,
|
||||
@@ -223,6 +277,26 @@ impl Channel for GatewayChannel {
|
||||
parameters: serde_json::to_string_pretty(¶meters)
|
||||
.unwrap_or_else(|_| parameters.to_string()),
|
||||
},
|
||||
StatusUpdate::AuthRequired {
|
||||
extension_name,
|
||||
instructions,
|
||||
auth_url,
|
||||
setup_url,
|
||||
} => SseEvent::AuthRequired {
|
||||
extension_name,
|
||||
instructions,
|
||||
auth_url,
|
||||
setup_url,
|
||||
},
|
||||
StatusUpdate::AuthCompleted {
|
||||
extension_name,
|
||||
success,
|
||||
message,
|
||||
} => SseEvent::AuthCompleted {
|
||||
extension_name,
|
||||
success,
|
||||
message,
|
||||
},
|
||||
};
|
||||
|
||||
self.state.sse.broadcast(event);
|
||||
|
||||
+1333
-96
File diff suppressed because it is too large
Load Diff
+12
-2
@@ -79,7 +79,15 @@ impl SseManager {
|
||||
SseEvent::StreamChunk { .. } => "stream_chunk",
|
||||
SseEvent::Status { .. } => "status",
|
||||
SseEvent::ApprovalNeeded { .. } => "approval_needed",
|
||||
SseEvent::AuthRequired { .. } => "auth_required",
|
||||
SseEvent::AuthCompleted { .. } => "auth_completed",
|
||||
SseEvent::Error { .. } => "error",
|
||||
SseEvent::JobStarted { .. } => "job_started",
|
||||
SseEvent::JobMessage { .. } => "job_message",
|
||||
SseEvent::JobToolUse { .. } => "job_tool_use",
|
||||
SseEvent::JobToolResult { .. } => "job_tool_result",
|
||||
SseEvent::JobStatus { .. } => "job_status",
|
||||
SseEvent::JobResult { .. } => "job_result",
|
||||
SseEvent::Heartbeat => "heartbeat",
|
||||
};
|
||||
Ok(Event::default().event(event_type).data(data))
|
||||
@@ -152,13 +160,14 @@ mod tests {
|
||||
|
||||
manager.broadcast(SseEvent::Status {
|
||||
message: "test".to_string(),
|
||||
thread_id: None,
|
||||
});
|
||||
|
||||
let event = rx.next().await;
|
||||
assert!(event.is_some());
|
||||
let event = event.unwrap().unwrap();
|
||||
match event {
|
||||
SseEvent::Status { message } => assert_eq!(message, "test"),
|
||||
SseEvent::Status { message, .. } => assert_eq!(message, "test"),
|
||||
_ => panic!("unexpected event type"),
|
||||
}
|
||||
}
|
||||
@@ -172,11 +181,12 @@ mod tests {
|
||||
|
||||
manager.broadcast(SseEvent::Thinking {
|
||||
message: "working".to_string(),
|
||||
thread_id: None,
|
||||
});
|
||||
|
||||
let event = stream.next().await.unwrap();
|
||||
match event {
|
||||
SseEvent::Thinking { message } => assert_eq!(message, "working"),
|
||||
SseEvent::Thinking { message, .. } => assert_eq!(message, "working"),
|
||||
_ => panic!("Expected Thinking event"),
|
||||
}
|
||||
}
|
||||
|
||||
+1358
-39
File diff suppressed because it is too large
Load Diff
@@ -10,12 +10,19 @@
|
||||
<body>
|
||||
<!-- Auth Screen -->
|
||||
<div id="auth-screen">
|
||||
<h1>IronClaw</h1>
|
||||
<div class="auth-form">
|
||||
<input type="password" id="token-input" placeholder="Auth token" autofocus>
|
||||
<button onclick="authenticate()">Connect</button>
|
||||
<div class="auth-card-login">
|
||||
<div class="auth-brand">
|
||||
<h1>IronClaw</h1>
|
||||
<p class="auth-tagline">Secure AI Assistant</p>
|
||||
</div>
|
||||
<div class="auth-form">
|
||||
<label for="token-input">Gateway Token</label>
|
||||
<input type="password" id="token-input" placeholder="Paste your auth token" autofocus>
|
||||
<button onclick="authenticate()">Connect</button>
|
||||
</div>
|
||||
<div id="auth-error"></div>
|
||||
<p class="auth-hint">Enter the GATEWAY_AUTH_TOKEN from your .env configuration.</p>
|
||||
</div>
|
||||
<div id="auth-error"></div>
|
||||
</div>
|
||||
|
||||
<!-- Main App (hidden until authenticated) -->
|
||||
@@ -26,16 +33,33 @@
|
||||
<button data-tab="memory">Memory</button>
|
||||
<button data-tab="jobs">Jobs</button>
|
||||
<button data-tab="logs">Logs</button>
|
||||
<button data-tab="routines">Routines</button>
|
||||
<button data-tab="extensions">Extensions</button>
|
||||
<div class="spacer"></div>
|
||||
<div class="status">
|
||||
<div class="status" id="gateway-status-trigger">
|
||||
<div class="dot" id="sse-dot"></div>
|
||||
<span id="sse-status">Connected</span>
|
||||
<div class="gateway-popover" id="gateway-popover"></div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Chat Tab -->
|
||||
<div class="tab-panel active" id="tab-chat">
|
||||
<div class="thread-sidebar" id="thread-sidebar">
|
||||
<div class="thread-sidebar-header">
|
||||
<span>Threads</span>
|
||||
<button class="thread-new-btn" onclick="createNewThread()" title="New thread (Ctrl/Cmd+N)">+</button>
|
||||
<button class="thread-toggle-btn" id="thread-toggle-btn" onclick="toggleThreadSidebar()" title="Toggle sidebar">«</button>
|
||||
</div>
|
||||
<div class="assistant-item" id="assistant-thread" onclick="switchToAssistant()">
|
||||
<span class="assistant-label">Assistant</span>
|
||||
<span class="assistant-meta" id="assistant-meta"></span>
|
||||
</div>
|
||||
<div class="threads-section-header">
|
||||
<span>Conversations</span>
|
||||
</div>
|
||||
<div class="thread-list" id="thread-list"></div>
|
||||
</div>
|
||||
<div class="chat-container">
|
||||
<div class="chat-messages" id="chat-messages"></div>
|
||||
<div class="chat-status" id="chat-status"></div>
|
||||
@@ -56,10 +80,20 @@
|
||||
<div class="memory-tree" id="memory-tree"></div>
|
||||
</div>
|
||||
<div class="memory-content">
|
||||
<div class="memory-breadcrumb" id="memory-breadcrumb">workspace /</div>
|
||||
<div class="memory-breadcrumb" id="memory-breadcrumb">
|
||||
<span id="memory-breadcrumb-path">workspace /</span>
|
||||
<button class="memory-edit-btn" id="memory-edit-btn" style="display:none" onclick="startMemoryEdit()">Edit</button>
|
||||
</div>
|
||||
<div class="memory-viewer" id="memory-viewer">
|
||||
<div class="empty">Select a file to view its contents</div>
|
||||
</div>
|
||||
<div class="memory-editor" id="memory-editor" style="display:none">
|
||||
<textarea id="memory-edit-textarea"></textarea>
|
||||
<div class="memory-editor-actions">
|
||||
<button class="btn-save" onclick="saveMemoryEdit()">Save</button>
|
||||
<button class="btn-cancel-edit" onclick="cancelMemoryEdit()">Cancel</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
@@ -73,6 +107,7 @@
|
||||
<tr>
|
||||
<th>ID</th>
|
||||
<th>Title</th>
|
||||
<th>Source</th>
|
||||
<th>Status</th>
|
||||
<th>Created</th>
|
||||
<th>Actions</th>
|
||||
@@ -104,9 +139,48 @@
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Routines Tab -->
|
||||
<div class="tab-panel" id="tab-routines">
|
||||
<div class="routines-container">
|
||||
<div class="routines-summary" id="routines-summary"></div>
|
||||
<table class="routines-table" id="routines-table">
|
||||
<thead>
|
||||
<tr>
|
||||
<th>Name</th>
|
||||
<th>Trigger</th>
|
||||
<th>Action</th>
|
||||
<th>Last Run</th>
|
||||
<th>Next Run</th>
|
||||
<th>Runs</th>
|
||||
<th>Status</th>
|
||||
<th>Actions</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody id="routines-tbody"></tbody>
|
||||
</table>
|
||||
<div class="empty-state" id="routines-empty" style="display:none">
|
||||
No routines configured. Ask the assistant to create one.
|
||||
</div>
|
||||
<div class="routine-detail" id="routine-detail" style="display:none"></div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Extensions Tab -->
|
||||
<div class="tab-panel" id="tab-extensions">
|
||||
<div class="extensions-container">
|
||||
<div class="extensions-section">
|
||||
<h3>Install Extension</h3>
|
||||
<div class="ext-install-form" id="ext-install-form">
|
||||
<input type="text" id="ext-install-name" placeholder="Extension name (required)">
|
||||
<input type="text" id="ext-install-url" placeholder="URL (optional)">
|
||||
<select id="ext-install-kind">
|
||||
<option value="mcp_server">MCP Server</option>
|
||||
<option value="wasm_tool">WASM Tool</option>
|
||||
<option value="wasm_channel">WASM Channel</option>
|
||||
</select>
|
||||
<button onclick="installExtension()">Install</button>
|
||||
</div>
|
||||
</div>
|
||||
<div class="extensions-section">
|
||||
<h3>Installed Extensions</h3>
|
||||
<div class="extensions-list" id="extensions-list">
|
||||
@@ -130,6 +204,7 @@
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div id="toasts"></div>
|
||||
<script src="/app.js"></script>
|
||||
</body>
|
||||
</html>
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+413
-9
@@ -24,10 +24,17 @@ pub struct ThreadInfo {
|
||||
pub turn_count: usize,
|
||||
pub created_at: String,
|
||||
pub updated_at: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub title: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub thread_type: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct ThreadListResponse {
|
||||
/// The pinned assistant thread (always present after first load).
|
||||
pub assistant_thread: Option<ThreadInfo>,
|
||||
/// Regular conversation threads.
|
||||
pub threads: Vec<ThreadInfo>,
|
||||
pub active_thread: Option<Uuid>,
|
||||
}
|
||||
@@ -54,6 +61,12 @@ pub struct ToolCallInfo {
|
||||
pub struct HistoryResponse {
|
||||
pub thread_id: Uuid,
|
||||
pub turns: Vec<TurnInfo>,
|
||||
/// Whether there are older messages available.
|
||||
#[serde(default)]
|
||||
pub has_more: bool,
|
||||
/// Cursor for the next page (ISO8601 timestamp of the oldest message returned).
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub oldest_timestamp: Option<String>,
|
||||
}
|
||||
|
||||
// --- Approval ---
|
||||
@@ -63,6 +76,8 @@ pub struct ApprovalRequest {
|
||||
pub request_id: String,
|
||||
/// "approve", "always", or "deny"
|
||||
pub action: String,
|
||||
/// Thread that owns the pending approval (so the agent loop finds the right session).
|
||||
pub thread_id: Option<String>,
|
||||
}
|
||||
|
||||
// --- SSE Event Types ---
|
||||
@@ -73,17 +88,49 @@ pub enum SseEvent {
|
||||
#[serde(rename = "response")]
|
||||
Response { content: String, thread_id: String },
|
||||
#[serde(rename = "thinking")]
|
||||
Thinking { message: String },
|
||||
Thinking {
|
||||
message: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
thread_id: Option<String>,
|
||||
},
|
||||
#[serde(rename = "tool_started")]
|
||||
ToolStarted { name: String },
|
||||
ToolStarted {
|
||||
name: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
thread_id: Option<String>,
|
||||
},
|
||||
#[serde(rename = "tool_completed")]
|
||||
ToolCompleted { name: String, success: bool },
|
||||
ToolCompleted {
|
||||
name: String,
|
||||
success: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
thread_id: Option<String>,
|
||||
},
|
||||
#[serde(rename = "tool_result")]
|
||||
ToolResult { name: String, preview: String },
|
||||
ToolResult {
|
||||
name: String,
|
||||
preview: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
thread_id: Option<String>,
|
||||
},
|
||||
#[serde(rename = "stream_chunk")]
|
||||
StreamChunk { content: String },
|
||||
StreamChunk {
|
||||
content: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
thread_id: Option<String>,
|
||||
},
|
||||
#[serde(rename = "status")]
|
||||
Status { message: String },
|
||||
Status {
|
||||
message: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
thread_id: Option<String>,
|
||||
},
|
||||
#[serde(rename = "job_started")]
|
||||
JobStarted {
|
||||
job_id: String,
|
||||
title: String,
|
||||
browse_url: String,
|
||||
},
|
||||
#[serde(rename = "approval_needed")]
|
||||
ApprovalNeeded {
|
||||
request_id: String,
|
||||
@@ -91,10 +138,59 @@ pub enum SseEvent {
|
||||
description: String,
|
||||
parameters: String,
|
||||
},
|
||||
#[serde(rename = "auth_required")]
|
||||
AuthRequired {
|
||||
extension_name: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
instructions: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
auth_url: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
setup_url: Option<String>,
|
||||
},
|
||||
#[serde(rename = "auth_completed")]
|
||||
AuthCompleted {
|
||||
extension_name: String,
|
||||
success: bool,
|
||||
message: String,
|
||||
},
|
||||
#[serde(rename = "error")]
|
||||
Error { message: String },
|
||||
Error {
|
||||
message: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
thread_id: Option<String>,
|
||||
},
|
||||
#[serde(rename = "heartbeat")]
|
||||
Heartbeat,
|
||||
|
||||
// Sandbox job streaming events (worker + Claude Code bridge)
|
||||
#[serde(rename = "job_message")]
|
||||
JobMessage {
|
||||
job_id: String,
|
||||
role: String,
|
||||
content: String,
|
||||
},
|
||||
#[serde(rename = "job_tool_use")]
|
||||
JobToolUse {
|
||||
job_id: String,
|
||||
tool_name: String,
|
||||
input: serde_json::Value,
|
||||
},
|
||||
#[serde(rename = "job_tool_result")]
|
||||
JobToolResult {
|
||||
job_id: String,
|
||||
tool_name: String,
|
||||
output: String,
|
||||
},
|
||||
#[serde(rename = "job_status")]
|
||||
JobStatus { job_id: String, message: String },
|
||||
#[serde(rename = "job_result")]
|
||||
JobResult {
|
||||
job_id: String,
|
||||
status: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
session_id: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
// --- Memory ---
|
||||
@@ -188,6 +284,54 @@ pub struct JobSummaryResponse {
|
||||
pub stuck: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct JobDetailResponse {
|
||||
pub id: Uuid,
|
||||
pub title: String,
|
||||
pub description: String,
|
||||
pub state: String,
|
||||
pub user_id: String,
|
||||
pub created_at: String,
|
||||
pub started_at: Option<String>,
|
||||
pub completed_at: Option<String>,
|
||||
pub elapsed_secs: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub project_dir: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub browse_url: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub job_mode: Option<String>,
|
||||
pub transitions: Vec<TransitionInfo>,
|
||||
}
|
||||
|
||||
// --- Project Files ---
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct ProjectFileEntry {
|
||||
pub name: String,
|
||||
pub path: String,
|
||||
pub is_dir: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct ProjectFilesResponse {
|
||||
pub entries: Vec<ProjectFileEntry>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct ProjectFileReadResponse {
|
||||
pub path: String,
|
||||
pub content: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct TransitionInfo {
|
||||
pub from: String,
|
||||
pub to: String,
|
||||
pub timestamp: String,
|
||||
pub reason: Option<String>,
|
||||
}
|
||||
|
||||
// --- Extensions ---
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
@@ -262,6 +406,21 @@ impl ActionResponse {
|
||||
}
|
||||
}
|
||||
|
||||
// --- Auth Token ---
|
||||
|
||||
/// Request to submit an auth token for an extension (dedicated endpoint).
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct AuthTokenRequest {
|
||||
pub extension_name: String,
|
||||
pub token: String,
|
||||
}
|
||||
|
||||
/// Request to cancel an in-progress auth flow.
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct AuthCancelRequest {
|
||||
pub extension_name: String,
|
||||
}
|
||||
|
||||
// --- WebSocket ---
|
||||
|
||||
/// Message sent by a WebSocket client to the server.
|
||||
@@ -280,7 +439,18 @@ pub enum WsClientMessage {
|
||||
request_id: String,
|
||||
/// "approve", "always", or "deny"
|
||||
action: String,
|
||||
/// Thread that owns the pending approval.
|
||||
thread_id: Option<String>,
|
||||
},
|
||||
/// Submit an auth token for an extension (bypasses message pipeline).
|
||||
#[serde(rename = "auth_token")]
|
||||
AuthToken {
|
||||
extension_name: String,
|
||||
token: String,
|
||||
},
|
||||
/// Cancel an in-progress auth flow.
|
||||
#[serde(rename = "auth_cancel")]
|
||||
AuthCancel { extension_name: String },
|
||||
/// Client heartbeat ping.
|
||||
#[serde(rename = "ping")]
|
||||
Ping,
|
||||
@@ -317,9 +487,17 @@ impl WsServerMessage {
|
||||
SseEvent::ToolResult { .. } => "tool_result",
|
||||
SseEvent::StreamChunk { .. } => "stream_chunk",
|
||||
SseEvent::Status { .. } => "status",
|
||||
SseEvent::JobStarted { .. } => "job_started",
|
||||
SseEvent::ApprovalNeeded { .. } => "approval_needed",
|
||||
SseEvent::AuthRequired { .. } => "auth_required",
|
||||
SseEvent::AuthCompleted { .. } => "auth_completed",
|
||||
SseEvent::Error { .. } => "error",
|
||||
SseEvent::Heartbeat => "heartbeat",
|
||||
SseEvent::JobMessage { .. } => "job_message",
|
||||
SseEvent::JobToolUse { .. } => "job_tool_use",
|
||||
SseEvent::JobToolResult { .. } => "job_tool_result",
|
||||
SseEvent::JobStatus { .. } => "job_status",
|
||||
SseEvent::JobResult { .. } => "job_result",
|
||||
};
|
||||
let data = serde_json::to_value(event).unwrap_or(serde_json::Value::Null);
|
||||
WsServerMessage::Event {
|
||||
@@ -329,6 +507,96 @@ impl WsServerMessage {
|
||||
}
|
||||
}
|
||||
|
||||
// --- Routines ---
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct RoutineInfo {
|
||||
pub id: Uuid,
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
pub enabled: bool,
|
||||
pub trigger_type: String,
|
||||
pub trigger_summary: String,
|
||||
pub action_type: String,
|
||||
pub last_run_at: Option<String>,
|
||||
pub next_fire_at: Option<String>,
|
||||
pub run_count: u64,
|
||||
pub consecutive_failures: u32,
|
||||
pub status: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct RoutineListResponse {
|
||||
pub routines: Vec<RoutineInfo>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct RoutineSummaryResponse {
|
||||
pub total: u64,
|
||||
pub enabled: u64,
|
||||
pub disabled: u64,
|
||||
pub failing: u64,
|
||||
pub runs_today: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct RoutineDetailResponse {
|
||||
pub id: Uuid,
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
pub enabled: bool,
|
||||
pub trigger: serde_json::Value,
|
||||
pub action: serde_json::Value,
|
||||
pub guardrails: serde_json::Value,
|
||||
pub notify: serde_json::Value,
|
||||
pub last_run_at: Option<String>,
|
||||
pub next_fire_at: Option<String>,
|
||||
pub run_count: u64,
|
||||
pub consecutive_failures: u32,
|
||||
pub created_at: String,
|
||||
pub recent_runs: Vec<RoutineRunInfo>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct RoutineRunInfo {
|
||||
pub id: Uuid,
|
||||
pub trigger_type: String,
|
||||
pub started_at: String,
|
||||
pub completed_at: Option<String>,
|
||||
pub status: String,
|
||||
pub result_summary: Option<String>,
|
||||
pub tokens_used: Option<i32>,
|
||||
}
|
||||
|
||||
// --- Settings ---
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct SettingResponse {
|
||||
pub key: String,
|
||||
pub value: serde_json::Value,
|
||||
pub updated_at: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct SettingsListResponse {
|
||||
pub settings: Vec<SettingResponse>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct SettingWriteRequest {
|
||||
pub value: serde_json::Value,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct SettingsImportRequest {
|
||||
pub settings: std::collections::HashMap<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct SettingsExportResponse {
|
||||
pub settings: std::collections::HashMap<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
// --- Health ---
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
@@ -371,12 +639,36 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_ws_client_approval_parse() {
|
||||
let json = r#"{"type":"approval","request_id":"abc-123","action":"approve"}"#;
|
||||
let json =
|
||||
r#"{"type":"approval","request_id":"abc-123","action":"approve","thread_id":"t1"}"#;
|
||||
let msg: WsClientMessage = serde_json::from_str(json).unwrap();
|
||||
match msg {
|
||||
WsClientMessage::Approval { request_id, action } => {
|
||||
WsClientMessage::Approval {
|
||||
request_id,
|
||||
action,
|
||||
thread_id,
|
||||
} => {
|
||||
assert_eq!(request_id, "abc-123");
|
||||
assert_eq!(action, "approve");
|
||||
assert_eq!(thread_id.as_deref(), Some("t1"));
|
||||
}
|
||||
_ => panic!("Expected Approval variant"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ws_client_approval_parse_no_thread() {
|
||||
let json = r#"{"type":"approval","request_id":"abc-123","action":"deny"}"#;
|
||||
let msg: WsClientMessage = serde_json::from_str(json).unwrap();
|
||||
match msg {
|
||||
WsClientMessage::Approval {
|
||||
request_id,
|
||||
action,
|
||||
thread_id,
|
||||
} => {
|
||||
assert_eq!(request_id, "abc-123");
|
||||
assert_eq!(action, "deny");
|
||||
assert!(thread_id.is_none());
|
||||
}
|
||||
_ => panic!("Expected Approval variant"),
|
||||
}
|
||||
@@ -437,6 +729,7 @@ mod tests {
|
||||
fn test_ws_server_from_sse_thinking() {
|
||||
let sse = SseEvent::Thinking {
|
||||
message: "reasoning...".to_string(),
|
||||
thread_id: None,
|
||||
};
|
||||
let ws = WsServerMessage::from_sse_event(&sse);
|
||||
match ws {
|
||||
@@ -477,4 +770,115 @@ mod tests {
|
||||
_ => panic!("Expected Event variant"),
|
||||
}
|
||||
}
|
||||
|
||||
// ---- Auth type tests ----
|
||||
|
||||
#[test]
|
||||
fn test_ws_client_auth_token_parse() {
|
||||
let json = r#"{"type":"auth_token","extension_name":"notion","token":"sk-123"}"#;
|
||||
let msg: WsClientMessage = serde_json::from_str(json).unwrap();
|
||||
match msg {
|
||||
WsClientMessage::AuthToken {
|
||||
extension_name,
|
||||
token,
|
||||
} => {
|
||||
assert_eq!(extension_name, "notion");
|
||||
assert_eq!(token, "sk-123");
|
||||
}
|
||||
_ => panic!("Expected AuthToken variant"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ws_client_auth_cancel_parse() {
|
||||
let json = r#"{"type":"auth_cancel","extension_name":"notion"}"#;
|
||||
let msg: WsClientMessage = serde_json::from_str(json).unwrap();
|
||||
match msg {
|
||||
WsClientMessage::AuthCancel { extension_name } => {
|
||||
assert_eq!(extension_name, "notion");
|
||||
}
|
||||
_ => panic!("Expected AuthCancel variant"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sse_auth_required_serialize() {
|
||||
let event = SseEvent::AuthRequired {
|
||||
extension_name: "notion".to_string(),
|
||||
instructions: Some("Get your token from...".to_string()),
|
||||
auth_url: None,
|
||||
setup_url: Some("https://notion.so/integrations".to_string()),
|
||||
};
|
||||
let json = serde_json::to_string(&event).unwrap();
|
||||
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(parsed["type"], "auth_required");
|
||||
assert_eq!(parsed["extension_name"], "notion");
|
||||
assert_eq!(parsed["instructions"], "Get your token from...");
|
||||
assert!(parsed.get("auth_url").is_none());
|
||||
assert_eq!(parsed["setup_url"], "https://notion.so/integrations");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sse_auth_completed_serialize() {
|
||||
let event = SseEvent::AuthCompleted {
|
||||
extension_name: "notion".to_string(),
|
||||
success: true,
|
||||
message: "notion authenticated (3 tools loaded)".to_string(),
|
||||
};
|
||||
let json = serde_json::to_string(&event).unwrap();
|
||||
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(parsed["type"], "auth_completed");
|
||||
assert_eq!(parsed["extension_name"], "notion");
|
||||
assert_eq!(parsed["success"], true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ws_server_from_sse_auth_required() {
|
||||
let sse = SseEvent::AuthRequired {
|
||||
extension_name: "openai".to_string(),
|
||||
instructions: Some("Enter API key".to_string()),
|
||||
auth_url: None,
|
||||
setup_url: None,
|
||||
};
|
||||
let ws = WsServerMessage::from_sse_event(&sse);
|
||||
match ws {
|
||||
WsServerMessage::Event { event_type, data } => {
|
||||
assert_eq!(event_type, "auth_required");
|
||||
assert_eq!(data["extension_name"], "openai");
|
||||
}
|
||||
_ => panic!("Expected Event variant"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ws_server_from_sse_auth_completed() {
|
||||
let sse = SseEvent::AuthCompleted {
|
||||
extension_name: "slack".to_string(),
|
||||
success: false,
|
||||
message: "Invalid token".to_string(),
|
||||
};
|
||||
let ws = WsServerMessage::from_sse_event(&sse);
|
||||
match ws {
|
||||
WsServerMessage::Event { event_type, data } => {
|
||||
assert_eq!(event_type, "auth_completed");
|
||||
assert_eq!(data["success"], false);
|
||||
}
|
||||
_ => panic!("Expected Event variant"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_auth_token_request_deserialize() {
|
||||
let json = r#"{"extension_name":"telegram","token":"bot12345"}"#;
|
||||
let req: AuthTokenRequest = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(req.extension_name, "telegram");
|
||||
assert_eq!(req.token, "bot12345");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_auth_cancel_request_deserialize() {
|
||||
let json = r#"{"extension_name":"telegram"}"#;
|
||||
let req: AuthCancelRequest = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(req.extension_name, "telegram");
|
||||
}
|
||||
}
|
||||
|
||||
+73
-3
@@ -170,7 +170,11 @@ async fn handle_client_message(
|
||||
.await;
|
||||
}
|
||||
}
|
||||
WsClientMessage::Approval { request_id, action } => {
|
||||
WsClientMessage::Approval {
|
||||
request_id,
|
||||
action,
|
||||
thread_id,
|
||||
} => {
|
||||
let (approved, always) = match action.as_str() {
|
||||
"approve" => (true, false),
|
||||
"always" => (true, true),
|
||||
@@ -214,12 +218,71 @@ async fn handle_client_message(
|
||||
}
|
||||
};
|
||||
|
||||
let msg = IncomingMessage::new("gateway", user_id, content);
|
||||
let mut msg = IncomingMessage::new("gateway", user_id, content);
|
||||
if let Some(ref tid) = thread_id {
|
||||
msg = msg.with_thread(tid);
|
||||
}
|
||||
let tx_guard = state.msg_tx.read().await;
|
||||
if let Some(ref tx) = *tx_guard {
|
||||
let _ = tx.send(msg).await;
|
||||
}
|
||||
}
|
||||
WsClientMessage::AuthToken {
|
||||
extension_name,
|
||||
token,
|
||||
} => {
|
||||
if let Some(ref ext_mgr) = state.extension_manager {
|
||||
match ext_mgr.auth(&extension_name, Some(&token)).await {
|
||||
Ok(result) if result.status == "authenticated" => {
|
||||
let msg = match ext_mgr.activate(&extension_name).await {
|
||||
Ok(r) => format!(
|
||||
"{} authenticated ({} tools loaded)",
|
||||
extension_name,
|
||||
r.tools_loaded.len()
|
||||
),
|
||||
Err(e) => format!(
|
||||
"{} authenticated but activation failed: {}",
|
||||
extension_name, e
|
||||
),
|
||||
};
|
||||
crate::channels::web::server::clear_auth_mode(state).await;
|
||||
state
|
||||
.sse
|
||||
.broadcast(crate::channels::web::types::SseEvent::AuthCompleted {
|
||||
extension_name,
|
||||
success: true,
|
||||
message: msg,
|
||||
});
|
||||
}
|
||||
Ok(result) => {
|
||||
state
|
||||
.sse
|
||||
.broadcast(crate::channels::web::types::SseEvent::AuthRequired {
|
||||
extension_name,
|
||||
instructions: result.instructions,
|
||||
auth_url: result.auth_url,
|
||||
setup_url: result.setup_url,
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = direct_tx
|
||||
.send(WsServerMessage::Error {
|
||||
message: format!("Auth failed: {}", e),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
let _ = direct_tx
|
||||
.send(WsServerMessage::Error {
|
||||
message: "Extension manager not available".to_string(),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
WsClientMessage::AuthCancel { .. } => {
|
||||
crate::channels::web::server::clear_auth_mode(state).await;
|
||||
}
|
||||
WsClientMessage::Ping => {
|
||||
let _ = direct_tx.send(WsServerMessage::Pong).await;
|
||||
}
|
||||
@@ -328,6 +391,7 @@ mod tests {
|
||||
WsClientMessage::Approval {
|
||||
request_id: request_id.to_string(),
|
||||
action: "approve".to_string(),
|
||||
thread_id: Some("thread-42".to_string()),
|
||||
},
|
||||
&state,
|
||||
"user1",
|
||||
@@ -338,6 +402,8 @@ mod tests {
|
||||
let incoming = agent_rx.recv().await.unwrap();
|
||||
// The content should be a serialized ExecApproval
|
||||
assert!(incoming.content.contains("ExecApproval"));
|
||||
// Thread should be forwarded onto the IncomingMessage.
|
||||
assert_eq!(incoming.thread_id.as_deref(), Some("thread-42"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -349,6 +415,7 @@ mod tests {
|
||||
WsClientMessage::Approval {
|
||||
request_id: Uuid::new_v4().to_string(),
|
||||
action: "maybe".to_string(),
|
||||
thread_id: None,
|
||||
},
|
||||
&state,
|
||||
"user1",
|
||||
@@ -374,6 +441,7 @@ mod tests {
|
||||
WsClientMessage::Approval {
|
||||
request_id: "not-a-uuid".to_string(),
|
||||
action: "approve".to_string(),
|
||||
thread_id: None,
|
||||
},
|
||||
&state,
|
||||
"user1",
|
||||
@@ -398,11 +466,13 @@ mod tests {
|
||||
msg_tx: tokio::sync::RwLock::new(msg_tx),
|
||||
sse: SseManager::new(),
|
||||
workspace: None,
|
||||
context_manager: None,
|
||||
session_manager: None,
|
||||
log_broadcaster: None,
|
||||
extension_manager: None,
|
||||
tool_registry: None,
|
||||
store: None,
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
user_id: "test".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
|
||||
+121
-54
@@ -1,6 +1,7 @@
|
||||
//! Configuration management CLI commands.
|
||||
//!
|
||||
//! Commands for viewing and modifying settings.
|
||||
//! Settings are stored in PostgreSQL (env > DB > default).
|
||||
|
||||
use clap::Subcommand;
|
||||
|
||||
@@ -36,41 +37,82 @@ pub enum ConfigCommand {
|
||||
path: String,
|
||||
},
|
||||
|
||||
/// Show the settings file path
|
||||
/// Show the settings storage info
|
||||
Path,
|
||||
}
|
||||
|
||||
/// Run a config command.
|
||||
pub fn run_config_command(cmd: ConfigCommand) -> anyhow::Result<()> {
|
||||
///
|
||||
/// Connects to the database to read/write settings. Falls back to disk
|
||||
/// if the database is not available.
|
||||
pub async fn run_config_command(cmd: ConfigCommand) -> anyhow::Result<()> {
|
||||
let _ = dotenvy::dotenv();
|
||||
|
||||
// Try to connect to the DB for settings access
|
||||
let store = match connect_store().await {
|
||||
Ok(s) => Some(s),
|
||||
Err(e) => {
|
||||
eprintln!(
|
||||
"Warning: Could not connect to database ({}), using disk fallback",
|
||||
e
|
||||
);
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
match cmd {
|
||||
ConfigCommand::List { filter } => list_settings(filter),
|
||||
ConfigCommand::Get { path } => get_setting(&path),
|
||||
ConfigCommand::Set { path, value } => set_setting(&path, &value),
|
||||
ConfigCommand::Reset { path } => reset_setting(&path),
|
||||
ConfigCommand::Path => show_path(),
|
||||
ConfigCommand::List { filter } => list_settings(store.as_ref(), filter).await,
|
||||
ConfigCommand::Get { path } => get_setting(store.as_ref(), &path).await,
|
||||
ConfigCommand::Set { path, value } => set_setting(store.as_ref(), &path, &value).await,
|
||||
ConfigCommand::Reset { path } => reset_setting(store.as_ref(), &path).await,
|
||||
ConfigCommand::Path => show_path(store.is_some()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Bootstrap a DB connection for config commands.
|
||||
async fn connect_store() -> anyhow::Result<crate::history::Store> {
|
||||
let config = crate::config::Config::from_env()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
let store = crate::history::Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
Ok(store)
|
||||
}
|
||||
|
||||
const DEFAULT_USER_ID: &str = "default";
|
||||
|
||||
/// Load settings: DB if available, else disk.
|
||||
async fn load_settings(store: Option<&crate::history::Store>) -> Settings {
|
||||
if let Some(store) = store {
|
||||
match store.get_all_settings(DEFAULT_USER_ID).await {
|
||||
Ok(map) if !map.is_empty() => return Settings::from_db_map(&map),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Settings::load()
|
||||
}
|
||||
|
||||
/// List all settings.
|
||||
fn list_settings(filter: Option<String>) -> anyhow::Result<()> {
|
||||
let settings = Settings::load();
|
||||
async fn list_settings(
|
||||
store: Option<&crate::history::Store>,
|
||||
filter: Option<String>,
|
||||
) -> anyhow::Result<()> {
|
||||
let settings = load_settings(store).await;
|
||||
let all = settings.list();
|
||||
|
||||
// Find the longest key for alignment
|
||||
let max_key_len = all.iter().map(|(k, _)| k.len()).max().unwrap_or(0);
|
||||
|
||||
println!("Settings:");
|
||||
let source = if store.is_some() { "database" } else { "disk" };
|
||||
println!("Settings (source: {}):", source);
|
||||
println!();
|
||||
|
||||
for (key, value) in all {
|
||||
// Skip if filter is set and doesn't match
|
||||
if let Some(ref f) = filter {
|
||||
if !key.starts_with(f) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
// Truncate long values for display
|
||||
let display_value = if value.len() > 60 {
|
||||
format!("{}...", &value[..57])
|
||||
} else {
|
||||
@@ -84,8 +126,8 @@ fn list_settings(filter: Option<String>) -> anyhow::Result<()> {
|
||||
}
|
||||
|
||||
/// Get a specific setting.
|
||||
fn get_setting(path: &str) -> anyhow::Result<()> {
|
||||
let settings = Settings::load();
|
||||
async fn get_setting(store: Option<&crate::history::Store>, path: &str) -> anyhow::Result<()> {
|
||||
let settings = load_settings(store).await;
|
||||
|
||||
match settings.get(path) {
|
||||
Some(value) => {
|
||||
@@ -99,67 +141,92 @@ fn get_setting(path: &str) -> anyhow::Result<()> {
|
||||
}
|
||||
|
||||
/// Set a setting value.
|
||||
fn set_setting(path: &str, value: &str) -> anyhow::Result<()> {
|
||||
let mut settings = Settings::load();
|
||||
async fn set_setting(
|
||||
store: Option<&crate::history::Store>,
|
||||
path: &str,
|
||||
value: &str,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut settings = load_settings(store).await;
|
||||
|
||||
// Try to set the value
|
||||
settings
|
||||
.set(path, value)
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
|
||||
// Save to disk
|
||||
settings.save()?;
|
||||
// Save to DB if available, otherwise disk
|
||||
if let Some(store) = store {
|
||||
let json_value = match serde_json::from_str::<serde_json::Value>(value) {
|
||||
Ok(v) => v,
|
||||
Err(_) => serde_json::Value::String(value.to_string()),
|
||||
};
|
||||
store
|
||||
.set_setting(DEFAULT_USER_ID, path, &json_value)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to save to database: {}", e))?;
|
||||
} else {
|
||||
settings.save()?;
|
||||
}
|
||||
|
||||
println!("Set {} = {}", path, value);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Reset a setting to default.
|
||||
fn reset_setting(path: &str) -> anyhow::Result<()> {
|
||||
let mut settings = Settings::load();
|
||||
|
||||
// Get the default value for display
|
||||
async fn reset_setting(store: Option<&crate::history::Store>, path: &str) -> anyhow::Result<()> {
|
||||
let default = Settings::default();
|
||||
let default_value = default
|
||||
.get(path)
|
||||
.ok_or_else(|| anyhow::anyhow!("Unknown setting: {}", path))?;
|
||||
|
||||
// Reset it
|
||||
settings.reset(path).map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
|
||||
// Save to disk
|
||||
settings.save()?;
|
||||
// Delete from DB (falling back to default) or reset on disk
|
||||
if let Some(store) = store {
|
||||
store
|
||||
.delete_setting(DEFAULT_USER_ID, path)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to delete setting from database: {}", e))?;
|
||||
} else {
|
||||
let mut settings = Settings::load();
|
||||
settings.reset(path).map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
settings.save()?;
|
||||
}
|
||||
|
||||
println!("Reset {} to default: {}", path, default_value);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Show the settings file path.
|
||||
fn show_path() -> anyhow::Result<()> {
|
||||
let path = Settings::default_path();
|
||||
println!("{}", path.display());
|
||||
|
||||
if path.exists() {
|
||||
let metadata = std::fs::metadata(&path)?;
|
||||
println!(" Size: {} bytes", metadata.len());
|
||||
if let Ok(modified) = metadata.modified() {
|
||||
use std::time::SystemTime;
|
||||
let duration = SystemTime::now()
|
||||
.duration_since(modified)
|
||||
.unwrap_or_default();
|
||||
let secs = duration.as_secs();
|
||||
if secs < 60 {
|
||||
println!(" Modified: {} seconds ago", secs);
|
||||
} else if secs < 3600 {
|
||||
println!(" Modified: {} minutes ago", secs / 60);
|
||||
} else if secs < 86400 {
|
||||
println!(" Modified: {} hours ago", secs / 3600);
|
||||
} else {
|
||||
println!(" Modified: {} days ago", secs / 86400);
|
||||
}
|
||||
}
|
||||
/// Show the settings storage info.
|
||||
fn show_path(has_db: bool) -> anyhow::Result<()> {
|
||||
if has_db {
|
||||
println!("Settings stored in: PostgreSQL (settings table)");
|
||||
println!(
|
||||
"Bootstrap config: {}",
|
||||
crate::bootstrap::BootstrapConfig::default_path().display()
|
||||
);
|
||||
} else {
|
||||
println!(" (does not exist, using defaults)");
|
||||
let path = Settings::default_path();
|
||||
println!("Settings stored in: {} (disk fallback)", path.display());
|
||||
|
||||
if path.exists() {
|
||||
let metadata = std::fs::metadata(&path)?;
|
||||
println!(" Size: {} bytes", metadata.len());
|
||||
if let Ok(modified) = metadata.modified() {
|
||||
use std::time::SystemTime;
|
||||
let duration = SystemTime::now()
|
||||
.duration_since(modified)
|
||||
.unwrap_or_default();
|
||||
let secs = duration.as_secs();
|
||||
if secs < 60 {
|
||||
println!(" Modified: {} seconds ago", secs);
|
||||
} else if secs < 3600 {
|
||||
println!(" Modified: {} minutes ago", secs / 60);
|
||||
} else if secs < 86400 {
|
||||
println!(" Modified: {} hours ago", secs / 3600);
|
||||
} else {
|
||||
println!(" Modified: {} days ago", secs / 86400);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
println!(" (does not exist, using defaults)");
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
|
||||
+61
-12
@@ -13,9 +13,7 @@ use crate::secrets::{PostgresSecretsStore, SecretsCrypto, SecretsStore};
|
||||
use crate::tools::mcp::{
|
||||
McpClient, McpServerConfig, McpSessionManager, OAuthConfig,
|
||||
auth::{authorize_mcp_server, is_authenticated},
|
||||
config::{
|
||||
add_mcp_server, get_mcp_server, load_mcp_servers, remove_mcp_server, save_mcp_servers,
|
||||
},
|
||||
config::{self, McpServersFile},
|
||||
};
|
||||
|
||||
#[derive(Subcommand, Debug, Clone)]
|
||||
@@ -173,8 +171,11 @@ async fn add_server(
|
||||
// Validate
|
||||
config.validate()?;
|
||||
|
||||
// Save
|
||||
add_mcp_server(config).await?;
|
||||
// Save (DB if available, else disk)
|
||||
let store = connect_store().await;
|
||||
let mut servers = load_servers(store.as_ref()).await?;
|
||||
servers.upsert(config);
|
||||
save_servers(store.as_ref(), &servers).await?;
|
||||
|
||||
println!();
|
||||
println!(" ✓ Added MCP server '{}'", name);
|
||||
@@ -192,7 +193,12 @@ async fn add_server(
|
||||
|
||||
/// Remove an MCP server.
|
||||
async fn remove_server(name: String) -> anyhow::Result<()> {
|
||||
remove_mcp_server(&name).await?;
|
||||
let store = connect_store().await;
|
||||
let mut servers = load_servers(store.as_ref()).await?;
|
||||
if !servers.remove(&name) {
|
||||
anyhow::bail!("Server '{}' not found", name);
|
||||
}
|
||||
save_servers(store.as_ref(), &servers).await?;
|
||||
|
||||
println!();
|
||||
println!(" ✓ Removed MCP server '{}'", name);
|
||||
@@ -203,7 +209,8 @@ async fn remove_server(name: String) -> anyhow::Result<()> {
|
||||
|
||||
/// List configured MCP servers.
|
||||
async fn list_servers(verbose: bool) -> anyhow::Result<()> {
|
||||
let servers = load_mcp_servers().await?;
|
||||
let store = connect_store().await;
|
||||
let servers = load_servers(store.as_ref()).await?;
|
||||
|
||||
if servers.servers.is_empty() {
|
||||
println!();
|
||||
@@ -261,7 +268,12 @@ async fn list_servers(verbose: bool) -> anyhow::Result<()> {
|
||||
/// Authenticate with an MCP server.
|
||||
async fn auth_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
// Get server config
|
||||
let server = get_mcp_server(&name).await?;
|
||||
let store = connect_store().await;
|
||||
let servers = load_servers(store.as_ref()).await?;
|
||||
let server = servers
|
||||
.get(&name)
|
||||
.cloned()
|
||||
.ok_or_else(|| anyhow::anyhow!("Server '{}' not found", name))?;
|
||||
|
||||
// Initialize secrets store
|
||||
let secrets = get_secrets_store().await?;
|
||||
@@ -329,7 +341,12 @@ async fn auth_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
/// Test connection to an MCP server.
|
||||
async fn test_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
// Get server config
|
||||
let server = get_mcp_server(&name).await?;
|
||||
let store = connect_store().await;
|
||||
let servers = load_servers(store.as_ref()).await?;
|
||||
let server = servers
|
||||
.get(&name)
|
||||
.cloned()
|
||||
.ok_or_else(|| anyhow::anyhow!("Server '{}' not found", name))?;
|
||||
|
||||
println!();
|
||||
println!(" Testing connection to '{}'...", name);
|
||||
@@ -420,7 +437,8 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
|
||||
/// Toggle server enabled/disabled state.
|
||||
async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Result<()> {
|
||||
let mut servers = load_mcp_servers().await?;
|
||||
let store = connect_store().await;
|
||||
let mut servers = load_servers(store.as_ref()).await?;
|
||||
|
||||
let server = servers
|
||||
.get_mut(&name)
|
||||
@@ -435,7 +453,7 @@ async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Res
|
||||
};
|
||||
|
||||
server.enabled = new_state;
|
||||
save_mcp_servers(&servers).await?;
|
||||
save_servers(store.as_ref(), &servers).await?;
|
||||
|
||||
let status = if new_state { "enabled" } else { "disabled" };
|
||||
println!();
|
||||
@@ -445,9 +463,40 @@ async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Res
|
||||
Ok(())
|
||||
}
|
||||
|
||||
const DEFAULT_USER_ID: &str = "default";
|
||||
|
||||
/// Try to connect to the database store for DB-backed config.
|
||||
async fn connect_store() -> Option<Store> {
|
||||
let config = Config::from_env().await.ok()?;
|
||||
let store = Store::new(&config.database).await.ok()?;
|
||||
store.run_migrations().await.ok()?;
|
||||
Some(store)
|
||||
}
|
||||
|
||||
/// Load MCP servers (DB if available, else disk).
|
||||
async fn load_servers(store: Option<&Store>) -> Result<McpServersFile, config::ConfigError> {
|
||||
if let Some(store) = store {
|
||||
config::load_mcp_servers_from_db(store, DEFAULT_USER_ID).await
|
||||
} else {
|
||||
config::load_mcp_servers().await
|
||||
}
|
||||
}
|
||||
|
||||
/// Save MCP servers (DB if available, else disk).
|
||||
async fn save_servers(
|
||||
store: Option<&Store>,
|
||||
servers: &McpServersFile,
|
||||
) -> Result<(), config::ConfigError> {
|
||||
if let Some(store) = store {
|
||||
config::save_mcp_servers_to_db(store, DEFAULT_USER_ID, servers).await
|
||||
} else {
|
||||
config::save_mcp_servers(servers).await
|
||||
}
|
||||
}
|
||||
|
||||
/// Initialize and return the secrets store.
|
||||
async fn get_secrets_store() -> anyhow::Result<Arc<dyn SecretsStore + Send + Sync>> {
|
||||
let config = Config::from_env()?;
|
||||
let config = Config::from_env().await?;
|
||||
|
||||
let master_key = config.secrets.master_key().ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
|
||||
@@ -12,12 +12,14 @@
|
||||
mod config;
|
||||
mod mcp;
|
||||
pub mod memory;
|
||||
mod pairing;
|
||||
pub mod status;
|
||||
mod tool;
|
||||
|
||||
pub use config::{ConfigCommand, run_config_command};
|
||||
pub use mcp::{McpCommand, run_mcp_command};
|
||||
pub use memory::{MemoryCommand, run_memory_command};
|
||||
pub use pairing::{PairingCommand, run_pairing_command, run_pairing_command_with_store};
|
||||
pub use status::run_status_command;
|
||||
pub use tool::{ToolCommand, run_tool_command};
|
||||
|
||||
@@ -86,8 +88,48 @@ pub enum Command {
|
||||
#[command(subcommand)]
|
||||
Memory(MemoryCommand),
|
||||
|
||||
/// DM pairing (approve inbound requests from unknown senders)
|
||||
#[command(subcommand)]
|
||||
Pairing(PairingCommand),
|
||||
|
||||
/// Show system health and diagnostics
|
||||
Status,
|
||||
|
||||
/// Run as a sandboxed worker inside a Docker container (internal use).
|
||||
/// This is invoked automatically by the orchestrator, not by users directly.
|
||||
Worker {
|
||||
/// Job ID to execute.
|
||||
#[arg(long)]
|
||||
job_id: uuid::Uuid,
|
||||
|
||||
/// URL of the orchestrator's internal API.
|
||||
#[arg(long, default_value = "http://host.docker.internal:50051")]
|
||||
orchestrator_url: String,
|
||||
|
||||
/// Maximum iterations before stopping.
|
||||
#[arg(long, default_value = "50")]
|
||||
max_iterations: u32,
|
||||
},
|
||||
|
||||
/// Run as a Claude Code bridge inside a Docker container (internal use).
|
||||
/// Spawns the `claude` CLI and streams output back to the orchestrator.
|
||||
ClaudeBridge {
|
||||
/// Job ID to execute.
|
||||
#[arg(long)]
|
||||
job_id: uuid::Uuid,
|
||||
|
||||
/// URL of the orchestrator's internal API.
|
||||
#[arg(long, default_value = "http://host.docker.internal:50051")]
|
||||
orchestrator_url: String,
|
||||
|
||||
/// Maximum agentic turns for Claude Code.
|
||||
#[arg(long, default_value = "50")]
|
||||
max_turns: u32,
|
||||
|
||||
/// Claude model to use (e.g. "sonnet", "opus").
|
||||
#[arg(long, default_value = "sonnet")]
|
||||
model: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl Cli {
|
||||
|
||||
@@ -0,0 +1,187 @@
|
||||
//! DM pairing CLI commands.
|
||||
//!
|
||||
//! Manage pairing requests for channels (Telegram, Slack, etc.).
|
||||
|
||||
use clap::Subcommand;
|
||||
|
||||
use crate::pairing::PairingStore;
|
||||
|
||||
/// Pairing subcommands.
|
||||
#[derive(Subcommand, Debug, Clone)]
|
||||
pub enum PairingCommand {
|
||||
/// List pending pairing requests
|
||||
List {
|
||||
/// Channel name (e.g., telegram, slack)
|
||||
#[arg(required = true)]
|
||||
channel: String,
|
||||
|
||||
/// Output as JSON
|
||||
#[arg(long)]
|
||||
json: bool,
|
||||
},
|
||||
|
||||
/// Approve a pairing request by code
|
||||
Approve {
|
||||
/// Channel name (e.g., telegram, slack)
|
||||
#[arg(required = true)]
|
||||
channel: String,
|
||||
|
||||
/// Pairing code (e.g., ABC12345)
|
||||
#[arg(required = true)]
|
||||
code: String,
|
||||
},
|
||||
}
|
||||
|
||||
/// Run pairing CLI command.
|
||||
pub fn run_pairing_command(cmd: PairingCommand) -> Result<(), String> {
|
||||
run_pairing_command_with_store(&PairingStore::new(), cmd)
|
||||
}
|
||||
|
||||
/// Run pairing CLI command with a given store (for testing).
|
||||
pub fn run_pairing_command_with_store(
|
||||
store: &PairingStore,
|
||||
cmd: PairingCommand,
|
||||
) -> Result<(), String> {
|
||||
match cmd {
|
||||
PairingCommand::List { channel, json } => run_list(store, &channel, json),
|
||||
PairingCommand::Approve { channel, code } => run_approve(store, &channel, &code),
|
||||
}
|
||||
}
|
||||
|
||||
fn run_list(store: &PairingStore, channel: &str, json: bool) -> Result<(), String> {
|
||||
let requests = store.list_pending(channel).map_err(|e| e.to_string())?;
|
||||
|
||||
if json {
|
||||
println!(
|
||||
"{}",
|
||||
serde_json::to_string_pretty(&requests).map_err(|e| e.to_string())?
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if requests.is_empty() {
|
||||
println!("No pending {} pairing requests.", channel);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
println!("Pairing requests ({}):", requests.len());
|
||||
for r in &requests {
|
||||
let meta = r
|
||||
.meta
|
||||
.as_ref()
|
||||
.and_then(|m| m.as_object())
|
||||
.map(|o| {
|
||||
o.iter()
|
||||
.filter_map(|(k, v)| v.as_str().map(|s| format!("{}={}", k, s)))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ")
|
||||
})
|
||||
.unwrap_or_default();
|
||||
println!(" {} {} {} {}", r.code, r.id, meta, r.created_at);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn run_approve(store: &PairingStore, channel: &str, code: &str) -> Result<(), String> {
|
||||
match store.approve(channel, code) {
|
||||
Ok(Some(entry)) => {
|
||||
println!("Approved {} sender {}.", channel, entry.id);
|
||||
Ok(())
|
||||
}
|
||||
Ok(None) => Err(format!(
|
||||
"No pending pairing request found for code: {}",
|
||||
code
|
||||
)),
|
||||
Err(crate::pairing::PairingStoreError::ApproveRateLimited) => Err(
|
||||
"Too many failed approve attempts. Wait a few minutes before trying again.".to_string(),
|
||||
),
|
||||
Err(e) => Err(e.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn test_store() -> (PairingStore, TempDir) {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let store = PairingStore::with_base_dir(dir.path().to_path_buf());
|
||||
(store, dir)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_list_empty_returns_ok() {
|
||||
let (store, _) = test_store();
|
||||
let result = run_pairing_command_with_store(
|
||||
&store,
|
||||
PairingCommand::List {
|
||||
channel: "telegram".to_string(),
|
||||
json: false,
|
||||
},
|
||||
);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_list_json_empty_returns_ok() {
|
||||
let (store, _) = test_store();
|
||||
let result = run_pairing_command_with_store(
|
||||
&store,
|
||||
PairingCommand::List {
|
||||
channel: "telegram".to_string(),
|
||||
json: true,
|
||||
},
|
||||
);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_approve_invalid_code_returns_err() {
|
||||
let (store, _) = test_store();
|
||||
// Create a pending request so the pairing file exists, then approve with wrong code
|
||||
store.upsert_request("telegram", "user1", None).unwrap();
|
||||
|
||||
let result = run_pairing_command_with_store(
|
||||
&store,
|
||||
PairingCommand::Approve {
|
||||
channel: "telegram".to_string(),
|
||||
code: "BADCODE1".to_string(),
|
||||
},
|
||||
);
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("No pending pairing request"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_approve_valid_code_returns_ok() {
|
||||
let (store, _) = test_store();
|
||||
let r = store.upsert_request("telegram", "user1", None).unwrap();
|
||||
assert!(r.created);
|
||||
|
||||
let result = run_pairing_command_with_store(
|
||||
&store,
|
||||
PairingCommand::Approve {
|
||||
channel: "telegram".to_string(),
|
||||
code: r.code,
|
||||
},
|
||||
);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_list_with_pending_returns_ok() {
|
||||
let (store, _) = test_store();
|
||||
store.upsert_request("telegram", "user1", None).unwrap();
|
||||
|
||||
let result = run_pairing_command_with_store(
|
||||
&store,
|
||||
PairingCommand::List {
|
||||
channel: "telegram".to_string(),
|
||||
json: false,
|
||||
},
|
||||
);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
}
|
||||
+1
-1
@@ -47,7 +47,7 @@ pub async fn run_status_command() -> anyhow::Result<()> {
|
||||
print!(" Secrets: ");
|
||||
let secrets_configured = settings.secrets_master_key_source != crate::settings::KeySource::None
|
||||
|| std::env::var("SECRETS_MASTER_KEY").is_ok()
|
||||
|| crate::secrets::keychain::has_master_key();
|
||||
|| crate::secrets::keychain::has_master_key().await;
|
||||
if secrets_configured {
|
||||
println!("configured ({:?})", settings.secrets_master_key_source);
|
||||
} else {
|
||||
|
||||
+1
-1
@@ -715,7 +715,7 @@ async fn auth_tool(name: String, dir: Option<PathBuf>, user_id: String) -> anyho
|
||||
println!();
|
||||
|
||||
// Initialize secrets store
|
||||
let config = Config::from_env()?;
|
||||
let config = Config::from_env().await?;
|
||||
let master_key = config.secrets.master_key().ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"SECRETS_MASTER_KEY not set. Run 'ironclaw onboard' first or set it in .env"
|
||||
|
||||
+390
-123
@@ -1,4 +1,9 @@
|
||||
//! Configuration for IronClaw.
|
||||
//!
|
||||
//! Settings are loaded with priority: env var > database > default.
|
||||
//! The database replaces the old `settings.json` file for all settings
|
||||
//! except the 4 bootstrap fields (database_url, pool_size, secrets key
|
||||
//! source, onboard_completed) which live in `~/.ironclaw/bootstrap.json`.
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::time::Duration;
|
||||
@@ -6,6 +11,7 @@ use std::time::Duration;
|
||||
use secrecy::{ExposeSecret, SecretString};
|
||||
|
||||
use crate::error::ConfigError;
|
||||
use crate::settings::Settings;
|
||||
|
||||
/// Main configuration for the agent.
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -21,28 +27,67 @@ pub struct Config {
|
||||
pub secrets: SecretsConfig,
|
||||
pub builder: BuilderModeConfig,
|
||||
pub heartbeat: HeartbeatConfig,
|
||||
pub routines: RoutineConfig,
|
||||
pub sandbox: SandboxModeConfig,
|
||||
pub claude_code: ClaudeCodeConfig,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
/// Load configuration from environment variables.
|
||||
pub fn from_env() -> Result<Self, ConfigError> {
|
||||
// Load .env file if present (ignore errors if not found)
|
||||
/// Load configuration from environment variables and the database.
|
||||
///
|
||||
/// Priority: env var > DB settings > default.
|
||||
/// This is the primary way to load config after DB is connected.
|
||||
pub async fn from_db(
|
||||
store: &crate::history::Store,
|
||||
user_id: &str,
|
||||
bootstrap: &crate::bootstrap::BootstrapConfig,
|
||||
) -> Result<Self, ConfigError> {
|
||||
let _ = dotenvy::dotenv();
|
||||
|
||||
// Load all settings from DB into a Settings struct
|
||||
let db_settings = match store.get_all_settings(user_id).await {
|
||||
Ok(map) => Settings::from_db_map(&map),
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to load settings from DB, using defaults: {}", e);
|
||||
Settings::default()
|
||||
}
|
||||
};
|
||||
|
||||
Self::build(bootstrap, &db_settings).await
|
||||
}
|
||||
|
||||
/// Load configuration from environment variables only (no database).
|
||||
///
|
||||
/// Used during early startup before the database is connected,
|
||||
/// and by CLI commands that don't have DB access.
|
||||
/// Falls back to legacy `settings.json` on disk if present.
|
||||
pub async fn from_env() -> Result<Self, ConfigError> {
|
||||
let _ = dotenvy::dotenv();
|
||||
let bootstrap = crate::bootstrap::BootstrapConfig::load();
|
||||
let settings = Settings::load();
|
||||
Self::build(&bootstrap, &settings).await
|
||||
}
|
||||
|
||||
/// Build config from bootstrap + settings (shared by from_env and from_db).
|
||||
async fn build(
|
||||
bootstrap: &crate::bootstrap::BootstrapConfig,
|
||||
settings: &Settings,
|
||||
) -> Result<Self, ConfigError> {
|
||||
Ok(Self {
|
||||
database: DatabaseConfig::from_env()?,
|
||||
llm: LlmConfig::from_env()?,
|
||||
embeddings: EmbeddingsConfig::from_env()?,
|
||||
tunnel: TunnelConfig::from_env()?,
|
||||
channels: ChannelsConfig::from_env()?,
|
||||
agent: AgentConfig::from_env()?,
|
||||
safety: SafetyConfig::from_env()?,
|
||||
wasm: WasmConfig::from_env()?,
|
||||
secrets: SecretsConfig::from_env()?,
|
||||
builder: BuilderModeConfig::from_env()?,
|
||||
heartbeat: HeartbeatConfig::from_env()?,
|
||||
sandbox: SandboxModeConfig::from_env()?,
|
||||
database: DatabaseConfig::resolve(bootstrap)?,
|
||||
llm: LlmConfig::resolve(settings)?,
|
||||
embeddings: EmbeddingsConfig::resolve(settings)?,
|
||||
tunnel: TunnelConfig::resolve(settings)?,
|
||||
channels: ChannelsConfig::resolve(settings)?,
|
||||
agent: AgentConfig::resolve(settings)?,
|
||||
safety: SafetyConfig::resolve()?,
|
||||
wasm: WasmConfig::resolve()?,
|
||||
secrets: SecretsConfig::resolve(bootstrap).await?,
|
||||
builder: BuilderModeConfig::resolve()?,
|
||||
heartbeat: HeartbeatConfig::resolve(settings)?,
|
||||
routines: RoutineConfig::resolve()?,
|
||||
sandbox: SandboxModeConfig::resolve()?,
|
||||
claude_code: ClaudeCodeConfig::resolve()?,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -51,48 +96,17 @@ impl Config {
|
||||
///
|
||||
/// Used by channels and tools that need public webhook endpoints.
|
||||
/// The tunnel URL is shared across all channels (Telegram, Slack, etc.).
|
||||
///
|
||||
/// # Security Notes
|
||||
///
|
||||
/// **Webhook endpoints** (e.g., `/webhook/telegram`) should NOT use tunnel-level
|
||||
/// authentication because webhook providers (Telegram, Slack, GitHub) need
|
||||
/// unauthenticated access to POST updates. Security for webhooks comes from:
|
||||
/// - Webhook signature verification (provider-specific secrets)
|
||||
/// - IP allowlisting (if supported by provider)
|
||||
///
|
||||
/// **Non-webhook endpoints** (admin APIs, health checks) CAN be protected using
|
||||
/// tunnel provider features:
|
||||
/// - ngrok: Basic Auth, OAuth, IP restrictions
|
||||
/// - Cloudflare: Access policies, mTLS
|
||||
///
|
||||
/// These protections are configured in the tunnel provider, not here.
|
||||
///
|
||||
/// # Supported Providers
|
||||
///
|
||||
/// - **ngrok**: `ngrok http 8080` -> `https://abc123.ngrok.io`
|
||||
/// - **Cloudflare Tunnel**: `cloudflared tunnel --url http://localhost:8080`
|
||||
/// - **localtunnel**: `lt --port 8080`
|
||||
/// - Any service that provides a public HTTPS URL to localhost
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct TunnelConfig {
|
||||
/// Public URL from tunnel provider (e.g., "https://abc123.ngrok.io").
|
||||
///
|
||||
/// When set, channels that support webhooks will register their endpoints
|
||||
/// with this base URL instead of using polling.
|
||||
pub public_url: Option<String>,
|
||||
}
|
||||
|
||||
impl TunnelConfig {
|
||||
fn from_env() -> Result<Self, ConfigError> {
|
||||
// Priority: env var > settings file
|
||||
let public_url = optional_env("TUNNEL_URL")?.or_else(|| {
|
||||
crate::settings::Settings::load()
|
||||
.tunnel
|
||||
.public_url
|
||||
.filter(|s| !s.is_empty())
|
||||
});
|
||||
fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
|
||||
let public_url = optional_env("TUNNEL_URL")?
|
||||
.or_else(|| settings.tunnel.public_url.clone().filter(|s| !s.is_empty()));
|
||||
|
||||
// Validate URL format if provided
|
||||
if let Some(ref url) = public_url {
|
||||
if !url.starts_with("https://") {
|
||||
return Err(ConfigError::InvalidValue {
|
||||
@@ -111,8 +125,6 @@ impl TunnelConfig {
|
||||
}
|
||||
|
||||
/// Get the webhook URL for a given path.
|
||||
///
|
||||
/// Returns `None` if no tunnel is configured.
|
||||
pub fn webhook_url(&self, path: &str) -> Option<String> {
|
||||
self.public_url.as_ref().map(|base| {
|
||||
let base = base.trim_end_matches('/');
|
||||
@@ -130,18 +142,14 @@ pub struct DatabaseConfig {
|
||||
}
|
||||
|
||||
impl DatabaseConfig {
|
||||
fn from_env() -> Result<Self, ConfigError> {
|
||||
let settings = crate::settings::Settings::load();
|
||||
|
||||
// Priority: env var > settings > error (required)
|
||||
fn resolve(bootstrap: &crate::bootstrap::BootstrapConfig) -> Result<Self, ConfigError> {
|
||||
let url = optional_env("DATABASE_URL")?
|
||||
.or(settings.database_url.clone())
|
||||
.or_else(|| bootstrap.database_url.clone())
|
||||
.ok_or_else(|| ConfigError::MissingRequired {
|
||||
key: "database_url".to_string(),
|
||||
hint: "Run 'ironclaw onboard' or set DATABASE_URL environment variable".to_string(),
|
||||
})?;
|
||||
|
||||
// Priority: env var > settings > default
|
||||
let pool_size = optional_env("DATABASE_POOL_SIZE")?
|
||||
.map(|s| s.parse())
|
||||
.transpose()
|
||||
@@ -149,7 +157,7 @@ impl DatabaseConfig {
|
||||
key: "DATABASE_POOL_SIZE".to_string(),
|
||||
message: format!("must be a positive integer: {e}"),
|
||||
})?
|
||||
.or(settings.database_pool_size)
|
||||
.or(bootstrap.database_pool_size)
|
||||
.unwrap_or(10);
|
||||
|
||||
Ok(Self {
|
||||
@@ -164,10 +172,102 @@ impl DatabaseConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// LLM provider configuration (NEAR AI only).
|
||||
/// Which LLM backend to use.
|
||||
///
|
||||
/// Defaults to `NearAi` to keep IronClaw close to the NEAR ecosystem.
|
||||
/// Users can override with `LLM_BACKEND` env var to use their own API keys.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum LlmBackend {
|
||||
/// NEAR AI proxy (default) -- session or API key auth
|
||||
#[default]
|
||||
NearAi,
|
||||
/// Direct OpenAI API
|
||||
OpenAi,
|
||||
/// Direct Anthropic API
|
||||
Anthropic,
|
||||
/// Local Ollama instance
|
||||
Ollama,
|
||||
/// Any OpenAI-compatible endpoint (e.g. vLLM, LiteLLM, Together)
|
||||
OpenAiCompatible,
|
||||
}
|
||||
|
||||
impl std::str::FromStr for LlmBackend {
|
||||
type Err = String;
|
||||
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
match s.to_lowercase().as_str() {
|
||||
"nearai" | "near_ai" | "near" => Ok(Self::NearAi),
|
||||
"openai" | "open_ai" => Ok(Self::OpenAi),
|
||||
"anthropic" | "claude" => Ok(Self::Anthropic),
|
||||
"ollama" => Ok(Self::Ollama),
|
||||
"openai_compatible" | "openai-compatible" | "compatible" => Ok(Self::OpenAiCompatible),
|
||||
_ => Err(format!(
|
||||
"invalid LLM backend '{}', expected one of: nearai, openai, anthropic, ollama, openai_compatible",
|
||||
s
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for LlmBackend {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::NearAi => write!(f, "nearai"),
|
||||
Self::OpenAi => write!(f, "openai"),
|
||||
Self::Anthropic => write!(f, "anthropic"),
|
||||
Self::Ollama => write!(f, "ollama"),
|
||||
Self::OpenAiCompatible => write!(f, "openai_compatible"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Configuration for direct OpenAI API access.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct OpenAiDirectConfig {
|
||||
pub api_key: SecretString,
|
||||
pub model: String,
|
||||
}
|
||||
|
||||
/// Configuration for direct Anthropic API access.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AnthropicDirectConfig {
|
||||
pub api_key: SecretString,
|
||||
pub model: String,
|
||||
}
|
||||
|
||||
/// Configuration for local Ollama.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct OllamaConfig {
|
||||
pub base_url: String,
|
||||
pub model: String,
|
||||
}
|
||||
|
||||
/// Configuration for any OpenAI-compatible endpoint.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct OpenAiCompatibleConfig {
|
||||
pub base_url: String,
|
||||
pub api_key: Option<SecretString>,
|
||||
pub model: String,
|
||||
}
|
||||
|
||||
/// LLM provider configuration.
|
||||
///
|
||||
/// NEAR AI remains the default backend. Users can switch to other providers
|
||||
/// by setting `LLM_BACKEND` (e.g. `openai`, `anthropic`, `ollama`).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct LlmConfig {
|
||||
/// Which backend to use (default: NearAi)
|
||||
pub backend: LlmBackend,
|
||||
/// NEAR AI config (always populated for NEAR AI embeddings, etc.)
|
||||
pub nearai: NearAiConfig,
|
||||
/// Direct OpenAI config (populated when backend=openai)
|
||||
pub openai: Option<OpenAiDirectConfig>,
|
||||
/// Direct Anthropic config (populated when backend=anthropic)
|
||||
pub anthropic: Option<AnthropicDirectConfig>,
|
||||
/// Ollama config (populated when backend=ollama)
|
||||
pub ollama: Option<OllamaConfig>,
|
||||
/// OpenAI-compatible config (populated when backend=openai_compatible)
|
||||
pub openai_compatible: Option<OpenAiCompatibleConfig>,
|
||||
}
|
||||
|
||||
/// API mode for NEAR AI.
|
||||
@@ -215,42 +315,110 @@ pub struct NearAiConfig {
|
||||
}
|
||||
|
||||
impl LlmConfig {
|
||||
fn from_env() -> Result<Self, ConfigError> {
|
||||
let api_key = optional_env("NEARAI_API_KEY")?.map(SecretString::from);
|
||||
fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
|
||||
// Determine backend (default: NearAi)
|
||||
let backend: LlmBackend = if let Some(b) = optional_env("LLM_BACKEND")? {
|
||||
b.parse().map_err(|e| ConfigError::InvalidValue {
|
||||
key: "LLM_BACKEND".to_string(),
|
||||
message: e,
|
||||
})?
|
||||
} else {
|
||||
LlmBackend::NearAi
|
||||
};
|
||||
|
||||
// Always resolve NEAR AI config (used as fallback and for embeddings)
|
||||
let nearai_api_key = optional_env("NEARAI_API_KEY")?.map(SecretString::from);
|
||||
|
||||
// Determine API mode: explicit setting, or infer from API key presence
|
||||
let api_mode = if let Some(mode_str) = optional_env("NEARAI_API_MODE")? {
|
||||
mode_str.parse().map_err(|e| ConfigError::InvalidValue {
|
||||
key: "NEARAI_API_MODE".to_string(),
|
||||
message: e,
|
||||
})?
|
||||
} else if api_key.is_some() {
|
||||
// If API key is provided, default to chat_completions mode
|
||||
} else if nearai_api_key.is_some() {
|
||||
NearAiApiMode::ChatCompletions
|
||||
} else {
|
||||
NearAiApiMode::Responses
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
nearai: NearAiConfig {
|
||||
// Load model from saved settings first, then env, then default
|
||||
model: crate::settings::Settings::load()
|
||||
.selected_model
|
||||
.or_else(|| optional_env("NEARAI_MODEL").ok().flatten())
|
||||
.unwrap_or_else(|| {
|
||||
"fireworks::accounts/fireworks/models/llama4-maverick-instruct-basic"
|
||||
.to_string()
|
||||
}),
|
||||
base_url: optional_env("NEARAI_BASE_URL")?
|
||||
.unwrap_or_else(|| "https://cloud-api.near.ai".to_string()),
|
||||
auth_base_url: optional_env("NEARAI_AUTH_URL")?
|
||||
.unwrap_or_else(|| "https://private.near.ai".to_string()),
|
||||
session_path: optional_env("NEARAI_SESSION_PATH")?
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(default_session_path),
|
||||
api_mode,
|
||||
let nearai = NearAiConfig {
|
||||
model: optional_env("NEARAI_MODEL")?
|
||||
.or_else(|| settings.selected_model.clone())
|
||||
.unwrap_or_else(|| {
|
||||
"fireworks::accounts/fireworks/models/llama4-maverick-instruct-basic"
|
||||
.to_string()
|
||||
}),
|
||||
base_url: optional_env("NEARAI_BASE_URL")?
|
||||
.unwrap_or_else(|| "https://cloud-api.near.ai".to_string()),
|
||||
auth_base_url: optional_env("NEARAI_AUTH_URL")?
|
||||
.unwrap_or_else(|| "https://private.near.ai".to_string()),
|
||||
session_path: optional_env("NEARAI_SESSION_PATH")?
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(default_session_path),
|
||||
api_mode,
|
||||
api_key: nearai_api_key,
|
||||
};
|
||||
|
||||
// Resolve provider-specific configs based on backend
|
||||
let openai = if backend == LlmBackend::OpenAi {
|
||||
let api_key = optional_env("OPENAI_API_KEY")?
|
||||
.map(SecretString::from)
|
||||
.ok_or_else(|| ConfigError::MissingRequired {
|
||||
key: "OPENAI_API_KEY".to_string(),
|
||||
hint: "Set OPENAI_API_KEY when LLM_BACKEND=openai".to_string(),
|
||||
})?;
|
||||
let model = optional_env("OPENAI_MODEL")?.unwrap_or_else(|| "gpt-4o".to_string());
|
||||
Some(OpenAiDirectConfig { api_key, model })
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let anthropic = if backend == LlmBackend::Anthropic {
|
||||
let api_key = optional_env("ANTHROPIC_API_KEY")?
|
||||
.map(SecretString::from)
|
||||
.ok_or_else(|| ConfigError::MissingRequired {
|
||||
key: "ANTHROPIC_API_KEY".to_string(),
|
||||
hint: "Set ANTHROPIC_API_KEY when LLM_BACKEND=anthropic".to_string(),
|
||||
})?;
|
||||
let model = optional_env("ANTHROPIC_MODEL")?
|
||||
.unwrap_or_else(|| "claude-sonnet-4-20250514".to_string());
|
||||
Some(AnthropicDirectConfig { api_key, model })
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let ollama = if backend == LlmBackend::Ollama {
|
||||
let base_url = optional_env("OLLAMA_BASE_URL")?
|
||||
.unwrap_or_else(|| "http://localhost:11434".to_string());
|
||||
let model = optional_env("OLLAMA_MODEL")?.unwrap_or_else(|| "llama3".to_string());
|
||||
Some(OllamaConfig { base_url, model })
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let openai_compatible = if backend == LlmBackend::OpenAiCompatible {
|
||||
let base_url =
|
||||
optional_env("LLM_BASE_URL")?.ok_or_else(|| ConfigError::MissingRequired {
|
||||
key: "LLM_BASE_URL".to_string(),
|
||||
hint: "Set LLM_BASE_URL when LLM_BACKEND=openai_compatible".to_string(),
|
||||
})?;
|
||||
let api_key = optional_env("LLM_API_KEY")?.map(SecretString::from);
|
||||
let model = optional_env("LLM_MODEL")?.unwrap_or_else(|| "default".to_string());
|
||||
Some(OpenAiCompatibleConfig {
|
||||
base_url,
|
||||
api_key,
|
||||
},
|
||||
model,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
backend,
|
||||
nearai,
|
||||
openai,
|
||||
anthropic,
|
||||
ollama,
|
||||
openai_compatible,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -265,8 +433,6 @@ pub struct EmbeddingsConfig {
|
||||
/// OpenAI API key (for OpenAI provider).
|
||||
pub openai_api_key: Option<SecretString>,
|
||||
/// Model to use for embeddings.
|
||||
/// For OpenAI: "text-embedding-3-small", "text-embedding-3-large", "text-embedding-ada-002"
|
||||
/// For NEAR AI: Uses the configured session for auth.
|
||||
pub model: String,
|
||||
}
|
||||
|
||||
@@ -282,18 +448,15 @@ impl Default for EmbeddingsConfig {
|
||||
}
|
||||
|
||||
impl EmbeddingsConfig {
|
||||
fn from_env() -> Result<Self, ConfigError> {
|
||||
let settings = crate::settings::Settings::load();
|
||||
fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
|
||||
let openai_api_key = optional_env("OPENAI_API_KEY")?.map(SecretString::from);
|
||||
|
||||
// Priority: env var > settings > default
|
||||
let provider = optional_env("EMBEDDING_PROVIDER")?
|
||||
.unwrap_or_else(|| settings.embeddings.provider.clone());
|
||||
|
||||
let model =
|
||||
optional_env("EMBEDDING_MODEL")?.unwrap_or_else(|| settings.embeddings.model.clone());
|
||||
|
||||
// Priority: env var > settings > auto-detect from API key
|
||||
let enabled = optional_env("EMBEDDING_ENABLED")?
|
||||
.map(|s| s.parse())
|
||||
.transpose()
|
||||
@@ -301,10 +464,7 @@ impl EmbeddingsConfig {
|
||||
key: "EMBEDDING_ENABLED".to_string(),
|
||||
message: format!("must be 'true' or 'false': {e}"),
|
||||
})?
|
||||
.unwrap_or_else(|| {
|
||||
// Check settings, or auto-enable if API key present
|
||||
settings.embeddings.enabled || openai_api_key.is_some()
|
||||
});
|
||||
.unwrap_or_else(|| settings.embeddings.enabled || openai_api_key.is_some());
|
||||
|
||||
Ok(Self {
|
||||
enabled,
|
||||
@@ -338,6 +498,8 @@ pub struct ChannelsConfig {
|
||||
pub wasm_channels_dir: std::path::PathBuf,
|
||||
/// Whether WASM channels are enabled.
|
||||
pub wasm_channels_enabled: bool,
|
||||
/// Telegram owner user ID. When set, the bot only responds to this user.
|
||||
pub telegram_owner_id: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -364,7 +526,7 @@ pub struct GatewayConfig {
|
||||
}
|
||||
|
||||
impl ChannelsConfig {
|
||||
fn from_env() -> Result<Self, ConfigError> {
|
||||
fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
|
||||
let http = if optional_env("HTTP_PORT")?.is_some() || optional_env("HTTP_HOST")?.is_some() {
|
||||
Some(HttpConfig {
|
||||
host: optional_env("HTTP_HOST")?.unwrap_or_else(|| "0.0.0.0".to_string()),
|
||||
@@ -385,7 +547,7 @@ impl ChannelsConfig {
|
||||
|
||||
let gateway = if optional_env("GATEWAY_ENABLED")?
|
||||
.map(|s| s.to_lowercase() == "true" || s == "1")
|
||||
.unwrap_or(false)
|
||||
.unwrap_or(true)
|
||||
{
|
||||
Some(GatewayConfig {
|
||||
host: optional_env("GATEWAY_HOST")?.unwrap_or_else(|| "127.0.0.1".to_string()),
|
||||
@@ -425,6 +587,14 @@ impl ChannelsConfig {
|
||||
message: format!("must be 'true' or 'false': {e}"),
|
||||
})?
|
||||
.unwrap_or(true),
|
||||
telegram_owner_id: optional_env("TELEGRAM_OWNER_ID")?
|
||||
.map(|s| s.parse())
|
||||
.transpose()
|
||||
.map_err(|e| ConfigError::InvalidValue {
|
||||
key: "TELEGRAM_OWNER_ID".to_string(),
|
||||
message: format!("must be an integer: {e}"),
|
||||
})?
|
||||
.or(settings.channels.telegram_owner_id),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -450,14 +620,13 @@ pub struct AgentConfig {
|
||||
pub use_planning: bool,
|
||||
/// Session idle timeout. Sessions inactive longer than this are pruned.
|
||||
pub session_idle_timeout: Duration,
|
||||
/// Allow chat to use filesystem/shell tools directly (bypass sandbox).
|
||||
pub allow_local_tools: bool,
|
||||
}
|
||||
|
||||
impl AgentConfig {
|
||||
fn from_env() -> Result<Self, ConfigError> {
|
||||
let settings = crate::settings::Settings::load();
|
||||
|
||||
fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
|
||||
Ok(Self {
|
||||
// Priority: env var > settings > default
|
||||
name: optional_env("AGENT_NAME")?.unwrap_or_else(|| settings.agent.name.clone()),
|
||||
max_parallel_jobs: optional_env("AGENT_MAX_PARALLEL_JOBS")?
|
||||
.map(|s| s.parse())
|
||||
@@ -523,6 +692,14 @@ impl AgentConfig {
|
||||
})?
|
||||
.unwrap_or(settings.agent.session_idle_timeout_secs),
|
||||
),
|
||||
allow_local_tools: optional_env("ALLOW_LOCAL_TOOLS")?
|
||||
.map(|s| s.parse())
|
||||
.transpose()
|
||||
.map_err(|e| ConfigError::InvalidValue {
|
||||
key: "ALLOW_LOCAL_TOOLS".to_string(),
|
||||
message: format!("must be 'true' or 'false': {e}"),
|
||||
})?
|
||||
.unwrap_or(false),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -535,7 +712,7 @@ pub struct SafetyConfig {
|
||||
}
|
||||
|
||||
impl SafetyConfig {
|
||||
fn from_env() -> Result<Self, ConfigError> {
|
||||
fn resolve() -> Result<Self, ConfigError> {
|
||||
Ok(Self {
|
||||
max_output_length: parse_optional_env("SAFETY_MAX_OUTPUT_LENGTH", 100_000)?,
|
||||
injection_check_enabled: optional_env("SAFETY_INJECTION_CHECK_ENABLED")?
|
||||
@@ -573,7 +750,6 @@ pub struct WasmConfig {
|
||||
#[derive(Clone, Default)]
|
||||
pub struct SecretsConfig {
|
||||
/// Master key for encrypting secrets.
|
||||
/// Source determined by KeySource in settings.
|
||||
pub master_key: Option<SecretString>,
|
||||
/// Whether secrets management is enabled.
|
||||
pub enabled: bool,
|
||||
@@ -592,20 +768,16 @@ impl std::fmt::Debug for SecretsConfig {
|
||||
}
|
||||
|
||||
impl SecretsConfig {
|
||||
fn from_env() -> Result<Self, ConfigError> {
|
||||
async fn resolve(bootstrap: &crate::bootstrap::BootstrapConfig) -> Result<Self, ConfigError> {
|
||||
use crate::settings::KeySource;
|
||||
|
||||
let settings = crate::settings::Settings::load();
|
||||
|
||||
// Priority: env var > keychain (based on settings) > disabled
|
||||
let (master_key, source) = if let Some(env_key) = optional_env("SECRETS_MASTER_KEY")? {
|
||||
// Env var takes priority (for CI/Docker)
|
||||
(Some(SecretString::from(env_key)), KeySource::Env)
|
||||
} else {
|
||||
match settings.secrets_master_key_source {
|
||||
match bootstrap.secrets_master_key_source {
|
||||
KeySource::Keychain => {
|
||||
// Try to load from OS keychain
|
||||
match crate::secrets::keychain::get_master_key() {
|
||||
// Try to load from OS keychain (async on Linux)
|
||||
match crate::secrets::keychain::get_master_key().await {
|
||||
Ok(key_bytes) => {
|
||||
let key_hex: String =
|
||||
key_bytes.iter().map(|b| format!("{:02x}", b)).collect();
|
||||
@@ -623,7 +795,6 @@ impl SecretsConfig {
|
||||
}
|
||||
}
|
||||
KeySource::Env => {
|
||||
// Settings say env, but no env var found
|
||||
tracing::warn!(
|
||||
"Secrets configured for env var but SECRETS_MASTER_KEY not set."
|
||||
);
|
||||
@@ -635,7 +806,6 @@ impl SecretsConfig {
|
||||
|
||||
let enabled = master_key.is_some();
|
||||
|
||||
// Validate master key length if provided
|
||||
if let Some(ref key) = master_key {
|
||||
if key.expose_secret().len() < 32 {
|
||||
return Err(ConfigError::InvalidValue {
|
||||
@@ -681,7 +851,7 @@ fn default_tools_dir() -> PathBuf {
|
||||
}
|
||||
|
||||
impl WasmConfig {
|
||||
fn from_env() -> Result<Self, ConfigError> {
|
||||
fn resolve() -> Result<Self, ConfigError> {
|
||||
Ok(Self {
|
||||
enabled: optional_env("WASM_ENABLED")?
|
||||
.map(|s| s.parse())
|
||||
@@ -752,7 +922,7 @@ pub struct BuilderModeConfig {
|
||||
impl Default for BuilderModeConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: true, // Builder enabled by default
|
||||
enabled: true,
|
||||
build_dir: None,
|
||||
max_iterations: 20,
|
||||
timeout_secs: 600,
|
||||
@@ -762,7 +932,7 @@ impl Default for BuilderModeConfig {
|
||||
}
|
||||
|
||||
impl BuilderModeConfig {
|
||||
fn from_env() -> Result<Self, ConfigError> {
|
||||
fn resolve() -> Result<Self, ConfigError> {
|
||||
Ok(Self {
|
||||
enabled: optional_env("BUILDER_ENABLED")?
|
||||
.map(|s| s.parse())
|
||||
@@ -771,7 +941,7 @@ impl BuilderModeConfig {
|
||||
key: "BUILDER_ENABLED".to_string(),
|
||||
message: format!("must be 'true' or 'false': {e}"),
|
||||
})?
|
||||
.unwrap_or(true), // Builder enabled by default
|
||||
.unwrap_or(true),
|
||||
build_dir: optional_env("BUILDER_DIR")?.map(PathBuf::from),
|
||||
max_iterations: parse_optional_env("BUILDER_MAX_ITERATIONS", 20)?,
|
||||
timeout_secs: parse_optional_env("BUILDER_TIMEOUT_SECS", 600)?,
|
||||
@@ -826,11 +996,8 @@ impl Default for HeartbeatConfig {
|
||||
}
|
||||
|
||||
impl HeartbeatConfig {
|
||||
fn from_env() -> Result<Self, ConfigError> {
|
||||
let settings = crate::settings::Settings::load();
|
||||
|
||||
fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
|
||||
Ok(Self {
|
||||
// Priority: env var > settings > default
|
||||
enabled: optional_env("HEARTBEAT_ENABLED")?
|
||||
.map(|s| s.parse())
|
||||
.transpose()
|
||||
@@ -848,9 +1015,55 @@ impl HeartbeatConfig {
|
||||
})?
|
||||
.unwrap_or(settings.heartbeat.interval_secs),
|
||||
notify_channel: optional_env("HEARTBEAT_NOTIFY_CHANNEL")?
|
||||
.or(settings.heartbeat.notify_channel.clone()),
|
||||
.or_else(|| settings.heartbeat.notify_channel.clone()),
|
||||
notify_user: optional_env("HEARTBEAT_NOTIFY_USER")?
|
||||
.or(settings.heartbeat.notify_user.clone()),
|
||||
.or_else(|| settings.heartbeat.notify_user.clone()),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Routines configuration.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RoutineConfig {
|
||||
/// Whether the routines system is enabled.
|
||||
pub enabled: bool,
|
||||
/// How often (seconds) to poll for cron routines that need firing.
|
||||
pub cron_check_interval_secs: u64,
|
||||
/// Max routines executing concurrently across all users.
|
||||
pub max_concurrent_routines: usize,
|
||||
/// Default cooldown between fires (seconds).
|
||||
pub default_cooldown_secs: u64,
|
||||
/// Max output tokens for lightweight routine LLM calls.
|
||||
pub max_lightweight_tokens: u32,
|
||||
}
|
||||
|
||||
impl Default for RoutineConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: true,
|
||||
cron_check_interval_secs: 15,
|
||||
max_concurrent_routines: 10,
|
||||
default_cooldown_secs: 300,
|
||||
max_lightweight_tokens: 4096,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RoutineConfig {
|
||||
fn resolve() -> Result<Self, ConfigError> {
|
||||
Ok(Self {
|
||||
enabled: optional_env("ROUTINES_ENABLED")?
|
||||
.map(|s| s.parse())
|
||||
.transpose()
|
||||
.map_err(|e| ConfigError::InvalidValue {
|
||||
key: "ROUTINES_ENABLED".to_string(),
|
||||
message: format!("must be 'true' or 'false': {e}"),
|
||||
})?
|
||||
.unwrap_or(true),
|
||||
cron_check_interval_secs: parse_optional_env("ROUTINES_CRON_INTERVAL", 15)?,
|
||||
max_concurrent_routines: parse_optional_env("ROUTINES_MAX_CONCURRENT", 10)?,
|
||||
default_cooldown_secs: parse_optional_env("ROUTINES_DEFAULT_COOLDOWN", 300)?,
|
||||
max_lightweight_tokens: parse_optional_env("ROUTINES_MAX_TOKENS", 4096)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -879,7 +1092,7 @@ pub struct SandboxModeConfig {
|
||||
impl Default for SandboxModeConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: true, // Enabled by default
|
||||
enabled: true,
|
||||
policy: "readonly".to_string(),
|
||||
timeout_secs: 120,
|
||||
memory_limit_mb: 2048,
|
||||
@@ -892,7 +1105,7 @@ impl Default for SandboxModeConfig {
|
||||
}
|
||||
|
||||
impl SandboxModeConfig {
|
||||
fn from_env() -> Result<Self, ConfigError> {
|
||||
fn resolve() -> Result<Self, ConfigError> {
|
||||
let extra_domains = optional_env("SANDBOX_EXTRA_DOMAINS")?
|
||||
.map(|s| s.split(',').map(|d| d.trim().to_string()).collect())
|
||||
.unwrap_or_default();
|
||||
@@ -948,6 +1161,60 @@ impl SandboxModeConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// Claude Code sandbox configuration.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ClaudeCodeConfig {
|
||||
/// Whether Claude Code sandbox mode is available.
|
||||
pub enabled: bool,
|
||||
/// Host directory containing Claude auth session (mounted read-only).
|
||||
pub config_dir: std::path::PathBuf,
|
||||
/// Claude model to use (e.g. "sonnet", "opus").
|
||||
pub model: String,
|
||||
/// Maximum agentic turns before stopping.
|
||||
pub max_turns: u32,
|
||||
/// Memory limit in MB for Claude Code containers (heavier than workers).
|
||||
pub memory_limit_mb: u64,
|
||||
}
|
||||
|
||||
impl Default for ClaudeCodeConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: false,
|
||||
config_dir: dirs::home_dir()
|
||||
.unwrap_or_else(|| std::path::PathBuf::from("."))
|
||||
.join(".claude"),
|
||||
model: "sonnet".to_string(),
|
||||
max_turns: 50,
|
||||
memory_limit_mb: 4096,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ClaudeCodeConfig {
|
||||
fn resolve() -> Result<Self, ConfigError> {
|
||||
let defaults = Self::default();
|
||||
Ok(Self {
|
||||
enabled: optional_env("CLAUDE_CODE_ENABLED")?
|
||||
.map(|s| s.parse())
|
||||
.transpose()
|
||||
.map_err(|e| ConfigError::InvalidValue {
|
||||
key: "CLAUDE_CODE_ENABLED".to_string(),
|
||||
message: format!("must be 'true' or 'false': {e}"),
|
||||
})?
|
||||
.unwrap_or(defaults.enabled),
|
||||
config_dir: optional_env("CLAUDE_CONFIG_DIR")?
|
||||
.map(std::path::PathBuf::from)
|
||||
.unwrap_or(defaults.config_dir),
|
||||
model: optional_env("CLAUDE_CODE_MODEL")?.unwrap_or(defaults.model),
|
||||
max_turns: parse_optional_env("CLAUDE_CODE_MAX_TURNS", defaults.max_turns)?,
|
||||
memory_limit_mb: parse_optional_env(
|
||||
"CLAUDE_CODE_MEMORY_LIMIT_MB",
|
||||
defaults.memory_limit_mb,
|
||||
)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Helper functions
|
||||
|
||||
fn optional_env(key: &str) -> Result<Option<String>, ConfigError> {
|
||||
|
||||
@@ -39,6 +39,12 @@ pub enum Error {
|
||||
|
||||
#[error("Workspace error: {0}")]
|
||||
Workspace(#[from] WorkspaceError),
|
||||
|
||||
#[error("Orchestrator error: {0}")]
|
||||
Orchestrator(#[from] OrchestratorError),
|
||||
|
||||
#[error("Worker error: {0}")]
|
||||
Worker(#[from] WorkerError),
|
||||
}
|
||||
|
||||
/// Configuration-related errors.
|
||||
@@ -308,5 +314,52 @@ pub enum WorkspaceError {
|
||||
HeartbeatError { reason: String },
|
||||
}
|
||||
|
||||
/// Orchestrator errors (internal API, container management).
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum OrchestratorError {
|
||||
#[error("Container creation failed for job {job_id}: {reason}")]
|
||||
ContainerCreationFailed { job_id: Uuid, reason: String },
|
||||
|
||||
#[error("Container not found for job {job_id}")]
|
||||
ContainerNotFound { job_id: Uuid },
|
||||
|
||||
#[error("Container for job {job_id} is in unexpected state: {state}")]
|
||||
InvalidContainerState { job_id: Uuid, state: String },
|
||||
|
||||
#[error("Worker authentication failed: {reason}")]
|
||||
AuthFailed { reason: String },
|
||||
|
||||
#[error("Internal API error: {reason}")]
|
||||
ApiError { reason: String },
|
||||
|
||||
#[error("Docker error: {reason}")]
|
||||
Docker { reason: String },
|
||||
|
||||
#[error("Job {job_id} timed out in container")]
|
||||
ContainerTimeout { job_id: Uuid },
|
||||
}
|
||||
|
||||
/// Worker errors (container-side execution).
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum WorkerError {
|
||||
#[error("Failed to connect to orchestrator at {url}: {reason}")]
|
||||
ConnectionFailed { url: String, reason: String },
|
||||
|
||||
#[error("LLM proxy request failed: {reason}")]
|
||||
LlmProxyFailed { reason: String },
|
||||
|
||||
#[error("Secret resolution failed for {secret_name}: {reason}")]
|
||||
SecretResolveFailed { secret_name: String, reason: String },
|
||||
|
||||
#[error("Orchestrator returned error for job {job_id}: {reason}")]
|
||||
OrchestratorRejected { job_id: Uuid, reason: String },
|
||||
|
||||
#[error("Worker execution failed: {reason}")]
|
||||
ExecutionFailed { reason: String },
|
||||
|
||||
#[error("Missing worker token (IRONCLAW_WORKER_TOKEN not set)")]
|
||||
MissingToken,
|
||||
}
|
||||
|
||||
/// Result type alias for the agent.
|
||||
pub type Result<T> = std::result::Result<T, Error>;
|
||||
|
||||
+62
-10
@@ -23,9 +23,7 @@ use crate::tools::mcp::auth::{
|
||||
PkceChallenge, authorize_mcp_server, build_authorization_url, discover_full_oauth_metadata,
|
||||
find_available_port, is_authenticated, register_client,
|
||||
};
|
||||
use crate::tools::mcp::config::{
|
||||
McpServerConfig, add_mcp_server, get_mcp_server, load_mcp_servers, remove_mcp_server,
|
||||
};
|
||||
use crate::tools::mcp::config::McpServerConfig;
|
||||
use crate::tools::mcp::session::McpSessionManager;
|
||||
use crate::tools::wasm::{WasmToolLoader, WasmToolRuntime, discover_tools};
|
||||
|
||||
@@ -58,6 +56,8 @@ pub struct ExtensionManager {
|
||||
/// Tunnel URL for remote OAuth callbacks (used in future iterations).
|
||||
_tunnel_url: Option<String>,
|
||||
user_id: String,
|
||||
/// Optional database store for DB-backed MCP config.
|
||||
store: Option<Arc<crate::history::Store>>,
|
||||
}
|
||||
|
||||
impl ExtensionManager {
|
||||
@@ -71,6 +71,7 @@ impl ExtensionManager {
|
||||
wasm_channels_dir: PathBuf,
|
||||
tunnel_url: Option<String>,
|
||||
user_id: String,
|
||||
store: Option<Arc<crate::history::Store>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
registry: ExtensionRegistry::new(),
|
||||
@@ -85,6 +86,7 @@ impl ExtensionManager {
|
||||
pending_auth: RwLock::new(HashMap::new()),
|
||||
_tunnel_url: tunnel_url,
|
||||
user_id,
|
||||
store,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -191,7 +193,7 @@ impl ExtensionManager {
|
||||
|
||||
// List MCP servers
|
||||
if kind_filter.is_none() || kind_filter == Some(ExtensionKind::McpServer) {
|
||||
match load_mcp_servers().await {
|
||||
match self.load_mcp_servers().await {
|
||||
Ok(servers) => {
|
||||
for server in &servers.servers {
|
||||
let authenticated =
|
||||
@@ -304,7 +306,7 @@ impl ExtensionManager {
|
||||
self.mcp_clients.write().await.remove(name);
|
||||
|
||||
// Remove from config
|
||||
remove_mcp_server(name)
|
||||
self.remove_mcp_server(name)
|
||||
.await
|
||||
.map_err(|e| ExtensionError::Config(e.to_string()))?;
|
||||
|
||||
@@ -342,6 +344,54 @@ impl ExtensionManager {
|
||||
}
|
||||
}
|
||||
|
||||
// ── MCP config helpers (DB with disk fallback) ─────────────────────
|
||||
|
||||
async fn load_mcp_servers(
|
||||
&self,
|
||||
) -> Result<crate::tools::mcp::config::McpServersFile, crate::tools::mcp::config::ConfigError>
|
||||
{
|
||||
if let Some(ref store) = self.store {
|
||||
crate::tools::mcp::config::load_mcp_servers_from_db(store, &self.user_id).await
|
||||
} else {
|
||||
crate::tools::mcp::config::load_mcp_servers().await
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_mcp_server(
|
||||
&self,
|
||||
name: &str,
|
||||
) -> Result<McpServerConfig, crate::tools::mcp::config::ConfigError> {
|
||||
let servers = self.load_mcp_servers().await?;
|
||||
servers.get(name).cloned().ok_or_else(|| {
|
||||
crate::tools::mcp::config::ConfigError::ServerNotFound {
|
||||
name: name.to_string(),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async fn add_mcp_server(
|
||||
&self,
|
||||
config: McpServerConfig,
|
||||
) -> Result<(), crate::tools::mcp::config::ConfigError> {
|
||||
config.validate()?;
|
||||
if let Some(ref store) = self.store {
|
||||
crate::tools::mcp::config::add_mcp_server_db(store, &self.user_id, config).await
|
||||
} else {
|
||||
crate::tools::mcp::config::add_mcp_server(config).await
|
||||
}
|
||||
}
|
||||
|
||||
async fn remove_mcp_server(
|
||||
&self,
|
||||
name: &str,
|
||||
) -> Result<(), crate::tools::mcp::config::ConfigError> {
|
||||
if let Some(ref store) = self.store {
|
||||
crate::tools::mcp::config::remove_mcp_server_db(store, &self.user_id, name).await
|
||||
} else {
|
||||
crate::tools::mcp::config::remove_mcp_server(name).await
|
||||
}
|
||||
}
|
||||
|
||||
// ── Private helpers ──────────────────────────────────────────────────
|
||||
|
||||
async fn install_from_entry(
|
||||
@@ -381,7 +431,7 @@ impl ExtensionManager {
|
||||
url: &str,
|
||||
) -> Result<InstallResult, ExtensionError> {
|
||||
// Check if already installed
|
||||
if get_mcp_server(name).await.is_ok() {
|
||||
if self.get_mcp_server(name).await.is_ok() {
|
||||
return Err(ExtensionError::AlreadyInstalled(name.to_string()));
|
||||
}
|
||||
|
||||
@@ -390,7 +440,7 @@ impl ExtensionManager {
|
||||
.validate()
|
||||
.map_err(|e| ExtensionError::InvalidUrl(e.to_string()))?;
|
||||
|
||||
add_mcp_server(config)
|
||||
self.add_mcp_server(config)
|
||||
.await
|
||||
.map_err(|e| ExtensionError::Config(e.to_string()))?;
|
||||
|
||||
@@ -465,7 +515,8 @@ impl ExtensionManager {
|
||||
name: &str,
|
||||
token: Option<&str>,
|
||||
) -> Result<AuthResult, ExtensionError> {
|
||||
let server = get_mcp_server(name)
|
||||
let server = self
|
||||
.get_mcp_server(name)
|
||||
.await
|
||||
.map_err(|e| ExtensionError::NotInstalled(e.to_string()))?;
|
||||
|
||||
@@ -784,7 +835,8 @@ impl ExtensionManager {
|
||||
}
|
||||
}
|
||||
|
||||
let server = get_mcp_server(name)
|
||||
let server = self
|
||||
.get_mcp_server(name)
|
||||
.await
|
||||
.map_err(|e| ExtensionError::NotInstalled(e.to_string()))?;
|
||||
|
||||
@@ -893,7 +945,7 @@ impl ExtensionManager {
|
||||
/// Determine what kind of installed extension this is.
|
||||
async fn determine_installed_kind(&self, name: &str) -> Result<ExtensionKind, ExtensionError> {
|
||||
// Check MCP servers first
|
||||
if get_mcp_server(name).await.is_ok() {
|
||||
if self.get_mcp_server(name).await.is_ok() {
|
||||
return Ok(ExtensionKind::McpServer);
|
||||
}
|
||||
|
||||
|
||||
+4
-1
@@ -9,4 +9,7 @@ mod analytics;
|
||||
mod store;
|
||||
|
||||
pub use analytics::{JobStats, ToolStats};
|
||||
pub use store::{LlmCallRecord, Store};
|
||||
pub use store::{
|
||||
ConversationMessage, ConversationSummary, JobEventRecord, LlmCallRecord, SandboxJobRecord,
|
||||
SandboxJobSummary, Store,
|
||||
};
|
||||
|
||||
+1117
-6
File diff suppressed because it is too large
Load Diff
@@ -39,6 +39,7 @@
|
||||
//! - **Continuous learning** - Improve estimates from historical data
|
||||
|
||||
pub mod agent;
|
||||
pub mod bootstrap;
|
||||
pub mod channels;
|
||||
pub mod cli;
|
||||
pub mod config;
|
||||
@@ -49,12 +50,15 @@ pub mod evaluation;
|
||||
pub mod extensions;
|
||||
pub mod history;
|
||||
pub mod llm;
|
||||
pub mod orchestrator;
|
||||
pub mod pairing;
|
||||
pub mod safety;
|
||||
pub mod sandbox;
|
||||
pub mod secrets;
|
||||
pub mod settings;
|
||||
pub mod setup;
|
||||
pub mod tools;
|
||||
pub mod worker;
|
||||
pub mod workspace;
|
||||
|
||||
pub use config::Config;
|
||||
|
||||
@@ -0,0 +1,124 @@
|
||||
//! Per-model cost lookup table for multi-provider LLM support.
|
||||
//!
|
||||
//! Returns (input_cost_per_token, output_cost_per_token) as Decimal pairs.
|
||||
//! Ollama and other local models return zero cost.
|
||||
|
||||
use rust_decimal::Decimal;
|
||||
use rust_decimal_macros::dec;
|
||||
|
||||
/// Look up known per-token costs for a model by its identifier.
|
||||
///
|
||||
/// Returns `Some((input_cost, output_cost))` for known models, `None` otherwise.
|
||||
pub fn model_cost(model_id: &str) -> Option<(Decimal, Decimal)> {
|
||||
// Normalize: strip provider prefixes (e.g., "openai/gpt-4o" -> "gpt-4o")
|
||||
let id = model_id
|
||||
.rsplit_once('/')
|
||||
.map(|(_, name)| name)
|
||||
.unwrap_or(model_id);
|
||||
|
||||
match id {
|
||||
// OpenAI models -- prices per token (USD)
|
||||
"gpt-4o" | "gpt-4o-2024-11-20" | "gpt-4o-2024-08-06" => {
|
||||
Some((dec!(0.0000025), dec!(0.00001)))
|
||||
}
|
||||
"gpt-4o-mini" | "gpt-4o-mini-2024-07-18" => Some((dec!(0.00000015), dec!(0.0000006))),
|
||||
"gpt-4-turbo" | "gpt-4-turbo-2024-04-09" => Some((dec!(0.00001), dec!(0.00003))),
|
||||
"gpt-4" | "gpt-4-0613" => Some((dec!(0.00003), dec!(0.00006))),
|
||||
"gpt-3.5-turbo" | "gpt-3.5-turbo-0125" => Some((dec!(0.0000005), dec!(0.0000015))),
|
||||
"o1" | "o1-2024-12-17" => Some((dec!(0.000015), dec!(0.00006))),
|
||||
"o1-mini" | "o1-mini-2024-09-12" => Some((dec!(0.000003), dec!(0.000012))),
|
||||
"o3-mini" | "o3-mini-2025-01-31" => Some((dec!(0.0000011), dec!(0.0000044))),
|
||||
|
||||
// Anthropic models
|
||||
"claude-3-5-sonnet-20241022" | "claude-3-5-sonnet-latest" | "claude-sonnet-4-20250514" => {
|
||||
Some((dec!(0.000003), dec!(0.000015)))
|
||||
}
|
||||
"claude-3-5-haiku-20241022" | "claude-3-5-haiku-latest" => {
|
||||
Some((dec!(0.0000008), dec!(0.000004)))
|
||||
}
|
||||
"claude-3-opus-20240229" | "claude-3-opus-latest" | "claude-opus-4-20250514" => {
|
||||
Some((dec!(0.000015), dec!(0.000075)))
|
||||
}
|
||||
"claude-3-haiku-20240307" => Some((dec!(0.00000025), dec!(0.00000125))),
|
||||
|
||||
// Ollama / local models -- free
|
||||
_ if is_local_model(id) => Some((Decimal::ZERO, Decimal::ZERO)),
|
||||
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Default cost for unknown models.
|
||||
pub fn default_cost() -> (Decimal, Decimal) {
|
||||
// Conservative estimate: roughly GPT-4o pricing
|
||||
(dec!(0.0000025), dec!(0.00001))
|
||||
}
|
||||
|
||||
/// Heuristic to detect local/self-hosted models (Ollama, llama.cpp, etc.).
|
||||
fn is_local_model(model_id: &str) -> bool {
|
||||
let lower = model_id.to_lowercase();
|
||||
lower.starts_with("llama")
|
||||
|| lower.starts_with("mistral")
|
||||
|| lower.starts_with("mixtral")
|
||||
|| lower.starts_with("phi")
|
||||
|| lower.starts_with("gemma")
|
||||
|| lower.starts_with("qwen")
|
||||
|| lower.starts_with("codellama")
|
||||
|| lower.starts_with("deepseek")
|
||||
|| lower.starts_with("starcoder")
|
||||
|| lower.starts_with("vicuna")
|
||||
|| lower.starts_with("yi")
|
||||
|| lower.contains(":latest")
|
||||
|| lower.contains(":instruct")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_known_model_costs() {
|
||||
let (input, output) = model_cost("gpt-4o").unwrap();
|
||||
assert!(input > Decimal::ZERO);
|
||||
assert!(output > input);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_claude_costs() {
|
||||
let (input, output) = model_cost("claude-3-5-sonnet-20241022").unwrap();
|
||||
assert!(input > Decimal::ZERO);
|
||||
assert!(output > input);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_local_model_free() {
|
||||
let (input, output) = model_cost("llama3").unwrap();
|
||||
assert_eq!(input, Decimal::ZERO);
|
||||
assert_eq!(output, Decimal::ZERO);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ollama_tagged_model_free() {
|
||||
let (input, output) = model_cost("mistral:latest").unwrap();
|
||||
assert_eq!(input, Decimal::ZERO);
|
||||
assert_eq!(output, Decimal::ZERO);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_unknown_model_returns_none() {
|
||||
assert!(model_cost("some-totally-unknown-model-xyz").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_default_cost_nonzero() {
|
||||
let (input, output) = default_cost();
|
||||
assert!(input > Decimal::ZERO);
|
||||
assert!(output > Decimal::ZERO);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_provider_prefix_stripped() {
|
||||
// "openai/gpt-4o" should resolve to same as "gpt-4o"
|
||||
assert_eq!(model_cost("openai/gpt-4o"), model_cost("gpt-4o"));
|
||||
}
|
||||
}
|
||||
+134
-10
@@ -1,48 +1,172 @@
|
||||
//! LLM integration for the agent.
|
||||
//!
|
||||
//! Supports two API modes:
|
||||
//! - **Responses API** (chat-api): Session-based auth, uses `/v1/responses` endpoint
|
||||
//! - **Chat Completions API** (cloud-api): API key auth, uses `/v1/chat/completions` endpoint
|
||||
//! Supports multiple backends:
|
||||
//! - **NEAR AI** (default): Session-based or API key auth via NEAR AI proxy
|
||||
//! - **OpenAI**: Direct API access with your own key
|
||||
//! - **Anthropic**: Direct API access with your own key
|
||||
//! - **Ollama**: Local model inference
|
||||
//! - **OpenAI-compatible**: Any endpoint that speaks the OpenAI API
|
||||
|
||||
mod costs;
|
||||
mod nearai;
|
||||
mod nearai_chat;
|
||||
mod provider;
|
||||
mod reasoning;
|
||||
mod rig_adapter;
|
||||
pub mod session;
|
||||
|
||||
pub use nearai::{ModelInfo, NearAiProvider};
|
||||
pub use nearai_chat::NearAiChatProvider;
|
||||
pub use provider::{
|
||||
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, Role, ToolCall,
|
||||
ToolCompletionRequest, ToolCompletionResponse, ToolDefinition, ToolResult,
|
||||
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, ModelMetadata,
|
||||
Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse, ToolDefinition, ToolResult,
|
||||
};
|
||||
pub use reasoning::{ActionPlan, Reasoning, ReasoningContext, RespondResult, ToolSelection};
|
||||
pub use rig_adapter::RigAdapter;
|
||||
pub use session::{SessionConfig, SessionManager, create_session_manager};
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::config::{LlmConfig, NearAiApiMode};
|
||||
use rig::client::CompletionClient;
|
||||
use secrecy::ExposeSecret;
|
||||
|
||||
use crate::config::{LlmBackend, LlmConfig, NearAiApiMode};
|
||||
use crate::error::LlmError;
|
||||
|
||||
/// Create an LLM provider based on configuration.
|
||||
///
|
||||
/// - For `Responses` mode: Requires a session manager for authentication
|
||||
/// - For `ChatCompletions` mode: Uses API key from config (session not needed)
|
||||
/// - `NearAi` backend: Uses session manager for authentication (Responses API)
|
||||
/// or API key (Chat Completions API)
|
||||
/// - Other backends: Use rig-core adapter with provider-specific clients
|
||||
pub fn create_llm_provider(
|
||||
config: &LlmConfig,
|
||||
session: Arc<SessionManager>,
|
||||
) -> Result<Arc<dyn LlmProvider>, LlmError> {
|
||||
match config.backend {
|
||||
LlmBackend::NearAi => create_nearai_provider(config, session),
|
||||
LlmBackend::OpenAi => create_openai_provider(config),
|
||||
LlmBackend::Anthropic => create_anthropic_provider(config),
|
||||
LlmBackend::Ollama => create_ollama_provider(config),
|
||||
LlmBackend::OpenAiCompatible => create_openai_compatible_provider(config),
|
||||
}
|
||||
}
|
||||
|
||||
fn create_nearai_provider(
|
||||
config: &LlmConfig,
|
||||
session: Arc<SessionManager>,
|
||||
) -> Result<Arc<dyn LlmProvider>, LlmError> {
|
||||
match config.nearai.api_mode {
|
||||
NearAiApiMode::Responses => {
|
||||
tracing::info!("Using Responses API (chat-api) with session auth");
|
||||
tracing::info!("Using NEAR AI Responses API (chat-api) with session auth");
|
||||
Ok(Arc::new(NearAiProvider::new(
|
||||
config.nearai.clone(),
|
||||
session,
|
||||
)))
|
||||
}
|
||||
NearAiApiMode::ChatCompletions => {
|
||||
tracing::info!("Using Chat Completions API (cloud-api) with API key auth");
|
||||
tracing::info!("Using NEAR AI Chat Completions API (cloud-api) with API key auth");
|
||||
Ok(Arc::new(NearAiChatProvider::new(config.nearai.clone())?))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn create_openai_provider(config: &LlmConfig) -> Result<Arc<dyn LlmProvider>, LlmError> {
|
||||
let oai = config.openai.as_ref().ok_or_else(|| LlmError::AuthFailed {
|
||||
provider: "openai".to_string(),
|
||||
})?;
|
||||
|
||||
use rig::providers::openai;
|
||||
|
||||
let client: openai::Client =
|
||||
openai::Client::new(oai.api_key.expose_secret()).map_err(|e| LlmError::RequestFailed {
|
||||
provider: "openai".to_string(),
|
||||
reason: format!("Failed to create OpenAI client: {}", e),
|
||||
})?;
|
||||
|
||||
let model = client.completion_model(&oai.model);
|
||||
tracing::info!("Using OpenAI direct API (model: {})", oai.model);
|
||||
Ok(Arc::new(RigAdapter::new(model, &oai.model)))
|
||||
}
|
||||
|
||||
fn create_anthropic_provider(config: &LlmConfig) -> Result<Arc<dyn LlmProvider>, LlmError> {
|
||||
let anth = config
|
||||
.anthropic
|
||||
.as_ref()
|
||||
.ok_or_else(|| LlmError::AuthFailed {
|
||||
provider: "anthropic".to_string(),
|
||||
})?;
|
||||
|
||||
use rig::providers::anthropic;
|
||||
|
||||
let client: anthropic::Client =
|
||||
anthropic::Client::new(anth.api_key.expose_secret()).map_err(|e| {
|
||||
LlmError::RequestFailed {
|
||||
provider: "anthropic".to_string(),
|
||||
reason: format!("Failed to create Anthropic client: {}", e),
|
||||
}
|
||||
})?;
|
||||
|
||||
let model = client.completion_model(&anth.model);
|
||||
tracing::info!("Using Anthropic direct API (model: {})", anth.model);
|
||||
Ok(Arc::new(RigAdapter::new(model, &anth.model)))
|
||||
}
|
||||
|
||||
fn create_ollama_provider(config: &LlmConfig) -> Result<Arc<dyn LlmProvider>, LlmError> {
|
||||
let oll = config.ollama.as_ref().ok_or_else(|| LlmError::AuthFailed {
|
||||
provider: "ollama".to_string(),
|
||||
})?;
|
||||
|
||||
use rig::client::Nothing;
|
||||
use rig::providers::ollama;
|
||||
|
||||
let client: ollama::Client = ollama::Client::builder()
|
||||
.base_url(&oll.base_url)
|
||||
.api_key(Nothing)
|
||||
.build()
|
||||
.map_err(|e| LlmError::RequestFailed {
|
||||
provider: "ollama".to_string(),
|
||||
reason: format!("Failed to create Ollama client: {}", e),
|
||||
})?;
|
||||
|
||||
let model = client.completion_model(&oll.model);
|
||||
tracing::info!(
|
||||
"Using Ollama (base_url: {}, model: {})",
|
||||
oll.base_url,
|
||||
oll.model
|
||||
);
|
||||
Ok(Arc::new(RigAdapter::new(model, &oll.model)))
|
||||
}
|
||||
|
||||
fn create_openai_compatible_provider(config: &LlmConfig) -> Result<Arc<dyn LlmProvider>, LlmError> {
|
||||
let compat = config
|
||||
.openai_compatible
|
||||
.as_ref()
|
||||
.ok_or_else(|| LlmError::AuthFailed {
|
||||
provider: "openai_compatible".to_string(),
|
||||
})?;
|
||||
|
||||
use rig::providers::openai;
|
||||
|
||||
let api_key = compat
|
||||
.api_key
|
||||
.as_ref()
|
||||
.map(|k| k.expose_secret().to_string())
|
||||
.unwrap_or_else(|| "no-key".to_string());
|
||||
|
||||
let client: openai::Client = openai::Client::builder()
|
||||
.base_url(&compat.base_url)
|
||||
.api_key(api_key)
|
||||
.build()
|
||||
.map_err(|e| LlmError::RequestFailed {
|
||||
provider: "openai_compatible".to_string(),
|
||||
reason: format!("Failed to create OpenAI-compatible client: {}", e),
|
||||
})?;
|
||||
|
||||
let model = client.completion_model(&compat.model);
|
||||
tracing::info!(
|
||||
"Using OpenAI-compatible endpoint (base_url: {}, model: {})",
|
||||
compat.base_url,
|
||||
compat.model
|
||||
);
|
||||
Ok(Arc::new(RigAdapter::new(model, &compat.model)))
|
||||
}
|
||||
|
||||
+308
-20
@@ -3,6 +3,7 @@
|
||||
//! This provider uses the NEAR AI chat-api which provides a unified interface
|
||||
//! to multiple LLM models (OpenAI, Anthropic, etc.) with user authentication.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
@@ -31,11 +32,23 @@ pub struct ModelInfo {
|
||||
pub provider: Option<String>,
|
||||
}
|
||||
|
||||
/// Per-thread chaining state: the last response ID and how many input
|
||||
/// messages were included in that request. This lets subsequent calls send
|
||||
/// only the delta (new messages since last call).
|
||||
struct ChainState {
|
||||
response_id: String,
|
||||
input_count: usize,
|
||||
}
|
||||
|
||||
/// NEAR AI Chat API provider.
|
||||
pub struct NearAiProvider {
|
||||
client: Client,
|
||||
config: NearAiConfig,
|
||||
session: Arc<SessionManager>,
|
||||
active_model: std::sync::RwLock<String>,
|
||||
/// Per-thread response ID chaining state.
|
||||
/// Key is thread_id from request metadata.
|
||||
response_chains: std::sync::RwLock<HashMap<String, ChainState>>,
|
||||
}
|
||||
|
||||
impl NearAiProvider {
|
||||
@@ -46,13 +59,64 @@ impl NearAiProvider {
|
||||
.build()
|
||||
.unwrap_or_else(|_| Client::new());
|
||||
|
||||
let active_model = std::sync::RwLock::new(config.model.clone());
|
||||
Self {
|
||||
client,
|
||||
config,
|
||||
session,
|
||||
active_model,
|
||||
response_chains: std::sync::RwLock::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Seed a response chain for a thread (e.g. when restoring from DB).
|
||||
pub fn seed_response_id(&self, thread_id: &str, response_id: String) {
|
||||
let mut chains = self
|
||||
.response_chains
|
||||
.write()
|
||||
.expect("response_chains lock poisoned");
|
||||
chains.insert(
|
||||
thread_id.to_string(),
|
||||
ChainState {
|
||||
response_id,
|
||||
input_count: 0,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
/// Get the last response ID for a thread (for persistence).
|
||||
pub fn get_response_id(&self, thread_id: &str) -> Option<String> {
|
||||
let chains = self
|
||||
.response_chains
|
||||
.read()
|
||||
.expect("response_chains lock poisoned");
|
||||
chains.get(thread_id).map(|c| c.response_id.clone())
|
||||
}
|
||||
|
||||
/// Store a response chain state after a successful call.
|
||||
fn store_chain(&self, thread_id: &str, response_id: String, input_count: usize) {
|
||||
let mut chains = self
|
||||
.response_chains
|
||||
.write()
|
||||
.expect("response_chains lock poisoned");
|
||||
chains.insert(
|
||||
thread_id.to_string(),
|
||||
ChainState {
|
||||
response_id,
|
||||
input_count,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
/// Clear the chain for a thread (on error / fallback).
|
||||
fn clear_chain(&self, thread_id: &str) {
|
||||
let mut chains = self
|
||||
.response_chains
|
||||
.write()
|
||||
.expect("response_chains lock poisoned");
|
||||
chains.remove(thread_id);
|
||||
}
|
||||
|
||||
fn api_url(&self, path: &str) -> String {
|
||||
format!(
|
||||
"{}/v1/{}",
|
||||
@@ -291,18 +355,34 @@ impl NearAiProvider {
|
||||
}
|
||||
}
|
||||
|
||||
/// Split messages into system instructions and non-system input messages.
|
||||
/// Split messages into system instructions and non-system input items.
|
||||
/// The OpenAI Responses API expects system prompts in an `instructions` field,
|
||||
/// not as a message with role "system" in the input array.
|
||||
fn split_messages(messages: Vec<ChatMessage>) -> (Option<String>, Vec<NearAiMessage>) {
|
||||
///
|
||||
/// When `chaining` is true, tool result messages (role=tool) are converted to
|
||||
/// `NearAiInputItem::FunctionCallOutput` for the Responses API protocol.
|
||||
fn split_messages(
|
||||
messages: Vec<ChatMessage>,
|
||||
chaining: bool,
|
||||
) -> (Option<String>, Vec<NearAiInputItem>) {
|
||||
let mut instructions: Vec<String> = Vec::new();
|
||||
let mut input: Vec<NearAiMessage> = Vec::new();
|
||||
let mut input: Vec<NearAiInputItem> = Vec::new();
|
||||
|
||||
for msg in messages {
|
||||
if msg.role == Role::System {
|
||||
instructions.push(msg.content);
|
||||
} else if chaining && msg.role == Role::Tool {
|
||||
if let Some(ref call_id) = msg.tool_call_id {
|
||||
input.push(NearAiInputItem::FunctionCallOutput {
|
||||
item_type: "function_call_output".to_string(),
|
||||
call_id: call_id.clone(),
|
||||
output: msg.content,
|
||||
});
|
||||
} else {
|
||||
input.push(NearAiInputItem::Message(msg.into()));
|
||||
}
|
||||
} else {
|
||||
input.push(msg.into());
|
||||
input.push(NearAiInputItem::Message(msg.into()));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -318,12 +398,14 @@ fn split_messages(messages: Vec<ChatMessage>) -> (Option<String>, Vec<NearAiMess
|
||||
#[async_trait]
|
||||
impl LlmProvider for NearAiProvider {
|
||||
async fn complete(&self, req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||
let (instructions, input) = split_messages(req.messages);
|
||||
let thread_id = req.metadata.get("thread_id").cloned();
|
||||
let (instructions, input) = split_messages(req.messages, false);
|
||||
|
||||
let request = NearAiRequest {
|
||||
model: self.config.model.clone(),
|
||||
model: self.active_model_name(),
|
||||
instructions,
|
||||
input,
|
||||
previous_response_id: None,
|
||||
temperature: req.temperature,
|
||||
max_output_tokens: req.max_tokens,
|
||||
stream: Some(false),
|
||||
@@ -350,6 +432,7 @@ impl LlmProvider for NearAiProvider {
|
||||
finish_reason: FinishReason::Stop,
|
||||
input_tokens: usage.input_tokens,
|
||||
output_tokens: usage.output_tokens,
|
||||
response_id: None,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -367,6 +450,7 @@ impl LlmProvider for NearAiProvider {
|
||||
finish_reason: FinishReason::Stop,
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
response_id: None,
|
||||
});
|
||||
}
|
||||
Err(e) => return Err(e),
|
||||
@@ -423,11 +507,17 @@ impl LlmProvider for NearAiProvider {
|
||||
);
|
||||
}
|
||||
|
||||
// Store response ID for chaining
|
||||
if let Some(ref tid) = thread_id {
|
||||
self.store_chain(tid, response.id.clone(), 0);
|
||||
}
|
||||
|
||||
Ok(CompletionResponse {
|
||||
content: text,
|
||||
finish_reason: FinishReason::Stop,
|
||||
input_tokens: response.usage.input_tokens,
|
||||
output_tokens: response.usage.output_tokens,
|
||||
response_id: Some(response.id),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -435,7 +525,33 @@ impl LlmProvider for NearAiProvider {
|
||||
&self,
|
||||
req: ToolCompletionRequest,
|
||||
) -> Result<ToolCompletionResponse, LlmError> {
|
||||
let (instructions, input) = split_messages(req.messages);
|
||||
let thread_id = req.metadata.get("thread_id").cloned();
|
||||
|
||||
// Look up chaining state for this thread
|
||||
let chain_state = thread_id.as_ref().and_then(|tid| {
|
||||
let chains = self
|
||||
.response_chains
|
||||
.read()
|
||||
.expect("response_chains lock poisoned");
|
||||
chains
|
||||
.get(tid)
|
||||
.map(|c| (c.response_id.clone(), c.input_count))
|
||||
});
|
||||
|
||||
let chaining = chain_state.is_some();
|
||||
let (previous_response_id, prev_input_count) = chain_state
|
||||
.map(|(rid, count)| (Some(rid), count))
|
||||
.unwrap_or((None, 0));
|
||||
|
||||
// When chaining, only send new messages (the delta since last call).
|
||||
// Tool results are converted to function_call_output items.
|
||||
let (instructions, all_input) = split_messages(req.messages, chaining);
|
||||
let input = if chaining && all_input.len() > prev_input_count {
|
||||
all_input[prev_input_count..].to_vec()
|
||||
} else {
|
||||
all_input.clone()
|
||||
};
|
||||
let total_input_count = all_input.len();
|
||||
|
||||
let tools: Vec<NearAiTool> = req
|
||||
.tools
|
||||
@@ -449,18 +565,58 @@ impl LlmProvider for NearAiProvider {
|
||||
.collect();
|
||||
|
||||
let request = NearAiRequest {
|
||||
model: self.config.model.clone(),
|
||||
instructions,
|
||||
model: self.active_model_name(),
|
||||
instructions: if chaining { None } else { instructions.clone() },
|
||||
input,
|
||||
previous_response_id: previous_response_id.clone(),
|
||||
temperature: req.temperature,
|
||||
max_output_tokens: req.max_tokens,
|
||||
stream: Some(false),
|
||||
tools: if tools.is_empty() { None } else { Some(tools) },
|
||||
tools: if tools.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(tools.clone())
|
||||
},
|
||||
};
|
||||
|
||||
// Try to get structured response, fall back to alternative formats
|
||||
// Try to get structured response, fall back to alternative formats.
|
||||
// If chaining fails (bad previous_response_id), retry with full history.
|
||||
let response: NearAiResponse = match self.send_request("responses", &request).await {
|
||||
Ok(r) => r,
|
||||
Err(ref e) if chaining && is_chain_error(e) => {
|
||||
tracing::warn!(
|
||||
"Response chaining failed, retrying with full history: {}",
|
||||
e
|
||||
);
|
||||
if let Some(ref tid) = thread_id {
|
||||
self.clear_chain(tid);
|
||||
}
|
||||
let (instructions_full, input_full) = split_messages(
|
||||
// Rebuild from the original input (non-chaining mode)
|
||||
{
|
||||
let mut msgs = Vec::new();
|
||||
if let Some(ref instr) = instructions {
|
||||
msgs.push(ChatMessage::system(instr.clone()));
|
||||
}
|
||||
for item in &all_input {
|
||||
msgs.push(item.to_chat_message());
|
||||
}
|
||||
msgs
|
||||
},
|
||||
false,
|
||||
);
|
||||
let retry_request = NearAiRequest {
|
||||
model: self.active_model_name(),
|
||||
instructions: instructions_full,
|
||||
input: input_full,
|
||||
previous_response_id: None,
|
||||
temperature: request.temperature,
|
||||
max_output_tokens: request.max_output_tokens,
|
||||
stream: Some(false),
|
||||
tools: request.tools.clone(),
|
||||
};
|
||||
self.send_request("responses", &retry_request).await?
|
||||
}
|
||||
Err(LlmError::InvalidResponse { reason, .. }) if reason.contains("Raw: ") => {
|
||||
let raw_text = reason.split("Raw: ").nth(1).unwrap_or("");
|
||||
|
||||
@@ -490,6 +646,7 @@ impl LlmProvider for NearAiProvider {
|
||||
finish_reason,
|
||||
input_tokens: usage.input_tokens,
|
||||
output_tokens: usage.output_tokens,
|
||||
response_id: None,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -507,6 +664,7 @@ impl LlmProvider for NearAiProvider {
|
||||
finish_reason: FinishReason::Stop,
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
response_id: None,
|
||||
});
|
||||
}
|
||||
Err(e) => return Err(e),
|
||||
@@ -560,12 +718,18 @@ impl LlmProvider for NearAiProvider {
|
||||
FinishReason::ToolUse
|
||||
};
|
||||
|
||||
// Store response ID for chaining on subsequent calls
|
||||
if let Some(ref tid) = thread_id {
|
||||
self.store_chain(tid, response.id.clone(), total_input_count);
|
||||
}
|
||||
|
||||
Ok(ToolCompletionResponse {
|
||||
content: if text.is_empty() { None } else { Some(text) },
|
||||
tool_calls,
|
||||
finish_reason,
|
||||
input_tokens: response.usage.input_tokens,
|
||||
output_tokens: response.usage.output_tokens,
|
||||
response_id: Some(response.id),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -584,6 +748,30 @@ impl LlmProvider for NearAiProvider {
|
||||
let models = NearAiProvider::list_models(self).await?;
|
||||
Ok(models.into_iter().map(|m| m.name).collect())
|
||||
}
|
||||
|
||||
fn active_model_name(&self) -> String {
|
||||
self.active_model
|
||||
.read()
|
||||
.expect("active_model lock poisoned")
|
||||
.clone()
|
||||
}
|
||||
|
||||
fn set_model(&self, model: &str) -> Result<(), LlmError> {
|
||||
let mut guard = self
|
||||
.active_model
|
||||
.write()
|
||||
.expect("active_model lock poisoned");
|
||||
*guard = model.to_string();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn seed_response_chain(&self, thread_id: &str, response_id: String) {
|
||||
self.seed_response_id(thread_id, response_id);
|
||||
}
|
||||
|
||||
fn get_response_chain_id(&self, thread_id: &str) -> Option<String> {
|
||||
self.get_response_id(thread_id)
|
||||
}
|
||||
}
|
||||
|
||||
// NEAR AI API types
|
||||
@@ -597,8 +785,11 @@ struct NearAiRequest {
|
||||
/// System instructions (replaces sending system role in input)
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
instructions: Option<String>,
|
||||
/// Input messages (user/assistant/tool only, NOT system)
|
||||
input: Vec<NearAiMessage>,
|
||||
/// Input items: messages and/or function_call_output entries.
|
||||
input: Vec<NearAiInputItem>,
|
||||
/// Chain this request to a previous response (avoids resending full context).
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
previous_response_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
temperature: Option<f32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
@@ -609,7 +800,7 @@ struct NearAiRequest {
|
||||
tools: Option<Vec<NearAiTool>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
struct NearAiMessage {
|
||||
role: String,
|
||||
content: String,
|
||||
@@ -630,7 +821,68 @@ impl From<ChatMessage> for NearAiMessage {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
/// Input item for the Responses API. Either a regular message or a
|
||||
/// function_call_output (for returning tool results when chaining).
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
#[serde(untagged)]
|
||||
enum NearAiInputItem {
|
||||
Message(NearAiMessage),
|
||||
FunctionCallOutput {
|
||||
#[serde(rename = "type")]
|
||||
item_type: String,
|
||||
call_id: String,
|
||||
output: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl NearAiInputItem {
|
||||
/// Convert back to a ChatMessage (used for fallback retry).
|
||||
fn to_chat_message(&self) -> ChatMessage {
|
||||
match self {
|
||||
NearAiInputItem::Message(msg) => {
|
||||
let role = match msg.role.as_str() {
|
||||
"system" => Role::System,
|
||||
"user" => Role::User,
|
||||
"assistant" => Role::Assistant,
|
||||
"tool" => Role::Tool,
|
||||
_ => Role::User,
|
||||
};
|
||||
ChatMessage {
|
||||
role,
|
||||
content: msg.content.clone(),
|
||||
tool_call_id: None,
|
||||
name: None,
|
||||
tool_calls: None,
|
||||
}
|
||||
}
|
||||
NearAiInputItem::FunctionCallOutput {
|
||||
call_id, output, ..
|
||||
} => ChatMessage {
|
||||
role: Role::Tool,
|
||||
content: output.clone(),
|
||||
tool_call_id: Some(call_id.clone()),
|
||||
name: None,
|
||||
tool_calls: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if an LLM error is likely caused by an invalid previous_response_id.
|
||||
fn is_chain_error(err: &LlmError) -> bool {
|
||||
match err {
|
||||
LlmError::RequestFailed { reason, .. } => {
|
||||
let lower = reason.to_lowercase();
|
||||
lower.contains("previous_response_id")
|
||||
|| lower.contains("previous response")
|
||||
|| lower.contains("not found")
|
||||
|| lower.contains("invalid response id")
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
struct NearAiTool {
|
||||
#[serde(rename = "type")]
|
||||
tool_type: String,
|
||||
@@ -833,14 +1085,17 @@ mod tests {
|
||||
ChatMessage::user("Hello"),
|
||||
ChatMessage::assistant("Hi there!"),
|
||||
];
|
||||
let (instructions, input) = split_messages(messages);
|
||||
let (instructions, input) = split_messages(messages, false);
|
||||
assert_eq!(
|
||||
instructions,
|
||||
Some("You are a helpful assistant".to_string())
|
||||
);
|
||||
assert_eq!(input.len(), 2);
|
||||
assert_eq!(input[0].role, "user");
|
||||
assert_eq!(input[1].role, "assistant");
|
||||
// Verify the input items are messages
|
||||
match &input[0] {
|
||||
NearAiInputItem::Message(m) => assert_eq!(m.role, "user"),
|
||||
_ => panic!("expected Message"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -849,7 +1104,7 @@ mod tests {
|
||||
ChatMessage::user("Hello"),
|
||||
ChatMessage::assistant("Hi there!"),
|
||||
];
|
||||
let (instructions, input) = split_messages(messages);
|
||||
let (instructions, input) = split_messages(messages, false);
|
||||
assert!(instructions.is_none());
|
||||
assert_eq!(input.len(), 2);
|
||||
}
|
||||
@@ -861,11 +1116,44 @@ mod tests {
|
||||
ChatMessage::system("Second instruction"),
|
||||
ChatMessage::user("Hello"),
|
||||
];
|
||||
let (instructions, input) = split_messages(messages);
|
||||
let (instructions, input) = split_messages(messages, false);
|
||||
assert_eq!(
|
||||
instructions,
|
||||
Some("First instruction\n\nSecond instruction".to_string())
|
||||
);
|
||||
assert_eq!(input.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_split_messages_chaining_converts_tool_results() {
|
||||
let messages = vec![
|
||||
ChatMessage::user("Hello"),
|
||||
ChatMessage::tool_result("call_123", "my_tool", "result data"),
|
||||
];
|
||||
let (_, input) = split_messages(messages, true);
|
||||
assert_eq!(input.len(), 2);
|
||||
match &input[1] {
|
||||
NearAiInputItem::FunctionCallOutput {
|
||||
call_id, output, ..
|
||||
} => {
|
||||
assert_eq!(call_id, "call_123");
|
||||
assert_eq!(output, "result data");
|
||||
}
|
||||
_ => panic!("expected FunctionCallOutput"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_split_messages_no_chaining_keeps_tool_as_message() {
|
||||
let messages = vec![
|
||||
ChatMessage::user("Hello"),
|
||||
ChatMessage::tool_result("call_123", "my_tool", "result data"),
|
||||
];
|
||||
let (_, input) = split_messages(messages, false);
|
||||
assert_eq!(input.len(), 2);
|
||||
match &input[1] {
|
||||
NearAiInputItem::Message(m) => assert_eq!(m.role, "tool"),
|
||||
_ => panic!("expected Message"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+69
-18
@@ -13,14 +13,15 @@ use serde::{Deserialize, Serialize};
|
||||
use crate::config::NearAiConfig;
|
||||
use crate::error::LlmError;
|
||||
use crate::llm::provider::{
|
||||
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, Role, ToolCall,
|
||||
ToolCompletionRequest, ToolCompletionResponse,
|
||||
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, ModelMetadata,
|
||||
Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse,
|
||||
};
|
||||
|
||||
/// NEAR AI Chat Completions API provider.
|
||||
pub struct NearAiChatProvider {
|
||||
client: Client,
|
||||
config: NearAiConfig,
|
||||
active_model: std::sync::RwLock<String>,
|
||||
}
|
||||
|
||||
impl NearAiChatProvider {
|
||||
@@ -37,7 +38,12 @@ impl NearAiChatProvider {
|
||||
.build()
|
||||
.unwrap_or_else(|_| Client::new());
|
||||
|
||||
Ok(Self { client, config })
|
||||
let active_model = std::sync::RwLock::new(config.model.clone());
|
||||
Ok(Self {
|
||||
client,
|
||||
config,
|
||||
active_model,
|
||||
})
|
||||
}
|
||||
|
||||
fn api_url(&self, path: &str) -> String {
|
||||
@@ -65,6 +71,11 @@ impl NearAiChatProvider {
|
||||
|
||||
tracing::debug!("Sending request to NEAR AI Chat: {}", url);
|
||||
|
||||
// Log the request body for debugging tool call issues
|
||||
if let Ok(json) = serde_json::to_string(body) {
|
||||
tracing::debug!("NEAR AI Chat request body: {}", json);
|
||||
}
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
@@ -111,8 +122,8 @@ impl NearAiChatProvider {
|
||||
})
|
||||
}
|
||||
|
||||
/// Fetch available models.
|
||||
pub async fn list_models(&self) -> Result<Vec<String>, LlmError> {
|
||||
/// Fetch available models with full metadata from the `/v1/models` endpoint.
|
||||
async fn fetch_models(&self) -> Result<Vec<ApiModelEntry>, LlmError> {
|
||||
let url = self.api_url("models");
|
||||
|
||||
let response = self
|
||||
@@ -138,12 +149,7 @@ impl NearAiChatProvider {
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ModelsResponse {
|
||||
data: Vec<ModelEntry>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ModelEntry {
|
||||
id: String,
|
||||
data: Vec<ApiModelEntry>,
|
||||
}
|
||||
|
||||
let resp: ModelsResponse =
|
||||
@@ -152,10 +158,18 @@ impl NearAiChatProvider {
|
||||
reason: format!("JSON parse error: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(resp.data.into_iter().map(|m| m.id).collect())
|
||||
Ok(resp.data)
|
||||
}
|
||||
}
|
||||
|
||||
/// Model entry as returned by the `/v1/models` API.
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ApiModelEntry {
|
||||
id: String,
|
||||
#[serde(default)]
|
||||
context_length: Option<u32>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LlmProvider for NearAiChatProvider {
|
||||
async fn complete(&self, req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||
@@ -163,7 +177,7 @@ impl LlmProvider for NearAiChatProvider {
|
||||
req.messages.into_iter().map(|m| m.into()).collect();
|
||||
|
||||
let request = ChatCompletionRequest {
|
||||
model: self.config.model.clone(),
|
||||
model: self.active_model_name(),
|
||||
messages,
|
||||
temperature: req.temperature,
|
||||
max_tokens: req.max_tokens,
|
||||
@@ -197,6 +211,7 @@ impl LlmProvider for NearAiChatProvider {
|
||||
finish_reason,
|
||||
input_tokens: response.usage.prompt_tokens,
|
||||
output_tokens: response.usage.completion_tokens,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -221,7 +236,7 @@ impl LlmProvider for NearAiChatProvider {
|
||||
.collect();
|
||||
|
||||
let request = ChatCompletionRequest {
|
||||
model: self.config.model.clone(),
|
||||
model: self.active_model_name(),
|
||||
messages,
|
||||
temperature: req.temperature,
|
||||
max_tokens: req.max_tokens,
|
||||
@@ -278,6 +293,7 @@ impl LlmProvider for NearAiChatProvider {
|
||||
finish_reason,
|
||||
input_tokens: response.usage.prompt_tokens,
|
||||
output_tokens: response.usage.completion_tokens,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -291,7 +307,34 @@ impl LlmProvider for NearAiChatProvider {
|
||||
}
|
||||
|
||||
async fn list_models(&self) -> Result<Vec<String>, LlmError> {
|
||||
NearAiChatProvider::list_models(self).await
|
||||
let models = self.fetch_models().await?;
|
||||
Ok(models.into_iter().map(|m| m.id).collect())
|
||||
}
|
||||
|
||||
async fn model_metadata(&self) -> Result<ModelMetadata, LlmError> {
|
||||
let active = self.active_model_name();
|
||||
let models = self.fetch_models().await?;
|
||||
let current = models.iter().find(|m| m.id == active);
|
||||
Ok(ModelMetadata {
|
||||
id: active,
|
||||
context_length: current.and_then(|m| m.context_length),
|
||||
})
|
||||
}
|
||||
|
||||
fn active_model_name(&self) -> String {
|
||||
self.active_model
|
||||
.read()
|
||||
.expect("active_model lock poisoned")
|
||||
.clone()
|
||||
}
|
||||
|
||||
fn set_model(&self, model: &str) -> Result<(), crate::error::LlmError> {
|
||||
let mut guard = self
|
||||
.active_model
|
||||
.write()
|
||||
.expect("active_model lock poisoned");
|
||||
*guard = model.to_string();
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -332,6 +375,7 @@ impl From<ChatMessage> for ChatCompletionMessage {
|
||||
Role::Assistant => "assistant",
|
||||
Role::Tool => "tool",
|
||||
};
|
||||
|
||||
let tool_calls = msg.tool_calls.map(|calls| {
|
||||
calls
|
||||
.into_iter()
|
||||
@@ -345,9 +389,16 @@ impl From<ChatMessage> for ChatCompletionMessage {
|
||||
})
|
||||
.collect()
|
||||
});
|
||||
|
||||
let content = if role == "assistant" && tool_calls.is_some() && msg.content.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(msg.content)
|
||||
};
|
||||
|
||||
Self {
|
||||
role: role.to_string(),
|
||||
content: Some(msg.content),
|
||||
content,
|
||||
tool_call_id: msg.tool_call_id,
|
||||
name: msg.name,
|
||||
tool_calls,
|
||||
@@ -454,7 +505,7 @@ mod tests {
|
||||
},
|
||||
];
|
||||
|
||||
let msg = ChatMessage::assistant_with_tool_calls("", tool_calls);
|
||||
let msg = ChatMessage::assistant_with_tool_calls(None, tool_calls);
|
||||
let chat_msg: ChatCompletionMessage = msg.into();
|
||||
|
||||
assert_eq!(chat_msg.role, "assistant");
|
||||
@@ -484,7 +535,7 @@ mod tests {
|
||||
name: "test".to_string(),
|
||||
arguments: serde_json::json!({"key": "value"}),
|
||||
};
|
||||
let msg = ChatMessage::assistant_with_tool_calls("", vec![tc]);
|
||||
let msg = ChatMessage::assistant_with_tool_calls(None, vec![tc]);
|
||||
let chat_msg: ChatCompletionMessage = msg.into();
|
||||
|
||||
let calls = chat_msg.tool_calls.unwrap();
|
||||
|
||||
+64
-11
@@ -27,9 +27,8 @@ pub struct ChatMessage {
|
||||
/// Name of the tool for tool results.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub name: Option<String>,
|
||||
/// Tool calls requested by the assistant (for conversation replay).
|
||||
/// OpenAI-compatible APIs require the assistant message to include
|
||||
/// tool_calls when followed by tool result messages.
|
||||
/// Tool calls made by the assistant (OpenAI protocol requires these
|
||||
/// to appear on the assistant message preceding tool result messages).
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tool_calls: Option<Vec<ToolCall>>,
|
||||
}
|
||||
@@ -68,17 +67,14 @@ impl ChatMessage {
|
||||
}
|
||||
}
|
||||
|
||||
/// Create an assistant message that requested tool calls.
|
||||
/// Create an assistant message that includes tool calls.
|
||||
///
|
||||
/// OpenAI-compatible APIs require the assistant message to carry the
|
||||
/// `tool_calls` array when followed by tool-result messages.
|
||||
pub fn assistant_with_tool_calls(
|
||||
content: impl Into<String>,
|
||||
tool_calls: Vec<ToolCall>,
|
||||
) -> Self {
|
||||
/// Per the OpenAI protocol, an assistant message with tool_calls must
|
||||
/// precede the corresponding tool result messages in the conversation.
|
||||
pub fn assistant_with_tool_calls(content: Option<String>, tool_calls: Vec<ToolCall>) -> Self {
|
||||
Self {
|
||||
role: Role::Assistant,
|
||||
content: content.into(),
|
||||
content: content.unwrap_or_default(),
|
||||
tool_call_id: None,
|
||||
name: None,
|
||||
tool_calls: if tool_calls.is_empty() {
|
||||
@@ -112,6 +108,8 @@ pub struct CompletionRequest {
|
||||
pub max_tokens: Option<u32>,
|
||||
pub temperature: Option<f32>,
|
||||
pub stop_sequences: Option<Vec<String>>,
|
||||
/// Opaque metadata passed through to the provider (e.g. thread_id for chaining).
|
||||
pub metadata: std::collections::HashMap<String, String>,
|
||||
}
|
||||
|
||||
impl CompletionRequest {
|
||||
@@ -122,6 +120,7 @@ impl CompletionRequest {
|
||||
max_tokens: None,
|
||||
temperature: None,
|
||||
stop_sequences: None,
|
||||
metadata: std::collections::HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -145,6 +144,8 @@ pub struct CompletionResponse {
|
||||
pub input_tokens: u32,
|
||||
pub output_tokens: u32,
|
||||
pub finish_reason: FinishReason,
|
||||
/// Provider-specific response ID (e.g. for NEAR AI response chaining).
|
||||
pub response_id: Option<String>,
|
||||
}
|
||||
|
||||
/// Why the completion finished.
|
||||
@@ -191,6 +192,8 @@ pub struct ToolCompletionRequest {
|
||||
pub temperature: Option<f32>,
|
||||
/// How to handle tool use: "auto", "required", or "none".
|
||||
pub tool_choice: Option<String>,
|
||||
/// Opaque metadata passed through to the provider (e.g. thread_id for chaining).
|
||||
pub metadata: std::collections::HashMap<String, String>,
|
||||
}
|
||||
|
||||
impl ToolCompletionRequest {
|
||||
@@ -202,6 +205,7 @@ impl ToolCompletionRequest {
|
||||
max_tokens: None,
|
||||
temperature: None,
|
||||
tool_choice: None,
|
||||
metadata: std::collections::HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -234,6 +238,16 @@ pub struct ToolCompletionResponse {
|
||||
pub input_tokens: u32,
|
||||
pub output_tokens: u32,
|
||||
pub finish_reason: FinishReason,
|
||||
/// Provider-specific response ID (e.g. for NEAR AI response chaining).
|
||||
pub response_id: Option<String>,
|
||||
}
|
||||
|
||||
/// Metadata about a model returned by the provider's API.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ModelMetadata {
|
||||
pub id: String,
|
||||
/// Total context window size in tokens.
|
||||
pub context_length: Option<u32>,
|
||||
}
|
||||
|
||||
/// Trait for LLM providers.
|
||||
@@ -260,6 +274,45 @@ pub trait LlmProvider: Send + Sync {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
/// Fetch metadata for the current model (context length, etc.).
|
||||
/// Default returns the model name with no size info.
|
||||
async fn model_metadata(&self) -> Result<ModelMetadata, LlmError> {
|
||||
Ok(ModelMetadata {
|
||||
id: self.model_name().to_string(),
|
||||
context_length: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// Get the currently active model name.
|
||||
///
|
||||
/// May differ from `model_name()` if the model was switched at runtime
|
||||
/// via `set_model()`. Default returns `model_name()`.
|
||||
fn active_model_name(&self) -> String {
|
||||
self.model_name().to_string()
|
||||
}
|
||||
|
||||
/// Switch the active model at runtime. Not all providers support this.
|
||||
fn set_model(&self, _model: &str) -> Result<(), LlmError> {
|
||||
Err(LlmError::RequestFailed {
|
||||
provider: "unknown".to_string(),
|
||||
reason: "Runtime model switching not supported by this provider".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Seed a response chain for a thread (e.g. restoring from DB).
|
||||
///
|
||||
/// Providers that support response chaining (e.g. NEAR AI `previous_response_id`)
|
||||
/// store this so subsequent calls send only delta messages.
|
||||
fn seed_response_chain(&self, _thread_id: &str, _response_id: String) {}
|
||||
|
||||
/// Get the last response chain ID for a thread.
|
||||
///
|
||||
/// Returns `None` if the provider doesn't support chaining or has no
|
||||
/// stored state for this thread.
|
||||
fn get_response_chain_id(&self, _thread_id: &str) -> Option<String> {
|
||||
None
|
||||
}
|
||||
|
||||
/// Calculate cost for a completion.
|
||||
fn calculate_cost(&self, input_tokens: u32, output_tokens: u32) -> Decimal {
|
||||
let (input_cost, output_cost) = self.cost_per_token();
|
||||
|
||||
+353
-29
@@ -21,6 +21,8 @@ pub struct ReasoningContext {
|
||||
pub job_description: Option<String>,
|
||||
/// Current state description.
|
||||
pub current_state: Option<String>,
|
||||
/// Opaque metadata forwarded to the LLM provider (e.g. thread_id for chaining).
|
||||
pub metadata: std::collections::HashMap<String, String>,
|
||||
}
|
||||
|
||||
impl ReasoningContext {
|
||||
@@ -31,6 +33,7 @@ impl ReasoningContext {
|
||||
available_tools: Vec::new(),
|
||||
job_description: None,
|
||||
current_state: None,
|
||||
metadata: std::collections::HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -57,6 +60,12 @@ impl ReasoningContext {
|
||||
self.job_description = Some(description.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// Set metadata (forwarded to the LLM provider).
|
||||
pub fn with_metadata(mut self, metadata: std::collections::HashMap<String, String>) -> Self {
|
||||
self.metadata = metadata;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ReasoningContext {
|
||||
@@ -114,7 +123,12 @@ pub enum RespondResult {
|
||||
/// A text response (no tools needed).
|
||||
Text(String),
|
||||
/// The model wants to call tools. Caller should execute them and call back.
|
||||
ToolCalls(Vec<ToolCall>),
|
||||
/// Includes the optional content from the assistant message (some models
|
||||
/// include explanatory text alongside tool calls).
|
||||
ToolCalls {
|
||||
tool_calls: Vec<ToolCall>,
|
||||
content: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Reasoning engine for the agent.
|
||||
@@ -192,10 +206,11 @@ impl Reasoning {
|
||||
return Ok(vec![]);
|
||||
}
|
||||
|
||||
let request =
|
||||
let mut request =
|
||||
ToolCompletionRequest::new(context.messages.clone(), context.available_tools.clone())
|
||||
.with_max_tokens(1024)
|
||||
.with_tool_choice("auto");
|
||||
request.metadata = context.metadata.clone();
|
||||
|
||||
let response = self.llm.complete_with_tools(request).await?;
|
||||
|
||||
@@ -271,7 +286,9 @@ Respond in JSON format:
|
||||
pub async fn respond(&self, context: &ReasoningContext) -> Result<String, LlmError> {
|
||||
match self.respond_with_tools(context).await? {
|
||||
RespondResult::Text(text) => Ok(text),
|
||||
RespondResult::ToolCalls(calls) => {
|
||||
RespondResult::ToolCalls {
|
||||
tool_calls: calls, ..
|
||||
} => {
|
||||
// Format tool calls as text (legacy behavior for non-agentic callers)
|
||||
let tool_info: Vec<String> = calls
|
||||
.iter()
|
||||
@@ -298,28 +315,49 @@ Respond in JSON format:
|
||||
|
||||
// If we have tools, use tool completion mode
|
||||
if !context.available_tools.is_empty() {
|
||||
let request = ToolCompletionRequest::new(messages, context.available_tools.clone())
|
||||
let mut request = ToolCompletionRequest::new(messages, context.available_tools.clone())
|
||||
.with_max_tokens(4096)
|
||||
.with_temperature(0.7)
|
||||
.with_tool_choice("auto");
|
||||
request.metadata = context.metadata.clone();
|
||||
|
||||
let response = self.llm.complete_with_tools(request).await?;
|
||||
|
||||
// If there were tool calls, return them for execution
|
||||
if !response.tool_calls.is_empty() {
|
||||
return Ok(RespondResult::ToolCalls(response.tool_calls));
|
||||
return Ok(RespondResult::ToolCalls {
|
||||
tool_calls: response.tool_calls,
|
||||
content: response.content,
|
||||
});
|
||||
}
|
||||
|
||||
let content = response
|
||||
.content
|
||||
.unwrap_or_else(|| "I'm not sure how to respond to that.".to_string());
|
||||
|
||||
// Some models (e.g. GLM-4.7) emit tool calls as XML tags in content
|
||||
// instead of using the structured tool_calls field. Try to recover
|
||||
// them before giving up and returning plain text.
|
||||
let recovered = recover_tool_calls_from_content(&content, &context.available_tools);
|
||||
if !recovered.is_empty() {
|
||||
let cleaned = clean_response(&content);
|
||||
return Ok(RespondResult::ToolCalls {
|
||||
tool_calls: recovered,
|
||||
content: if cleaned.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(cleaned)
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
Ok(RespondResult::Text(clean_response(&content)))
|
||||
} else {
|
||||
// No tools, use simple completion
|
||||
let request = CompletionRequest::new(messages)
|
||||
let mut request = CompletionRequest::new(messages)
|
||||
.with_max_tokens(4096)
|
||||
.with_temperature(0.7);
|
||||
request.metadata = context.metadata.clone();
|
||||
|
||||
let response = self.llm.complete(request).await?;
|
||||
Ok(RespondResult::Text(clean_response(&response.content)))
|
||||
@@ -462,47 +500,178 @@ fn extract_json(text: &str) -> Option<&str> {
|
||||
}
|
||||
}
|
||||
|
||||
/// Clean up LLM response by stripping thinking tags and reasoning patterns.
|
||||
/// Clean up LLM response by stripping model-internal tags and reasoning patterns.
|
||||
///
|
||||
/// Some models (GLM-4.7, etc.) emit XML-tagged internal state like
|
||||
/// Try to extract tool calls from content text where the model emitted them
|
||||
/// as XML tags instead of using the structured tool_calls field.
|
||||
///
|
||||
/// Handles these formats:
|
||||
/// - `<tool_call>tool_name</tool_call>` (bare name)
|
||||
/// - `<tool_call>{"name":"x","arguments":{}}</tool_call>` (JSON)
|
||||
/// - `<|tool_call|>...<|/tool_call|>` (pipe-delimited variant)
|
||||
/// - `<function_call>...</function_call>` (function_call variant)
|
||||
///
|
||||
/// Only returns calls whose name matches an available tool.
|
||||
fn recover_tool_calls_from_content(
|
||||
content: &str,
|
||||
available_tools: &[ToolDefinition],
|
||||
) -> Vec<ToolCall> {
|
||||
let tool_names: std::collections::HashSet<&str> =
|
||||
available_tools.iter().map(|t| t.name.as_str()).collect();
|
||||
let mut calls = Vec::new();
|
||||
|
||||
for (open, close) in &[
|
||||
("<tool_call>", "</tool_call>"),
|
||||
("<|tool_call|>", "<|/tool_call|>"),
|
||||
("<function_call>", "</function_call>"),
|
||||
("<|function_call|>", "<|/function_call|>"),
|
||||
] {
|
||||
let mut remaining = content;
|
||||
while let Some(start) = remaining.find(open) {
|
||||
let inner_start = start + open.len();
|
||||
let after = &remaining[inner_start..];
|
||||
let Some(end) = after.find(close) else {
|
||||
break;
|
||||
};
|
||||
let inner = after[..end].trim();
|
||||
remaining = &after[end + close.len()..];
|
||||
|
||||
if inner.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Try JSON first: {"name":"x","arguments":{}}
|
||||
if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(inner) {
|
||||
if let Some(name) = parsed.get("name").and_then(|v| v.as_str()) {
|
||||
if tool_names.contains(name) {
|
||||
let arguments = parsed
|
||||
.get("arguments")
|
||||
.cloned()
|
||||
.unwrap_or(serde_json::Value::Object(Default::default()));
|
||||
calls.push(ToolCall {
|
||||
id: format!("recovered_{}", calls.len()),
|
||||
name: name.to_string(),
|
||||
arguments,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Bare tool name (e.g. "<tool_call>tool_list</tool_call>")
|
||||
let name = inner.trim();
|
||||
if tool_names.contains(name) {
|
||||
calls.push(ToolCall {
|
||||
id: format!("recovered_{}", calls.len()),
|
||||
name: name.to_string(),
|
||||
arguments: serde_json::Value::Object(Default::default()),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
calls
|
||||
}
|
||||
|
||||
/// `<tool_call>tool_list</tool_call>` or `<|tool_call|>` in the content field
|
||||
/// instead of using the standard OpenAI tool_calls array. We strip all of
|
||||
/// these before the response reaches channels/users.
|
||||
fn clean_response(text: &str) -> String {
|
||||
let text = strip_thinking_tags(text);
|
||||
let text = strip_internal_tags(text);
|
||||
strip_reasoning_patterns(&text)
|
||||
}
|
||||
|
||||
/// Strip `<thinking>...</thinking>` blocks from LLM output.
|
||||
/// Tags that are model-internal and should never reach users.
|
||||
const INTERNAL_TAGS: &[&str] = &["thinking", "tool_call", "function_call", "tool_calls"];
|
||||
|
||||
/// Strip all model-internal XML tags from LLM output.
|
||||
///
|
||||
/// Some models (especially Claude with extended thinking) include internal
|
||||
/// reasoning in thinking tags. We strip these before showing to users.
|
||||
fn strip_thinking_tags(text: &str) -> String {
|
||||
/// Handles standard XML tags (`<tag>...</tag>`) and pipe-delimited variants
|
||||
/// (`<|tag|>...<|/tag|>`) used by some models (e.g. GLM-4.7).
|
||||
fn strip_internal_tags(text: &str) -> String {
|
||||
let mut result = text.to_string();
|
||||
for tag in INTERNAL_TAGS {
|
||||
result = strip_xml_tag(&result, tag);
|
||||
result = strip_pipe_tag(&result, tag);
|
||||
}
|
||||
// Collapse triple+ newlines left behind by removed blocks
|
||||
while result.contains("\n\n\n") {
|
||||
result = result.replace("\n\n\n", "\n\n");
|
||||
}
|
||||
result.trim().to_string()
|
||||
}
|
||||
|
||||
/// Strip `<tag>...</tag>` and `<tag ...>...</tag>` blocks from text.
|
||||
fn strip_xml_tag(text: &str, tag: &str) -> String {
|
||||
let open_exact = format!("<{}>", tag);
|
||||
let open_prefix = format!("<{} ", tag); // for <tag attr="...">
|
||||
let close = format!("</{}>", tag);
|
||||
|
||||
let mut result = String::with_capacity(text.len());
|
||||
let mut remaining = text;
|
||||
|
||||
while let Some(start) = remaining.find("<thinking>") {
|
||||
loop {
|
||||
// Find the next opening tag (exact or with attributes)
|
||||
let exact_pos = remaining.find(&open_exact);
|
||||
let prefix_pos = remaining.find(&open_prefix);
|
||||
let start = match (exact_pos, prefix_pos) {
|
||||
(Some(a), Some(b)) => a.min(b),
|
||||
(Some(a), None) => a,
|
||||
(None, Some(b)) => b,
|
||||
(None, None) => break,
|
||||
};
|
||||
|
||||
// Add everything before the tag
|
||||
result.push_str(&remaining[..start]);
|
||||
|
||||
// Find the end of the opening tag (the closing >)
|
||||
let after_open = &remaining[start..];
|
||||
let open_end = match after_open.find('>') {
|
||||
Some(pos) => start + pos + 1,
|
||||
None => break, // malformed, stop
|
||||
};
|
||||
|
||||
// Find the closing tag
|
||||
if let Some(end_offset) = remaining[start..].find("</thinking>") {
|
||||
// Skip past the closing tag (start + offset + tag length)
|
||||
let end = start + end_offset + "</thinking>".len();
|
||||
if let Some(close_offset) = remaining[open_end..].find(&close) {
|
||||
let end = open_end + close_offset + close.len();
|
||||
remaining = &remaining[end..];
|
||||
} else {
|
||||
// No closing tag found, discard everything from here
|
||||
// (malformed, but handle gracefully by not including the unclosed tag)
|
||||
// No closing tag, discard from here (malformed)
|
||||
remaining = "";
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Add any remaining content after the last thinking block
|
||||
result.push_str(remaining);
|
||||
result
|
||||
}
|
||||
|
||||
// Clean up any double newlines left behind
|
||||
let mut cleaned = result.trim().to_string();
|
||||
while cleaned.contains("\n\n\n") {
|
||||
cleaned = cleaned.replace("\n\n\n", "\n\n");
|
||||
/// Strip `<|tag|>...<|/tag|>` pipe-delimited blocks from text.
|
||||
///
|
||||
/// Some models (e.g. certain Chinese LLMs) use this format instead of
|
||||
/// standard XML tags.
|
||||
fn strip_pipe_tag(text: &str, tag: &str) -> String {
|
||||
let open = format!("<|{}|>", tag);
|
||||
let close = format!("<|/{}|>", tag);
|
||||
|
||||
let mut result = String::with_capacity(text.len());
|
||||
let mut remaining = text;
|
||||
|
||||
while let Some(start) = remaining.find(&open) {
|
||||
result.push_str(&remaining[..start]);
|
||||
|
||||
if let Some(close_offset) = remaining[start..].find(&close) {
|
||||
let end = start + close_offset + close.len();
|
||||
remaining = &remaining[end..];
|
||||
} else {
|
||||
remaining = "";
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
cleaned
|
||||
result.push_str(remaining);
|
||||
result
|
||||
}
|
||||
|
||||
/// Strip any remaining reasoning that wasn't in proper <thinking> tags.
|
||||
@@ -574,7 +743,7 @@ That's my plan."#;
|
||||
#[test]
|
||||
fn test_strip_thinking_tags_basic() {
|
||||
let input = "<thinking>Let me think about this...</thinking>Hello, user!";
|
||||
let output = strip_thinking_tags(input);
|
||||
let output = strip_internal_tags(input);
|
||||
assert_eq!(output, "Hello, user!");
|
||||
}
|
||||
|
||||
@@ -582,7 +751,7 @@ That's my plan."#;
|
||||
fn test_strip_thinking_tags_multiple() {
|
||||
let input =
|
||||
"<thinking>First thought</thinking>Hello<thinking>Second thought</thinking> world!";
|
||||
let output = strip_thinking_tags(input);
|
||||
let output = strip_internal_tags(input);
|
||||
assert_eq!(output, "Hello world!");
|
||||
}
|
||||
|
||||
@@ -594,14 +763,14 @@ I need to consider:
|
||||
2. How to respond
|
||||
</thinking>
|
||||
Here is my response to your question."#;
|
||||
let output = strip_thinking_tags(input);
|
||||
let output = strip_internal_tags(input);
|
||||
assert_eq!(output, "Here is my response to your question.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strip_thinking_tags_no_tags() {
|
||||
let input = "Just a normal response without thinking tags.";
|
||||
let output = strip_thinking_tags(input);
|
||||
let output = strip_internal_tags(input);
|
||||
assert_eq!(output, "Just a normal response without thinking tags.");
|
||||
}
|
||||
|
||||
@@ -609,10 +778,77 @@ Here is my response to your question."#;
|
||||
fn test_strip_thinking_tags_unclosed() {
|
||||
// Malformed: unclosed tag should strip from there to end
|
||||
let input = "Hello <thinking>this never closes";
|
||||
let output = strip_thinking_tags(input);
|
||||
let output = strip_internal_tags(input);
|
||||
assert_eq!(output, "Hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strip_tool_call_tags() {
|
||||
// GLM-4.7 emits this garbage instead of using the tool_calls array
|
||||
let input = "<tool_call>tool_list</tool_call>";
|
||||
let output = strip_internal_tags(input);
|
||||
assert_eq!(output, "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strip_tool_call_with_surrounding_text() {
|
||||
let input = "Here is my answer.\n\n<tool_call>\n{\"name\": \"search\", \"arguments\": {}}\n</tool_call>";
|
||||
let output = strip_internal_tags(input);
|
||||
assert_eq!(output, "Here is my answer.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strip_multiple_internal_tags() {
|
||||
let input = "<thinking>Let me think</thinking>Hello!\n<tool_call>some_tool</tool_call>";
|
||||
let output = strip_internal_tags(input);
|
||||
assert_eq!(output, "Hello!");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strip_function_call_tags() {
|
||||
let input = "Response text<function_call>{\"name\": \"foo\"}</function_call>";
|
||||
let output = strip_internal_tags(input);
|
||||
assert_eq!(output, "Response text");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strip_tool_calls_plural() {
|
||||
let input = "<tool_calls>[{\"id\": \"1\"}]</tool_calls>Actual response.";
|
||||
let output = strip_internal_tags(input);
|
||||
assert_eq!(output, "Actual response.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strip_pipe_delimited_tags() {
|
||||
let input = "<|tool_call|>{\"name\": \"search\"}<|/tool_call|>Hello!";
|
||||
let output = strip_internal_tags(input);
|
||||
assert_eq!(output, "Hello!");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strip_pipe_delimited_thinking() {
|
||||
let input = "<|thinking|>reasoning here<|/thinking|>The answer is 42.";
|
||||
let output = strip_internal_tags(input);
|
||||
assert_eq!(output, "The answer is 42.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strip_xml_tag_with_attributes() {
|
||||
let input = "<tool_call type=\"function\">search()</tool_call>Done.";
|
||||
let output = strip_internal_tags(input);
|
||||
assert_eq!(output, "Done.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_clean_response_preserves_normal_content() {
|
||||
let input = "The function tool_call_handler works great. No tags here!";
|
||||
let output = clean_response(input);
|
||||
assert_eq!(
|
||||
output,
|
||||
"The function tool_call_handler works great. No tags here!"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strip_reasoning_paragraph_break() {
|
||||
// Content after paragraph break with "here" marker
|
||||
@@ -660,4 +896,92 @@ Here is my response to your question."#;
|
||||
let output = clean_response(input);
|
||||
assert_eq!(output, "Here's the answer.");
|
||||
}
|
||||
|
||||
// -- recover_tool_calls_from_content tests --
|
||||
|
||||
fn make_tools(names: &[&str]) -> Vec<ToolDefinition> {
|
||||
names
|
||||
.iter()
|
||||
.map(|n| ToolDefinition {
|
||||
name: n.to_string(),
|
||||
description: String::new(),
|
||||
parameters: serde_json::json!({}),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recover_bare_tool_name() {
|
||||
let tools = make_tools(&["tool_list", "tool_auth"]);
|
||||
let content = "<tool_call>tool_list</tool_call>";
|
||||
let calls = recover_tool_calls_from_content(content, &tools);
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].name, "tool_list");
|
||||
assert_eq!(calls[0].arguments, serde_json::json!({}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recover_json_tool_call() {
|
||||
let tools = make_tools(&["memory_search"]);
|
||||
let content =
|
||||
r#"<tool_call>{"name": "memory_search", "arguments": {"query": "test"}}</tool_call>"#;
|
||||
let calls = recover_tool_calls_from_content(content, &tools);
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].name, "memory_search");
|
||||
assert_eq!(calls[0].arguments, serde_json::json!({"query": "test"}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recover_pipe_delimited() {
|
||||
let tools = make_tools(&["tool_list"]);
|
||||
let content = "<|tool_call|>tool_list<|/tool_call|>";
|
||||
let calls = recover_tool_calls_from_content(content, &tools);
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].name, "tool_list");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recover_unknown_tool_ignored() {
|
||||
let tools = make_tools(&["tool_list"]);
|
||||
let content = "<tool_call>nonexistent_tool</tool_call>";
|
||||
let calls = recover_tool_calls_from_content(content, &tools);
|
||||
assert!(calls.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recover_no_tags() {
|
||||
let tools = make_tools(&["tool_list"]);
|
||||
let content = "Just a normal response.";
|
||||
let calls = recover_tool_calls_from_content(content, &tools);
|
||||
assert!(calls.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recover_multiple_tool_calls() {
|
||||
let tools = make_tools(&["tool_list", "tool_auth"]);
|
||||
let content = "<tool_call>tool_list</tool_call>\n<tool_call>tool_auth</tool_call>";
|
||||
let calls = recover_tool_calls_from_content(content, &tools);
|
||||
assert_eq!(calls.len(), 2);
|
||||
assert_eq!(calls[0].name, "tool_list");
|
||||
assert_eq!(calls[1].name, "tool_auth");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recover_function_call_variant() {
|
||||
let tools = make_tools(&["shell"]);
|
||||
let content =
|
||||
r#"<function_call>{"name": "shell", "arguments": {"cmd": "ls"}}</function_call>"#;
|
||||
let calls = recover_tool_calls_from_content(content, &tools);
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].name, "shell");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recover_with_surrounding_text() {
|
||||
let tools = make_tools(&["tool_list"]);
|
||||
let content = "Let me check.\n\n<tool_call>tool_list</tool_call>\n\nDone.";
|
||||
let calls = recover_tool_calls_from_content(content, &tools);
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].name, "tool_list");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,451 @@
|
||||
//! Generic adapter that bridges rig-core's `CompletionModel` trait to IronClaw's `LlmProvider`.
|
||||
//!
|
||||
//! This lets us use any rig-core provider (OpenAI, Anthropic, Ollama, etc.) as an
|
||||
//! `Arc<dyn LlmProvider>` without changing any of the agent, reasoning, or tool code.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use rig::OneOrMany;
|
||||
use rig::completion::{
|
||||
AssistantContent, CompletionModel, CompletionRequest as RigRequest,
|
||||
ToolDefinition as RigToolDefinition, Usage as RigUsage,
|
||||
};
|
||||
use rig::message::{
|
||||
Message as RigMessage, ToolChoice as RigToolChoice, ToolFunction, ToolResult as RigToolResult,
|
||||
ToolResultContent, UserContent,
|
||||
};
|
||||
use rust_decimal::Decimal;
|
||||
use serde::Serialize;
|
||||
use serde::de::DeserializeOwned;
|
||||
|
||||
use crate::error::LlmError;
|
||||
use crate::llm::costs;
|
||||
use crate::llm::provider::{
|
||||
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider,
|
||||
ToolCall as IronToolCall, ToolCompletionRequest, ToolCompletionResponse,
|
||||
ToolDefinition as IronToolDefinition,
|
||||
};
|
||||
|
||||
/// Adapter that wraps a rig-core `CompletionModel` and implements `LlmProvider`.
|
||||
pub struct RigAdapter<M: CompletionModel> {
|
||||
model: M,
|
||||
model_name: String,
|
||||
input_cost: Decimal,
|
||||
output_cost: Decimal,
|
||||
}
|
||||
|
||||
impl<M: CompletionModel> RigAdapter<M> {
|
||||
/// Create a new adapter wrapping the given rig-core model.
|
||||
pub fn new(model: M, model_name: impl Into<String>) -> Self {
|
||||
let name = model_name.into();
|
||||
let (input_cost, output_cost) =
|
||||
costs::model_cost(&name).unwrap_or_else(costs::default_cost);
|
||||
Self {
|
||||
model,
|
||||
model_name: name,
|
||||
input_cost,
|
||||
output_cost,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// -- Type conversion helpers --
|
||||
|
||||
/// Convert IronClaw messages to rig-core format.
|
||||
///
|
||||
/// Returns `(preamble, chat_history)` where preamble is extracted from
|
||||
/// any System message and chat_history contains the rest.
|
||||
fn convert_messages(messages: &[ChatMessage]) -> (Option<String>, Vec<RigMessage>) {
|
||||
let mut preamble: Option<String> = None;
|
||||
let mut history = Vec::new();
|
||||
|
||||
for msg in messages {
|
||||
match msg.role {
|
||||
crate::llm::Role::System => {
|
||||
// Concatenate system messages into preamble
|
||||
match preamble {
|
||||
Some(ref mut p) => {
|
||||
p.push('\n');
|
||||
p.push_str(&msg.content);
|
||||
}
|
||||
None => preamble = Some(msg.content.clone()),
|
||||
}
|
||||
}
|
||||
crate::llm::Role::User => {
|
||||
history.push(RigMessage::user(&msg.content));
|
||||
}
|
||||
crate::llm::Role::Assistant => {
|
||||
if let Some(ref tool_calls) = msg.tool_calls {
|
||||
// Assistant message with tool calls
|
||||
let mut contents: Vec<AssistantContent> = Vec::new();
|
||||
if !msg.content.is_empty() {
|
||||
contents.push(AssistantContent::text(&msg.content));
|
||||
}
|
||||
for tc in tool_calls {
|
||||
contents.push(AssistantContent::ToolCall(rig::message::ToolCall::new(
|
||||
tc.id.clone(),
|
||||
ToolFunction::new(tc.name.clone(), tc.arguments.clone()),
|
||||
)));
|
||||
}
|
||||
if let Ok(many) = OneOrMany::many(contents) {
|
||||
history.push(RigMessage::Assistant {
|
||||
id: None,
|
||||
content: many,
|
||||
});
|
||||
} else {
|
||||
// Shouldn't happen but fall back to text
|
||||
history.push(RigMessage::assistant(&msg.content));
|
||||
}
|
||||
} else {
|
||||
history.push(RigMessage::assistant(&msg.content));
|
||||
}
|
||||
}
|
||||
crate::llm::Role::Tool => {
|
||||
// Tool result message: wrap as User { ToolResult }
|
||||
let tool_id = msg.tool_call_id.clone().unwrap_or_default();
|
||||
history.push(RigMessage::User {
|
||||
content: OneOrMany::one(UserContent::ToolResult(RigToolResult {
|
||||
id: tool_id,
|
||||
call_id: None,
|
||||
content: OneOrMany::one(ToolResultContent::text(&msg.content)),
|
||||
})),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
(preamble, history)
|
||||
}
|
||||
|
||||
/// Convert IronClaw tool definitions to rig-core format.
|
||||
fn convert_tools(tools: &[IronToolDefinition]) -> Vec<RigToolDefinition> {
|
||||
tools
|
||||
.iter()
|
||||
.map(|t| RigToolDefinition {
|
||||
name: t.name.clone(),
|
||||
description: t.description.clone(),
|
||||
parameters: t.parameters.clone(),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Convert IronClaw tool_choice string to rig-core ToolChoice.
|
||||
fn convert_tool_choice(choice: Option<&str>) -> Option<RigToolChoice> {
|
||||
match choice.map(|s| s.to_lowercase()).as_deref() {
|
||||
Some("auto") => Some(RigToolChoice::Auto),
|
||||
Some("required") => Some(RigToolChoice::Required),
|
||||
Some("none") => Some(RigToolChoice::None),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract text and tool calls from a rig-core completion response.
|
||||
fn extract_response(
|
||||
choice: &OneOrMany<AssistantContent>,
|
||||
_usage: &RigUsage,
|
||||
) -> (Option<String>, Vec<IronToolCall>, FinishReason) {
|
||||
let mut text_parts: Vec<String> = Vec::new();
|
||||
let mut tool_calls: Vec<IronToolCall> = Vec::new();
|
||||
|
||||
for content in choice.iter() {
|
||||
match content {
|
||||
AssistantContent::Text(t) => {
|
||||
if !t.text.is_empty() {
|
||||
text_parts.push(t.text.clone());
|
||||
}
|
||||
}
|
||||
AssistantContent::ToolCall(tc) => {
|
||||
tool_calls.push(IronToolCall {
|
||||
id: tc.id.clone(),
|
||||
name: tc.function.name.clone(),
|
||||
arguments: tc.function.arguments.clone(),
|
||||
});
|
||||
}
|
||||
// Reasoning and Image variants are not mapped to IronClaw types
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
let text = if text_parts.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(text_parts.join(""))
|
||||
};
|
||||
|
||||
let finish = if !tool_calls.is_empty() {
|
||||
FinishReason::ToolUse
|
||||
} else {
|
||||
FinishReason::Stop
|
||||
};
|
||||
|
||||
(text, tool_calls, finish)
|
||||
}
|
||||
|
||||
/// Saturate u64 to u32 for token counts.
|
||||
fn saturate_u32(val: u64) -> u32 {
|
||||
val.min(u32::MAX as u64) as u32
|
||||
}
|
||||
|
||||
/// Build a rig-core CompletionRequest from our internal types.
|
||||
fn build_rig_request(
|
||||
preamble: Option<String>,
|
||||
mut history: Vec<RigMessage>,
|
||||
tools: Vec<RigToolDefinition>,
|
||||
tool_choice: Option<RigToolChoice>,
|
||||
temperature: Option<f32>,
|
||||
max_tokens: Option<u32>,
|
||||
) -> Result<RigRequest, LlmError> {
|
||||
// rig-core requires at least one message in chat_history
|
||||
if history.is_empty() {
|
||||
history.push(RigMessage::user("Hello"));
|
||||
}
|
||||
|
||||
let chat_history = OneOrMany::many(history).map_err(|e| LlmError::RequestFailed {
|
||||
provider: "rig".to_string(),
|
||||
reason: format!("Failed to build chat history: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(RigRequest {
|
||||
preamble,
|
||||
chat_history,
|
||||
documents: Vec::new(),
|
||||
tools,
|
||||
temperature: temperature.map(|t| t as f64),
|
||||
max_tokens: max_tokens.map(|t| t as u64),
|
||||
tool_choice,
|
||||
additional_params: None,
|
||||
})
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<M> LlmProvider for RigAdapter<M>
|
||||
where
|
||||
M: CompletionModel + Send + Sync + 'static,
|
||||
M::Response: Send + Sync + Serialize + DeserializeOwned,
|
||||
{
|
||||
fn model_name(&self) -> &str {
|
||||
&self.model_name
|
||||
}
|
||||
|
||||
fn cost_per_token(&self) -> (Decimal, Decimal) {
|
||||
(self.input_cost, self.output_cost)
|
||||
}
|
||||
|
||||
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||
let (preamble, history) = convert_messages(&request.messages);
|
||||
|
||||
let rig_req = build_rig_request(
|
||||
preamble,
|
||||
history,
|
||||
Vec::new(),
|
||||
None,
|
||||
request.temperature,
|
||||
request.max_tokens,
|
||||
)?;
|
||||
|
||||
let response =
|
||||
self.model
|
||||
.completion(rig_req)
|
||||
.await
|
||||
.map_err(|e| LlmError::RequestFailed {
|
||||
provider: self.model_name.clone(),
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
let (text, _tool_calls, finish) = extract_response(&response.choice, &response.usage);
|
||||
|
||||
Ok(CompletionResponse {
|
||||
content: text.unwrap_or_default(),
|
||||
input_tokens: saturate_u32(response.usage.input_tokens),
|
||||
output_tokens: saturate_u32(response.usage.output_tokens),
|
||||
finish_reason: finish,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
async fn complete_with_tools(
|
||||
&self,
|
||||
request: ToolCompletionRequest,
|
||||
) -> Result<ToolCompletionResponse, LlmError> {
|
||||
let (preamble, history) = convert_messages(&request.messages);
|
||||
let tools = convert_tools(&request.tools);
|
||||
let tool_choice = convert_tool_choice(request.tool_choice.as_deref());
|
||||
|
||||
let rig_req = build_rig_request(
|
||||
preamble,
|
||||
history,
|
||||
tools,
|
||||
tool_choice,
|
||||
request.temperature,
|
||||
request.max_tokens,
|
||||
)?;
|
||||
|
||||
let response =
|
||||
self.model
|
||||
.completion(rig_req)
|
||||
.await
|
||||
.map_err(|e| LlmError::RequestFailed {
|
||||
provider: self.model_name.clone(),
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
let (text, tool_calls, finish) = extract_response(&response.choice, &response.usage);
|
||||
|
||||
Ok(ToolCompletionResponse {
|
||||
content: text,
|
||||
tool_calls,
|
||||
input_tokens: saturate_u32(response.usage.input_tokens),
|
||||
output_tokens: saturate_u32(response.usage.output_tokens),
|
||||
finish_reason: finish,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn active_model_name(&self) -> String {
|
||||
self.model_name.clone()
|
||||
}
|
||||
|
||||
fn set_model(&self, _model: &str) -> Result<(), LlmError> {
|
||||
// rig-core models are baked at construction time.
|
||||
// Switching requires creating a new adapter.
|
||||
Err(LlmError::RequestFailed {
|
||||
provider: self.model_name.clone(),
|
||||
reason: "Runtime model switching not supported for rig-core providers. \
|
||||
Restart with a different model configured."
|
||||
.to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_convert_messages_system_to_preamble() {
|
||||
let messages = vec![
|
||||
ChatMessage::system("You are a helpful assistant."),
|
||||
ChatMessage::user("Hello"),
|
||||
];
|
||||
let (preamble, history) = convert_messages(&messages);
|
||||
assert_eq!(preamble, Some("You are a helpful assistant.".to_string()));
|
||||
assert_eq!(history.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_convert_messages_multiple_systems_concatenated() {
|
||||
let messages = vec![
|
||||
ChatMessage::system("System 1"),
|
||||
ChatMessage::system("System 2"),
|
||||
ChatMessage::user("Hi"),
|
||||
];
|
||||
let (preamble, history) = convert_messages(&messages);
|
||||
assert_eq!(preamble, Some("System 1\nSystem 2".to_string()));
|
||||
assert_eq!(history.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_convert_messages_tool_result() {
|
||||
let messages = vec![ChatMessage::tool_result(
|
||||
"call_123",
|
||||
"search",
|
||||
"result text",
|
||||
)];
|
||||
let (preamble, history) = convert_messages(&messages);
|
||||
assert!(preamble.is_none());
|
||||
assert_eq!(history.len(), 1);
|
||||
// Tool results become User messages in rig-core
|
||||
match &history[0] {
|
||||
RigMessage::User { .. } => {}
|
||||
other => panic!("Expected User message, got: {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_convert_messages_assistant_with_tool_calls() {
|
||||
let tc = IronToolCall {
|
||||
id: "call_1".to_string(),
|
||||
name: "search".to_string(),
|
||||
arguments: serde_json::json!({"query": "test"}),
|
||||
};
|
||||
let msg = ChatMessage::assistant_with_tool_calls(Some("thinking".to_string()), vec![tc]);
|
||||
let messages = vec![msg];
|
||||
let (_preamble, history) = convert_messages(&messages);
|
||||
assert_eq!(history.len(), 1);
|
||||
match &history[0] {
|
||||
RigMessage::Assistant { content, .. } => {
|
||||
// Should have both text and tool call
|
||||
assert!(content.iter().count() >= 2);
|
||||
}
|
||||
other => panic!("Expected Assistant message, got: {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_convert_tools() {
|
||||
let tools = vec![IronToolDefinition {
|
||||
name: "search".to_string(),
|
||||
description: "Search the web".to_string(),
|
||||
parameters: serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {"type": "string"}
|
||||
}
|
||||
}),
|
||||
}];
|
||||
let rig_tools = convert_tools(&tools);
|
||||
assert_eq!(rig_tools.len(), 1);
|
||||
assert_eq!(rig_tools[0].name, "search");
|
||||
assert_eq!(rig_tools[0].description, "Search the web");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_convert_tool_choice() {
|
||||
assert!(matches!(
|
||||
convert_tool_choice(Some("auto")),
|
||||
Some(RigToolChoice::Auto)
|
||||
));
|
||||
assert!(matches!(
|
||||
convert_tool_choice(Some("required")),
|
||||
Some(RigToolChoice::Required)
|
||||
));
|
||||
assert!(matches!(
|
||||
convert_tool_choice(Some("none")),
|
||||
Some(RigToolChoice::None)
|
||||
));
|
||||
assert!(matches!(
|
||||
convert_tool_choice(Some("AUTO")),
|
||||
Some(RigToolChoice::Auto)
|
||||
));
|
||||
assert!(convert_tool_choice(None).is_none());
|
||||
assert!(convert_tool_choice(Some("unknown")).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_response_text_only() {
|
||||
let content = OneOrMany::one(AssistantContent::text("Hello world"));
|
||||
let usage = RigUsage::new();
|
||||
let (text, calls, finish) = extract_response(&content, &usage);
|
||||
assert_eq!(text, Some("Hello world".to_string()));
|
||||
assert!(calls.is_empty());
|
||||
assert_eq!(finish, FinishReason::Stop);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_response_tool_call() {
|
||||
let tc = AssistantContent::tool_call("call_1", "search", serde_json::json!({"q": "test"}));
|
||||
let content = OneOrMany::one(tc);
|
||||
let usage = RigUsage::new();
|
||||
let (text, calls, finish) = extract_response(&content, &usage);
|
||||
assert!(text.is_none());
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].name, "search");
|
||||
assert_eq!(finish, FinishReason::ToolUse);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_saturate_u32() {
|
||||
assert_eq!(saturate_u32(100), 100);
|
||||
assert_eq!(saturate_u32(u64::MAX), u32::MAX);
|
||||
assert_eq!(saturate_u32(u32::MAX as u64), u32::MAX);
|
||||
}
|
||||
}
|
||||
+77
-2
@@ -61,6 +61,10 @@ pub struct SessionManager {
|
||||
token: RwLock<Option<SecretString>>,
|
||||
/// Prevents thundering herd during concurrent 401s.
|
||||
renewal_lock: Mutex<()>,
|
||||
/// Optional database store for persisting session to the settings table.
|
||||
store: RwLock<Option<Arc<crate::history::Store>>>,
|
||||
/// User ID for DB settings (default: "default").
|
||||
user_id: RwLock<String>,
|
||||
}
|
||||
|
||||
impl SessionManager {
|
||||
@@ -74,6 +78,8 @@ impl SessionManager {
|
||||
.unwrap_or_else(|_| Client::new()),
|
||||
token: RwLock::new(None),
|
||||
renewal_lock: Mutex::new(()),
|
||||
store: RwLock::new(None),
|
||||
user_id: RwLock::new("default".to_string()),
|
||||
};
|
||||
|
||||
// Try to load existing session synchronously during construction
|
||||
@@ -103,6 +109,8 @@ impl SessionManager {
|
||||
.unwrap_or_else(|_| Client::new()),
|
||||
token: RwLock::new(None),
|
||||
renewal_lock: Mutex::new(()),
|
||||
store: RwLock::new(None),
|
||||
user_id: RwLock::new("default".to_string()),
|
||||
};
|
||||
|
||||
if let Err(e) = manager.load_session().await {
|
||||
@@ -112,6 +120,21 @@ impl SessionManager {
|
||||
manager
|
||||
}
|
||||
|
||||
/// Attach a database store for persisting session tokens.
|
||||
///
|
||||
/// When a store is attached, session tokens are saved to the `settings`
|
||||
/// table (key: `nearai.session_token`) in addition to the disk file.
|
||||
/// On load, DB is preferred over disk.
|
||||
pub async fn attach_store(&self, store: Arc<crate::history::Store>, user_id: &str) {
|
||||
*self.store.write().await = Some(store);
|
||||
*self.user_id.write().await = user_id.to_string();
|
||||
|
||||
// Try to load from DB (may have been saved by a previous run)
|
||||
if let Err(e) = self.load_session_from_db().await {
|
||||
tracing::debug!("No session in DB: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the current session token, returning an error if not authenticated.
|
||||
pub async fn get_token(&self) -> Result<SecretString, LlmError> {
|
||||
let guard = self.token.read().await;
|
||||
@@ -460,7 +483,7 @@ impl SessionManager {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Save session data to disk.
|
||||
/// Save session data to disk and (if available) to the database.
|
||||
async fn save_session(&self, token: &str, auth_provider: Option<&str>) -> Result<(), LlmError> {
|
||||
let session = SessionData {
|
||||
session_token: token.to_string(),
|
||||
@@ -468,7 +491,7 @@ impl SessionManager {
|
||||
auth_provider: auth_provider.map(String::from),
|
||||
};
|
||||
|
||||
// Ensure parent directory exists
|
||||
// Save to disk (always, as bootstrap fallback)
|
||||
if let Some(parent) = self.config.session_path.parent() {
|
||||
tokio::fs::create_dir_all(parent).await.map_err(|e| {
|
||||
LlmError::Io(std::io::Error::new(
|
||||
@@ -498,6 +521,58 @@ impl SessionManager {
|
||||
})?;
|
||||
|
||||
tracing::debug!("Session saved to {}", self.config.session_path.display());
|
||||
|
||||
// Also save to DB if a store is attached
|
||||
if let Some(ref store) = *self.store.read().await {
|
||||
let user_id = self.user_id.read().await.clone();
|
||||
let session_json = serde_json::to_value(&session)
|
||||
.unwrap_or(serde_json::Value::String(token.to_string()));
|
||||
if let Err(e) = store
|
||||
.set_setting(&user_id, "nearai.session_token", &session_json)
|
||||
.await
|
||||
{
|
||||
tracing::warn!("Failed to save session to DB: {}", e);
|
||||
} else {
|
||||
tracing::debug!("Session also saved to DB settings");
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Try to load session from the database.
|
||||
async fn load_session_from_db(&self) -> Result<(), LlmError> {
|
||||
let store_guard = self.store.read().await;
|
||||
let store = store_guard
|
||||
.as_ref()
|
||||
.ok_or_else(|| LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: "No DB store attached".to_string(),
|
||||
})?;
|
||||
|
||||
let user_id = self.user_id.read().await.clone();
|
||||
let value = store
|
||||
.get_setting(&user_id, "nearai.session_token")
|
||||
.await
|
||||
.map_err(|e| LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("DB query failed: {}", e),
|
||||
})?
|
||||
.ok_or_else(|| LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: "No session in DB".to_string(),
|
||||
})?;
|
||||
|
||||
let session: SessionData =
|
||||
serde_json::from_value(value).map_err(|e| LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("Failed to parse DB session: {}", e),
|
||||
})?;
|
||||
|
||||
let mut guard = self.token.write().await;
|
||||
*guard = Some(SecretString::from(session.session_token));
|
||||
tracing::info!("Loaded session from DB settings");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
+315
-29
@@ -17,21 +17,26 @@ use ironclaw::{
|
||||
web::log_layer::{LogBroadcaster, WebLogLayer},
|
||||
},
|
||||
cli::{
|
||||
Cli, Command, run_mcp_command, run_memory_command, run_status_command, run_tool_command,
|
||||
Cli, Command, run_mcp_command, run_memory_command, run_pairing_command, run_status_command,
|
||||
run_tool_command,
|
||||
},
|
||||
config::Config,
|
||||
context::ContextManager,
|
||||
extensions::ExtensionManager,
|
||||
history::Store,
|
||||
llm::{SessionConfig, create_llm_provider, create_session_manager},
|
||||
orchestrator::{
|
||||
ContainerJobConfig, ContainerJobManager, OrchestratorApi, TokenStore,
|
||||
api::OrchestratorState,
|
||||
},
|
||||
pairing::PairingStore,
|
||||
safety::SafetyLayer,
|
||||
secrets::{PostgresSecretsStore, SecretsCrypto, SecretsStore},
|
||||
settings::Settings,
|
||||
setup::{SetupConfig, SetupWizard},
|
||||
tools::{
|
||||
ToolRegistry,
|
||||
mcp::{McpClient, McpSessionManager, config::load_mcp_servers, is_authenticated},
|
||||
wasm::{WasmToolLoader, WasmToolRuntime},
|
||||
mcp::{McpClient, McpSessionManager, config::load_mcp_servers_from_db, is_authenticated},
|
||||
wasm::{WasmToolLoader, WasmToolRuntime, load_dev_tools},
|
||||
},
|
||||
workspace::{EmbeddingProvider, NearAiEmbeddings, OpenAiEmbeddings, Workspace},
|
||||
};
|
||||
@@ -53,9 +58,14 @@ async fn main() -> anyhow::Result<()> {
|
||||
return run_tool_command(tool_cmd.clone()).await;
|
||||
}
|
||||
Some(Command::Config(config_cmd)) => {
|
||||
// Config commands don't need logging setup
|
||||
return ironclaw::cli::run_config_command(config_cmd.clone())
|
||||
.map_err(|e| anyhow::anyhow!("{}", e));
|
||||
// Config commands need DB access for settings
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")),
|
||||
)
|
||||
.init();
|
||||
|
||||
return ironclaw::cli::run_config_command(config_cmd.clone()).await;
|
||||
}
|
||||
Some(Command::Mcp(mcp_cmd)) => {
|
||||
// Simple logging for MCP commands
|
||||
@@ -76,7 +86,9 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
// Memory commands need database (and optionally embeddings)
|
||||
let _ = dotenvy::dotenv();
|
||||
let config = Config::from_env().map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
let config = Config::from_env()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
let store = ironclaw::history::Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
|
||||
@@ -120,6 +132,15 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
return run_memory_command(mem_cmd.clone(), store.pool(), embeddings).await;
|
||||
}
|
||||
Some(Command::Pairing(pairing_cmd)) => {
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")),
|
||||
)
|
||||
.init();
|
||||
|
||||
return run_pairing_command(pairing_cmd.clone()).map_err(|e| anyhow::anyhow!("{}", e));
|
||||
}
|
||||
Some(Command::Status) => {
|
||||
let _ = dotenvy::dotenv();
|
||||
tracing_subscriber::fmt()
|
||||
@@ -130,6 +151,83 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
return run_status_command().await;
|
||||
}
|
||||
Some(Command::Worker {
|
||||
job_id,
|
||||
orchestrator_url,
|
||||
max_iterations,
|
||||
}) => {
|
||||
// Worker mode: runs inside a Docker container.
|
||||
// Simple logging (no TUI, no DB, no channels).
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| EnvFilter::new("ironclaw=info")),
|
||||
)
|
||||
.init();
|
||||
|
||||
tracing::info!(
|
||||
"Starting worker for job {} (orchestrator: {})",
|
||||
job_id,
|
||||
orchestrator_url
|
||||
);
|
||||
|
||||
let config = ironclaw::worker::runtime::WorkerConfig {
|
||||
job_id: *job_id,
|
||||
orchestrator_url: orchestrator_url.clone(),
|
||||
max_iterations: *max_iterations,
|
||||
timeout: std::time::Duration::from_secs(600),
|
||||
};
|
||||
|
||||
let runtime = ironclaw::worker::WorkerRuntime::new(config)
|
||||
.map_err(|e| anyhow::anyhow!("Worker init failed: {}", e))?;
|
||||
|
||||
runtime
|
||||
.run()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Worker failed: {}", e))?;
|
||||
|
||||
return Ok(());
|
||||
}
|
||||
Some(Command::ClaudeBridge {
|
||||
job_id,
|
||||
orchestrator_url,
|
||||
max_turns,
|
||||
model,
|
||||
}) => {
|
||||
// Claude Code bridge mode: runs inside a Docker container.
|
||||
// Spawns the `claude` CLI and streams output to the orchestrator.
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| EnvFilter::new("ironclaw=info")),
|
||||
)
|
||||
.init();
|
||||
|
||||
tracing::info!(
|
||||
"Starting Claude Code bridge for job {} (orchestrator: {}, model: {})",
|
||||
job_id,
|
||||
orchestrator_url,
|
||||
model
|
||||
);
|
||||
|
||||
let config = ironclaw::worker::claude_bridge::ClaudeBridgeConfig {
|
||||
job_id: *job_id,
|
||||
orchestrator_url: orchestrator_url.clone(),
|
||||
max_turns: *max_turns,
|
||||
model: model.clone(),
|
||||
timeout: std::time::Duration::from_secs(1800),
|
||||
};
|
||||
|
||||
let runtime = ironclaw::worker::ClaudeBridgeRuntime::new(config)
|
||||
.map_err(|e| anyhow::anyhow!("Claude bridge init failed: {}", e))?;
|
||||
|
||||
runtime
|
||||
.run()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Claude bridge failed: {}", e))?;
|
||||
|
||||
return Ok(());
|
||||
}
|
||||
Some(Command::Onboard {
|
||||
skip_auth,
|
||||
channels_only,
|
||||
@@ -155,7 +253,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
// Enhanced first-run detection
|
||||
if !cli.no_onboard {
|
||||
if let Some(reason) = check_onboard_needed() {
|
||||
if let Some(reason) = check_onboard_needed().await {
|
||||
println!("Onboarding needed: {}", reason);
|
||||
println!();
|
||||
let mut wizard = SetupWizard::new();
|
||||
@@ -163,8 +261,11 @@ async fn main() -> anyhow::Result<()> {
|
||||
}
|
||||
}
|
||||
|
||||
// Load configuration (after potential setup)
|
||||
let config = match Config::from_env() {
|
||||
// Load bootstrap config (4 fields that must live on disk)
|
||||
let bootstrap = ironclaw::bootstrap::BootstrapConfig::load();
|
||||
|
||||
// Load initial config from env + disk (before DB is available)
|
||||
let mut config = match Config::from_env().await {
|
||||
Ok(c) => c,
|
||||
Err(ironclaw::error::ConfigError::MissingRequired { key, hint }) => {
|
||||
eprintln!("Configuration error: Missing required setting '{}'", key);
|
||||
@@ -186,8 +287,10 @@ async fn main() -> anyhow::Result<()> {
|
||||
};
|
||||
let session = create_session_manager(session_config).await;
|
||||
|
||||
// Ensure we're authenticated before proceeding (may trigger login flow)
|
||||
session.ensure_authenticated().await?;
|
||||
// Ensure we're authenticated before proceeding (only needed for NEAR AI backend)
|
||||
if config.llm.backend == ironclaw::config::LlmBackend::NearAi {
|
||||
session.ensure_authenticated().await?;
|
||||
}
|
||||
|
||||
// Initialize tracing
|
||||
let env_filter = EnvFilter::try_from_default_env()
|
||||
@@ -214,7 +317,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
tracing::info!("Starting IronClaw...");
|
||||
tracing::info!("Loaded configuration for agent: {}", config.agent.name);
|
||||
tracing::info!("NEAR AI session authenticated");
|
||||
tracing::info!("LLM backend: {}", config.llm.backend);
|
||||
|
||||
// Initialize database store (optional for testing)
|
||||
let store = if cli.no_db {
|
||||
@@ -224,7 +327,38 @@ async fn main() -> anyhow::Result<()> {
|
||||
let store = Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
tracing::info!("Database connected and migrations applied");
|
||||
Some(Arc::new(store))
|
||||
|
||||
// One-time migration: move disk config files into the DB settings table.
|
||||
if let Err(e) = ironclaw::bootstrap::migrate_disk_to_db(&store, "default").await {
|
||||
tracing::warn!("Disk-to-DB settings migration failed: {}", e);
|
||||
}
|
||||
|
||||
// Reload config from DB now that we have a connection.
|
||||
// Priority: env var > DB setting > default.
|
||||
match Config::from_db(&store, "default", &bootstrap).await {
|
||||
Ok(db_config) => {
|
||||
config = db_config;
|
||||
tracing::info!("Configuration reloaded from database");
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Failed to reload config from DB, keeping env-based config: {}",
|
||||
e
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let store = Arc::new(store);
|
||||
|
||||
// Attach store to session manager so tokens save to DB too
|
||||
session.attach_store(Arc::clone(&store), "default").await;
|
||||
|
||||
// Mark any jobs left in "running" or "creating" state as "interrupted".
|
||||
if let Err(e) = store.cleanup_stale_sandbox_jobs().await {
|
||||
tracing::warn!("Failed to cleanup stale sandbox jobs: {}", e);
|
||||
}
|
||||
|
||||
Some(store)
|
||||
};
|
||||
|
||||
// Initialize LLM provider (clone session so we can reuse it for embeddings)
|
||||
@@ -289,8 +423,11 @@ async fn main() -> anyhow::Result<()> {
|
||||
tools.register_memory_tools(workspace);
|
||||
}
|
||||
|
||||
// Register builder tool if enabled
|
||||
if config.builder.enabled {
|
||||
// Register builder tool if enabled.
|
||||
// When sandbox is enabled and allow_local_tools is false, skip builder registration
|
||||
// because register_builder_tool also registers dev tools (shell, file ops) that would
|
||||
// bypass the sandbox. The builder runs inside containers instead.
|
||||
if config.builder.enabled && (config.agent.allow_local_tools || !config.sandbox.enabled) {
|
||||
tools
|
||||
.register_builder_tool(
|
||||
llm.clone(),
|
||||
@@ -339,6 +476,8 @@ async fn main() -> anyhow::Result<()> {
|
||||
let wasm_tools_future = async {
|
||||
if let Some(ref runtime) = wasm_tool_runtime {
|
||||
let loader = WasmToolLoader::new(Arc::clone(runtime), Arc::clone(&tools));
|
||||
|
||||
// Load installed tools from ~/.ironclaw/tools/
|
||||
match loader.load_from_dir(&config.wasm.tools_dir).await {
|
||||
Ok(results) => {
|
||||
if !results.loaded.is_empty() {
|
||||
@@ -356,12 +495,32 @@ async fn main() -> anyhow::Result<()> {
|
||||
tracing::warn!("Failed to scan WASM tools directory: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
// Load dev tools from build artifacts (overrides installed if newer)
|
||||
match load_dev_tools(&loader, &config.wasm.tools_dir).await {
|
||||
Ok(results) => {
|
||||
if !results.loaded.is_empty() {
|
||||
tracing::info!(
|
||||
"Loaded {} dev WASM tools from build artifacts",
|
||||
results.loaded.len()
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::debug!("No dev WASM tools found: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let mcp_servers_future = async {
|
||||
if let Some(ref secrets) = secrets_store {
|
||||
match load_mcp_servers().await {
|
||||
let servers_result = if let Some(ref s) = store {
|
||||
load_mcp_servers_from_db(s, "default").await
|
||||
} else {
|
||||
ironclaw::tools::mcp::config::load_mcp_servers().await
|
||||
};
|
||||
match servers_result {
|
||||
Ok(servers) => {
|
||||
let enabled: Vec<_> = servers.enabled_servers().cloned().collect();
|
||||
if !enabled.is_empty() {
|
||||
@@ -470,6 +629,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
config.channels.wasm_channels_dir.clone(),
|
||||
config.tunnel.public_url.clone(),
|
||||
"default".to_string(),
|
||||
store.clone(),
|
||||
));
|
||||
tools.register_extension_tools(Arc::clone(&manager));
|
||||
tracing::info!("Extension manager initialized with in-chat discovery tools");
|
||||
@@ -482,6 +642,77 @@ async fn main() -> anyhow::Result<()> {
|
||||
None
|
||||
};
|
||||
|
||||
// Set up orchestrator for sandboxed job execution
|
||||
// When allow_local_tools is false (default), the LLM uses create_job for FS/shell work.
|
||||
// When allow_local_tools is true, dev tools are also registered directly (current behavior).
|
||||
if config.agent.allow_local_tools {
|
||||
tools.register_dev_tools();
|
||||
tracing::info!(
|
||||
"Local tools enabled (allow_local_tools=true), dev tools registered directly"
|
||||
);
|
||||
}
|
||||
|
||||
// Shared state for job events (used by both orchestrator and web gateway)
|
||||
let job_event_tx: Option<
|
||||
tokio::sync::broadcast::Sender<(uuid::Uuid, ironclaw::channels::web::types::SseEvent)>,
|
||||
> = if config.sandbox.enabled {
|
||||
let (tx, _) = tokio::sync::broadcast::channel(256);
|
||||
Some(tx)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let prompt_queue = Arc::new(tokio::sync::Mutex::new(std::collections::HashMap::<
|
||||
uuid::Uuid,
|
||||
std::collections::VecDeque<ironclaw::orchestrator::api::PendingPrompt>,
|
||||
>::new()));
|
||||
|
||||
let container_job_manager: Option<Arc<ContainerJobManager>> = if config.sandbox.enabled {
|
||||
let token_store = TokenStore::new();
|
||||
let job_config = ContainerJobConfig {
|
||||
image: config.sandbox.image.clone(),
|
||||
memory_limit_mb: config.sandbox.memory_limit_mb,
|
||||
cpu_shares: config.sandbox.cpu_shares,
|
||||
orchestrator_port: 50051,
|
||||
claude_config_dir: if config.claude_code.enabled {
|
||||
Some(config.claude_code.config_dir.clone())
|
||||
} else {
|
||||
None
|
||||
},
|
||||
claude_code_model: config.claude_code.model.clone(),
|
||||
claude_code_max_turns: config.claude_code.max_turns,
|
||||
claude_code_memory_limit_mb: config.claude_code.memory_limit_mb,
|
||||
};
|
||||
let jm = Arc::new(ContainerJobManager::new(job_config, token_store.clone()));
|
||||
|
||||
// Start the orchestrator internal API in the background
|
||||
let orchestrator_state = OrchestratorState {
|
||||
llm: llm.clone(),
|
||||
job_manager: Arc::clone(&jm),
|
||||
token_store,
|
||||
job_event_tx: job_event_tx.clone(),
|
||||
prompt_queue: Arc::clone(&prompt_queue),
|
||||
store: store.clone(),
|
||||
};
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = OrchestratorApi::start(orchestrator_state, 50051).await {
|
||||
tracing::error!("Orchestrator API failed: {}", e);
|
||||
}
|
||||
});
|
||||
|
||||
tracing::info!("Orchestrator API started on :50051, sandbox delegation enabled");
|
||||
if config.claude_code.enabled {
|
||||
tracing::info!(
|
||||
"Claude Code sandbox mode available (model: {}, max_turns: {})",
|
||||
config.claude_code.model,
|
||||
config.claude_code.max_turns
|
||||
);
|
||||
}
|
||||
Some(jm)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
tracing::info!(
|
||||
"Tool registry initialized with {} total tools",
|
||||
tools.count()
|
||||
@@ -507,7 +738,8 @@ async fn main() -> anyhow::Result<()> {
|
||||
match WasmChannelRuntime::new(WasmChannelRuntimeConfig::default()) {
|
||||
Ok(runtime) => {
|
||||
let runtime = Arc::new(runtime);
|
||||
let loader = WasmChannelLoader::new(Arc::clone(&runtime));
|
||||
let pairing_store = Arc::new(PairingStore::new());
|
||||
let loader = WasmChannelLoader::new(Arc::clone(&runtime), pairing_store);
|
||||
|
||||
match loader
|
||||
.load_from_dir(&config.channels.wasm_channels_dir)
|
||||
@@ -563,6 +795,17 @@ async fn main() -> anyhow::Result<()> {
|
||||
);
|
||||
}
|
||||
|
||||
// Inject owner_id for Telegram so the bot only responds
|
||||
// to the bound user account.
|
||||
if channel_name == "telegram" {
|
||||
if let Some(owner_id) = config.channels.telegram_owner_id {
|
||||
config_updates.insert(
|
||||
"owner_id".to_string(),
|
||||
serde_json::json!(owner_id),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
if !config_updates.is_empty() {
|
||||
channel_arc.update_config(config_updates).await;
|
||||
tracing::info!(
|
||||
@@ -621,7 +864,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
channels.add(Box::new(SharedWasmChannel::new(channel_arc)));
|
||||
}
|
||||
|
||||
if has_webhook_channels && config.tunnel.public_url.is_some() {
|
||||
if has_webhook_channels {
|
||||
webhook_routes.push(create_wasm_channel_router(
|
||||
wasm_router,
|
||||
extension_manager.as_ref().map(Arc::clone),
|
||||
@@ -693,6 +936,19 @@ async fn main() -> anyhow::Result<()> {
|
||||
Arc::new(ws)
|
||||
});
|
||||
|
||||
// Seed workspace with core identity files on first boot
|
||||
if let Some(ref ws) = workspace {
|
||||
match ws.seed_if_empty().await {
|
||||
Ok(count) if count > 0 => {
|
||||
tracing::info!("Workspace seeded with {} core files", count);
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to seed workspace: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Backfill embeddings if we just enabled the provider
|
||||
if let (Some(ws), Some(_)) = (&workspace, &embeddings) {
|
||||
match ws.backfill_embeddings().await {
|
||||
@@ -712,8 +968,12 @@ async fn main() -> anyhow::Result<()> {
|
||||
// Create session manager (shared between agent and web gateway)
|
||||
let session_manager = Arc::new(SessionManager::new());
|
||||
|
||||
// Register job tools
|
||||
tools.register_job_tools(Arc::clone(&context_manager));
|
||||
// Register job tools (sandbox deps auto-injected when container_job_manager is available)
|
||||
tools.register_job_tools(
|
||||
Arc::clone(&context_manager),
|
||||
container_job_manager.clone(),
|
||||
store.clone(),
|
||||
);
|
||||
|
||||
// Add web gateway channel if configured
|
||||
if let Some(ref gw_config) = config.channels.gateway {
|
||||
@@ -721,19 +981,44 @@ async fn main() -> anyhow::Result<()> {
|
||||
if let Some(ref ws) = workspace {
|
||||
gw = gw.with_workspace(Arc::clone(ws));
|
||||
}
|
||||
gw = gw.with_context_manager(Arc::clone(&context_manager));
|
||||
gw = gw.with_session_manager(Arc::clone(&session_manager));
|
||||
gw = gw.with_log_broadcaster(Arc::clone(&log_broadcaster));
|
||||
gw = gw.with_tool_registry(Arc::clone(&tools));
|
||||
if let Some(ref ext_mgr) = extension_manager {
|
||||
gw = gw.with_extension_manager(Arc::clone(ext_mgr));
|
||||
}
|
||||
if let Some(ref s) = store {
|
||||
gw = gw.with_store(Arc::clone(s));
|
||||
}
|
||||
if let Some(ref jm) = container_job_manager {
|
||||
gw = gw.with_job_manager(Arc::clone(jm));
|
||||
}
|
||||
if config.sandbox.enabled {
|
||||
gw = gw.with_prompt_queue(Arc::clone(&prompt_queue));
|
||||
|
||||
// Spawn a task to forward job events from the broadcast channel to SSE
|
||||
if let Some(ref tx) = job_event_tx {
|
||||
let mut rx = tx.subscribe();
|
||||
let gw_state = Arc::clone(gw.state());
|
||||
tokio::spawn(async move {
|
||||
while let Ok((_job_id, event)) = rx.recv().await {
|
||||
gw_state.sse.broadcast(event);
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
"Web gateway enabled on {}:{}",
|
||||
gw_config.host,
|
||||
gw_config.port
|
||||
);
|
||||
tracing::info!(
|
||||
"Web UI: http://{}:{}/?token={}",
|
||||
gw_config.host,
|
||||
gw_config.port,
|
||||
gw.auth_token()
|
||||
);
|
||||
|
||||
channels.add(Box::new(gw));
|
||||
}
|
||||
@@ -752,6 +1037,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
deps,
|
||||
channels,
|
||||
Some(config.heartbeat.clone()),
|
||||
Some(config.routines.clone()),
|
||||
Some(context_manager),
|
||||
Some(session_manager),
|
||||
);
|
||||
@@ -773,18 +1059,18 @@ async fn main() -> anyhow::Result<()> {
|
||||
/// Check if onboarding is needed and return the reason.
|
||||
///
|
||||
/// Returns `Some(reason)` if onboarding should be triggered, `None` otherwise.
|
||||
fn check_onboard_needed() -> Option<&'static str> {
|
||||
let settings = Settings::load();
|
||||
async fn check_onboard_needed() -> Option<&'static str> {
|
||||
let bootstrap = ironclaw::bootstrap::BootstrapConfig::load();
|
||||
|
||||
// Database not configured (and not in env)
|
||||
if settings.database_url.is_none() && std::env::var("DATABASE_URL").is_err() {
|
||||
if bootstrap.database_url.is_none() && std::env::var("DATABASE_URL").is_err() {
|
||||
return Some("Database not configured");
|
||||
}
|
||||
|
||||
// Secrets not configured (and not in env)
|
||||
if settings.secrets_master_key_source == ironclaw::settings::KeySource::None
|
||||
if bootstrap.secrets_master_key_source == ironclaw::settings::KeySource::None
|
||||
&& std::env::var("SECRETS_MASTER_KEY").is_err()
|
||||
&& !ironclaw::secrets::keychain::has_master_key()
|
||||
&& !ironclaw::secrets::keychain::has_master_key().await
|
||||
{
|
||||
// Only require secrets setup if user hasn't explicitly disabled it
|
||||
// For now, we don't require it for first run
|
||||
@@ -792,7 +1078,7 @@ fn check_onboard_needed() -> Option<&'static str> {
|
||||
|
||||
// First run (onboarding never completed and no session)
|
||||
let session_path = ironclaw::llm::session::default_session_path();
|
||||
if !settings.onboard_completed && !session_path.exists() {
|
||||
if !bootstrap.onboard_completed && !session_path.exists() {
|
||||
return Some("First run");
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,511 @@
|
||||
//! Internal HTTP API for worker-to-orchestrator communication.
|
||||
//!
|
||||
//! This runs on a separate port (default 50051) from the web gateway.
|
||||
//! All endpoints are authenticated via per-job bearer tokens.
|
||||
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::extract::{Path, State};
|
||||
use axum::http::StatusCode;
|
||||
use axum::routing::{get, post};
|
||||
use axum::{Json, Router};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::sync::{Mutex, broadcast};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::channels::web::types::SseEvent;
|
||||
use crate::history::Store;
|
||||
use crate::llm::{CompletionRequest, LlmProvider, ToolCompletionRequest};
|
||||
use crate::orchestrator::auth::{TokenStore, worker_auth_middleware};
|
||||
use crate::orchestrator::job_manager::ContainerJobManager;
|
||||
use crate::worker::api::JobEventPayload;
|
||||
use crate::worker::api::{
|
||||
CompletionReport, JobDescription, ProxyCompletionRequest, ProxyCompletionResponse,
|
||||
ProxyToolCompletionRequest, ProxyToolCompletionResponse, StatusUpdate,
|
||||
};
|
||||
|
||||
/// A follow-up prompt queued for a Claude Code bridge.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PendingPrompt {
|
||||
pub content: String,
|
||||
pub done: bool,
|
||||
}
|
||||
|
||||
/// Shared state for the orchestrator API.
|
||||
#[derive(Clone)]
|
||||
pub struct OrchestratorState {
|
||||
pub llm: Arc<dyn LlmProvider>,
|
||||
pub job_manager: Arc<ContainerJobManager>,
|
||||
pub token_store: TokenStore,
|
||||
/// Broadcast channel for job events (consumed by the web gateway SSE).
|
||||
pub job_event_tx: Option<broadcast::Sender<(Uuid, SseEvent)>>,
|
||||
/// Buffered follow-up prompts for sandbox jobs, keyed by job_id.
|
||||
pub prompt_queue: Arc<Mutex<HashMap<Uuid, VecDeque<PendingPrompt>>>>,
|
||||
/// Database handle for persisting job events.
|
||||
pub store: Option<Arc<Store>>,
|
||||
}
|
||||
|
||||
/// The orchestrator's internal API server.
|
||||
pub struct OrchestratorApi;
|
||||
|
||||
impl OrchestratorApi {
|
||||
/// Build the axum router for the internal API.
|
||||
pub fn router(state: OrchestratorState) -> Router {
|
||||
Router::new()
|
||||
// Worker routes: authenticated via route_layer middleware.
|
||||
.route("/worker/{job_id}/job", get(get_job))
|
||||
.route("/worker/{job_id}/llm/complete", post(llm_complete))
|
||||
.route(
|
||||
"/worker/{job_id}/llm/complete_with_tools",
|
||||
post(llm_complete_with_tools),
|
||||
)
|
||||
.route("/worker/{job_id}/status", post(report_status))
|
||||
.route("/worker/{job_id}/complete", post(report_complete))
|
||||
.route("/worker/{job_id}/event", post(job_event_handler))
|
||||
.route("/worker/{job_id}/prompt", get(get_prompt_handler))
|
||||
.route_layer(axum::middleware::from_fn_with_state(
|
||||
state.token_store.clone(),
|
||||
worker_auth_middleware,
|
||||
))
|
||||
// Unauthenticated routes (added after the layer).
|
||||
.route("/health", get(health_check))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
/// Start the internal API server on the given port.
|
||||
///
|
||||
/// On macOS/Windows (Docker Desktop), binds to loopback only because
|
||||
/// Docker Desktop routes `host.docker.internal` through its VM to the
|
||||
/// host's `127.0.0.1`.
|
||||
///
|
||||
/// On Linux, containers reach the host via the docker bridge gateway
|
||||
/// (`172.17.0.1`), which is NOT loopback. Binding to `127.0.0.1`
|
||||
/// would reject container traffic. We bind to all interfaces instead
|
||||
/// and rely on `worker_auth_middleware` (applied as a route_layer on
|
||||
/// every `/worker/` endpoint) to reject unauthenticated requests.
|
||||
pub async fn start(
|
||||
state: OrchestratorState,
|
||||
port: u16,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
let router = Self::router(state);
|
||||
let addr = if cfg!(target_os = "linux") {
|
||||
std::net::SocketAddr::from(([0, 0, 0, 0], port))
|
||||
} else {
|
||||
std::net::SocketAddr::from(([127, 0, 0, 1], port))
|
||||
};
|
||||
|
||||
tracing::info!("Orchestrator internal API listening on {}", addr);
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(addr).await?;
|
||||
axum::serve(listener, router).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// -- Handlers --
|
||||
//
|
||||
// All /worker/ handlers below are behind the worker_auth_middleware route_layer,
|
||||
// so they don't need to validate tokens themselves.
|
||||
|
||||
async fn health_check() -> &'static str {
|
||||
"ok"
|
||||
}
|
||||
|
||||
async fn get_job(
|
||||
State(state): State<OrchestratorState>,
|
||||
Path(job_id): Path<Uuid>,
|
||||
) -> Result<Json<JobDescription>, StatusCode> {
|
||||
let handle = state
|
||||
.job_manager
|
||||
.get_handle(job_id)
|
||||
.await
|
||||
.ok_or(StatusCode::NOT_FOUND)?;
|
||||
|
||||
Ok(Json(JobDescription {
|
||||
title: format!("Job {}", job_id),
|
||||
description: handle.task_description,
|
||||
project_dir: handle.project_dir.map(|p| p.display().to_string()),
|
||||
}))
|
||||
}
|
||||
|
||||
async fn llm_complete(
|
||||
State(state): State<OrchestratorState>,
|
||||
Path(job_id): Path<Uuid>,
|
||||
Json(req): Json<ProxyCompletionRequest>,
|
||||
) -> Result<Json<ProxyCompletionResponse>, StatusCode> {
|
||||
let completion_req = CompletionRequest {
|
||||
messages: req.messages,
|
||||
max_tokens: req.max_tokens,
|
||||
temperature: req.temperature,
|
||||
stop_sequences: req.stop_sequences,
|
||||
metadata: std::collections::HashMap::new(),
|
||||
};
|
||||
|
||||
let resp = state.llm.complete(completion_req).await.map_err(|e| {
|
||||
tracing::error!("LLM completion failed for job {}: {}", job_id, e);
|
||||
StatusCode::BAD_GATEWAY
|
||||
})?;
|
||||
|
||||
Ok(Json(ProxyCompletionResponse {
|
||||
content: resp.content,
|
||||
input_tokens: resp.input_tokens,
|
||||
output_tokens: resp.output_tokens,
|
||||
finish_reason: format_finish_reason(resp.finish_reason),
|
||||
}))
|
||||
}
|
||||
|
||||
async fn llm_complete_with_tools(
|
||||
State(state): State<OrchestratorState>,
|
||||
Path(job_id): Path<Uuid>,
|
||||
Json(req): Json<ProxyToolCompletionRequest>,
|
||||
) -> Result<Json<ProxyToolCompletionResponse>, StatusCode> {
|
||||
let tool_req = ToolCompletionRequest {
|
||||
messages: req.messages,
|
||||
tools: req.tools,
|
||||
max_tokens: req.max_tokens,
|
||||
temperature: req.temperature,
|
||||
tool_choice: req.tool_choice,
|
||||
metadata: std::collections::HashMap::new(),
|
||||
};
|
||||
|
||||
let resp = state.llm.complete_with_tools(tool_req).await.map_err(|e| {
|
||||
tracing::error!("LLM tool completion failed for job {}: {}", job_id, e);
|
||||
StatusCode::BAD_GATEWAY
|
||||
})?;
|
||||
|
||||
Ok(Json(ProxyToolCompletionResponse {
|
||||
content: resp.content,
|
||||
tool_calls: resp.tool_calls,
|
||||
input_tokens: resp.input_tokens,
|
||||
output_tokens: resp.output_tokens,
|
||||
finish_reason: format_finish_reason(resp.finish_reason),
|
||||
}))
|
||||
}
|
||||
|
||||
async fn report_status(
|
||||
Path(job_id): Path<Uuid>,
|
||||
Json(update): Json<StatusUpdate>,
|
||||
) -> Result<StatusCode, StatusCode> {
|
||||
tracing::debug!(
|
||||
job_id = %job_id,
|
||||
state = %update.state,
|
||||
iteration = update.iteration,
|
||||
"Worker status update"
|
||||
);
|
||||
|
||||
Ok(StatusCode::OK)
|
||||
}
|
||||
|
||||
async fn report_complete(
|
||||
State(state): State<OrchestratorState>,
|
||||
Path(job_id): Path<Uuid>,
|
||||
Json(report): Json<CompletionReport>,
|
||||
) -> Result<StatusCode, StatusCode> {
|
||||
if report.success {
|
||||
tracing::info!(
|
||||
job_id = %job_id,
|
||||
"Worker reported job complete"
|
||||
);
|
||||
} else {
|
||||
tracing::warn!(
|
||||
job_id = %job_id,
|
||||
message = ?report.message,
|
||||
"Worker reported job failure"
|
||||
);
|
||||
}
|
||||
|
||||
// Store the result and clean up the container
|
||||
let result = crate::orchestrator::job_manager::CompletionResult {
|
||||
success: report.success,
|
||||
message: report.message.clone(),
|
||||
};
|
||||
let _ = state.job_manager.complete_job(job_id, result).await;
|
||||
|
||||
Ok(StatusCode::OK)
|
||||
}
|
||||
|
||||
// -- Sandbox job event handlers --
|
||||
|
||||
/// Receive a job event from a worker or Claude Code bridge and broadcast + persist it.
|
||||
async fn job_event_handler(
|
||||
State(state): State<OrchestratorState>,
|
||||
Path(job_id): Path<Uuid>,
|
||||
Json(payload): Json<JobEventPayload>,
|
||||
) -> Result<StatusCode, StatusCode> {
|
||||
tracing::debug!(
|
||||
job_id = %job_id,
|
||||
event_type = %payload.event_type,
|
||||
"Job event received"
|
||||
);
|
||||
|
||||
// Persist to DB (fire-and-forget)
|
||||
if let Some(ref store) = state.store {
|
||||
let store = Arc::clone(store);
|
||||
let event_type = payload.event_type.clone();
|
||||
let data = payload.data.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = store.save_job_event(job_id, &event_type, &data).await {
|
||||
tracing::warn!(job_id = %job_id, "Failed to persist job event: {}", e);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// Convert to SSE event and broadcast
|
||||
let job_id_str = job_id.to_string();
|
||||
let sse_event = match payload.event_type.as_str() {
|
||||
"message" => SseEvent::JobMessage {
|
||||
job_id: job_id_str,
|
||||
role: payload
|
||||
.data
|
||||
.get("role")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("assistant")
|
||||
.to_string(),
|
||||
content: payload
|
||||
.data
|
||||
.get("content")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
},
|
||||
"tool_use" => SseEvent::JobToolUse {
|
||||
job_id: job_id_str,
|
||||
tool_name: payload
|
||||
.data
|
||||
.get("tool_name")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("unknown")
|
||||
.to_string(),
|
||||
input: payload
|
||||
.data
|
||||
.get("input")
|
||||
.cloned()
|
||||
.unwrap_or(serde_json::Value::Null),
|
||||
},
|
||||
"tool_result" => SseEvent::JobToolResult {
|
||||
job_id: job_id_str,
|
||||
tool_name: payload
|
||||
.data
|
||||
.get("tool_name")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("unknown")
|
||||
.to_string(),
|
||||
output: payload
|
||||
.data
|
||||
.get("output")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
},
|
||||
"result" => SseEvent::JobResult {
|
||||
job_id: job_id_str,
|
||||
status: payload
|
||||
.data
|
||||
.get("status")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("unknown")
|
||||
.to_string(),
|
||||
session_id: payload
|
||||
.data
|
||||
.get("session_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string()),
|
||||
},
|
||||
_ => SseEvent::JobStatus {
|
||||
job_id: job_id_str,
|
||||
message: payload
|
||||
.data
|
||||
.get("message")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
},
|
||||
};
|
||||
|
||||
// Broadcast via the channel (if configured)
|
||||
if let Some(ref tx) = state.job_event_tx {
|
||||
let _ = tx.send((job_id, sse_event));
|
||||
}
|
||||
|
||||
Ok(StatusCode::OK)
|
||||
}
|
||||
|
||||
/// Return the next queued follow-up prompt for a Claude Code bridge.
|
||||
/// Returns 204 No Content if no prompt is available.
|
||||
async fn get_prompt_handler(
|
||||
State(state): State<OrchestratorState>,
|
||||
Path(job_id): Path<Uuid>,
|
||||
) -> Result<(StatusCode, Json<serde_json::Value>), StatusCode> {
|
||||
let mut queue = state.prompt_queue.lock().await;
|
||||
if let Some(prompts) = queue.get_mut(&job_id) {
|
||||
if let Some(prompt) = prompts.pop_front() {
|
||||
return Ok((
|
||||
StatusCode::OK,
|
||||
Json(serde_json::json!({
|
||||
"content": prompt.content,
|
||||
"done": prompt.done,
|
||||
})),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
// Return 204 with an empty body. The Json wrapper requires some value
|
||||
// but the status code signals "nothing here".
|
||||
Ok((StatusCode::NO_CONTENT, Json(serde_json::Value::Null)))
|
||||
}
|
||||
|
||||
fn format_finish_reason(reason: crate::llm::FinishReason) -> String {
|
||||
match reason {
|
||||
crate::llm::FinishReason::Stop => "stop".to_string(),
|
||||
crate::llm::FinishReason::Length => "length".to_string(),
|
||||
crate::llm::FinishReason::ToolUse => "tool_use".to_string(),
|
||||
crate::llm::FinishReason::ContentFilter => "content_filter".to_string(),
|
||||
crate::llm::FinishReason::Unknown => "unknown".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::http::Request;
|
||||
use tower::ServiceExt;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::error::LlmError;
|
||||
use crate::llm::{
|
||||
CompletionRequest, CompletionResponse, ToolCompletionRequest, ToolCompletionResponse,
|
||||
};
|
||||
use crate::orchestrator::auth::TokenStore;
|
||||
use crate::orchestrator::job_manager::{ContainerJobConfig, ContainerJobManager};
|
||||
|
||||
use super::*;
|
||||
|
||||
/// Stub LLM provider that panics if called (tests only exercise routing/auth).
|
||||
struct StubLlm;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl crate::llm::LlmProvider for StubLlm {
|
||||
fn model_name(&self) -> &str {
|
||||
"stub"
|
||||
}
|
||||
fn cost_per_token(&self) -> (rust_decimal::Decimal, rust_decimal::Decimal) {
|
||||
(rust_decimal::Decimal::ZERO, rust_decimal::Decimal::ZERO)
|
||||
}
|
||||
async fn complete(&self, _req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||
Err(LlmError::RequestFailed {
|
||||
provider: "stub".into(),
|
||||
reason: "not implemented".into(),
|
||||
})
|
||||
}
|
||||
async fn complete_with_tools(
|
||||
&self,
|
||||
_req: ToolCompletionRequest,
|
||||
) -> Result<ToolCompletionResponse, LlmError> {
|
||||
Err(LlmError::RequestFailed {
|
||||
provider: "stub".into(),
|
||||
reason: "not implemented".into(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn test_state() -> OrchestratorState {
|
||||
let token_store = TokenStore::new();
|
||||
let jm = ContainerJobManager::new(ContainerJobConfig::default(), token_store.clone());
|
||||
OrchestratorState {
|
||||
llm: Arc::new(StubLlm),
|
||||
job_manager: Arc::new(jm),
|
||||
token_store,
|
||||
job_event_tx: None,
|
||||
prompt_queue: Arc::new(Mutex::new(HashMap::new())),
|
||||
store: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn health_requires_no_auth() {
|
||||
let state = test_state();
|
||||
let router = OrchestratorApi::router(state);
|
||||
|
||||
let req = Request::builder()
|
||||
.uri("/health")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
|
||||
let resp = router.oneshot(req).await.unwrap();
|
||||
assert_eq!(resp.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn worker_route_rejects_missing_token() {
|
||||
let state = test_state();
|
||||
let router = OrchestratorApi::router(state);
|
||||
|
||||
let job_id = Uuid::new_v4();
|
||||
let req = Request::builder()
|
||||
.uri(format!("/worker/{}/job", job_id))
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
|
||||
let resp = router.oneshot(req).await.unwrap();
|
||||
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn worker_route_rejects_wrong_token() {
|
||||
let state = test_state();
|
||||
let router = OrchestratorApi::router(state);
|
||||
|
||||
let job_id = Uuid::new_v4();
|
||||
let req = Request::builder()
|
||||
.uri(format!("/worker/{}/job", job_id))
|
||||
.header("Authorization", "Bearer totally-bogus")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
|
||||
let resp = router.oneshot(req).await.unwrap();
|
||||
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn worker_route_accepts_valid_token() {
|
||||
let state = test_state();
|
||||
let job_id = Uuid::new_v4();
|
||||
let token = state.token_store.create_token(job_id).await;
|
||||
|
||||
let router = OrchestratorApi::router(state);
|
||||
|
||||
let req = Request::builder()
|
||||
.uri(format!("/worker/{}/job", job_id))
|
||||
.header("Authorization", format!("Bearer {}", token))
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
|
||||
let resp = router.oneshot(req).await.unwrap();
|
||||
// 404 because no container exists for this job_id, but NOT 401.
|
||||
assert_eq!(resp.status(), StatusCode::NOT_FOUND);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn token_for_job_a_rejected_on_job_b() {
|
||||
let state = test_state();
|
||||
let job_a = Uuid::new_v4();
|
||||
let job_b = Uuid::new_v4();
|
||||
let token_a = state.token_store.create_token(job_a).await;
|
||||
|
||||
let router = OrchestratorApi::router(state);
|
||||
|
||||
// Use job_a's token to hit job_b's endpoint
|
||||
let req = Request::builder()
|
||||
.uri(format!("/worker/{}/job", job_b))
|
||||
.header("Authorization", format!("Bearer {}", token_a))
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
|
||||
let resp = router.oneshot(req).await.unwrap();
|
||||
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,162 @@
|
||||
//! Per-job bearer token authentication for worker-to-orchestrator communication.
|
||||
//!
|
||||
//! Security properties:
|
||||
//! - Tokens are cryptographically random (32 bytes, hex-encoded)
|
||||
//! - Tokens are scoped to a specific job_id
|
||||
//! - Tokens are ephemeral (in-memory only, never persisted)
|
||||
//! - A token for Job A cannot access endpoints for Job B
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::extract::{Request, State};
|
||||
use axum::http::StatusCode;
|
||||
use axum::middleware::Next;
|
||||
use axum::response::Response;
|
||||
use rand::Rng;
|
||||
use tokio::sync::RwLock;
|
||||
use uuid::Uuid;
|
||||
|
||||
/// In-memory store for per-job authentication tokens.
|
||||
#[derive(Clone)]
|
||||
pub struct TokenStore {
|
||||
/// Maps job_id -> bearer token. Never logged or persisted.
|
||||
tokens: Arc<RwLock<HashMap<Uuid, String>>>,
|
||||
}
|
||||
|
||||
impl TokenStore {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
tokens: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Generate and store a new token for a job.
|
||||
pub async fn create_token(&self, job_id: Uuid) -> String {
|
||||
let token = generate_token();
|
||||
self.tokens.write().await.insert(job_id, token.clone());
|
||||
token
|
||||
}
|
||||
|
||||
/// Validate a token for a specific job.
|
||||
pub async fn validate(&self, job_id: Uuid, token: &str) -> bool {
|
||||
self.tokens
|
||||
.read()
|
||||
.await
|
||||
.get(&job_id)
|
||||
.map(|stored| stored == token)
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
/// Remove a token (on container cleanup).
|
||||
pub async fn revoke(&self, job_id: Uuid) {
|
||||
self.tokens.write().await.remove(&job_id);
|
||||
}
|
||||
|
||||
/// Get the number of active tokens (for diagnostics).
|
||||
pub async fn active_count(&self) -> usize {
|
||||
self.tokens.read().await.len()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for TokenStore {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Generate a cryptographically random token (32 bytes, hex-encoded = 64 chars).
|
||||
fn generate_token() -> String {
|
||||
let mut bytes = [0u8; 32];
|
||||
rand::thread_rng().fill(&mut bytes);
|
||||
hex_encode(&bytes)
|
||||
}
|
||||
|
||||
fn hex_encode(bytes: &[u8]) -> String {
|
||||
bytes.iter().map(|b| format!("{:02x}", b)).collect()
|
||||
}
|
||||
|
||||
/// Axum middleware that validates worker bearer tokens.
|
||||
///
|
||||
/// Extracts the job_id from the path (`/worker/{job_id}/...`) and validates
|
||||
/// the `Authorization: Bearer <token>` header against the token store.
|
||||
///
|
||||
/// Wire up with `axum::middleware::from_fn_with_state(token_store, worker_auth_middleware)`.
|
||||
pub async fn worker_auth_middleware(
|
||||
State(token_store): State<TokenStore>,
|
||||
request: Request,
|
||||
next: Next,
|
||||
) -> Result<Response, StatusCode> {
|
||||
let path = request.uri().path().to_string();
|
||||
let job_id = extract_job_id_from_path(&path).ok_or(StatusCode::BAD_REQUEST)?;
|
||||
|
||||
let token = request
|
||||
.headers()
|
||||
.get("authorization")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|v| v.strip_prefix("Bearer "))
|
||||
.ok_or(StatusCode::UNAUTHORIZED)?;
|
||||
|
||||
if !token_store.validate(job_id, token).await {
|
||||
return Err(StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
|
||||
Ok(next.run(request).await)
|
||||
}
|
||||
|
||||
/// Extract job UUID from a path like `/worker/{uuid}/...`
|
||||
fn extract_job_id_from_path(path: &str) -> Option<Uuid> {
|
||||
let parts: Vec<&str> = path.trim_start_matches('/').split('/').collect();
|
||||
if parts.len() >= 2 && parts[0] == "worker" {
|
||||
Uuid::parse_str(parts[1]).ok()
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_token_create_and_validate() {
|
||||
let store = TokenStore::new();
|
||||
let job_id = Uuid::new_v4();
|
||||
|
||||
let token = store.create_token(job_id).await;
|
||||
assert_eq!(token.len(), 64); // 32 bytes hex = 64 chars
|
||||
|
||||
assert!(store.validate(job_id, &token).await);
|
||||
assert!(!store.validate(job_id, "wrong-token").await);
|
||||
assert!(!store.validate(Uuid::new_v4(), &token).await);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_token_revoke() {
|
||||
let store = TokenStore::new();
|
||||
let job_id = Uuid::new_v4();
|
||||
|
||||
let token = store.create_token(job_id).await;
|
||||
assert!(store.validate(job_id, &token).await);
|
||||
|
||||
store.revoke(job_id).await;
|
||||
assert!(!store.validate(job_id, &token).await);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_job_id() {
|
||||
let id = Uuid::new_v4();
|
||||
let path = format!("/worker/{}/llm/complete", id);
|
||||
assert_eq!(extract_job_id_from_path(&path), Some(id));
|
||||
|
||||
assert_eq!(extract_job_id_from_path("/other/path"), None);
|
||||
assert_eq!(extract_job_id_from_path("/worker/not-a-uuid/foo"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_token_is_random() {
|
||||
let t1 = generate_token();
|
||||
let t2 = generate_token();
|
||||
assert_ne!(t1, t2);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,474 @@
|
||||
//! Container lifecycle management for sandboxed jobs.
|
||||
//!
|
||||
//! Extends the existing `SandboxManager` infrastructure to support persistent
|
||||
//! containers with their own agent loops (as opposed to ephemeral per-command containers).
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use tokio::sync::RwLock;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::error::OrchestratorError;
|
||||
use crate::orchestrator::auth::TokenStore;
|
||||
use crate::sandbox::connect_docker;
|
||||
|
||||
/// Which mode a sandbox container runs in.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum JobMode {
|
||||
/// Standard IronClaw worker with proxied LLM calls.
|
||||
Worker,
|
||||
/// Claude Code bridge that spawns the `claude` CLI directly.
|
||||
ClaudeCode,
|
||||
}
|
||||
|
||||
impl JobMode {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Worker => "worker",
|
||||
Self::ClaudeCode => "claude_code",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for JobMode {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}", self.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
/// Configuration for the container job manager.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ContainerJobConfig {
|
||||
/// Docker image for worker containers.
|
||||
pub image: String,
|
||||
/// Default memory limit in MB.
|
||||
pub memory_limit_mb: u64,
|
||||
/// Default CPU shares.
|
||||
pub cpu_shares: u32,
|
||||
/// Port the orchestrator internal API listens on.
|
||||
pub orchestrator_port: u16,
|
||||
/// Host directory containing Claude auth config (mounted read-only for ClaudeCode mode).
|
||||
pub claude_config_dir: Option<PathBuf>,
|
||||
/// Claude model to use in ClaudeCode mode.
|
||||
pub claude_code_model: String,
|
||||
/// Maximum turns for Claude Code.
|
||||
pub claude_code_max_turns: u32,
|
||||
/// Memory limit in MB for Claude Code containers (heavier than workers).
|
||||
pub claude_code_memory_limit_mb: u64,
|
||||
}
|
||||
|
||||
impl Default for ContainerJobConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
image: "ironclaw-worker:latest".to_string(),
|
||||
memory_limit_mb: 2048,
|
||||
cpu_shares: 1024,
|
||||
orchestrator_port: 50051,
|
||||
claude_config_dir: None,
|
||||
claude_code_model: "sonnet".to_string(),
|
||||
claude_code_max_turns: 50,
|
||||
claude_code_memory_limit_mb: 4096,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// State of a container.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ContainerState {
|
||||
Creating,
|
||||
Running,
|
||||
Stopped,
|
||||
Failed,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for ContainerState {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Creating => write!(f, "creating"),
|
||||
Self::Running => write!(f, "running"),
|
||||
Self::Stopped => write!(f, "stopped"),
|
||||
Self::Failed => write!(f, "failed"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle to a running container job.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ContainerHandle {
|
||||
pub job_id: Uuid,
|
||||
pub container_id: String,
|
||||
pub state: ContainerState,
|
||||
pub mode: JobMode,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub project_dir: Option<PathBuf>,
|
||||
pub task_description: String,
|
||||
/// Completion result from the worker (set when the worker reports done).
|
||||
pub completion_result: Option<CompletionResult>,
|
||||
// NOTE: auth_token is intentionally NOT in this struct.
|
||||
// It lives only in the TokenStore (never logged, serialized, or persisted).
|
||||
}
|
||||
|
||||
/// Result reported by a worker on completion.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CompletionResult {
|
||||
pub success: bool,
|
||||
pub message: Option<String>,
|
||||
}
|
||||
|
||||
/// Manages the lifecycle of Docker containers for sandboxed job execution.
|
||||
pub struct ContainerJobManager {
|
||||
config: ContainerJobConfig,
|
||||
token_store: TokenStore,
|
||||
containers: Arc<RwLock<HashMap<Uuid, ContainerHandle>>>,
|
||||
}
|
||||
|
||||
impl ContainerJobManager {
|
||||
pub fn new(config: ContainerJobConfig, token_store: TokenStore) -> Self {
|
||||
Self {
|
||||
config,
|
||||
token_store,
|
||||
containers: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create and start a new container for a job.
|
||||
///
|
||||
/// The caller provides the `job_id` so it can be persisted to the database
|
||||
/// before the container is created. Returns the auth token for the worker.
|
||||
pub async fn create_job(
|
||||
&self,
|
||||
job_id: Uuid,
|
||||
task: &str,
|
||||
project_dir: Option<PathBuf>,
|
||||
mode: JobMode,
|
||||
) -> Result<String, OrchestratorError> {
|
||||
// Generate auth token (stored in TokenStore, never logged)
|
||||
let token = self.token_store.create_token(job_id).await;
|
||||
|
||||
// Record the handle
|
||||
let handle = ContainerHandle {
|
||||
job_id,
|
||||
container_id: String::new(), // set after container creation
|
||||
state: ContainerState::Creating,
|
||||
mode,
|
||||
created_at: Utc::now(),
|
||||
project_dir: project_dir.clone(),
|
||||
task_description: task.to_string(),
|
||||
completion_result: None,
|
||||
};
|
||||
self.containers.write().await.insert(job_id, handle);
|
||||
|
||||
// Connect to Docker
|
||||
let docker = connect_docker()
|
||||
.await
|
||||
.map_err(|e| OrchestratorError::Docker {
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
// Build container configuration
|
||||
let orchestrator_host = if cfg!(target_os = "linux") {
|
||||
"172.17.0.1"
|
||||
} else {
|
||||
"host.docker.internal"
|
||||
};
|
||||
|
||||
let orchestrator_url = format!(
|
||||
"http://{}:{}",
|
||||
orchestrator_host, self.config.orchestrator_port
|
||||
);
|
||||
|
||||
let mut env_vec = vec![
|
||||
format!("IRONCLAW_WORKER_TOKEN={}", token),
|
||||
format!("IRONCLAW_JOB_ID={}", job_id),
|
||||
format!("IRONCLAW_ORCHESTRATOR_URL={}", orchestrator_url),
|
||||
];
|
||||
|
||||
// Build volume mounts (validate project_dir stays within ~/.ironclaw/projects/)
|
||||
let mut binds = Vec::new();
|
||||
if let Some(ref dir) = project_dir {
|
||||
let canonical =
|
||||
dir.canonicalize()
|
||||
.map_err(|e| OrchestratorError::ContainerCreationFailed {
|
||||
job_id,
|
||||
reason: format!(
|
||||
"failed to canonicalize project dir {}: {}",
|
||||
dir.display(),
|
||||
e
|
||||
),
|
||||
})?;
|
||||
let projects_base = dirs::home_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join(".ironclaw")
|
||||
.join("projects");
|
||||
if let Ok(canonical_base) = projects_base.canonicalize() {
|
||||
if !canonical.starts_with(&canonical_base) {
|
||||
return Err(OrchestratorError::ContainerCreationFailed {
|
||||
job_id,
|
||||
reason: format!(
|
||||
"project directory {} is outside allowed base {}",
|
||||
canonical.display(),
|
||||
canonical_base.display()
|
||||
),
|
||||
});
|
||||
}
|
||||
}
|
||||
binds.push(format!("{}:/workspace:rw", canonical.display()));
|
||||
env_vec.push("IRONCLAW_WORKSPACE=/workspace".to_string());
|
||||
}
|
||||
|
||||
// Claude Code mode: mount host ~/.claude read-only for auth
|
||||
if mode == JobMode::ClaudeCode {
|
||||
if let Some(ref claude_dir) = self.config.claude_config_dir {
|
||||
binds.push(format!("{}:/home/sandbox/.claude:ro", claude_dir.display()));
|
||||
}
|
||||
}
|
||||
|
||||
// Memory limit: Claude Code gets more memory
|
||||
let memory_mb = match mode {
|
||||
JobMode::ClaudeCode => self.config.claude_code_memory_limit_mb,
|
||||
JobMode::Worker => self.config.memory_limit_mb,
|
||||
};
|
||||
|
||||
// Create the container
|
||||
use bollard::container::{Config, CreateContainerOptions};
|
||||
use bollard::models::HostConfig;
|
||||
|
||||
let host_config = HostConfig {
|
||||
binds: if binds.is_empty() { None } else { Some(binds) },
|
||||
memory: Some((memory_mb * 1024 * 1024) as i64),
|
||||
cpu_shares: Some(self.config.cpu_shares as i64),
|
||||
network_mode: Some("bridge".to_string()),
|
||||
extra_hosts: Some(vec!["host.docker.internal:host-gateway".to_string()]),
|
||||
cap_drop: Some(vec!["ALL".to_string()]),
|
||||
cap_add: Some(vec![
|
||||
"CHOWN".to_string(),
|
||||
"SETUID".to_string(),
|
||||
"SETGID".to_string(),
|
||||
]),
|
||||
security_opt: Some(vec!["no-new-privileges:true".to_string()]),
|
||||
tmpfs: Some(
|
||||
[("/tmp".to_string(), "size=512M".to_string())]
|
||||
.into_iter()
|
||||
.collect(),
|
||||
),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
// Build CMD based on mode
|
||||
let cmd = match mode {
|
||||
JobMode::Worker => vec![
|
||||
"worker".to_string(),
|
||||
"--job-id".to_string(),
|
||||
job_id.to_string(),
|
||||
"--orchestrator-url".to_string(),
|
||||
orchestrator_url,
|
||||
],
|
||||
JobMode::ClaudeCode => vec![
|
||||
"claude-bridge".to_string(),
|
||||
"--job-id".to_string(),
|
||||
job_id.to_string(),
|
||||
"--orchestrator-url".to_string(),
|
||||
orchestrator_url,
|
||||
"--max-turns".to_string(),
|
||||
self.config.claude_code_max_turns.to_string(),
|
||||
"--model".to_string(),
|
||||
self.config.claude_code_model.clone(),
|
||||
],
|
||||
};
|
||||
|
||||
let container_config = Config {
|
||||
image: Some(self.config.image.clone()),
|
||||
cmd: Some(cmd),
|
||||
env: Some(env_vec),
|
||||
host_config: Some(host_config),
|
||||
user: Some("1000:1000".to_string()),
|
||||
working_dir: Some("/workspace".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let container_name = match mode {
|
||||
JobMode::Worker => format!("ironclaw-worker-{}", job_id),
|
||||
JobMode::ClaudeCode => format!("ironclaw-claude-{}", job_id),
|
||||
};
|
||||
let options = CreateContainerOptions {
|
||||
name: container_name,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let response = docker
|
||||
.create_container(Some(options), container_config)
|
||||
.await
|
||||
.map_err(|e| OrchestratorError::ContainerCreationFailed {
|
||||
job_id,
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
let container_id = response.id;
|
||||
|
||||
// Start the container
|
||||
docker
|
||||
.start_container::<String>(&container_id, None)
|
||||
.await
|
||||
.map_err(|e| OrchestratorError::ContainerCreationFailed {
|
||||
job_id,
|
||||
reason: format!("failed to start container: {}", e),
|
||||
})?;
|
||||
|
||||
// Update handle with container ID
|
||||
if let Some(handle) = self.containers.write().await.get_mut(&job_id) {
|
||||
handle.container_id = container_id;
|
||||
handle.state = ContainerState::Running;
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
job_id = %job_id,
|
||||
"Created and started worker container"
|
||||
);
|
||||
|
||||
Ok(token)
|
||||
}
|
||||
|
||||
/// Stop a running container job.
|
||||
pub async fn stop_job(&self, job_id: Uuid) -> Result<(), OrchestratorError> {
|
||||
let container_id = {
|
||||
let containers = self.containers.read().await;
|
||||
containers
|
||||
.get(&job_id)
|
||||
.map(|h| h.container_id.clone())
|
||||
.ok_or(OrchestratorError::ContainerNotFound { job_id })?
|
||||
};
|
||||
|
||||
if container_id.is_empty() {
|
||||
return Err(OrchestratorError::InvalidContainerState {
|
||||
job_id,
|
||||
state: "creating (no container ID yet)".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
let docker = connect_docker()
|
||||
.await
|
||||
.map_err(|e| OrchestratorError::Docker {
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
// Stop the container (10 second grace period)
|
||||
let _ = docker
|
||||
.stop_container(
|
||||
&container_id,
|
||||
Some(bollard::container::StopContainerOptions { t: 10 }),
|
||||
)
|
||||
.await;
|
||||
|
||||
// Remove the container
|
||||
let _ = docker
|
||||
.remove_container(
|
||||
&container_id,
|
||||
Some(bollard::container::RemoveContainerOptions {
|
||||
force: true,
|
||||
..Default::default()
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
// Update state
|
||||
if let Some(handle) = self.containers.write().await.get_mut(&job_id) {
|
||||
handle.state = ContainerState::Stopped;
|
||||
}
|
||||
|
||||
// Revoke the auth token
|
||||
self.token_store.revoke(job_id).await;
|
||||
|
||||
tracing::info!(job_id = %job_id, "Stopped worker container");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Mark a job as complete with a result. The container is stopped but the
|
||||
/// handle is kept so `CreateJobTool` can read the completion message.
|
||||
pub async fn complete_job(
|
||||
&self,
|
||||
job_id: Uuid,
|
||||
result: CompletionResult,
|
||||
) -> Result<(), OrchestratorError> {
|
||||
// Store the result before stopping
|
||||
{
|
||||
let mut containers = self.containers.write().await;
|
||||
if let Some(handle) = containers.get_mut(&job_id) {
|
||||
handle.completion_result = Some(result);
|
||||
handle.state = ContainerState::Stopped;
|
||||
}
|
||||
}
|
||||
|
||||
// Stop container and revoke token (but keep handle in map)
|
||||
let container_id = {
|
||||
let containers = self.containers.read().await;
|
||||
containers.get(&job_id).map(|h| h.container_id.clone())
|
||||
};
|
||||
if let Some(cid) = container_id {
|
||||
if !cid.is_empty() {
|
||||
if let Ok(docker) = connect_docker().await {
|
||||
let _ = docker
|
||||
.stop_container(
|
||||
&cid,
|
||||
Some(bollard::container::StopContainerOptions { t: 5 }),
|
||||
)
|
||||
.await;
|
||||
let _ = docker
|
||||
.remove_container(
|
||||
&cid,
|
||||
Some(bollard::container::RemoveContainerOptions {
|
||||
force: true,
|
||||
..Default::default()
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
self.token_store.revoke(job_id).await;
|
||||
|
||||
tracing::info!(job_id = %job_id, "Completed worker container");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Remove a completed job handle from memory (called after result is read).
|
||||
pub async fn cleanup_job(&self, job_id: Uuid) {
|
||||
self.containers.write().await.remove(&job_id);
|
||||
}
|
||||
|
||||
/// Get the handle for a job.
|
||||
pub async fn get_handle(&self, job_id: Uuid) -> Option<ContainerHandle> {
|
||||
self.containers.read().await.get(&job_id).cloned()
|
||||
}
|
||||
|
||||
/// List all active container jobs.
|
||||
pub async fn list_jobs(&self) -> Vec<ContainerHandle> {
|
||||
self.containers.read().await.values().cloned().collect()
|
||||
}
|
||||
|
||||
/// Get a reference to the token store.
|
||||
pub fn token_store(&self) -> &TokenStore {
|
||||
&self.token_store
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_container_job_config_default() {
|
||||
let config = ContainerJobConfig::default();
|
||||
assert_eq!(config.orchestrator_port, 50051);
|
||||
assert_eq!(config.memory_limit_mb, 2048);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_container_state_display() {
|
||||
assert_eq!(ContainerState::Running.to_string(), "running");
|
||||
assert_eq!(ContainerState::Stopped.to_string(), "stopped");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
//! Orchestrator for managing sandboxed worker containers.
|
||||
//!
|
||||
//! The orchestrator runs in the main agent process and provides:
|
||||
//! - An internal HTTP API for worker communication (LLM proxy, status, secrets)
|
||||
//! - Per-job bearer token authentication
|
||||
//! - Container lifecycle management (create, monitor, stop)
|
||||
//!
|
||||
//! ```text
|
||||
//! ┌───────────────────────────────────────────────┐
|
||||
//! │ Orchestrator │
|
||||
//! │ │
|
||||
//! │ Internal API (:50051) │
|
||||
//! │ POST /worker/{id}/llm/complete │
|
||||
//! │ POST /worker/{id}/llm/complete_with_tools │
|
||||
//! │ GET /worker/{id}/job │
|
||||
//! │ POST /worker/{id}/status │
|
||||
//! │ POST /worker/{id}/complete │
|
||||
//! │ │
|
||||
//! │ ContainerJobManager │
|
||||
//! │ create_job() -> container + token │
|
||||
//! │ stop_job() │
|
||||
//! │ list_jobs() │
|
||||
//! │ │
|
||||
//! │ TokenStore │
|
||||
//! │ per-job bearer tokens (in-memory only) │
|
||||
//! └───────────────────────────────────────────────┘
|
||||
//! ```
|
||||
|
||||
pub mod api;
|
||||
pub mod auth;
|
||||
pub mod job_manager;
|
||||
|
||||
pub use api::OrchestratorApi;
|
||||
pub use auth::TokenStore;
|
||||
pub use job_manager::{
|
||||
CompletionResult, ContainerHandle, ContainerJobConfig, ContainerJobManager, JobMode,
|
||||
};
|
||||
@@ -0,0 +1,10 @@
|
||||
//! DM pairing for channels.
|
||||
//!
|
||||
//! Gates DMs from unknown senders. Only approved senders can message the agent.
|
||||
//! Unknown senders receive a pairing code and must be approved via `ironclaw pairing approve`.
|
||||
//!
|
||||
//! OpenClaw reference: src/pairing/pairing-store.ts
|
||||
|
||||
mod store;
|
||||
|
||||
pub use store::{PairingRequest, PairingStore, PairingStoreError};
|
||||
@@ -0,0 +1,699 @@
|
||||
//! Pairing store: pending requests and allowFrom list.
|
||||
//!
|
||||
//! Stored in ~/.ironclaw/{channel}-pairing.json and {channel}-allowFrom.json.
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::fs;
|
||||
use std::io::{Seek, SeekFrom, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use fs4::FileExt;
|
||||
use rand::Rng;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
const PAIRING_CODE_LENGTH: usize = 8;
|
||||
const PAIRING_ALPHABET: &[u8] = b"ABCDEFGHJKLMNPQRSTUVWXYZ23456789";
|
||||
/// TTL for pending pairing requests (minutes, not hours — reduces brute-force window).
|
||||
const PAIRING_PENDING_TTL_SECS: u64 = 15 * 60;
|
||||
const PAIRING_PENDING_MAX: usize = 3;
|
||||
/// Max failed approve attempts per channel before rate limit kicks in.
|
||||
const PAIRING_APPROVE_RATE_LIMIT: usize = 10;
|
||||
/// Time window for rate limit (seconds).
|
||||
const PAIRING_APPROVE_RATE_WINDOW_SECS: u64 = 5 * 60;
|
||||
|
||||
/// Error from pairing store operations.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum PairingStoreError {
|
||||
#[error("Invalid channel: {0}")]
|
||||
InvalidChannel(String),
|
||||
|
||||
#[error("IO error: {0}")]
|
||||
Io(#[from] std::io::Error),
|
||||
|
||||
#[error("JSON error: {0}")]
|
||||
Json(#[from] serde_json::Error),
|
||||
|
||||
#[error("Rate limit: too many failed approve attempts; try again later")]
|
||||
ApproveRateLimited,
|
||||
}
|
||||
|
||||
/// Result of upserting a pairing request.
|
||||
#[derive(Debug)]
|
||||
pub struct UpsertResult {
|
||||
pub code: String,
|
||||
pub created: bool,
|
||||
}
|
||||
|
||||
/// A pending pairing request.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PairingRequest {
|
||||
pub id: String,
|
||||
pub code: String,
|
||||
pub created_at: String,
|
||||
pub last_seen_at: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub meta: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
struct PairingStoreFile {
|
||||
version: u8,
|
||||
requests: Vec<PairingRequest>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
struct AllowFromStoreFile {
|
||||
version: u8,
|
||||
#[serde(rename = "allowFrom")]
|
||||
allow_from: Vec<String>,
|
||||
}
|
||||
|
||||
fn default_pairing_dir() -> PathBuf {
|
||||
dirs::home_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join(".ironclaw")
|
||||
}
|
||||
|
||||
fn safe_channel_key(channel: &str) -> Result<String, PairingStoreError> {
|
||||
let raw = channel.trim().to_lowercase();
|
||||
if raw.is_empty() {
|
||||
return Err(PairingStoreError::InvalidChannel("empty".to_string()));
|
||||
}
|
||||
let safe = raw
|
||||
.chars()
|
||||
.map(|c| match c {
|
||||
'\\' | '/' | ':' | '*' | '?' | '"' | '<' | '>' | '|' => '_',
|
||||
_ => c,
|
||||
})
|
||||
.collect::<String>()
|
||||
.replace("..", "_");
|
||||
if safe.is_empty() || safe == "_" {
|
||||
return Err(PairingStoreError::InvalidChannel(channel.to_string()));
|
||||
}
|
||||
Ok(safe)
|
||||
}
|
||||
|
||||
fn pairing_path(base_dir: &Path, channel: &str) -> Result<PathBuf, PairingStoreError> {
|
||||
let key = safe_channel_key(channel)?;
|
||||
Ok(base_dir.join(format!("{}-pairing.json", key)))
|
||||
}
|
||||
|
||||
fn allow_from_path(base_dir: &Path, channel: &str) -> Result<PathBuf, PairingStoreError> {
|
||||
let key = safe_channel_key(channel)?;
|
||||
Ok(base_dir.join(format!("{}-allowFrom.json", key)))
|
||||
}
|
||||
|
||||
fn approve_attempts_path(base_dir: &Path, channel: &str) -> Result<PathBuf, PairingStoreError> {
|
||||
let key = safe_channel_key(channel)?;
|
||||
Ok(base_dir.join(format!("{}-approve-attempts.json", key)))
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Serialize, Deserialize)]
|
||||
struct ApproveAttemptsFile {
|
||||
failed_at: Vec<u64>,
|
||||
}
|
||||
|
||||
fn now_iso() -> String {
|
||||
let now = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default();
|
||||
#[allow(clippy::cast_possible_wrap)]
|
||||
chrono::DateTime::from_timestamp(now.as_secs() as i64, 0)
|
||||
.map(|dt| dt.to_rfc3339())
|
||||
.unwrap_or_else(|| now.as_secs().to_string())
|
||||
}
|
||||
|
||||
fn now_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn parse_timestamp(value: &str) -> Option<u64> {
|
||||
chrono::DateTime::parse_from_rfc3339(value)
|
||||
.ok()
|
||||
.map(|dt| dt.timestamp() as u64)
|
||||
.or_else(|| value.parse::<u64>().ok())
|
||||
}
|
||||
|
||||
fn is_expired(req: &PairingRequest, now_secs: u64) -> bool {
|
||||
let created = parse_timestamp(&req.created_at).unwrap_or(0);
|
||||
now_secs.saturating_sub(created) > PAIRING_PENDING_TTL_SECS
|
||||
}
|
||||
|
||||
fn random_code() -> String {
|
||||
let mut rng = rand::thread_rng();
|
||||
(0..PAIRING_CODE_LENGTH)
|
||||
.map(|_| {
|
||||
let idx = rng.gen_range(0..PAIRING_ALPHABET.len());
|
||||
PAIRING_ALPHABET[idx] as char
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn generate_unique_code(existing: &HashSet<String>) -> String {
|
||||
let mut rng = rand::thread_rng();
|
||||
for _ in 0..500 {
|
||||
let code = random_code();
|
||||
if !existing.contains(&code) {
|
||||
return code;
|
||||
}
|
||||
}
|
||||
// Fallback: add suffix
|
||||
format!("{}{:04}", random_code(), rng.gen_range(0..10000))
|
||||
}
|
||||
|
||||
/// Pairing store for a channel.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PairingStore {
|
||||
base_dir: PathBuf,
|
||||
}
|
||||
|
||||
impl PairingStore {
|
||||
/// Create a new pairing store using default directory (~/.ironclaw).
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
base_dir: default_pairing_dir(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a pairing store with a custom base directory (for testing).
|
||||
pub fn with_base_dir(base_dir: PathBuf) -> Self {
|
||||
Self { base_dir }
|
||||
}
|
||||
|
||||
/// List pending pairing requests for a channel.
|
||||
pub fn list_pending(&self, channel: &str) -> Result<Vec<PairingRequest>, PairingStoreError> {
|
||||
let path = pairing_path(&self.base_dir, channel)?;
|
||||
let content = match fs::read_to_string(&path) {
|
||||
Ok(c) => c,
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
Err(e) => return Err(e.into()),
|
||||
};
|
||||
|
||||
let file: PairingStoreFile = serde_json::from_str(&content).unwrap_or(PairingStoreFile {
|
||||
version: 1,
|
||||
requests: Vec::new(),
|
||||
});
|
||||
|
||||
let now = now_secs();
|
||||
let original_len = file.requests.len();
|
||||
let mut requests: Vec<_> = file
|
||||
.requests
|
||||
.into_iter()
|
||||
.filter(|r| !is_expired(r, now))
|
||||
.collect();
|
||||
|
||||
if requests.len() != original_len {
|
||||
self.write_pairing_file(channel, &requests)?;
|
||||
}
|
||||
|
||||
requests.sort_by(|a, b| a.created_at.cmp(&b.created_at));
|
||||
Ok(requests)
|
||||
}
|
||||
|
||||
/// Upsert a pairing request. Returns (code, created).
|
||||
pub fn upsert_request(
|
||||
&self,
|
||||
channel: &str,
|
||||
id: &str,
|
||||
meta: Option<serde_json::Value>,
|
||||
) -> Result<UpsertResult, PairingStoreError> {
|
||||
let path = pairing_path(&self.base_dir, channel)?;
|
||||
fs::create_dir_all(path.parent().unwrap())?;
|
||||
|
||||
let mut file = fs::OpenOptions::new()
|
||||
.read(true)
|
||||
.write(true)
|
||||
.create(true)
|
||||
.truncate(false)
|
||||
.open(&path)?;
|
||||
|
||||
file.lock_exclusive()?;
|
||||
|
||||
let content = fs::read_to_string(&path).unwrap_or_default();
|
||||
let mut store: PairingStoreFile =
|
||||
serde_json::from_str(&content).unwrap_or(PairingStoreFile {
|
||||
version: 1,
|
||||
requests: Vec::new(),
|
||||
});
|
||||
|
||||
let now = now_iso();
|
||||
let now_secs = now_secs();
|
||||
let id = id.trim().to_string();
|
||||
if id.is_empty() {
|
||||
fs4::FileExt::unlock(&file)?;
|
||||
return Err(PairingStoreError::InvalidChannel("empty id".to_string()));
|
||||
}
|
||||
|
||||
store.requests.retain(|r| !is_expired(r, now_secs));
|
||||
let existing_codes: HashSet<String> = store
|
||||
.requests
|
||||
.iter()
|
||||
.map(|r| r.code.to_uppercase())
|
||||
.collect();
|
||||
|
||||
if let Some(idx) = store.requests.iter().position(|r| r.id == id) {
|
||||
let req = &mut store.requests[idx];
|
||||
let code = if req.code.is_empty() {
|
||||
generate_unique_code(&existing_codes)
|
||||
} else {
|
||||
req.code.clone()
|
||||
};
|
||||
req.last_seen_at = now.clone();
|
||||
req.code = code.clone();
|
||||
if let Some(m) = meta {
|
||||
req.meta = Some(m);
|
||||
}
|
||||
self.write_pairing_file_locked(&mut file, channel, &store.requests)?;
|
||||
fs4::FileExt::unlock(&file)?;
|
||||
return Ok(UpsertResult {
|
||||
code,
|
||||
created: false,
|
||||
});
|
||||
}
|
||||
|
||||
if store.requests.len() >= PAIRING_PENDING_MAX {
|
||||
fs4::FileExt::unlock(&file)?;
|
||||
return Ok(UpsertResult {
|
||||
code: String::new(),
|
||||
created: false,
|
||||
});
|
||||
}
|
||||
|
||||
let code = generate_unique_code(&existing_codes);
|
||||
store.requests.push(PairingRequest {
|
||||
id: id.clone(),
|
||||
code: code.clone(),
|
||||
created_at: now.clone(),
|
||||
last_seen_at: now,
|
||||
meta,
|
||||
});
|
||||
|
||||
self.write_pairing_file_locked(&mut file, channel, &store.requests)?;
|
||||
fs4::FileExt::unlock(&file)?;
|
||||
|
||||
Ok(UpsertResult {
|
||||
code,
|
||||
created: true,
|
||||
})
|
||||
}
|
||||
|
||||
fn is_approve_rate_limited(&self, channel: &str) -> Result<bool, PairingStoreError> {
|
||||
let path = approve_attempts_path(&self.base_dir, channel)?;
|
||||
let content = match fs::read_to_string(&path) {
|
||||
Ok(c) => c,
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(false),
|
||||
Err(e) => return Err(e.into()),
|
||||
};
|
||||
let mut data: ApproveAttemptsFile = serde_json::from_str(&content).unwrap_or_default();
|
||||
let now = now_secs();
|
||||
let cutoff = now.saturating_sub(PAIRING_APPROVE_RATE_WINDOW_SECS);
|
||||
data.failed_at.retain(|&t| t >= cutoff);
|
||||
Ok(data.failed_at.len() >= PAIRING_APPROVE_RATE_LIMIT)
|
||||
}
|
||||
|
||||
fn record_failed_approve(&self, channel: &str) -> Result<(), PairingStoreError> {
|
||||
let path = approve_attempts_path(&self.base_dir, channel)?;
|
||||
fs::create_dir_all(path.parent().unwrap())?;
|
||||
let file = fs::OpenOptions::new()
|
||||
.read(true)
|
||||
.write(true)
|
||||
.create(true)
|
||||
.truncate(true)
|
||||
.open(&path)?;
|
||||
file.lock_exclusive()?;
|
||||
let content = fs::read_to_string(&path).unwrap_or_default();
|
||||
let mut data: ApproveAttemptsFile = serde_json::from_str(&content).unwrap_or_default();
|
||||
let now = now_secs();
|
||||
data.failed_at.push(now);
|
||||
let cutoff = now.saturating_sub(PAIRING_APPROVE_RATE_WINDOW_SECS);
|
||||
data.failed_at.retain(|&t| t >= cutoff);
|
||||
let json = serde_json::to_string_pretty(&data)?;
|
||||
fs::write(&path, json)?;
|
||||
fs4::FileExt::unlock(&file)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Approve a pairing code and add the sender to allowFrom.
|
||||
pub fn approve(
|
||||
&self,
|
||||
channel: &str,
|
||||
code: &str,
|
||||
) -> Result<Option<PairingRequest>, PairingStoreError> {
|
||||
let code = code.trim().to_uppercase();
|
||||
if code.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if self.is_approve_rate_limited(channel)? {
|
||||
return Err(PairingStoreError::ApproveRateLimited);
|
||||
}
|
||||
|
||||
let path = pairing_path(&self.base_dir, channel)?;
|
||||
let mut file = fs::OpenOptions::new()
|
||||
.read(true)
|
||||
.write(true)
|
||||
.create(false)
|
||||
.open(&path)
|
||||
.map_err(|e| {
|
||||
if e.kind() == std::io::ErrorKind::NotFound {
|
||||
PairingStoreError::InvalidChannel("no pairing file".to_string())
|
||||
} else {
|
||||
PairingStoreError::Io(e)
|
||||
}
|
||||
})?;
|
||||
|
||||
file.lock_exclusive()?;
|
||||
|
||||
let content = fs::read_to_string(&path).unwrap_or_default();
|
||||
let mut store: PairingStoreFile =
|
||||
serde_json::from_str(&content).unwrap_or(PairingStoreFile {
|
||||
version: 1,
|
||||
requests: Vec::new(),
|
||||
});
|
||||
|
||||
let now_secs = now_secs();
|
||||
store.requests.retain(|r| !is_expired(r, now_secs));
|
||||
|
||||
let idx = store
|
||||
.requests
|
||||
.iter()
|
||||
.position(|r| r.code.to_uppercase() == code);
|
||||
|
||||
let entry = match idx {
|
||||
Some(i) => store.requests.remove(i),
|
||||
None => {
|
||||
fs4::FileExt::unlock(&file)?;
|
||||
self.record_failed_approve(channel)?;
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
self.write_pairing_file_locked(&mut file, channel, &store.requests)?;
|
||||
fs4::FileExt::unlock(&file)?;
|
||||
|
||||
self.add_allow_from(channel, &entry.id)?;
|
||||
|
||||
Ok(Some(entry))
|
||||
}
|
||||
|
||||
/// Read the allowFrom list for a channel.
|
||||
pub fn read_allow_from(&self, channel: &str) -> Result<Vec<String>, PairingStoreError> {
|
||||
let path = allow_from_path(&self.base_dir, channel)?;
|
||||
let content = match fs::read_to_string(&path) {
|
||||
Ok(c) => c,
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
Err(e) => return Err(e.into()),
|
||||
};
|
||||
|
||||
let file: AllowFromStoreFile =
|
||||
serde_json::from_str(&content).unwrap_or(AllowFromStoreFile {
|
||||
version: 1,
|
||||
allow_from: Vec::new(),
|
||||
});
|
||||
|
||||
Ok(file.allow_from)
|
||||
}
|
||||
|
||||
/// Check if a sender is allowed (by id or username).
|
||||
pub fn is_sender_allowed(
|
||||
&self,
|
||||
channel: &str,
|
||||
id: &str,
|
||||
username: Option<&str>,
|
||||
) -> Result<bool, PairingStoreError> {
|
||||
let allow = self.read_allow_from(channel)?;
|
||||
let id = id.trim();
|
||||
let id_ok = allow.iter().any(|e| e.trim() == id);
|
||||
if id_ok {
|
||||
return Ok(true);
|
||||
}
|
||||
if let Some(u) = username {
|
||||
let u = u.trim().to_lowercase();
|
||||
let u_norm = u.strip_prefix('@').unwrap_or(&u);
|
||||
if allow.iter().any(|e| {
|
||||
e.trim().to_lowercase() == u || e.trim().to_lowercase() == format!("@{}", u_norm)
|
||||
}) {
|
||||
return Ok(true);
|
||||
}
|
||||
}
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
fn add_allow_from(&self, channel: &str, entry: &str) -> Result<(), PairingStoreError> {
|
||||
let entry = entry.trim().to_string();
|
||||
if entry.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let path = allow_from_path(&self.base_dir, channel)?;
|
||||
fs::create_dir_all(path.parent().unwrap())?;
|
||||
|
||||
let file = fs::OpenOptions::new()
|
||||
.read(true)
|
||||
.write(true)
|
||||
.create(true)
|
||||
.truncate(true)
|
||||
.open(&path)?;
|
||||
|
||||
file.lock_exclusive()?;
|
||||
|
||||
let content = fs::read_to_string(&path).unwrap_or_default();
|
||||
let mut store: AllowFromStoreFile =
|
||||
serde_json::from_str(&content).unwrap_or(AllowFromStoreFile {
|
||||
version: 1,
|
||||
allow_from: Vec::new(),
|
||||
});
|
||||
|
||||
let normalized = entry.to_lowercase();
|
||||
if store
|
||||
.allow_from
|
||||
.iter()
|
||||
.any(|e| e.to_lowercase() == normalized)
|
||||
{
|
||||
fs4::FileExt::unlock(&file)?;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
store.allow_from.push(entry);
|
||||
let json = serde_json::to_string_pretty(&store)?;
|
||||
fs::write(&path, json)?;
|
||||
|
||||
fs4::FileExt::unlock(&file)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn write_pairing_file(
|
||||
&self,
|
||||
channel: &str,
|
||||
requests: &[PairingRequest],
|
||||
) -> Result<(), PairingStoreError> {
|
||||
let path = pairing_path(&self.base_dir, channel)?;
|
||||
let mut file = fs::OpenOptions::new()
|
||||
.write(true)
|
||||
.create(true)
|
||||
.truncate(true)
|
||||
.open(&path)?;
|
||||
file.lock_exclusive()?;
|
||||
self.write_pairing_file_locked(&mut file, channel, requests)?;
|
||||
fs4::FileExt::unlock(&file)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn write_pairing_file_locked(
|
||||
&self,
|
||||
file: &mut fs::File,
|
||||
_channel: &str,
|
||||
requests: &[PairingRequest],
|
||||
) -> Result<(), PairingStoreError> {
|
||||
let store = PairingStoreFile {
|
||||
version: 1,
|
||||
requests: requests.to_vec(),
|
||||
};
|
||||
let json = serde_json::to_string_pretty(&store)?;
|
||||
file.set_len(0)?;
|
||||
file.seek(SeekFrom::Start(0))?;
|
||||
file.write_all(json.as_bytes())?;
|
||||
file.sync_all()?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for PairingStore {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[test]
|
||||
fn test_safe_channel_key() {
|
||||
assert_eq!(safe_channel_key("telegram").unwrap(), "telegram");
|
||||
assert_eq!(safe_channel_key("Telegram").unwrap(), "telegram");
|
||||
safe_channel_key("").unwrap_err();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_random_code() {
|
||||
let c = random_code();
|
||||
assert_eq!(c.len(), PAIRING_CODE_LENGTH);
|
||||
assert!(c.chars().all(|c| PAIRING_ALPHABET.contains(&(c as u8))));
|
||||
}
|
||||
|
||||
fn test_store() -> (PairingStore, TempDir) {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let store = PairingStore::with_base_dir(dir.path().to_path_buf());
|
||||
(store, dir)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_list_pending_empty() {
|
||||
let (store, _) = test_store();
|
||||
let requests = store.list_pending("telegram").unwrap();
|
||||
assert!(requests.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_upsert_request_creates_new() {
|
||||
let (store, _) = test_store();
|
||||
let result = store
|
||||
.upsert_request(
|
||||
"telegram",
|
||||
"user123",
|
||||
Some(serde_json::json!({"chat_id": 456})),
|
||||
)
|
||||
.unwrap();
|
||||
assert!(result.created);
|
||||
assert_eq!(result.code.len(), PAIRING_CODE_LENGTH);
|
||||
assert!(
|
||||
result
|
||||
.code
|
||||
.chars()
|
||||
.all(|c| PAIRING_ALPHABET.contains(&(c as u8)))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_upsert_request_updates_existing() {
|
||||
let (store, _) = test_store();
|
||||
let r1 = store.upsert_request("telegram", "user123", None).unwrap();
|
||||
assert!(r1.created);
|
||||
let r2 = store
|
||||
.upsert_request("telegram", "user123", Some(serde_json::json!({"x": 1})))
|
||||
.unwrap();
|
||||
assert!(!r2.created);
|
||||
assert_eq!(r1.code, r2.code);
|
||||
|
||||
let pending = store.list_pending("telegram").unwrap();
|
||||
assert_eq!(pending.len(), 1);
|
||||
assert_eq!(pending[0].id, "user123");
|
||||
assert_eq!(pending[0].meta, Some(serde_json::json!({"x": 1})));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_approve_adds_to_allow_from() {
|
||||
let (store, _) = test_store();
|
||||
let r = store.upsert_request("telegram", "user456", None).unwrap();
|
||||
assert!(r.created);
|
||||
|
||||
let approved = store.approve("telegram", &r.code).unwrap();
|
||||
assert!(approved.is_some());
|
||||
assert_eq!(approved.unwrap().id, "user456");
|
||||
|
||||
let allow = store.read_allow_from("telegram").unwrap();
|
||||
assert_eq!(allow, vec!["user456"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_approve_case_insensitive_code() {
|
||||
let (store, _) = test_store();
|
||||
let r = store.upsert_request("telegram", "user789", None).unwrap();
|
||||
let code_lower = r.code.to_lowercase();
|
||||
let approved = store.approve("telegram", &code_lower).unwrap();
|
||||
assert!(approved.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_approve_invalid_code_returns_none() {
|
||||
let (store, _) = test_store();
|
||||
store.upsert_request("telegram", "user123", None).unwrap();
|
||||
let approved = store.approve("telegram", "BADCODE1").unwrap();
|
||||
assert!(approved.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_approve_rate_limited_after_many_failures() {
|
||||
let (store, _) = test_store();
|
||||
store.upsert_request("telegram", "user123", None).unwrap();
|
||||
for _ in 0..PAIRING_APPROVE_RATE_LIMIT {
|
||||
let _ = store.approve("telegram", "WRONG01");
|
||||
}
|
||||
let err = store.approve("telegram", "WRONG02").unwrap_err();
|
||||
assert!(matches!(err, PairingStoreError::ApproveRateLimited));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_sender_allowed_by_id() {
|
||||
let (store, _) = test_store();
|
||||
let r = store.upsert_request("telegram", "user999", None).unwrap();
|
||||
store.approve("telegram", &r.code).unwrap();
|
||||
|
||||
assert!(
|
||||
store
|
||||
.is_sender_allowed("telegram", "user999", None)
|
||||
.unwrap()
|
||||
);
|
||||
assert!(!store.is_sender_allowed("telegram", "other", None).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_sender_allowed_by_username() {
|
||||
let (store, _) = test_store();
|
||||
store
|
||||
.upsert_request(
|
||||
"telegram",
|
||||
"alice",
|
||||
Some(serde_json::json!({"username": "alice"})),
|
||||
)
|
||||
.unwrap();
|
||||
let pending = store.list_pending("telegram").unwrap();
|
||||
store.approve("telegram", &pending[0].code).unwrap();
|
||||
|
||||
// approve adds id to allow_from. For username we need to add it manually.
|
||||
// Actually approve adds entry.id which is "alice". So is_sender_allowed("telegram", "alice", None) would work.
|
||||
assert!(store.is_sender_allowed("telegram", "alice", None).unwrap());
|
||||
assert!(
|
||||
store
|
||||
.is_sender_allowed("telegram", "alice", Some("alice"))
|
||||
.unwrap()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_channel_normalization() {
|
||||
let (store, _) = test_store();
|
||||
store.upsert_request("Telegram", "u1", None).unwrap();
|
||||
let pending = store.list_pending("telegram").unwrap();
|
||||
assert_eq!(pending.len(), 1);
|
||||
assert_eq!(pending[0].id, "u1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_channel_rejected() {
|
||||
let (store, _) = test_store();
|
||||
store.upsert_request("telegram", "u1", None).unwrap();
|
||||
store.list_pending("").unwrap_err();
|
||||
store.upsert_request("", "u1", None).unwrap_err();
|
||||
}
|
||||
}
|
||||
@@ -491,9 +491,36 @@ impl ContainerRunner {
|
||||
}
|
||||
|
||||
/// Connect to the Docker daemon.
|
||||
///
|
||||
/// Tries these locations in order:
|
||||
/// 1. `DOCKER_HOST` env var (bollard default)
|
||||
/// 2. `/var/run/docker.sock` (Linux default)
|
||||
/// 3. `~/.docker/run/docker.sock` (Docker Desktop on macOS)
|
||||
pub async fn connect_docker() -> Result<Docker> {
|
||||
Docker::connect_with_local_defaults().map_err(|e| SandboxError::DockerNotAvailable {
|
||||
reason: e.to_string(),
|
||||
// First try bollard defaults (checks DOCKER_HOST, then /var/run/docker.sock)
|
||||
if let Ok(docker) = Docker::connect_with_local_defaults() {
|
||||
if docker.ping().await.is_ok() {
|
||||
return Ok(docker);
|
||||
}
|
||||
}
|
||||
|
||||
// Try Docker Desktop socket (macOS)
|
||||
if let Some(home) = std::env::var_os("HOME") {
|
||||
let desktop_sock = std::path::Path::new(&home).join(".docker/run/docker.sock");
|
||||
if desktop_sock.exists() {
|
||||
let sock_str = desktop_sock.to_string_lossy();
|
||||
if let Ok(docker) =
|
||||
Docker::connect_with_socket(&sock_str, 120, bollard::API_DEFAULT_VERSION)
|
||||
{
|
||||
if docker.ping().await.is_ok() {
|
||||
return Ok(docker);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Err(SandboxError::DockerNotAvailable {
|
||||
reason: "Socket not found: /var/run/docker.sock".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user