mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
Compare commits
47
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
dcd8bc3697 | ||
|
|
b125832787 | ||
|
|
294449065b | ||
|
|
8a3f0c0870 | ||
|
|
6685f2ca9d | ||
|
|
fb8ac7455e | ||
|
|
b5b9934df0 | ||
|
|
0e4d50454b | ||
|
|
ab37518008 | ||
|
|
e55b160a2e | ||
|
|
fdfc5fbccc | ||
|
|
402c71d642 | ||
|
|
831484abab | ||
|
|
ecc026843c | ||
|
|
87674420c5 | ||
|
|
c474972a49 | ||
|
|
fcb152e408 | ||
|
|
e86b372fa6 | ||
|
|
63f140d391 | ||
|
|
ab0a2e05de | ||
|
|
290d925c7f | ||
|
|
d73e35cfb0 | ||
|
|
30d81fcdee | ||
|
|
d8dcc34319 | ||
|
|
652f30a826 | ||
|
|
98e9a40762 | ||
|
|
553c306c52 | ||
|
|
7fb2f47999 | ||
|
|
02f85a8ad5 | ||
|
|
9401ab0d58 | ||
|
|
7d1461fc74 | ||
|
|
605a4ba46e | ||
|
|
fe91ba2ab4 | ||
|
|
da2569bb77 | ||
|
|
732b3ecfeb | ||
|
|
461d7712e8 | ||
|
|
1c5117eded | ||
|
|
33b02eabb7 | ||
|
|
068ad2d4b7 | ||
|
|
56b7218897 | ||
|
|
200aed16cd | ||
|
|
4c0275bcdc | ||
|
|
272d31797e | ||
|
|
edff54b0b1 | ||
|
|
4d61d3eedf | ||
|
|
df3635d6be | ||
|
|
a20e19ab16 |
@@ -5,6 +5,7 @@ DATABASE_POOL_SIZE=10
|
||||
# LLM Provider
|
||||
# LLM_BACKEND=nearai # default
|
||||
# Possible values: nearai, ollama, openai_compatible, openai, anthropic, tinfoil
|
||||
# LLM_REQUEST_TIMEOUT_SECS=120 # Increase for local LLMs (Ollama, vLLM, LM Studio)
|
||||
|
||||
# === Anthropic Direct ===
|
||||
# Two auth modes:
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
name: Claude Code Review
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
types: [opened, labeled]
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: write
|
||||
issues: write
|
||||
id-token: write
|
||||
|
||||
concurrency:
|
||||
group: claude-review-${{ github.event.pull_request.number || github.run_id }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
review:
|
||||
name: Claude Code Review
|
||||
if: contains(github.event.pull_request.labels.*.name, 'staging-promotion')
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Run Claude Code review
|
||||
uses: anthropics/claude-code-action@v1
|
||||
with:
|
||||
anthropic_api_key: ${{ secrets.ANTHROPIC_API_KEY }}
|
||||
claude_args: "--max-turns 50 --model claude-haiku-4-5-20251001 --allowedTools 'Bash(gh pr comment:*),Bash(gh pr diff:*),Bash(gh pr view:*),Bash(gh pr list:*),Bash(gh issue view:*),Bash(gh issue list:*),Bash(gh search:*),Bash(git blame:*),Bash(git log:*),Bash(git diff:*)'"
|
||||
prompt: |
|
||||
Code review this pull request. Follow these steps precisely:
|
||||
|
||||
1. Use a Haiku agent to find relevant CLAUDE.md files: the root CLAUDE.md
|
||||
and any CLAUDE.md files in directories whose files this PR modifies.
|
||||
|
||||
2. Use a Haiku agent to summarize the PR change (use `gh pr diff`).
|
||||
|
||||
3. Launch 4 parallel agents to review the change independently. Each agent should
|
||||
read the PR diff with `gh pr diff` and the full source files for changed
|
||||
code, then return a list of issues found:
|
||||
|
||||
Agent 1 — Security & Safety
|
||||
Check for: command injection, path traversal, SSRF, XSS, auth bypass,
|
||||
secrets in logs, .unwrap()/.expect() in production code (not tests),
|
||||
race conditions, TOCTOU, unsafe blocks, panics in async, unbounded allocations.
|
||||
|
||||
Agent 2 — Architecture & Patterns
|
||||
Check for: extensible design (traits/enums over nested conditionals),
|
||||
clean abstractions, proper error types (thiserror), CLAUDE.md compliance,
|
||||
type-driven design over stringly-typed code, DRY violations.
|
||||
|
||||
Agent 3 — Bug Scan
|
||||
Shallow diff-only scan for obvious bugs: logic errors, off-by-one,
|
||||
missing error handling, division by zero, incorrect return values.
|
||||
Ignore nitpicks and likely false positives. Do NOT read extra context
|
||||
beyond the diff — focus only on the changes.
|
||||
|
||||
Agent 4 — Performance & Production
|
||||
Check for: blocking in async, N+1 queries, unbounded loops, missing
|
||||
timeouts, resource leaks (file handles, connections), large allocations
|
||||
in hot paths.
|
||||
|
||||
4. For each issue found, launch a parallel Haiku agent to:
|
||||
a. Assign a severity:
|
||||
- CRITICAL: security vulns, panics in prod (.unwrap/.expect), data exfiltration, race conditions
|
||||
- HIGH: logic bugs, missing error handling, breaking API/schema changes
|
||||
- MEDIUM: missing tests, unnecessary complexity, performance issues
|
||||
- LOW: documentation gaps, naming suggestions
|
||||
b. Score confidence 0-100 (give this rubric verbatim):
|
||||
0: False positive, doesn't stand up to scrutiny, or pre-existing issue.
|
||||
25: Might be real, but may be false positive. Stylistic issues not in CLAUDE.md.
|
||||
50: Real issue but nitpick or rare in practice. Not very important.
|
||||
75: Verified real issue, will be hit in practice. Directly impacts functionality
|
||||
or explicitly mentioned in CLAUDE.md.
|
||||
100: Certain, confirmed, will happen frequently. Evidence directly confirms.
|
||||
|
||||
5. Post a single comment on the PR using `gh pr comment` with this format.
|
||||
If no issues were found, post "No issues found." instead:
|
||||
|
||||
### Code review
|
||||
|
||||
Found N issues:
|
||||
|
||||
1. [SEVERITY:CONFIDENCE] <brief description>
|
||||
|
||||
<permalink to file:line using full SHA, eg https://github.com/owner/repo/blob/abc123def/src/file.rs#L10-L15>
|
||||
|
||||
Example: [CRITICAL:92] `.unwrap()` can panic in production when config is missing
|
||||
|
||||
You MUST use the full git SHA in links (not HEAD or branch name).
|
||||
Provide 1 line of context before and after each linked range.
|
||||
|
||||
Notes:
|
||||
- Use `gh` for all GitHub interactions, not web fetch
|
||||
- Do NOT check build signal or attempt to build/test the code
|
||||
- Ignore pre-existing issues not introduced by this PR
|
||||
- Ignore issues a linter/compiler would catch (formatting, imports, types)
|
||||
@@ -1,5 +1,6 @@
|
||||
name: E2E Tests
|
||||
on:
|
||||
workflow_call:
|
||||
schedule:
|
||||
- cron: "0 6 * * 1" # Weekly Monday 6 AM UTC
|
||||
workflow_dispatch:
|
||||
|
||||
@@ -0,0 +1,473 @@
|
||||
name: Staging CI (Batched)
|
||||
|
||||
on:
|
||||
schedule:
|
||||
- cron: "0 * * * *" # Every 60 minutes
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
force:
|
||||
description: "Force run even if no new commits"
|
||||
type: boolean
|
||||
default: false
|
||||
skip_claude_gate:
|
||||
description: "Skip Claude review gate (bypass blocking findings)"
|
||||
type: boolean
|
||||
default: false
|
||||
|
||||
permissions:
|
||||
contents: write
|
||||
issues: write
|
||||
pull-requests: write
|
||||
checks: read
|
||||
|
||||
concurrency:
|
||||
group: staging-ci
|
||||
cancel-in-progress: false # Let running suites finish
|
||||
|
||||
jobs:
|
||||
# ── Check for new commits ──────────────────────────────────────
|
||||
check-changes:
|
||||
name: Check for new commits
|
||||
runs-on: ubuntu-latest
|
||||
outputs:
|
||||
has_changes: ${{ steps.check.outputs.has_changes }}
|
||||
current_head: ${{ steps.check.outputs.current_head }}
|
||||
diff_range: ${{ steps.check.outputs.diff_range }}
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: staging
|
||||
fetch-depth: 0
|
||||
fetch-tags: true
|
||||
|
||||
- name: Check for changes since last tested
|
||||
id: check
|
||||
env:
|
||||
FORCE_RUN: ${{ inputs.force }}
|
||||
run: |
|
||||
CURRENT_HEAD=$(git rev-parse HEAD)
|
||||
echo "current_head=${CURRENT_HEAD}" >> "$GITHUB_OUTPUT"
|
||||
|
||||
if git rev-parse staging-tested >/dev/null 2>&1; then
|
||||
LAST_TESTED=$(git rev-parse staging-tested)
|
||||
else
|
||||
LAST_TESTED=""
|
||||
fi
|
||||
|
||||
DIFF_RANGE=""
|
||||
if [ -n "$LAST_TESTED" ] && [ "$LAST_TESTED" = "$CURRENT_HEAD" ]; then
|
||||
echo "No new commits since last tested (${CURRENT_HEAD})"
|
||||
HAS_CHANGES=false
|
||||
else
|
||||
HAS_CHANGES=true
|
||||
if [ -n "$LAST_TESTED" ]; then
|
||||
COMMIT_COUNT=$(git rev-list --count "${LAST_TESTED}..HEAD")
|
||||
echo "Found ${COMMIT_COUNT} new commit(s) since last tested"
|
||||
DIFF_RANGE="${LAST_TESTED}..${CURRENT_HEAD}"
|
||||
else
|
||||
git fetch origin main
|
||||
MERGE_BASE=$(git merge-base origin/main HEAD)
|
||||
echo "First run -- reviewing from merge-base ${MERGE_BASE}"
|
||||
DIFF_RANGE="${MERGE_BASE}..${CURRENT_HEAD}"
|
||||
fi
|
||||
fi
|
||||
|
||||
# Force override from workflow_dispatch
|
||||
if [ "$FORCE_RUN" = "true" ]; then
|
||||
echo "Force run requested"
|
||||
HAS_CHANGES=true
|
||||
if [ -z "$DIFF_RANGE" ]; then
|
||||
DIFF_RANGE="${CURRENT_HEAD}..${CURRENT_HEAD}"
|
||||
fi
|
||||
fi
|
||||
|
||||
echo "has_changes=${HAS_CHANGES}" >> "$GITHUB_OUTPUT"
|
||||
echo "diff_range=${DIFF_RANGE}" >> "$GITHUB_OUTPUT"
|
||||
|
||||
# ── Run full test suite ──────────────────────────────────────────
|
||||
tests:
|
||||
name: Test Suite
|
||||
needs: check-changes
|
||||
if: needs.check-changes.outputs.has_changes == 'true'
|
||||
uses: ./.github/workflows/test.yml
|
||||
|
||||
# ── Run E2E browser tests ────────────────────────────────────────
|
||||
e2e:
|
||||
name: E2E Browser Tests
|
||||
needs: check-changes
|
||||
if: needs.check-changes.outputs.has_changes == 'true'
|
||||
uses: ./.github/workflows/e2e.yml
|
||||
|
||||
# ── Create promotion PR (triggers claude-review.yml on the PR) ──
|
||||
create-promotion-pr:
|
||||
name: Create Promotion PR
|
||||
needs: check-changes
|
||||
if: needs.check-changes.outputs.has_changes == 'true'
|
||||
runs-on: ubuntu-latest
|
||||
outputs:
|
||||
pr_number: ${{ steps.create-pr.outputs.pr_number }}
|
||||
promotion_branch: ${{ steps.branch.outputs.branch }}
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: staging
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Generate GitHub App token
|
||||
id: app-token
|
||||
continue-on-error: true
|
||||
uses: actions/create-github-app-token@v2
|
||||
with:
|
||||
app-id: ${{ secrets.GH_RELEASES_MANAGER_APP_ID }}
|
||||
private-key: ${{ secrets.GH_RELEASES_MANAGER_APP_PRIVATE_KEY }}
|
||||
|
||||
- name: Set token
|
||||
id: token
|
||||
run: |
|
||||
if [ -n "${{ steps.app-token.outputs.token }}" ]; then
|
||||
echo "token=${{ steps.app-token.outputs.token }}" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "token=${{ github.token }}" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
|
||||
- name: Check if staging is ahead of main
|
||||
id: ahead-check
|
||||
env:
|
||||
GH_TOKEN: ${{ steps.token.outputs.token }}
|
||||
run: |
|
||||
git fetch origin main
|
||||
AHEAD=$(git rev-list --count origin/main..origin/staging)
|
||||
echo "commits_ahead=${AHEAD}" >> "$GITHUB_OUTPUT"
|
||||
if [ "$AHEAD" -eq 0 ]; then
|
||||
echo "Staging is not ahead of main. Nothing to promote."
|
||||
else
|
||||
echo "Staging is ${AHEAD} commits ahead of main."
|
||||
fi
|
||||
|
||||
- name: Create promotion branch
|
||||
id: branch
|
||||
if: steps.ahead-check.outputs.commits_ahead != '0'
|
||||
run: |
|
||||
SHORT_SHA=$(echo "${{ needs.check-changes.outputs.current_head }}" | cut -c1-8)
|
||||
BRANCH="staging-promote/${SHORT_SHA}-${{ github.run_id }}"
|
||||
git checkout -b "$BRANCH"
|
||||
git push origin "$BRANCH"
|
||||
echo "branch=${BRANCH}" >> "$GITHUB_OUTPUT"
|
||||
echo "Created promotion branch: ${BRANCH}"
|
||||
|
||||
- name: Find base branch
|
||||
id: find-base
|
||||
if: steps.ahead-check.outputs.commits_ahead != '0'
|
||||
env:
|
||||
GH_TOKEN: ${{ steps.token.outputs.token }}
|
||||
run: |
|
||||
# Find the newest open promotion PR with a staging-promote/* head branch
|
||||
LATEST=$(gh pr list --label staging-promotion --state open \
|
||||
--json headRefName,createdAt \
|
||||
--jq '[.[] | select(.headRefName | startswith("staging-promote/"))] | sort_by(.createdAt) | last | .headRefName // empty')
|
||||
if [ -n "$LATEST" ]; then
|
||||
echo "base=${LATEST}" >> "$GITHUB_OUTPUT"
|
||||
echo "Chaining onto existing promotion branch: ${LATEST}"
|
||||
else
|
||||
echo "base=main" >> "$GITHUB_OUTPUT"
|
||||
echo "No existing promotion PR — targeting main"
|
||||
fi
|
||||
|
||||
- name: Create promotion PR
|
||||
id: create-pr
|
||||
if: steps.ahead-check.outputs.commits_ahead != '0'
|
||||
env:
|
||||
GH_TOKEN: ${{ steps.token.outputs.token }}
|
||||
run: |
|
||||
RANGE="${{ needs.check-changes.outputs.diff_range }}"
|
||||
TIMESTAMP=$(date -u +"%Y-%m-%d %H:%M UTC")
|
||||
BRANCH="${{ steps.branch.outputs.branch }}"
|
||||
BASE="${{ steps.find-base.outputs.base }}"
|
||||
|
||||
PR_URL=$(gh pr create \
|
||||
--base "$BASE" \
|
||||
--head "$BRANCH" \
|
||||
--title "chore: promote staging to main (${TIMESTAMP})" \
|
||||
--body "## Auto-promotion from staging CI
|
||||
|
||||
**Batch range:** \`${RANGE}\`
|
||||
**Promotion branch:** \`${BRANCH}\`
|
||||
**Base:** \`${BASE}\`
|
||||
**Triggered by:** Staging CI batch at ${TIMESTAMP}
|
||||
|
||||
Waiting for gates:
|
||||
- Tests: pending
|
||||
- E2E: pending
|
||||
- Claude Code review: pending (will post comments on this PR)
|
||||
|
||||
---
|
||||
*Auto-created by staging-ci workflow*" \
|
||||
--label "staging-promotion")
|
||||
|
||||
PR_NUM=$(echo "$PR_URL" | grep -oE '[0-9]+$')
|
||||
echo "pr_number=${PR_NUM}" >> "$GITHUB_OUTPUT"
|
||||
echo "Created promotion PR #${PR_NUM}"
|
||||
|
||||
# ── Gate: wait for review, process findings, merge or block ─────
|
||||
gate:
|
||||
name: Staging Gate
|
||||
needs: [check-changes, tests, e2e, create-promotion-pr]
|
||||
if: >
|
||||
always() &&
|
||||
needs.check-changes.outputs.has_changes == 'true' &&
|
||||
needs.tests.result == 'success' &&
|
||||
needs.e2e.result == 'success' &&
|
||||
needs.create-promotion-pr.result == 'success'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 25
|
||||
outputs:
|
||||
gate_passed: ${{ steps.evaluate.outputs.passed }}
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: staging
|
||||
fetch-depth: 1
|
||||
|
||||
- name: Generate GitHub App token
|
||||
id: app-token
|
||||
continue-on-error: true
|
||||
uses: actions/create-github-app-token@v2
|
||||
with:
|
||||
app-id: ${{ secrets.GH_RELEASES_MANAGER_APP_ID }}
|
||||
private-key: ${{ secrets.GH_RELEASES_MANAGER_APP_PRIVATE_KEY }}
|
||||
|
||||
- name: Set token
|
||||
id: token
|
||||
run: |
|
||||
if [ -n "${{ steps.app-token.outputs.token }}" ]; then
|
||||
echo "token=${{ steps.app-token.outputs.token }}" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "token=${{ github.token }}" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
|
||||
- name: Wait for Claude review job
|
||||
env:
|
||||
GH_TOKEN: ${{ steps.token.outputs.token }}
|
||||
PR_NUMBER: ${{ needs.create-promotion-pr.outputs.pr_number }}
|
||||
REPO: ${{ github.repository }}
|
||||
run: |
|
||||
if [ -z "$PR_NUMBER" ]; then
|
||||
echo "No PR number — skipping wait"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
PR_SHA=$(gh pr view "$PR_NUMBER" --json headRefOid --jq '.headRefOid' || echo "")
|
||||
if [ -z "$PR_SHA" ]; then
|
||||
echo "::warning::Could not get PR head SHA"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
echo "Polling for Claude Code Review job on PR #${PR_NUMBER} (SHA: ${PR_SHA})..."
|
||||
TIMEOUT=1200 # 20 minutes
|
||||
ELAPSED=0
|
||||
INTERVAL=30
|
||||
|
||||
while [ "$ELAPSED" -lt "$TIMEOUT" ]; do
|
||||
STATUS=$(gh api "repos/${REPO}/commits/${PR_SHA}/check-runs" \
|
||||
--jq '[.check_runs[] | select(.name == "Claude Code Review") | .conclusion // .status] | first // "pending"' 2>/dev/null || echo "pending")
|
||||
|
||||
if [ "$STATUS" = "success" ] || [ "$STATUS" = "failure" ] || [ "$STATUS" = "cancelled" ]; then
|
||||
echo "Claude review job completed with status: ${STATUS} (${ELAPSED}s)"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
echo "Claude review status: ${STATUS} (${ELAPSED}s elapsed)"
|
||||
sleep "$INTERVAL"
|
||||
ELAPSED=$((ELAPSED + INTERVAL))
|
||||
done
|
||||
|
||||
echo "::warning::Claude review job not completed after ${TIMEOUT}s"
|
||||
|
||||
- name: Process Claude review comments and create issues
|
||||
id: process-findings
|
||||
env:
|
||||
GH_TOKEN: ${{ steps.token.outputs.token }}
|
||||
PR_NUMBER: ${{ needs.create-promotion-pr.outputs.pr_number }}
|
||||
REPO: ${{ github.repository }}
|
||||
run: |
|
||||
HAS_BLOCKING=false
|
||||
ISSUES_CREATED=0
|
||||
|
||||
if [ -z "$PR_NUMBER" ]; then
|
||||
echo "No PR — skipping finding processing"
|
||||
echo "has_blocking=false" >> "$GITHUB_OUTPUT"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
# Check for "No issues found" first (clean pass)
|
||||
NO_ISSUES=$(gh api "repos/${REPO}/issues/${PR_NUMBER}/comments" \
|
||||
--jq '[.[] | select(.user.login == "claude[bot]") | select(.body | test("No issues found"))] | length' 2>/dev/null || echo "0")
|
||||
if [ "$NO_ISSUES" -gt 0 ]; then
|
||||
echo "Claude review found no issues — gate passes"
|
||||
echo "has_blocking=false" >> "$GITHUB_OUTPUT"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
# Get the last Claude comment that contains findings
|
||||
JQ_FILTER='[.[] | select(.user.login == "claude[bot]") | select(.body | test("Found [0-9]+ issue"))] | last'
|
||||
BODY=$(gh api "repos/${REPO}/issues/${PR_NUMBER}/comments" \
|
||||
--jq "${JQ_FILTER} | .body // empty" 2>/dev/null || echo "")
|
||||
COMMENT_URL=$(gh api "repos/${REPO}/issues/${PR_NUMBER}/comments" \
|
||||
--jq "${JQ_FILTER} | .html_url // empty" 2>/dev/null || echo "")
|
||||
|
||||
if [ -z "$BODY" ]; then
|
||||
echo "::warning::No Claude review comment found for PR #${PR_NUMBER} — treating as blocking"
|
||||
echo "has_blocking=true" >> "$GITHUB_OUTPUT"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
# Parse [SEVERITY:CONFIDENCE] tags from each numbered finding
|
||||
# Matrix: CRITICAL always→issue, ≥80→block. HIGH ≥50→issue. MEDIUM ≥80→issue. LOW ≥80→issue.
|
||||
# Use process substitution so variables propagate to parent shell
|
||||
while read -r line; do
|
||||
TAG=$(echo "$line" | grep -oE '^\[(CRITICAL|HIGH|MEDIUM|LOW):[0-9]+\]')
|
||||
SEVERITY=$(echo "$TAG" | sed 's/\[\(.*\):\(.*\)\]/\1/')
|
||||
CONFIDENCE=$(echo "$TAG" | sed 's/\[\(.*\):\(.*\)\]/\2/')
|
||||
DESC=$(echo "$line" | sed "s/\[${SEVERITY}:${CONFIDENCE}\] *//" | head -1)
|
||||
|
||||
echo "Found: [${SEVERITY}:${CONFIDENCE}] ${DESC}"
|
||||
|
||||
# Check if blocking (CRITICAL ≥80)
|
||||
if [ "$SEVERITY" = "CRITICAL" ] && [ "$CONFIDENCE" -ge 80 ]; then
|
||||
HAS_BLOCKING=true
|
||||
fi
|
||||
|
||||
# Determine if this should create an issue
|
||||
CREATE_ISSUE=false
|
||||
case "$SEVERITY" in
|
||||
CRITICAL) CREATE_ISSUE=true ;;
|
||||
HIGH) [ "$CONFIDENCE" -ge 50 ] && CREATE_ISSUE=true ;;
|
||||
MEDIUM) [ "$CONFIDENCE" -ge 80 ] && CREATE_ISSUE=true ;;
|
||||
LOW) [ "$CONFIDENCE" -ge 80 ] && CREATE_ISSUE=true ;;
|
||||
esac
|
||||
|
||||
if [ "$CREATE_ISSUE" = "true" ]; then
|
||||
case "$SEVERITY" in
|
||||
CRITICAL) LABELS="bug,risk: high,staging-ci-review" ;;
|
||||
HIGH) LABELS="bug,risk: medium,staging-ci-review" ;;
|
||||
MEDIUM) LABELS="risk: medium,staging-ci-review" ;;
|
||||
LOW) LABELS="risk: low,staging-ci-review" ;;
|
||||
esac
|
||||
|
||||
TITLE=$(echo "$DESC" | cut -c1-80)
|
||||
{
|
||||
echo "## [${SEVERITY}:${CONFIDENCE}] Issue Found by Staging CI Review"
|
||||
echo ""
|
||||
echo "**Severity:** ${SEVERITY}"
|
||||
echo "**Confidence:** ${CONFIDENCE}/100"
|
||||
echo "**PR comment:** ${COMMENT_URL}"
|
||||
echo ""
|
||||
echo "### Description"
|
||||
echo "$DESC"
|
||||
echo ""
|
||||
echo "---"
|
||||
echo "*Auto-created by staging-ci Claude Code review*"
|
||||
} > /tmp/issue-body.md
|
||||
|
||||
if gh issue create \
|
||||
--title "[${SEVERITY}] ${TITLE}" \
|
||||
--body-file /tmp/issue-body.md \
|
||||
--label "${LABELS}"; then
|
||||
ISSUES_CREATED=$((ISSUES_CREATED + 1))
|
||||
else
|
||||
echo "::warning::Failed to create issue for ${SEVERITY} finding"
|
||||
fi
|
||||
fi
|
||||
done < <(echo "$BODY" | grep -oE '\[(CRITICAL|HIGH|MEDIUM|LOW):[0-9]+\].*')
|
||||
|
||||
echo "Created ${ISSUES_CREATED} issues"
|
||||
echo "has_blocking=${HAS_BLOCKING}" >> "$GITHUB_OUTPUT"
|
||||
|
||||
- name: Evaluate gate
|
||||
id: evaluate
|
||||
env:
|
||||
PR_NUMBER: ${{ needs.create-promotion-pr.outputs.pr_number }}
|
||||
SKIP_GATE: ${{ inputs.skip_claude_gate }}
|
||||
HAS_BLOCKING: ${{ steps.process-findings.outputs.has_blocking }}
|
||||
run: |
|
||||
SKIP_INPUT="$SKIP_GATE"
|
||||
|
||||
if [ "$HAS_BLOCKING" = "true" ]; then
|
||||
echo "::warning::Claude review found blocking issues (CRITICAL ≥80 confidence)"
|
||||
if [ "$SKIP_INPUT" = "true" ]; then
|
||||
echo "::warning::Gate overridden by skip_claude_gate workflow input"
|
||||
echo "passed=true" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "::error::Blocking promotion due to CRITICAL findings (≥80 confidence)"
|
||||
echo "::error::PR #${PR_NUMBER} left open with review comments"
|
||||
echo "passed=false" >> "$GITHUB_OUTPUT"
|
||||
exit 1
|
||||
fi
|
||||
else
|
||||
echo "No blocking findings. Gate passed."
|
||||
echo "passed=true" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
|
||||
- name: Merge promotion PR
|
||||
id: merge
|
||||
if: steps.evaluate.outputs.passed == 'true'
|
||||
env:
|
||||
GH_TOKEN: ${{ steps.token.outputs.token }}
|
||||
PR_NUMBER: ${{ needs.create-promotion-pr.outputs.pr_number }}
|
||||
run: |
|
||||
if [ -n "$PR_NUMBER" ]; then
|
||||
echo "Merging promotion PR #${PR_NUMBER}"
|
||||
# Do NOT use --delete-branch: deleting a promotion branch closes
|
||||
# any chained PRs that use it as their base (verified in ironclaw-ci-test).
|
||||
# Stale promotion branches are cleaned up separately.
|
||||
gh pr merge "$PR_NUMBER" --merge
|
||||
echo "merged=true" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
|
||||
# ── Update tested tag (always, so next batch covers only new commits) ──
|
||||
update-tag:
|
||||
name: Update staging-tested tag
|
||||
needs: [check-changes, tests, e2e, create-promotion-pr, gate]
|
||||
if: >
|
||||
always() &&
|
||||
needs.check-changes.outputs.has_changes == 'true' &&
|
||||
needs.tests.result == 'success' &&
|
||||
needs.e2e.result == 'success' &&
|
||||
needs.create-promotion-pr.result == 'success'
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: staging
|
||||
fetch-depth: 1
|
||||
|
||||
- name: Update staging-tested tag
|
||||
run: |
|
||||
git tag -f staging-tested "${{ needs.check-changes.outputs.current_head }}"
|
||||
git push origin staging-tested --force
|
||||
echo "Updated staging-tested tag to ${{ needs.check-changes.outputs.current_head }}"
|
||||
|
||||
# ── Report ───────────────────────────────────────────────────────
|
||||
report:
|
||||
name: Staging CI Summary
|
||||
needs: [check-changes, tests, e2e, create-promotion-pr, gate, update-tag]
|
||||
if: always() && needs.check-changes.outputs.has_changes == 'true'
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Summary
|
||||
run: |
|
||||
echo "## Staging CI Batch Results" >> "$GITHUB_STEP_SUMMARY"
|
||||
echo "" >> "$GITHUB_STEP_SUMMARY"
|
||||
echo "| Check | Result |" >> "$GITHUB_STEP_SUMMARY"
|
||||
echo "|-------|--------|" >> "$GITHUB_STEP_SUMMARY"
|
||||
echo "| Tests | ${{ needs.tests.result }} |" >> "$GITHUB_STEP_SUMMARY"
|
||||
echo "| E2E | ${{ needs.e2e.result }} |" >> "$GITHUB_STEP_SUMMARY"
|
||||
echo "| Promotion PR | ${{ needs.create-promotion-pr.result }} |" >> "$GITHUB_STEP_SUMMARY"
|
||||
echo "| Gate | ${{ needs.gate.result }} |" >> "$GITHUB_STEP_SUMMARY"
|
||||
echo "| Tag Updated | ${{ needs.update-tag.result }} |" >> "$GITHUB_STEP_SUMMARY"
|
||||
echo "" >> "$GITHUB_STEP_SUMMARY"
|
||||
echo "Range: ${{ needs.check-changes.outputs.diff_range }}" >> "$GITHUB_STEP_SUMMARY"
|
||||
PR_NUM="${{ needs.create-promotion-pr.outputs.pr_number }}"
|
||||
if [ -n "$PR_NUM" ]; then
|
||||
echo "Promotion PR: #${PR_NUM}" >> "$GITHUB_STEP_SUMMARY"
|
||||
fi
|
||||
@@ -1,5 +1,6 @@
|
||||
name: Run Tests
|
||||
on:
|
||||
workflow_call:
|
||||
pull_request:
|
||||
push:
|
||||
branches:
|
||||
@@ -38,6 +39,9 @@ jobs:
|
||||
|
||||
telegram-tests:
|
||||
name: Telegram Channel Tests
|
||||
if: >
|
||||
github.event_name == 'push' ||
|
||||
(github.event_name == 'pull_request' && github.base_ref != 'staging')
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
@@ -50,6 +54,9 @@ jobs:
|
||||
|
||||
windows-build:
|
||||
name: Windows Build (${{ matrix.name }})
|
||||
if: >
|
||||
github.event_name == 'push' ||
|
||||
(github.event_name == 'pull_request' && github.base_ref != 'staging')
|
||||
runs-on: windows-latest
|
||||
strategy:
|
||||
fail-fast: false
|
||||
@@ -74,6 +81,9 @@ jobs:
|
||||
|
||||
wasm-wit-compat:
|
||||
name: WASM WIT Compatibility
|
||||
if: >
|
||||
github.event_name == 'push' ||
|
||||
(github.event_name == 'pull_request' && github.base_ref != 'staging')
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
@@ -94,6 +104,9 @@ jobs:
|
||||
|
||||
docker-build:
|
||||
name: Docker Build
|
||||
if: >
|
||||
github.event_name == 'push' ||
|
||||
(github.event_name == 'pull_request' && github.base_ref != 'staging')
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
@@ -123,12 +136,22 @@ jobs:
|
||||
needs: [tests, telegram-tests, wasm-wit-compat, docker-build, windows-build, version-check]
|
||||
steps:
|
||||
- run: |
|
||||
if [[ "${{ needs.tests.result }}" != "success" || "${{ needs.telegram-tests.result }}" != "success" || "${{ needs.wasm-wit-compat.result }}" != "success" || "${{ needs.docker-build.result }}" != "success" || "${{ needs.windows-build.result }}" != "success" ]]; then
|
||||
echo "One or more jobs failed"
|
||||
exit 1
|
||||
fi
|
||||
# version-check only runs on PRs, so skip/success are both acceptable
|
||||
if [[ "${{ needs.version-check.result }}" == "failure" ]]; then
|
||||
echo "Version bump check failed"
|
||||
# Unit tests must always pass
|
||||
if [[ "${{ needs.tests.result }}" != "success" ]]; then
|
||||
echo "Unit tests failed"
|
||||
exit 1
|
||||
fi
|
||||
# Gated jobs: must pass on promotion PRs / push, skipped on developer PRs
|
||||
for job in telegram-tests wasm-wit-compat docker-build windows-build version-check; do
|
||||
case "$job" in
|
||||
telegram-tests) result="${{ needs.telegram-tests.result }}" ;;
|
||||
wasm-wit-compat) result="${{ needs.wasm-wit-compat.result }}" ;;
|
||||
docker-build) result="${{ needs.docker-build.result }}" ;;
|
||||
windows-build) result="${{ needs.windows-build.result }}" ;;
|
||||
version-check) result="${{ needs.version-check.result }}" ;;
|
||||
esac
|
||||
if [[ "$result" == "failure" || "$result" == "cancelled" ]]; then
|
||||
echo "$job failed"
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
|
||||
+3
-1
@@ -4,8 +4,9 @@
|
||||
.env.*
|
||||
!.env.example
|
||||
|
||||
# Claude Code worktrees
|
||||
# Claude Code worktrees and lock files
|
||||
.claude/worktrees/
|
||||
.claude/scheduled_tasks.lock
|
||||
|
||||
# Sidecar tool data
|
||||
.sidecar/
|
||||
@@ -22,4 +23,5 @@ bench-results/
|
||||
# WASM build artifacts (loaded from disk, not bundled)
|
||||
*.wasm
|
||||
|
||||
# Traces
|
||||
trace_*.json
|
||||
|
||||
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
### Added
|
||||
|
||||
- AWS Bedrock LLM provider via native Converse API with IAM and SSO auth support (feature-gated: `--features bedrock`)
|
||||
|
||||
## [0.16.1](https://github.com/nearai/ironclaw/compare/v0.16.0...v0.16.1) - 2026-03-06
|
||||
|
||||
### Fixed
|
||||
|
||||
@@ -484,6 +484,13 @@ SKILLS_AUTO_DISCOVER=true # Scan skill directories on startup
|
||||
TINFOIL_API_KEY=... # Required when LLM_BACKEND=tinfoil
|
||||
TINFOIL_MODEL=kimi-k2-5 # Default model
|
||||
|
||||
# AWS Bedrock (native Converse API, requires --features bedrock)
|
||||
# LLM_BACKEND=bedrock
|
||||
# BEDROCK_REGION=us-east-1 # AWS region
|
||||
# BEDROCK_MODEL=anthropic.claude-opus-4-6-v1 # Required model ID
|
||||
# BEDROCK_CROSS_REGION=us # Cross-region prefix (us/eu/apac/global)
|
||||
# AWS_PROFILE=my-profile # Named profile (SSO/assume-role)
|
||||
|
||||
# Tunnel (public internet exposure for webhooks)
|
||||
TUNNEL_URL=https://abc123.ngrok.io # Static public URL (manual tunnel)
|
||||
# Or use a managed tunnel provider:
|
||||
@@ -500,7 +507,9 @@ OBSERVABILITY_BACKEND=none # none/noop (default) or log
|
||||
|
||||
### LLM Providers
|
||||
|
||||
Backends: `nearai` (default), `openai`, `anthropic`, `ollama`, `openai_compatible`, `tinfoil` — set via `LLM_BACKEND`. See [src/llm/CLAUDE.md](src/llm/CLAUDE.md) for per-provider auth and configuration details.
|
||||
Backends: `nearai` (default), `openai`, `anthropic`, `ollama`, `openai_compatible`, `tinfoil`, `bedrock` (requires `--features bedrock`) — set via `LLM_BACKEND`. See [src/llm/CLAUDE.md](src/llm/CLAUDE.md) for per-provider auth and configuration details.
|
||||
|
||||
**AWS Bedrock** -- Uses the native Converse API via `aws-sdk-bedrockruntime`. Requires `--features bedrock` at build time (not included in default features due to heavy AWS SDK dependencies). Supports standard AWS auth methods: IAM credentials (`AWS_ACCESS_KEY_ID`/`AWS_SECRET_ACCESS_KEY`), SSO profiles (`AWS_PROFILE`), and instance roles. Configure with `BEDROCK_REGION` (default: `us-east-1`), `BEDROCK_MODEL` (required, e.g., `anthropic.claude-opus-4-6-v1`), and `BEDROCK_CROSS_REGION` (optional: `us`, `eu`, `apac`, `global` for cross-region inference profiles). The SDK credential chain resolves auth automatically from the environment.
|
||||
|
||||
## Database
|
||||
|
||||
|
||||
Generated
+707
-23
File diff suppressed because it is too large
Load Diff
@@ -73,6 +73,8 @@ toml = "0.8"
|
||||
# Core types
|
||||
uuid = { version = "1", features = ["v4", "v5", "serde"] }
|
||||
chrono = { version = "0.4", features = ["serde"] }
|
||||
chrono-tz = "0.10"
|
||||
iana-time-zone = "0.1"
|
||||
rust_decimal = { version = "1", features = ["serde", "serde-with-str", "maths"] }
|
||||
rust_decimal_macros = "1"
|
||||
|
||||
@@ -140,6 +142,11 @@ subtle = "2" # Constant-time comparisons for token validation
|
||||
# Multi-provider LLM support
|
||||
rig-core = "0.30"
|
||||
|
||||
# AWS Bedrock (native Converse API, opt-in via --features bedrock)
|
||||
aws-config = { version = "1", features = ["behavior-version-latest"], optional = true }
|
||||
aws-sdk-bedrockruntime = { version = "1", optional = true }
|
||||
aws-smithy-types = { version = "1", optional = true }
|
||||
|
||||
# Docker sandbox
|
||||
bollard = "0.18"
|
||||
|
||||
@@ -201,6 +208,7 @@ postgres = [
|
||||
libsql = ["dep:libsql"]
|
||||
integration = []
|
||||
html-to-markdown = ["dep:html-to-markdown-rs", "dep:readabilityrs"]
|
||||
bedrock = ["dep:aws-config", "dep:aws-sdk-bedrockruntime", "dep:aws-smithy-types"]
|
||||
|
||||
[[test]]
|
||||
name = "html_to_markdown"
|
||||
|
||||
+2
-2
@@ -39,7 +39,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
| Network modes (loopback/LAN/remote) | ✅ | 🚧 | HTTP only |
|
||||
| OpenAI-compatible HTTP API | ✅ | ✅ | /v1/chat/completions, per-request `model` override |
|
||||
| Canvas hosting | ✅ | ❌ | Agent-driven UI |
|
||||
| Gateway lock (PID-based) | ✅ | ❌ | |
|
||||
| Gateway lock (PID-based) | ✅ | ✅ | `fs4` flock-based, acquired in `main.rs` before agent startup |
|
||||
| launchd/systemd integration | ✅ | ❌ | |
|
||||
| Bonjour/mDNS discovery | ✅ | ❌ | |
|
||||
| Tailscale integration | ✅ | ❌ | |
|
||||
@@ -215,7 +215,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
| NEAR AI | ✅ | ✅ | - | Primary provider |
|
||||
| Anthropic (Claude) | ✅ | 🚧 | - | Via NEAR AI proxy; Opus 4.5, Sonnet 4, Sonnet 4.6 |
|
||||
| OpenAI | ✅ | 🚧 | - | Via NEAR AI proxy |
|
||||
| AWS Bedrock | ✅ | ✅ | P3 | Via `openai_compatible` adapter (e.g. LiteLLM) |
|
||||
| AWS Bedrock | ✅ | ✅ | - | Native Converse API via aws-sdk-bedrockruntime (requires `--features bedrock`) |
|
||||
| Google Gemini | ✅ | ✅ | P3 | Via `gemini` adapter |
|
||||
| io.net | ✅ | ✅ | P3 | Via `ionet` adapter |
|
||||
| Mistral | ✅ | ✅ | P3 | Via `mistral` adapter |
|
||||
|
||||
@@ -14,6 +14,11 @@
|
||||
<a href="https://www.reddit.com/r/ironclawAI/"><img src="https://img.shields.io/badge/Reddit-r%2FironclawAI-FF4500?style=flat&logo=reddit&logoColor=white" alt="Reddit: r/ironclawAI" /></a>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a href="README.md">English</a> |
|
||||
<a href="README.zh-CN.md">简体中文</a>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a href="#philosophy">Philosophy</a> •
|
||||
<a href="#features">Features</a> •
|
||||
|
||||
+319
@@ -0,0 +1,319 @@
|
||||
<p align="center">
|
||||
<img src="ironclaw.png?v=2" alt="IronClaw" width="200"/>
|
||||
</p>
|
||||
|
||||
<h1 align="center">IronClaw</h1>
|
||||
|
||||
<p align="center">
|
||||
<strong>安全可靠的个人 AI 助手,始终站在你这边</strong>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a href="#license"><img src="https://img.shields.io/badge/license-MIT%20OR%20Apache%202.0-blue.svg" alt="License: MIT OR Apache-2.0" /></a>
|
||||
<a href="https://t.me/ironclawAI"><img src="https://img.shields.io/badge/Telegram-%40ironclawAI-26A5E4?style=flat&logo=telegram&logoColor=white" alt="Telegram: @ironclawAI" /></a>
|
||||
<a href="https://www.reddit.com/r/ironclawAI/"><img src="https://img.shields.io/badge/Reddit-r%2FironclawAI-FF4500?style=flat&logo=reddit&logoColor=white" alt="Reddit: r/ironclawAI" /></a>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a href="README.md">English</a> |
|
||||
<a href="README.zh-CN.md">简体中文</a>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a href="#设计理念">设计理念</a> •
|
||||
<a href="#功能特性">功能特性</a> •
|
||||
<a href="#安装">安装</a> •
|
||||
<a href="#配置">配置</a> •
|
||||
<a href="#安全机制">安全机制</a> •
|
||||
<a href="#系统架构">系统架构</a>
|
||||
</p>
|
||||
|
||||
---
|
||||
|
||||
## 设计理念
|
||||
|
||||
IronClaw 基于一个简单的原则:**你的 AI 助手应该为你服务,而不是与你为敌。**
|
||||
|
||||
在 AI 系统对数据处理日益不透明、与企业利益捆绑的今天,IronClaw 选择了一条不同的路:
|
||||
|
||||
- **数据归你所有** — 所有信息存储在本地,加密保护,始终在你掌控之下
|
||||
- **透明至上** — 完全开源,可审计,没有隐藏的遥测或数据收集
|
||||
- **自主扩展** — 随时构建新工具,无需等待供应商更新
|
||||
- **纵深防御** — 多层安全机制抵御提示注入和数据泄露
|
||||
|
||||
IronClaw 是一个你真正可以信赖的 AI 助手,无论是个人生活还是工作。
|
||||
|
||||
## 功能特性
|
||||
|
||||
### 安全优先
|
||||
|
||||
- **WASM 沙箱** — 不受信任的工具在隔离的 WebAssembly 容器中运行,采用基于能力的权限模型
|
||||
- **凭据保护** — 密钥永远不会暴露给工具;在宿主边界注入并进行泄露检测
|
||||
- **提示注入防御** — 模式检测、内容清理和策略执行
|
||||
- **端点白名单** — HTTP 请求仅限于明确批准的主机和路径
|
||||
|
||||
### 随时可用
|
||||
|
||||
- **多渠道接入** — REPL、HTTP webhook、WASM 渠道(Telegram、Slack)和 Web 网关
|
||||
- **Docker 沙箱** — 隔离的容器执行,支持每任务令牌和编排器/工作器模式
|
||||
- **Web 网关** — 浏览器 UI,支持实时 SSE/WebSocket 流式传输
|
||||
- **定时任务** — Cron 调度、事件触发器、Webhook 处理器,实现后台自动化
|
||||
- **心跳系统** — 主动后台执行,用于监控和维护任务
|
||||
- **并行任务** — 使用隔离上下文同时处理多个请求
|
||||
- **自修复** — 自动检测并恢复卡住的操作
|
||||
|
||||
### 自主扩展
|
||||
|
||||
- **动态工具构建** — 描述你的需求,IronClaw 会将其构建为 WASM 工具
|
||||
- **MCP 协议** — 连接模型上下文协议(Model Context Protocol)服务器以获取额外能力
|
||||
- **插件架构** — 无需重启即可加载新的 WASM 工具和渠道
|
||||
|
||||
### 持久记忆
|
||||
|
||||
- **混合搜索** — 全文搜索 + 向量搜索,采用倒数排名融合(Reciprocal Rank Fusion)
|
||||
- **工作空间文件系统** — 灵活的基于路径的存储,用于笔记、日志和上下文
|
||||
- **身份文件** — 跨会话保持一致的个性和偏好设置
|
||||
|
||||
## 安装
|
||||
|
||||
### 前置要求
|
||||
|
||||
- Rust 1.85+
|
||||
- PostgreSQL 15+,需安装 [pgvector](https://github.com/pgvector/pgvector) 扩展
|
||||
- NEAR AI 账户(通过设置向导进行身份验证)
|
||||
|
||||
## 下载或编译
|
||||
|
||||
访问 [Releases 页面](https://github.com/nearai/ironclaw/releases/) 查看最新版本。
|
||||
|
||||
<details>
|
||||
<summary>通过 Windows 安装程序安装 (Windows)</summary>
|
||||
|
||||
下载 [Windows 安装程序](https://github.com/nearai/ironclaw/releases/latest/download/ironclaw-x86_64-pc-windows-msvc.msi) 并运行。
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>通过 PowerShell 脚本安装 (Windows)</summary>
|
||||
|
||||
```sh
|
||||
irm https://github.com/nearai/ironclaw/releases/latest/download/ironclaw-installer.ps1 | iex
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>通过 Shell 脚本安装 (macOS、Linux、Windows/WSL)</summary>
|
||||
|
||||
```sh
|
||||
curl --proto '=https' --tlsv1.2 -LsSf https://github.com/nearai/ironclaw/releases/latest/download/ironclaw-installer.sh | sh
|
||||
```
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>通过 Homebrew 安装 (macOS/Linux)</summary>
|
||||
|
||||
```sh
|
||||
brew install ironclaw
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>从源码编译 (Windows、Linux、macOS 上使用 Cargo)</summary>
|
||||
|
||||
确保你已安装 [Rust](https://rustup.rs)。
|
||||
|
||||
```bash
|
||||
# 克隆仓库
|
||||
git clone https://github.com/nearai/ironclaw.git
|
||||
cd ironclaw
|
||||
|
||||
# 编译
|
||||
cargo build --release
|
||||
|
||||
# 运行测试
|
||||
cargo test
|
||||
```
|
||||
|
||||
如需进行**完整发布构建**(修改了渠道源码后),先运行 `./scripts/build-all.sh` 重新编译渠道。
|
||||
|
||||
</details>
|
||||
|
||||
### 数据库设置
|
||||
|
||||
```bash
|
||||
# 创建数据库
|
||||
createdb ironclaw
|
||||
|
||||
# 启用 pgvector 扩展
|
||||
psql ironclaw -c "CREATE EXTENSION IF NOT EXISTS vector;"
|
||||
```
|
||||
|
||||
## 配置
|
||||
|
||||
运行设置向导来配置 IronClaw:
|
||||
|
||||
```bash
|
||||
ironclaw onboard
|
||||
```
|
||||
|
||||
向导将引导你完成数据库连接、NEAR AI 身份验证(通过浏览器 OAuth)和密钥加密(使用系统钥匙串)。设置会保存在数据库中;引导变量(如 `DATABASE_URL`、`LLM_BACKEND`)写入 `~/.ironclaw/.env`,以便在数据库连接前可用。
|
||||
|
||||
### 替代 LLM 提供商
|
||||
|
||||
IronClaw 默认使用 NEAR AI,但兼容任何 OpenAI 兼容的端点。
|
||||
常用选项包括 **OpenRouter**(300+ 模型)、**Together AI**、**Fireworks AI**、**Ollama**(本地部署)以及自托管服务器如 **vLLM** 或 **LiteLLM**。
|
||||
|
||||
在向导中选择 *"OpenAI-compatible"*,或直接设置环境变量:
|
||||
|
||||
```env
|
||||
LLM_BACKEND=openai_compatible
|
||||
LLM_BASE_URL=https://openrouter.ai/api/v1
|
||||
LLM_API_KEY=sk-or-...
|
||||
LLM_MODEL=anthropic/claude-sonnet-4
|
||||
```
|
||||
|
||||
详见 [docs/LLM_PROVIDERS.md](docs/LLM_PROVIDERS.md) 获取完整的提供商指南。
|
||||
|
||||
## 安全机制
|
||||
|
||||
IronClaw 实现了纵深防御策略来保护你的数据并防止滥用。
|
||||
|
||||
### WASM 沙箱
|
||||
|
||||
所有不受信任的工具都在隔离的 WebAssembly 容器中运行:
|
||||
|
||||
- **基于能力的权限** — 明确授权 HTTP、密钥、工具调用等能力
|
||||
- **端点白名单** — HTTP 请求仅限已批准的主机和路径
|
||||
- **凭据注入** — 密钥在宿主边界注入,永远不会暴露给 WASM 代码
|
||||
- **泄露检测** — 扫描请求和响应以防止密钥外泄
|
||||
- **速率限制** — 每个工具独立的请求限制,防止滥用
|
||||
- **资源限制** — 内存、CPU 和执行时间约束
|
||||
|
||||
```
|
||||
WASM ──► 白名单 ──► 泄露扫描 ──► 凭据 ──► 执行 ──► 泄露扫描 ──► WASM
|
||||
验证器 (请求) 注入器 请求 (响应)
|
||||
```
|
||||
|
||||
### 提示注入防御
|
||||
|
||||
外部内容需通过多个安全层:
|
||||
|
||||
- 基于模式的注入尝试检测
|
||||
- 内容清理和转义
|
||||
- 带严重级别的策略规则(阻止/警告/审核/清理)
|
||||
- 工具输出包装,确保安全的 LLM 上下文注入
|
||||
|
||||
### 数据保护
|
||||
|
||||
- 所有数据存储在本地 PostgreSQL 数据库中
|
||||
- 密钥使用 AES-256-GCM 加密
|
||||
- 无遥测、无分析、无数据共享
|
||||
- 所有工具执行的完整审计日志
|
||||
|
||||
## 系统架构
|
||||
|
||||
```
|
||||
┌────────────────────────────────────────────────────────────────┐
|
||||
│ 渠道 │
|
||||
│ ┌──────┐ ┌──────┐ ┌─────────────┐ ┌─────────────┐ │
|
||||
│ │ REPL │ │ HTTP │ │ WASM 渠道 │ │ Web 网关 │ │
|
||||
│ └──┬───┘ └──┬───┘ └──────┬──────┘ │ (SSE + WS) │ │
|
||||
│ │ │ │ └──────┬──────┘ │
|
||||
│ └─────────┴──────────────┴────────────────┘ │
|
||||
│ │ │
|
||||
│ ┌─────────▼─────────┐ │
|
||||
│ │ 代理循环 │ 意图路由 │
|
||||
│ └────┬──────────┬───┘ │
|
||||
│ │ │ │
|
||||
│ ┌──────────▼────┐ ┌──▼───────────────┐ │
|
||||
│ │ 调度器 │ │ 定时任务引擎 │ │
|
||||
│ │ (并行任务) │ │(cron, 事件, wh) │ │
|
||||
│ └──────┬────────┘ └────────┬─────────┘ │
|
||||
│ │ │ │
|
||||
│ ┌─────────────┼────────────────────┘ │
|
||||
│ │ │ │
|
||||
│ ┌───▼─────┐ ┌────▼────────────────┐ │
|
||||
│ │ 本地 │ │ 编排器 │ │
|
||||
│ │ 工作器 │ │ ┌───────────────┐ │ │
|
||||
│ │(进程内) │ │ │ Docker 沙箱 │ │ │
|
||||
│ └───┬─────┘ │ │ 容器 │ │ │
|
||||
│ │ │ │ ┌───────────┐ │ │ │
|
||||
│ │ │ │ │工作器/CC │ │ │ │
|
||||
│ │ │ │ └───────────┘ │ │ │
|
||||
│ │ │ └───────────────┘ │ │
|
||||
│ │ └─────────┬───────────┘ │
|
||||
│ └──────────────────┤ │
|
||||
│ │ │
|
||||
│ ┌───────────▼──────────┐ │
|
||||
│ │ 工具注册表 │ │
|
||||
│ │ 内置、MCP、WASM │ │
|
||||
│ └──────────────────────┘ │
|
||||
└────────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### 核心组件
|
||||
|
||||
| 组件 | 用途 |
|
||||
|------|------|
|
||||
| **代理循环** | 主消息处理和任务协调 |
|
||||
| **路由器** | 分类用户意图(命令、查询、任务) |
|
||||
| **调度器** | 管理带优先级的并行任务执行 |
|
||||
| **工作器** | 执行包含 LLM 推理和工具调用的任务 |
|
||||
| **编排器** | 容器生命周期、LLM 代理、每任务认证 |
|
||||
| **Web 网关** | 浏览器 UI,含聊天、记忆、任务、日志、扩展、定时任务 |
|
||||
| **定时任务引擎** | 定时(cron)和响应式(事件、webhook)后台任务 |
|
||||
| **工作空间** | 带混合搜索的持久记忆 |
|
||||
| **安全层** | 提示注入防御和内容清理 |
|
||||
|
||||
## 使用方式
|
||||
|
||||
```bash
|
||||
# 首次设置(配置数据库、认证等)
|
||||
ironclaw onboard
|
||||
|
||||
# 启动交互式 REPL
|
||||
cargo run
|
||||
|
||||
# 启用调试日志
|
||||
RUST_LOG=ironclaw=debug cargo run
|
||||
```
|
||||
|
||||
## 开发
|
||||
|
||||
```bash
|
||||
# 格式化代码
|
||||
cargo fmt
|
||||
|
||||
# 代码检查
|
||||
cargo clippy --all --benches --tests --examples --all-features
|
||||
|
||||
# 运行测试
|
||||
createdb ironclaw_test
|
||||
cargo test
|
||||
|
||||
# 运行指定测试
|
||||
cargo test test_name
|
||||
```
|
||||
|
||||
- **Telegram 渠道**:参见 [docs/TELEGRAM_SETUP.md](docs/TELEGRAM_SETUP.md) 了解设置和私信配对。
|
||||
- **修改渠道源码**:在 `cargo build` 之前运行 `./channels-src/telegram/build.sh` 以便打包更新后的 WASM。
|
||||
|
||||
## OpenClaw 传承
|
||||
|
||||
IronClaw 是受 [OpenClaw](https://github.com/openclaw/openclaw) 启发的 Rust 重新实现。参见 [FEATURE_PARITY.md](FEATURE_PARITY.md) 了解完整的功能追踪矩阵。
|
||||
|
||||
主要差异:
|
||||
|
||||
- **Rust vs TypeScript** — 原生性能、内存安全、单一二进制文件
|
||||
- **WASM 沙箱 vs Docker** — 轻量级、基于能力的安全机制
|
||||
- **PostgreSQL vs SQLite** — 生产级持久化存储
|
||||
- **安全优先设计** — 多层防御、凭据保护
|
||||
|
||||
## 许可证
|
||||
|
||||
可选择以下任一许可证:
|
||||
|
||||
- Apache License, Version 2.0 ([LICENSE-APACHE](LICENSE-APACHE))
|
||||
- MIT License ([LICENSE-MIT](LICENSE-MIT))
|
||||
Generated
+1
-1
@@ -267,7 +267,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "slack-channel"
|
||||
version = "0.1.0"
|
||||
version = "0.2.1"
|
||||
dependencies = [
|
||||
"hex",
|
||||
"hmac",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "slack-channel"
|
||||
version = "0.2.0"
|
||||
version = "0.2.1"
|
||||
edition = "2021"
|
||||
description = "Slack Events API channel for IronClaw"
|
||||
license = "MIT OR Apache-2.0"
|
||||
|
||||
@@ -357,10 +357,108 @@ fn extract_slack_attachments(files: &Option<Vec<SlackFile>>) -> Vec<InboundAttac
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Download a file from Slack using the url_private endpoint.
|
||||
///
|
||||
/// Slack file downloads require Bearer auth with the bot token, which is
|
||||
/// injected by the host credential system via `channel_host::http_request`.
|
||||
fn download_slack_file(url: &str) -> Result<Vec<u8>, String> {
|
||||
let headers = serde_json::json!({});
|
||||
|
||||
let result = channel_host::http_request("GET", url, &headers.to_string(), None, None);
|
||||
|
||||
let response = result.map_err(|e| format!("Slack file download failed: {}", e))?;
|
||||
|
||||
if response.status != 200 {
|
||||
let body_str = String::from_utf8_lossy(&response.body);
|
||||
return Err(format!(
|
||||
"Slack file download returned {}: {}",
|
||||
response.status, body_str
|
||||
));
|
||||
}
|
||||
|
||||
Ok(response.body)
|
||||
}
|
||||
|
||||
/// Download file bytes and store them via the host for processing.
|
||||
///
|
||||
/// Downloads all file types (images, documents, etc.) so the host-side
|
||||
/// middleware can process them (vision pipeline for images, text extraction
|
||||
/// for documents, transcription for audio, etc.).
|
||||
/// Maximum file size to download (20 MB). Files larger than this are skipped
|
||||
/// to avoid excessive memory use and slow downloads in the WASM runtime.
|
||||
const MAX_DOWNLOAD_SIZE_BYTES: u64 = 20 * 1024 * 1024;
|
||||
|
||||
fn download_and_store_slack_files(attachments: &[InboundAttachment]) {
|
||||
for att in attachments {
|
||||
let Some(ref url) = att.source_url else {
|
||||
continue;
|
||||
};
|
||||
|
||||
// Skip files that exceed the size limit
|
||||
if let Some(size) = att.size_bytes {
|
||||
if size > MAX_DOWNLOAD_SIZE_BYTES {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Warn,
|
||||
&format!(
|
||||
"Skipping Slack file download: {} bytes exceeds {} MB limit (id={})",
|
||||
size,
|
||||
MAX_DOWNLOAD_SIZE_BYTES / (1024 * 1024),
|
||||
att.id
|
||||
),
|
||||
);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
match download_slack_file(url) {
|
||||
Ok(bytes) => {
|
||||
// Post-download size guard: metadata size_bytes is optional,
|
||||
// so a file with no size info could bypass the pre-download check.
|
||||
if bytes.len() as u64 > MAX_DOWNLOAD_SIZE_BYTES {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Warn,
|
||||
&format!(
|
||||
"Discarding Slack file after download: {} bytes exceeds {} MB limit (id={})",
|
||||
bytes.len(),
|
||||
MAX_DOWNLOAD_SIZE_BYTES / (1024 * 1024),
|
||||
att.id
|
||||
),
|
||||
);
|
||||
continue;
|
||||
}
|
||||
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Info,
|
||||
&format!(
|
||||
"Downloaded Slack file: {} bytes, mime={}",
|
||||
bytes.len(),
|
||||
att.mime_type
|
||||
),
|
||||
);
|
||||
if let Err(e) = channel_host::store_attachment_data(&att.id, &bytes) {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Error,
|
||||
&format!("Failed to store Slack file data: {}", e),
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Error,
|
||||
&format!("Failed to download Slack file: {}", e),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle a Slack event and emit message if applicable.
|
||||
fn handle_slack_event(event: SlackEvent, team_id: Option<String>, _event_id: Option<String>) {
|
||||
let attachments = extract_slack_attachments(&event.files);
|
||||
|
||||
// Download and store file attachments for host-side processing
|
||||
download_and_store_slack_files(&attachments);
|
||||
|
||||
match event.event_type.as_str() {
|
||||
// Direct mention of the bot (always in a channel, not a DM)
|
||||
"app_mention" => {
|
||||
@@ -722,4 +820,10 @@ mod tests {
|
||||
let event: SlackEvent = serde_json::from_str(json).unwrap();
|
||||
assert!(event.files.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_max_download_size_constant() {
|
||||
// Verify the constant is 20 MB
|
||||
assert_eq!(MAX_DOWNLOAD_SIZE_BYTES, 20 * 1024 * 1024);
|
||||
}
|
||||
}
|
||||
|
||||
Generated
+1
-1
@@ -212,7 +212,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "telegram-channel"
|
||||
version = "0.2.0"
|
||||
version = "0.2.1"
|
||||
dependencies = [
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "telegram-channel"
|
||||
version = "0.2.0"
|
||||
version = "0.2.1"
|
||||
edition = "2021"
|
||||
description = "Telegram Bot API channel for IronClaw"
|
||||
license = "MIT OR Apache-2.0"
|
||||
|
||||
@@ -878,10 +878,6 @@ fn send_message(
|
||||
// Voice File Download
|
||||
// ============================================================================
|
||||
|
||||
/// Download a voice file from Telegram by file_id.
|
||||
///
|
||||
/// 1. Call getFile to get the file_path.
|
||||
/// 2. Download the file bytes from /file/bot{TOKEN}/{file_path}.
|
||||
/// Percent-encode a string for safe use as a URL query parameter value.
|
||||
fn percent_encode(s: &str) -> String {
|
||||
let mut out = String::with_capacity(s.len());
|
||||
@@ -898,6 +894,10 @@ fn percent_encode(s: &str) -> String {
|
||||
out
|
||||
}
|
||||
|
||||
/// Maximum file size to download (20 MB). Files larger than this are discarded
|
||||
/// to avoid excessive memory use and slow downloads in the WASM runtime.
|
||||
const MAX_DOWNLOAD_SIZE_BYTES: u64 = 20 * 1024 * 1024;
|
||||
|
||||
fn download_telegram_file(file_id: &str) -> Result<Vec<u8>, String> {
|
||||
// Reject file_id containing curly braces to prevent credential placeholder injection
|
||||
if file_id.contains('{') || file_id.contains('}') {
|
||||
@@ -965,6 +965,16 @@ fn download_telegram_file(file_id: &str) -> Result<Vec<u8>, String> {
|
||||
));
|
||||
}
|
||||
|
||||
// Post-download size guard: Telegram metadata file_size is optional,
|
||||
// so enforce the limit on actual downloaded bytes.
|
||||
if response.body.len() as u64 > MAX_DOWNLOAD_SIZE_BYTES {
|
||||
return Err(format!(
|
||||
"Downloaded file exceeds {} MB limit ({} bytes)",
|
||||
MAX_DOWNLOAD_SIZE_BYTES / (1024 * 1024),
|
||||
response.body.len()
|
||||
));
|
||||
}
|
||||
|
||||
Ok(response.body)
|
||||
}
|
||||
|
||||
@@ -1535,6 +1545,39 @@ fn download_and_store_voice(attachments: &[InboundAttachment]) {
|
||||
}
|
||||
}
|
||||
|
||||
/// Download image file bytes and store them via the host for the vision pipeline.
|
||||
///
|
||||
/// Separated from `extract_attachments` so that function stays pure (no host
|
||||
/// calls) and remains testable in native unit tests.
|
||||
fn download_and_store_images(attachments: &[InboundAttachment]) {
|
||||
for att in attachments {
|
||||
if !att.mime_type.starts_with("image/") {
|
||||
continue;
|
||||
}
|
||||
|
||||
match download_telegram_file(&att.id) {
|
||||
Ok(bytes) => {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Info,
|
||||
&format!("Downloaded image file: {} bytes", bytes.len()),
|
||||
);
|
||||
if let Err(e) = channel_host::store_attachment_data(&att.id, &bytes) {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Error,
|
||||
&format!("Failed to store image data: {}", e),
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
channel_host::log(
|
||||
channel_host::LogLevel::Error,
|
||||
&format!("Failed to download image file: {}", e),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns true if the attachment should be downloaded for document text extraction.
|
||||
///
|
||||
/// Excludes voice (handled by transcription), image (vision pipeline),
|
||||
@@ -1608,6 +1651,9 @@ fn handle_message(message: TelegramMessage) {
|
||||
// Download and store voice attachments for host-side transcription
|
||||
download_and_store_voice(&attachments);
|
||||
|
||||
// Download and store image attachments for host-side vision pipeline
|
||||
download_and_store_images(&attachments);
|
||||
|
||||
// Download and store document attachments for host-side text extraction
|
||||
download_and_store_documents(&mut attachments);
|
||||
|
||||
@@ -1681,7 +1727,7 @@ fn handle_message(message: TelegramMessage) {
|
||||
let username_opt = from.username.as_deref();
|
||||
let is_allowed = allowed.contains(&"*".to_string())
|
||||
|| allowed.contains(&id_str)
|
||||
|| username_opt.map_or(false, |u| allowed.contains(&u.to_string()));
|
||||
|| username_opt.is_some_and(|u| allowed.contains(&u.to_string()));
|
||||
|
||||
if !is_allowed {
|
||||
if is_private && dm_policy == "pairing" {
|
||||
@@ -2605,4 +2651,10 @@ mod tests {
|
||||
assert!(!is_downloadable_document(&make("audio/mpeg", Some("song.mp3"))));
|
||||
assert!(!is_downloadable_document(&make("video/mp4", Some("clip.mp4"))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_max_download_size_constant() {
|
||||
// Verify the constant is 20 MB, matching the Slack channel limit
|
||||
assert_eq!(MAX_DOWNLOAD_SIZE_BYTES, 20 * 1024 * 1024);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,9 +1,17 @@
|
||||
{
|
||||
"version": "0.2.0",
|
||||
"version": "0.2.2",
|
||||
"wit_version": "0.3.0",
|
||||
"type": "channel",
|
||||
"name": "telegram",
|
||||
"description": "Telegram Bot API channel for receiving and responding to Telegram messages",
|
||||
"auth": {
|
||||
"secret_name": "telegram_bot_token",
|
||||
"display_name": "Telegram",
|
||||
"instructions": "Get your bot token from @BotFather on Telegram (https://t.me/BotFather). Send /newbot or /token to get it.",
|
||||
"setup_url": "https://t.me/BotFather",
|
||||
"token_hint": "Looks like 123456789:AABBccDDeeFFgg...",
|
||||
"env_var": "TELEGRAM_BOT_TOKEN"
|
||||
},
|
||||
"setup": {
|
||||
"required_secrets": [
|
||||
{
|
||||
|
||||
+1
-1
@@ -3,7 +3,7 @@ services:
|
||||
postgres:
|
||||
image: pgvector/pgvector:pg16
|
||||
ports:
|
||||
- "5432:5432"
|
||||
- "127.0.0.1:5432:5432"
|
||||
environment:
|
||||
POSTGRES_DB: ironclaw
|
||||
POSTGRES_USER: ironclaw
|
||||
|
||||
+50
-1
@@ -12,12 +12,12 @@ configurations.
|
||||
| Anthropic | `anthropic` | `ANTHROPIC_API_KEY` | Claude models |
|
||||
| OpenAI | `openai` | `OPENAI_API_KEY` | GPT models |
|
||||
| Google Gemini | `gemini` | `GEMINI_API_KEY` | Gemini models |
|
||||
| AWS Bedrock | `bedrock` | `BEDROCK_ACCESS_KEY` | Requires OpenAI proxy (e.g. LiteLLM) |
|
||||
| io.net | `ionet` | `IONET_API_KEY` | Intelligence API |
|
||||
| Mistral | `mistral` | `MISTRAL_API_KEY` | Mistral models |
|
||||
| Yandex AI Studio | `yandex` | `YANDEX_API_KEY` | YandexGPT models |
|
||||
| Cloudflare Workers AI | `cloudflare` | `CLOUDFLARE_API_KEY` | Access to Workers AI |
|
||||
| Ollama | `ollama` | No | Local inference |
|
||||
| AWS Bedrock | `bedrock` | AWS credentials | Native Converse API |
|
||||
| OpenRouter | `openai_compatible` | `LLM_API_KEY` | 300+ models |
|
||||
| Together AI | `openai_compatible` | `LLM_API_KEY` | Fast inference |
|
||||
| Fireworks AI | `openai_compatible` | `LLM_API_KEY` | Fast inference |
|
||||
@@ -74,6 +74,55 @@ Pull a model first: `ollama pull llama3.2`
|
||||
|
||||
---
|
||||
|
||||
## AWS Bedrock (requires `--features bedrock`)
|
||||
|
||||
Uses the native AWS Converse API via `aws-sdk-bedrockruntime`. Supports standard AWS
|
||||
authentication methods: IAM credentials, SSO profiles, and instance roles.
|
||||
|
||||
> **Build prerequisite:** The `aws-lc-sys` crate (transitive dependency via AWS SDK)
|
||||
> requires **CMake** to compile. Install it before building with `--features bedrock`:
|
||||
> - macOS: `brew install cmake`
|
||||
> - Ubuntu/Debian: `sudo apt install cmake`
|
||||
> - Fedora: `sudo dnf install cmake`
|
||||
|
||||
### With AWS credentials (IAM, SSO, instance roles)
|
||||
|
||||
```env
|
||||
LLM_BACKEND=bedrock
|
||||
BEDROCK_MODEL=anthropic.claude-opus-4-6-v1
|
||||
BEDROCK_REGION=us-east-1
|
||||
BEDROCK_CROSS_REGION=us
|
||||
# AWS_PROFILE=my-sso-profile # optional, for named profiles
|
||||
```
|
||||
|
||||
The AWS SDK credential chain automatically resolves credentials from environment
|
||||
variables (`AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`), shared credentials file
|
||||
(`~/.aws/credentials`), SSO profiles, and EC2/ECS instance roles.
|
||||
|
||||
### Cross-region inference
|
||||
|
||||
Set `BEDROCK_CROSS_REGION` to route requests across AWS regions for capacity:
|
||||
|
||||
| Prefix | Routing |
|
||||
|---|---|
|
||||
| `us` | US regions (us-east-1, us-east-2, us-west-2) |
|
||||
| `eu` | European regions |
|
||||
| `apac` | Asia-Pacific regions |
|
||||
| `global` | All commercial AWS regions |
|
||||
| _(unset)_ | Single-region only |
|
||||
|
||||
### Popular Bedrock model IDs
|
||||
|
||||
| Model | ID |
|
||||
|---|---|
|
||||
| Claude Opus 4.6 | `anthropic.claude-opus-4-6-v1` |
|
||||
| Claude Sonnet 4.5 | `anthropic.claude-sonnet-4-5-20250929-v1:0` |
|
||||
| Claude Haiku 4.5 | `anthropic.claude-haiku-4-5-20251001-v1:0` |
|
||||
| Amazon Nova Pro | `amazon.nova-pro-v1:0` |
|
||||
| Llama 4 Maverick | `meta.llama4-maverick-17b-instruct-v1:0` |
|
||||
|
||||
---
|
||||
|
||||
## OpenAI-Compatible Endpoints
|
||||
|
||||
All providers below use `LLM_BACKEND=openai_compatible`. Set `LLM_BASE_URL` to the
|
||||
|
||||
@@ -295,26 +295,6 @@
|
||||
"can_list_models": true
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "bedrock",
|
||||
"aliases": [
|
||||
"aws_bedrock",
|
||||
"aws"
|
||||
],
|
||||
"protocol": "open_ai_completions",
|
||||
"api_key_env": "BEDROCK_ACCESS_KEY",
|
||||
"api_key_required": false,
|
||||
"base_url_env": "BEDROCK_BASE_URL",
|
||||
"model_env": "BEDROCK_MODEL",
|
||||
"default_model": "anthropic.claude-3-5-sonnet-20241022-v2:0",
|
||||
"description": "AWS Bedrock (requires LiteLLM or OpenAI-compatible proxy)",
|
||||
"setup": {
|
||||
"kind": "open_ai_compatible",
|
||||
"secret_name": "llm_bedrock_api_key",
|
||||
"display_name": "AWS Bedrock",
|
||||
"can_list_models": false
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "ionet",
|
||||
"aliases": [
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
"name": "slack",
|
||||
"display_name": "Slack Channel",
|
||||
"kind": "channel",
|
||||
"version": "0.2.0",
|
||||
"version": "0.2.1",
|
||||
"wit_version": "0.3.0",
|
||||
"description": "Talk to your agent in Slack",
|
||||
"keywords": [
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
"name": "telegram",
|
||||
"display_name": "Telegram Channel",
|
||||
"kind": "channel",
|
||||
"version": "0.2.0",
|
||||
"version": "0.2.2",
|
||||
"wit_version": "0.3.0",
|
||||
"description": "Talk to your agent through a Telegram bot",
|
||||
"keywords": [
|
||||
|
||||
@@ -356,6 +356,12 @@ impl Agent {
|
||||
if let Some(workspace) = self.workspace() {
|
||||
let mut config = AgentHeartbeatConfig::default()
|
||||
.with_interval(std::time::Duration::from_secs(hb_config.interval_secs));
|
||||
config.quiet_hours_start = hb_config.quiet_hours_start;
|
||||
config.quiet_hours_end = hb_config.quiet_hours_end;
|
||||
config.timezone = hb_config
|
||||
.timezone
|
||||
.clone()
|
||||
.or_else(|| Some(self.config.default_timezone.clone()));
|
||||
if let (Some(user), Some(channel)) =
|
||||
(&hb_config.notify_user, &hb_config.notify_channel)
|
||||
{
|
||||
@@ -411,7 +417,6 @@ impl Agent {
|
||||
hygiene,
|
||||
workspace.clone(),
|
||||
self.cheap_llm().clone(),
|
||||
self.safety().clone(),
|
||||
Some(notify_tx),
|
||||
self.store().map(Arc::clone),
|
||||
))
|
||||
|
||||
+48
-7
@@ -345,7 +345,6 @@ impl Agent {
|
||||
crate::workspace::hygiene::HygieneConfig::default(),
|
||||
workspace.clone(),
|
||||
self.llm().clone(),
|
||||
self.safety().clone(),
|
||||
);
|
||||
|
||||
match runner.check_heartbeat().await {
|
||||
@@ -406,7 +405,7 @@ impl Agent {
|
||||
.with_max_tokens(512)
|
||||
.with_temperature(0.3);
|
||||
|
||||
let reasoning = Reasoning::new(self.llm().clone(), self.safety().clone());
|
||||
let reasoning = Reasoning::new(self.llm().clone());
|
||||
match reasoning.complete(request).await {
|
||||
Ok((text, _usage)) => Ok(SubmissionResult::response(format!(
|
||||
"Thread Summary:\n\n{}",
|
||||
@@ -454,7 +453,7 @@ impl Agent {
|
||||
.with_max_tokens(512)
|
||||
.with_temperature(0.5);
|
||||
|
||||
let reasoning = Reasoning::new(self.llm().clone(), self.safety().clone());
|
||||
let reasoning = Reasoning::new(self.llm().clone());
|
||||
match reasoning.complete(request).await {
|
||||
Ok((text, _usage)) => Ok(SubmissionResult::response(format!(
|
||||
"Suggested Next Steps:\n\n{}",
|
||||
@@ -663,10 +662,14 @@ impl Agent {
|
||||
}
|
||||
|
||||
match self.llm().set_model(requested) {
|
||||
Ok(()) => Ok(SubmissionResult::response(format!(
|
||||
"Switched model to: {}",
|
||||
requested
|
||||
))),
|
||||
Ok(()) => {
|
||||
// Persist the model choice so it survives restarts.
|
||||
self.persist_selected_model(requested).await;
|
||||
Ok(SubmissionResult::response(format!(
|
||||
"Switched model to: {}",
|
||||
requested
|
||||
)))
|
||||
}
|
||||
Err(e) => Ok(SubmissionResult::error(format!(
|
||||
"Failed to switch model: {}",
|
||||
e
|
||||
@@ -822,4 +825,42 @@ impl Agent {
|
||||
_ => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
/// Persist the selected model to the settings store (DB and/or TOML config).
|
||||
///
|
||||
/// Best-effort: logs warnings on failure but does not propagate errors,
|
||||
/// since the in-memory model switch already succeeded.
|
||||
async fn persist_selected_model(&self, model: &str) {
|
||||
// 1. Persist to DB if available.
|
||||
if let Some(store) = self.store() {
|
||||
let value = serde_json::Value::String(model.to_string());
|
||||
if let Err(e) = store.set_setting("default", "selected_model", &value).await {
|
||||
tracing::warn!("Failed to persist model to DB: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
// 2. Update TOML config file if it exists (sync I/O in spawn_blocking).
|
||||
let model_owned = model.to_string();
|
||||
if let Err(e) = tokio::task::spawn_blocking(move || {
|
||||
let toml_path = crate::settings::Settings::default_toml_path();
|
||||
match crate::settings::Settings::load_toml(&toml_path) {
|
||||
Ok(Some(mut settings)) => {
|
||||
settings.selected_model = Some(model_owned);
|
||||
if let Err(e) = settings.save_toml(&toml_path) {
|
||||
tracing::warn!("Failed to persist model to config.toml: {}", e);
|
||||
}
|
||||
}
|
||||
Ok(None) => {
|
||||
// No config file on disk; nothing to update.
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to load config.toml for model persistence: {}", e);
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
{
|
||||
tracing::warn!("Model TOML persistence task failed: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+112
-37
@@ -13,7 +13,6 @@ use crate::agent::context_monitor::{CompactionStrategy, ContextBreakdown};
|
||||
use crate::agent::session::Thread;
|
||||
use crate::error::Error;
|
||||
use crate::llm::{ChatMessage, CompletionRequest, LlmProvider, Reasoning};
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::workspace::Workspace;
|
||||
|
||||
/// Result of a compaction operation.
|
||||
@@ -34,13 +33,12 @@ pub struct CompactionResult {
|
||||
/// Compacts conversation context to stay within limits.
|
||||
pub struct ContextCompactor {
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
}
|
||||
|
||||
impl ContextCompactor {
|
||||
/// Create a new context compactor.
|
||||
pub fn new(llm: Arc<dyn LlmProvider>, safety: Arc<SafetyLayer>) -> Self {
|
||||
Self { llm, safety }
|
||||
pub fn new(llm: Arc<dyn LlmProvider>) -> Self {
|
||||
Self { llm }
|
||||
}
|
||||
|
||||
/// Compact a thread's context using the given strategy.
|
||||
@@ -105,27 +103,26 @@ impl ContextCompactor {
|
||||
// Generate summary
|
||||
let summary = self.generate_summary(&to_summarize).await?;
|
||||
|
||||
// Write to workspace if available
|
||||
let summary_written = if let Some(ws) = workspace {
|
||||
// Write to workspace if available.
|
||||
// If archival fails, preserve turns to avoid context loss.
|
||||
let (summary_written, turns_removed) = if let Some(ws) = workspace {
|
||||
match self.write_summary_to_workspace(ws, &summary).await {
|
||||
Ok(()) => true,
|
||||
Ok(()) => {
|
||||
thread.truncate_turns(keep_recent);
|
||||
(true, turns_to_remove)
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Compaction summary write failed (turns will still be truncated): {}",
|
||||
e
|
||||
);
|
||||
false
|
||||
tracing::warn!("Compaction summary write failed (turns preserved): {}", e);
|
||||
(false, 0)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
false
|
||||
thread.truncate_turns(keep_recent);
|
||||
(false, turns_to_remove)
|
||||
};
|
||||
|
||||
// Truncate thread
|
||||
thread.truncate_turns(keep_recent);
|
||||
|
||||
Ok(CompactionPartial {
|
||||
turns_removed: turns_to_remove,
|
||||
turns_removed,
|
||||
summary_written,
|
||||
summary: Some(summary),
|
||||
})
|
||||
@@ -167,23 +164,20 @@ impl ContextCompactor {
|
||||
// Format turns for storage
|
||||
let content = format_turns_for_storage(old_turns);
|
||||
|
||||
// Write to workspace
|
||||
let written = match self.write_context_to_workspace(ws, &content).await {
|
||||
Ok(()) => true,
|
||||
// Write to workspace. If archival fails, preserve turns.
|
||||
let (written, turns_removed) = match self.write_context_to_workspace(ws, &content).await {
|
||||
Ok(()) => {
|
||||
thread.truncate_turns(keep_recent);
|
||||
(true, turns_to_remove)
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Compaction context write failed (turns will still be truncated): {}",
|
||||
e
|
||||
);
|
||||
false
|
||||
tracing::warn!("Compaction context write failed (turns preserved): {}", e);
|
||||
(false, 0)
|
||||
}
|
||||
};
|
||||
|
||||
// Truncate
|
||||
thread.truncate_turns(keep_recent);
|
||||
|
||||
Ok(CompactionPartial {
|
||||
turns_removed: turns_to_remove,
|
||||
turns_removed,
|
||||
summary_written: written,
|
||||
summary: None,
|
||||
})
|
||||
@@ -233,7 +227,7 @@ Be brief but capture all important details. Use bullet points."#,
|
||||
.with_max_tokens(1024)
|
||||
.with_temperature(0.3);
|
||||
|
||||
let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone());
|
||||
let reasoning = Reasoning::new(self.llm.clone());
|
||||
let (text, _) = reasoning.complete(request).await?;
|
||||
Ok(text)
|
||||
}
|
||||
@@ -346,17 +340,11 @@ mod tests {
|
||||
// === QA Plan - Compaction strategy tests ===
|
||||
|
||||
use crate::agent::context_monitor::CompactionStrategy;
|
||||
use crate::config::SafetyConfig;
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::testing::StubLlm;
|
||||
|
||||
/// Helper: build a `ContextCompactor` with the given `StubLlm`.
|
||||
fn make_compactor(llm: Arc<StubLlm>) -> ContextCompactor {
|
||||
let safety = Arc::new(SafetyLayer::new(&SafetyConfig {
|
||||
max_output_length: 100_000,
|
||||
injection_check_enabled: false,
|
||||
}));
|
||||
ContextCompactor::new(llm, safety)
|
||||
ContextCompactor::new(llm)
|
||||
}
|
||||
|
||||
/// Helper: build a thread with `n` completed turns.
|
||||
@@ -370,6 +358,19 @@ mod tests {
|
||||
thread
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
async fn make_unmigrated_workspace() -> crate::workspace::Workspace {
|
||||
use crate::db::Database;
|
||||
use crate::db::libsql::LibSqlBackend;
|
||||
|
||||
// Intentionally skip migrations so workspace append operations fail.
|
||||
let backend = LibSqlBackend::new_memory()
|
||||
.await
|
||||
.expect("should create in-memory libsql backend");
|
||||
let db: Arc<dyn Database> = Arc::new(backend);
|
||||
crate::workspace::Workspace::new_with_db("compaction-test", db)
|
||||
}
|
||||
|
||||
// ------------------------------------------------------------------
|
||||
// 1. compact_truncate keeps last N turns
|
||||
// ------------------------------------------------------------------
|
||||
@@ -568,6 +569,43 @@ mod tests {
|
||||
assert_eq!(llm.calls(), 0);
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
#[tokio::test]
|
||||
async fn test_compact_with_summary_preserves_turns_when_workspace_write_fails() {
|
||||
let llm = Arc::new(StubLlm::new("summary"));
|
||||
let compactor = make_compactor(llm.clone());
|
||||
let mut thread = make_thread(8);
|
||||
let original_inputs: Vec<String> =
|
||||
thread.turns.iter().map(|t| t.user_input.clone()).collect();
|
||||
let workspace = make_unmigrated_workspace().await;
|
||||
|
||||
let result = compactor
|
||||
.compact(
|
||||
&mut thread,
|
||||
CompactionStrategy::Summarize { keep_recent: 3 },
|
||||
Some(&workspace),
|
||||
)
|
||||
.await
|
||||
.expect("compact should succeed even when workspace write fails");
|
||||
|
||||
// On archival failure, no turns should be removed.
|
||||
assert_eq!(thread.turns.len(), 8);
|
||||
assert_eq!(
|
||||
thread
|
||||
.turns
|
||||
.iter()
|
||||
.map(|t| t.user_input.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
original_inputs
|
||||
.iter()
|
||||
.map(|s| s.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
);
|
||||
assert_eq!(result.turns_removed, 0);
|
||||
assert!(!result.summary_written);
|
||||
assert_eq!(llm.calls(), 1);
|
||||
}
|
||||
|
||||
// ------------------------------------------------------------------
|
||||
// 7. compact_to_workspace without workspace falls back to truncation
|
||||
// ------------------------------------------------------------------
|
||||
@@ -616,6 +654,43 @@ mod tests {
|
||||
assert_eq!(result.turns_removed, 0);
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
#[tokio::test]
|
||||
async fn test_compact_to_workspace_preserves_turns_when_workspace_write_fails() {
|
||||
let llm = Arc::new(StubLlm::new("unused"));
|
||||
let compactor = make_compactor(llm.clone());
|
||||
let mut thread = make_thread(20);
|
||||
let original_inputs: Vec<String> =
|
||||
thread.turns.iter().map(|t| t.user_input.clone()).collect();
|
||||
let workspace = make_unmigrated_workspace().await;
|
||||
|
||||
let result = compactor
|
||||
.compact(
|
||||
&mut thread,
|
||||
CompactionStrategy::MoveToWorkspace,
|
||||
Some(&workspace),
|
||||
)
|
||||
.await
|
||||
.expect("compact should succeed even when workspace write fails");
|
||||
|
||||
// On archival failure, no turns should be removed.
|
||||
assert_eq!(thread.turns.len(), 20);
|
||||
assert_eq!(
|
||||
thread
|
||||
.turns
|
||||
.iter()
|
||||
.map(|t| t.user_input.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
original_inputs
|
||||
.iter()
|
||||
.map(|s| s.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
);
|
||||
assert_eq!(result.turns_removed, 0);
|
||||
assert!(!result.summary_written);
|
||||
assert_eq!(llm.calls(), 0);
|
||||
}
|
||||
|
||||
// ------------------------------------------------------------------
|
||||
// 9. format_turns_for_storage includes tool calls
|
||||
// ------------------------------------------------------------------
|
||||
|
||||
+128
-31
@@ -50,8 +50,18 @@ impl Agent {
|
||||
|
||||
// Load workspace system prompt (identity files: AGENTS.md, SOUL.md, etc.)
|
||||
// In group chats, MEMORY.md is excluded to prevent leaking personal context.
|
||||
// Resolve the user's timezone
|
||||
let user_tz = crate::timezone::resolve_timezone(
|
||||
message.timezone.as_deref(),
|
||||
None, // user setting lookup can be added later
|
||||
&self.config.default_timezone,
|
||||
);
|
||||
|
||||
let system_prompt = if let Some(ws) = self.workspace() {
|
||||
match ws.system_prompt_for_context(is_group_chat).await {
|
||||
match ws
|
||||
.system_prompt_for_context_tz(is_group_chat, user_tz)
|
||||
.await
|
||||
{
|
||||
Ok(prompt) if !prompt.is_empty() => Some(prompt),
|
||||
Ok(_) => None,
|
||||
Err(e) => {
|
||||
@@ -103,7 +113,7 @@ impl Agent {
|
||||
None
|
||||
};
|
||||
|
||||
let mut reasoning = Reasoning::new(self.llm().clone(), self.safety().clone())
|
||||
let mut reasoning = Reasoning::new(self.llm().clone())
|
||||
.with_channel(message.channel.clone())
|
||||
.with_model_name(self.llm().active_model_name())
|
||||
.with_group_chat(is_group_chat);
|
||||
@@ -130,6 +140,7 @@ impl Agent {
|
||||
let mut job_ctx =
|
||||
JobContext::with_user(&message.user_id, "chat", "Interactive chat session");
|
||||
job_ctx.http_interceptor = self.deps.http_interceptor.clone();
|
||||
job_ctx.user_timezone = user_tz.name().to_string();
|
||||
|
||||
// Build system prompts once for this turn. Two variants: with tools
|
||||
// (normal iterations) and without (force_text final iteration).
|
||||
@@ -670,8 +681,53 @@ impl Agent {
|
||||
.into())
|
||||
});
|
||||
|
||||
// Send ToolResult preview
|
||||
if let Ok(ref output) = tool_result
|
||||
// Detect image generation sentinel in tool output
|
||||
// (only from image tools — avoids parsing all tool outputs)
|
||||
let is_image_sentinel = if let Ok(ref output) = tool_result
|
||||
&& matches!(tc.name.as_str(), "image_generate" | "image_edit")
|
||||
{
|
||||
if let Ok(sentinel) =
|
||||
serde_json::from_str::<serde_json::Value>(output)
|
||||
&& sentinel.get("type").and_then(|v| v.as_str())
|
||||
== Some("image_generated")
|
||||
{
|
||||
let data_url = sentinel
|
||||
.get("data")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
let path = sentinel
|
||||
.get("path")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from);
|
||||
// Skip broadcasting if data_url is empty to avoid
|
||||
// sending a broken ImageGenerated SSE event.
|
||||
if data_url.is_empty() {
|
||||
tracing::warn!(
|
||||
"Image generation sentinel has empty data URL, skipping broadcast"
|
||||
);
|
||||
} else {
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::ImageGenerated { data_url, path },
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
} else {
|
||||
false
|
||||
};
|
||||
|
||||
// Send ToolResult preview (skip for image sentinels to avoid
|
||||
// broadcasting multi-MB base64 data as a preview)
|
||||
if !is_image_sentinel
|
||||
&& let Ok(ref output) = tool_result
|
||||
&& !output.is_empty()
|
||||
{
|
||||
let _ = self
|
||||
@@ -687,23 +743,6 @@ impl Agent {
|
||||
.await;
|
||||
}
|
||||
|
||||
// Record result in thread
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id)
|
||||
&& let Some(turn) = thread.last_turn_mut()
|
||||
{
|
||||
match &tool_result {
|
||||
Ok(output) => {
|
||||
turn.record_tool_result(serde_json::json!(output));
|
||||
}
|
||||
Err(e) => {
|
||||
turn.record_tool_error(e.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check for auth awaiting — defer the return
|
||||
// until all results are recorded.
|
||||
if deferred_auth.is_none()
|
||||
@@ -743,6 +782,7 @@ impl Agent {
|
||||
}
|
||||
|
||||
// Sanitize and add tool result to context
|
||||
let is_tool_error = tool_result.is_err();
|
||||
let result_content = match tool_result {
|
||||
Ok(output) => {
|
||||
let sanitized =
|
||||
@@ -756,6 +796,23 @@ impl Agent {
|
||||
Err(e) => format!("Tool '{}' failed: {}", tc.name, e),
|
||||
};
|
||||
|
||||
// Record sanitized result in thread so messages()
|
||||
// and persist_tool_calls() use cleaned content.
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id)
|
||||
&& let Some(turn) = thread.last_turn_mut()
|
||||
{
|
||||
if is_tool_error {
|
||||
turn.record_tool_error(result_content.clone());
|
||||
} else {
|
||||
turn.record_tool_result(serde_json::json!(
|
||||
result_content
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
context_messages.push(ChatMessage::tool_result(
|
||||
&tc.id,
|
||||
&tc.name,
|
||||
@@ -785,6 +842,7 @@ impl Agent {
|
||||
tool_call_id: tc.id.clone(),
|
||||
context_messages: context_messages.clone(),
|
||||
deferred_tool_calls: tool_calls[approval_idx + 1..].to_vec(),
|
||||
user_timezone: Some(user_tz.name().to_string()),
|
||||
};
|
||||
|
||||
return Ok(AgenticLoopResult::NeedApproval { pending });
|
||||
@@ -1146,6 +1204,7 @@ mod tests {
|
||||
max_actions_per_hour: None,
|
||||
max_tool_iterations: 50,
|
||||
auto_approve_tools: false,
|
||||
default_timezone: "UTC".to_string(),
|
||||
},
|
||||
deps,
|
||||
Arc::new(ChannelManager::new()),
|
||||
@@ -1248,6 +1307,7 @@ mod tests {
|
||||
arguments: serde_json::json!({"message": "done"}),
|
||||
},
|
||||
],
|
||||
user_timezone: None,
|
||||
};
|
||||
|
||||
let json = serde_json::to_string(&pending).expect("serialize");
|
||||
@@ -1595,12 +1655,8 @@ mod tests {
|
||||
use crate::testing::StubLlm;
|
||||
|
||||
let stub = Arc::new(StubLlm::failing_non_transient("ctx-bomb"));
|
||||
let safety = Arc::new(SafetyLayer::new(&SafetyConfig {
|
||||
max_output_length: 100_000,
|
||||
injection_check_enabled: false,
|
||||
}));
|
||||
|
||||
let reasoning = Reasoning::new(stub.clone(), safety);
|
||||
let reasoning = Reasoning::new(stub.clone());
|
||||
|
||||
// Build a fat context with lots of history.
|
||||
let messages = vec![
|
||||
@@ -1710,11 +1766,7 @@ mod tests {
|
||||
use crate::llm::{Reasoning, ReasoningContext, RespondResult, ToolDefinition};
|
||||
|
||||
let provider = Arc::new(AlwaysToolCallProvider);
|
||||
let safety = Arc::new(SafetyLayer::new(&SafetyConfig {
|
||||
max_output_length: 100_000,
|
||||
injection_check_enabled: false,
|
||||
}));
|
||||
let reasoning = Reasoning::new(provider, safety);
|
||||
let reasoning = Reasoning::new(provider);
|
||||
|
||||
let tool_def = ToolDefinition {
|
||||
name: "echo".to_string(),
|
||||
@@ -1900,6 +1952,7 @@ mod tests {
|
||||
max_actions_per_hour: None,
|
||||
max_tool_iterations,
|
||||
auto_approve_tools: true,
|
||||
default_timezone: "UTC".to_string(),
|
||||
},
|
||||
deps,
|
||||
Arc::new(ChannelManager::new()),
|
||||
@@ -2015,6 +2068,7 @@ mod tests {
|
||||
max_actions_per_hour: None,
|
||||
max_tool_iterations: max_iter,
|
||||
auto_approve_tools: true,
|
||||
default_timezone: "UTC".to_string(),
|
||||
},
|
||||
deps,
|
||||
Arc::new(ChannelManager::new()),
|
||||
@@ -2116,4 +2170,47 @@ mod tests {
|
||||
"Error should include the underlying reason, got: {formatted}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_image_sentinel_empty_data_url_should_be_skipped() {
|
||||
// Regression: unwrap_or_default() on missing "data" field produces an empty
|
||||
// string. Broadcasting an empty data_url would send a broken SSE event.
|
||||
let sentinel = serde_json::json!({
|
||||
"type": "image_generated",
|
||||
"path": "/tmp/image.png"
|
||||
// "data" field is missing
|
||||
});
|
||||
|
||||
let data_url = sentinel
|
||||
.get("data")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
|
||||
assert!(
|
||||
data_url.is_empty(),
|
||||
"Missing 'data' field should produce empty string"
|
||||
);
|
||||
// The fix: empty data_url means we skip broadcasting
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_image_sentinel_present_data_url_is_valid() {
|
||||
let sentinel = serde_json::json!({
|
||||
"type": "image_generated",
|
||||
"data": "data:image/png;base64,abc123",
|
||||
"path": "/tmp/image.png"
|
||||
});
|
||||
|
||||
let data_url = sentinel
|
||||
.get("data")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
|
||||
assert!(
|
||||
!data_url.is_empty(),
|
||||
"Present 'data' field should produce non-empty string"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
+114
-8
@@ -31,7 +31,6 @@ use tokio::sync::mpsc;
|
||||
use crate::channels::OutgoingResponse;
|
||||
use crate::db::Database;
|
||||
use crate::llm::{ChatMessage, CompletionRequest, LlmProvider, Reasoning};
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::workspace::Workspace;
|
||||
use crate::workspace::hygiene::HygieneConfig;
|
||||
|
||||
@@ -48,6 +47,12 @@ pub struct HeartbeatConfig {
|
||||
pub notify_user_id: Option<String>,
|
||||
/// Channel to notify on heartbeat findings.
|
||||
pub notify_channel: Option<String>,
|
||||
/// Hour (0-23) when quiet hours start.
|
||||
pub quiet_hours_start: Option<u32>,
|
||||
/// Hour (0-23) when quiet hours end.
|
||||
pub quiet_hours_end: Option<u32>,
|
||||
/// Timezone for quiet hours evaluation (IANA name).
|
||||
pub timezone: Option<String>,
|
||||
}
|
||||
|
||||
impl Default for HeartbeatConfig {
|
||||
@@ -58,6 +63,9 @@ impl Default for HeartbeatConfig {
|
||||
max_failures: 3,
|
||||
notify_user_id: None,
|
||||
notify_channel: None,
|
||||
quiet_hours_start: None,
|
||||
quiet_hours_end: None,
|
||||
timezone: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -75,6 +83,26 @@ impl HeartbeatConfig {
|
||||
self
|
||||
}
|
||||
|
||||
/// Check whether the current time falls within configured quiet hours.
|
||||
pub fn is_quiet_hours(&self) -> bool {
|
||||
use chrono::Timelike;
|
||||
let (Some(start), Some(end)) = (self.quiet_hours_start, self.quiet_hours_end) else {
|
||||
return false;
|
||||
};
|
||||
let tz = self
|
||||
.timezone
|
||||
.as_deref()
|
||||
.and_then(crate::timezone::parse_timezone)
|
||||
.unwrap_or(chrono_tz::UTC);
|
||||
let now_hour = crate::timezone::now_in_tz(tz).hour();
|
||||
if start <= end {
|
||||
now_hour >= start && now_hour < end
|
||||
} else {
|
||||
// Wraps midnight, e.g. 22..06
|
||||
now_hour >= start || now_hour < end
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the notification target.
|
||||
pub fn with_notify(mut self, user_id: impl Into<String>, channel: impl Into<String>) -> Self {
|
||||
self.notify_user_id = Some(user_id.into());
|
||||
@@ -102,7 +130,6 @@ pub struct HeartbeatRunner {
|
||||
hygiene_config: HygieneConfig,
|
||||
workspace: Arc<Workspace>,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
|
||||
store: Option<Arc<dyn Database>>,
|
||||
consecutive_failures: u32,
|
||||
@@ -115,14 +142,12 @@ impl HeartbeatRunner {
|
||||
hygiene_config: HygieneConfig,
|
||||
workspace: Arc<Workspace>,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config,
|
||||
hygiene_config,
|
||||
workspace,
|
||||
llm,
|
||||
safety,
|
||||
response_tx: None,
|
||||
store: None,
|
||||
consecutive_failures: 0,
|
||||
@@ -162,6 +187,12 @@ impl HeartbeatRunner {
|
||||
loop {
|
||||
interval.tick().await;
|
||||
|
||||
// Skip during quiet hours
|
||||
if self.config.is_quiet_hours() {
|
||||
tracing::debug!("Heartbeat skipped: quiet hours");
|
||||
continue;
|
||||
}
|
||||
|
||||
// Run memory hygiene in the background so it never delays the
|
||||
// heartbeat checklist. Failures are logged inside run_if_due.
|
||||
let hygiene_workspace = Arc::clone(&self.workspace);
|
||||
@@ -272,7 +303,7 @@ impl HeartbeatRunner {
|
||||
.with_max_tokens(max_tokens)
|
||||
.with_temperature(0.3);
|
||||
|
||||
let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone());
|
||||
let reasoning = Reasoning::new(self.llm.clone());
|
||||
let (content, _usage) = match reasoning.complete(request).await {
|
||||
Ok(r) => r,
|
||||
Err(e) => return HeartbeatResult::Failed(format!("LLM call failed: {}", e)),
|
||||
@@ -386,11 +417,10 @@ pub fn spawn_heartbeat(
|
||||
hygiene_config: HygieneConfig,
|
||||
workspace: Arc<Workspace>,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
|
||||
store: Option<Arc<dyn Database>>,
|
||||
) -> tokio::task::JoinHandle<()> {
|
||||
let mut runner = HeartbeatRunner::new(config, hygiene_config, workspace, llm, safety);
|
||||
let mut runner = HeartbeatRunner::new(config, hygiene_config, workspace, llm);
|
||||
if let Some(tx) = response_tx {
|
||||
runner = runner.with_response_channel(tx);
|
||||
}
|
||||
@@ -532,6 +562,83 @@ mod tests {
|
||||
assert!(!is_effectively_empty(content));
|
||||
}
|
||||
|
||||
// ==================== quiet hours ====================
|
||||
|
||||
#[test]
|
||||
fn test_quiet_hours_inside() {
|
||||
use chrono::{Timelike, Utc};
|
||||
|
||||
let now_utc = Utc::now();
|
||||
let hour = now_utc.hour();
|
||||
let start = hour;
|
||||
let end = (hour + 1) % 24;
|
||||
|
||||
let config = HeartbeatConfig {
|
||||
quiet_hours_start: Some(start),
|
||||
quiet_hours_end: Some(end),
|
||||
timezone: Some("UTC".to_string()),
|
||||
..HeartbeatConfig::default()
|
||||
};
|
||||
// Current UTC hour is inside [start, end) by construction
|
||||
assert!(config.is_quiet_hours());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_quiet_hours_outside() {
|
||||
use chrono::{Timelike, Utc};
|
||||
|
||||
let now_utc = Utc::now();
|
||||
let hour = now_utc.hour();
|
||||
let start = (hour + 1) % 24;
|
||||
let end = (hour + 2) % 24;
|
||||
|
||||
let config = HeartbeatConfig {
|
||||
quiet_hours_start: Some(start),
|
||||
quiet_hours_end: Some(end),
|
||||
timezone: Some("UTC".to_string()),
|
||||
..HeartbeatConfig::default()
|
||||
};
|
||||
// Current UTC hour is outside [start, end) by construction
|
||||
assert!(!config.is_quiet_hours());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_quiet_hours_wraparound_excludes_now() {
|
||||
use chrono::{Timelike, Utc};
|
||||
|
||||
let now_utc = Utc::now();
|
||||
let hour = now_utc.hour();
|
||||
// Window covers all hours except the current one
|
||||
let start = (hour + 1) % 24;
|
||||
let end = hour;
|
||||
|
||||
let config = HeartbeatConfig {
|
||||
quiet_hours_start: Some(start),
|
||||
quiet_hours_end: Some(end),
|
||||
timezone: Some("UTC".to_string()),
|
||||
..HeartbeatConfig::default()
|
||||
};
|
||||
assert!(!config.is_quiet_hours());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_quiet_hours_none_configured() {
|
||||
let config = HeartbeatConfig::default();
|
||||
assert!(!config.is_quiet_hours());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_quiet_hours_same_start_end() {
|
||||
let config = HeartbeatConfig {
|
||||
quiet_hours_start: Some(10),
|
||||
quiet_hours_end: Some(10),
|
||||
timezone: Some("UTC".to_string()),
|
||||
..HeartbeatConfig::default()
|
||||
};
|
||||
// start == end means zero-width window, should be false
|
||||
assert!(!config.is_quiet_hours());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_spawn_heartbeat_accepts_store_param() {
|
||||
// Regression: spawn_heartbeat must accept an optional Database store
|
||||
@@ -543,7 +650,6 @@ mod tests {
|
||||
HygieneConfig,
|
||||
Arc<crate::workspace::Workspace>,
|
||||
Arc<dyn crate::llm::LlmProvider>,
|
||||
Arc<crate::safety::SafetyLayer>,
|
||||
Option<tokio::sync::mpsc::Sender<crate::channels::OutgoingResponse>>,
|
||||
Option<Arc<dyn crate::db::Database>>,
|
||||
) -> tokio::task::JoinHandle<()> = spawn_heartbeat;
|
||||
|
||||
+87
-9
@@ -57,7 +57,11 @@ pub struct Routine {
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum Trigger {
|
||||
/// Fire on a cron schedule (e.g. "0 9 * * MON-FRI" or "every 2h").
|
||||
Cron { schedule: String },
|
||||
Cron {
|
||||
schedule: String,
|
||||
#[serde(default)]
|
||||
timezone: Option<String>,
|
||||
},
|
||||
/// Fire when a channel message matches a pattern.
|
||||
Event {
|
||||
/// Optional channel filter (e.g. "telegram", "slack").
|
||||
@@ -99,7 +103,21 @@ impl Trigger {
|
||||
field: "schedule".into(),
|
||||
})?
|
||||
.to_string();
|
||||
Ok(Trigger::Cron { schedule })
|
||||
let timezone = config
|
||||
.get("timezone")
|
||||
.and_then(|v| v.as_str())
|
||||
.and_then(|tz| {
|
||||
if crate::timezone::parse_timezone(tz).is_some() {
|
||||
Some(tz.to_string())
|
||||
} else {
|
||||
tracing::warn!(
|
||||
"Ignoring invalid timezone '{}' from DB for cron trigger",
|
||||
tz
|
||||
);
|
||||
None
|
||||
}
|
||||
});
|
||||
Ok(Trigger::Cron { schedule, timezone })
|
||||
}
|
||||
"event" => {
|
||||
let pattern = config
|
||||
@@ -137,7 +155,10 @@ impl Trigger {
|
||||
/// Serialize trigger-specific config to JSON for DB storage.
|
||||
pub fn to_config_json(&self) -> serde_json::Value {
|
||||
match self {
|
||||
Trigger::Cron { schedule } => serde_json::json!({ "schedule": schedule }),
|
||||
Trigger::Cron { schedule, timezone } => serde_json::json!({
|
||||
"schedule": schedule,
|
||||
"timezone": timezone,
|
||||
}),
|
||||
Trigger::Event { channel, pattern } => serde_json::json!({
|
||||
"pattern": pattern,
|
||||
"channel": channel,
|
||||
@@ -415,12 +436,25 @@ pub fn content_hash(content: &str) -> u64 {
|
||||
}
|
||||
|
||||
/// Parse a cron expression and compute the next fire time from now.
|
||||
pub fn next_cron_fire(schedule: &str) -> Result<Option<DateTime<Utc>>, RoutineError> {
|
||||
///
|
||||
/// When `timezone` is provided and valid, the schedule is evaluated in that
|
||||
/// timezone and the result is converted back to UTC. Otherwise UTC is used.
|
||||
pub fn next_cron_fire(
|
||||
schedule: &str,
|
||||
timezone: Option<&str>,
|
||||
) -> Result<Option<DateTime<Utc>>, RoutineError> {
|
||||
let cron_schedule =
|
||||
cron::Schedule::from_str(schedule).map_err(|e| RoutineError::InvalidCron {
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
Ok(cron_schedule.upcoming(Utc).next())
|
||||
if let Some(tz) = timezone.and_then(crate::timezone::parse_timezone) {
|
||||
Ok(cron_schedule
|
||||
.upcoming(tz)
|
||||
.next()
|
||||
.map(|dt| dt.with_timezone(&Utc)))
|
||||
} else {
|
||||
Ok(cron_schedule.upcoming(Utc).next())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -433,10 +467,11 @@ mod tests {
|
||||
fn test_trigger_roundtrip() {
|
||||
let trigger = Trigger::Cron {
|
||||
schedule: "0 9 * * MON-FRI".to_string(),
|
||||
timezone: None,
|
||||
};
|
||||
let json = trigger.to_config_json();
|
||||
let parsed = Trigger::from_db("cron", json).expect("parse cron");
|
||||
assert!(matches!(parsed, Trigger::Cron { schedule } if schedule == "0 9 * * MON-FRI"));
|
||||
assert!(matches!(parsed, Trigger::Cron { schedule, .. } if schedule == "0 9 * * MON-FRI"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -509,16 +544,58 @@ mod tests {
|
||||
#[test]
|
||||
fn test_next_cron_fire_valid() {
|
||||
// Every minute should always have a next fire
|
||||
let next = next_cron_fire("* * * * * *").expect("valid cron");
|
||||
let next = next_cron_fire("* * * * * *", None).expect("valid cron");
|
||||
assert!(next.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_next_cron_fire_invalid() {
|
||||
let result = next_cron_fire("not a cron");
|
||||
let result = next_cron_fire("not a cron", None);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_trigger_cron_timezone_roundtrip() {
|
||||
let trigger = Trigger::Cron {
|
||||
schedule: "0 9 * * MON-FRI".to_string(),
|
||||
timezone: Some("America/New_York".to_string()),
|
||||
};
|
||||
let json = trigger.to_config_json();
|
||||
let parsed = Trigger::from_db("cron", json).expect("parse cron");
|
||||
assert!(matches!(parsed, Trigger::Cron { schedule, timezone }
|
||||
if schedule == "0 9 * * MON-FRI"
|
||||
&& timezone.as_deref() == Some("America/New_York")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_trigger_cron_no_timezone_backward_compat() {
|
||||
let json = serde_json::json!({"schedule": "0 9 * * *"});
|
||||
let parsed = Trigger::from_db("cron", json).expect("parse cron");
|
||||
assert!(matches!(parsed, Trigger::Cron { timezone, .. } if timezone.is_none()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_trigger_cron_invalid_timezone_coerced_to_none() {
|
||||
let json = serde_json::json!({"schedule": "0 9 * * *", "timezone": "Fake/Zone"});
|
||||
let parsed = Trigger::from_db("cron", json).expect("parse cron");
|
||||
assert!(
|
||||
matches!(parsed, Trigger::Cron { timezone, .. } if timezone.is_none()),
|
||||
"invalid timezone should be coerced to None"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_next_cron_fire_with_timezone() {
|
||||
let next_utc = next_cron_fire("0 0 9 * * * *", None)
|
||||
.expect("valid cron")
|
||||
.expect("has next");
|
||||
let next_est = next_cron_fire("0 0 9 * * * *", Some("America/New_York"))
|
||||
.expect("valid cron")
|
||||
.expect("has next");
|
||||
// EST is UTC-5 (or EDT UTC-4), so the UTC result should differ
|
||||
assert_ne!(next_utc, next_est, "timezone should shift the fire time");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_guardrails_default() {
|
||||
let g = RoutineGuardrails::default();
|
||||
@@ -531,7 +608,8 @@ mod tests {
|
||||
fn test_trigger_type_tag() {
|
||||
assert_eq!(
|
||||
Trigger::Cron {
|
||||
schedule: String::new()
|
||||
schedule: String::new(),
|
||||
timezone: None,
|
||||
}
|
||||
.type_tag(),
|
||||
"cron"
|
||||
|
||||
+12
-13
@@ -170,7 +170,7 @@ impl RoutineEngine {
|
||||
continue;
|
||||
}
|
||||
|
||||
let detail = if let Trigger::Cron { ref schedule } = routine.trigger {
|
||||
let detail = if let Trigger::Cron { ref schedule, .. } = routine.trigger {
|
||||
Some(schedule.clone())
|
||||
} else {
|
||||
None
|
||||
@@ -380,8 +380,12 @@ async fn execute_routine(ctx: EngineContext, routine: Routine, run: RoutineRun)
|
||||
|
||||
// Update routine runtime state
|
||||
let now = Utc::now();
|
||||
let next_fire = if let Trigger::Cron { ref schedule } = routine.trigger {
|
||||
next_cron_fire(schedule).unwrap_or(None)
|
||||
let next_fire = if let Trigger::Cron {
|
||||
ref schedule,
|
||||
ref timezone,
|
||||
} = routine.trigger
|
||||
{
|
||||
next_cron_fire(schedule, timezone.as_deref()).unwrap_or(None)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
@@ -488,18 +492,13 @@ async fn execute_full_job(
|
||||
reason: "scheduler not available".to_string(),
|
||||
})?;
|
||||
|
||||
// Set the message tool's default channel/target from the routine's notify config
|
||||
// so the LLM can send results without triggering cross-channel approval.
|
||||
// TODO: This mutates shared global state and can race with concurrent jobs.
|
||||
// Move notify config into JobContext metadata and apply per-job instead.
|
||||
let mut metadata = serde_json::json!({ "max_iterations": max_iterations });
|
||||
// Carry the routine's notify config in job metadata so the message tool
|
||||
// can resolve channel/target per-job without global state mutation.
|
||||
if let Some(channel) = &routine.notify.channel {
|
||||
scheduler
|
||||
.tools()
|
||||
.set_message_tool_context(Some(channel.clone()), Some(routine.notify.user.clone()))
|
||||
.await;
|
||||
metadata["notify_channel"] = serde_json::json!(channel);
|
||||
}
|
||||
|
||||
let metadata = serde_json::json!({ "max_iterations": max_iterations });
|
||||
metadata["notify_user"] = serde_json::json!(&routine.notify.user);
|
||||
|
||||
// Build approval context: UnlessAutoApproved tools are auto-approved for routines;
|
||||
// Always tools require explicit listing in tool_permissions.
|
||||
|
||||
+324
-11
@@ -16,6 +16,7 @@ use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::channels::web::util::truncate_preview;
|
||||
use crate::llm::{ChatMessage, ToolCall};
|
||||
|
||||
/// A session containing one or more threads.
|
||||
@@ -164,6 +165,10 @@ pub struct PendingApproval {
|
||||
/// executed yet when approval was requested.
|
||||
#[serde(default)]
|
||||
pub deferred_tool_calls: Vec<ToolCall>,
|
||||
/// User timezone at the time the approval was requested, so it persists
|
||||
/// through the approval flow even if the approval message lacks timezone.
|
||||
#[serde(default)]
|
||||
pub user_timezone: Option<String>,
|
||||
}
|
||||
|
||||
/// A conversation thread within a session.
|
||||
@@ -316,7 +321,13 @@ impl Thread {
|
||||
}
|
||||
}
|
||||
|
||||
/// Get all messages for context building.
|
||||
/// Get all messages for context building, including tool call history.
|
||||
///
|
||||
/// Emits the full LLM-compatible message sequence per turn:
|
||||
/// `user → [assistant_with_tool_calls → tool_result*] → assistant`
|
||||
///
|
||||
/// This ensures the LLM sees prior tool executions and won't re-attempt
|
||||
/// completed actions in subsequent turns.
|
||||
pub fn messages(&self) -> Vec<ChatMessage> {
|
||||
let mut messages = Vec::new();
|
||||
for turn in &self.turns {
|
||||
@@ -328,6 +339,42 @@ impl Thread {
|
||||
turn.image_content_parts.clone(),
|
||||
));
|
||||
}
|
||||
|
||||
if !turn.tool_calls.is_empty() {
|
||||
// Build ToolCall objects with synthetic stable IDs
|
||||
let tool_calls: Vec<ToolCall> = turn
|
||||
.tool_calls
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, tc)| ToolCall {
|
||||
id: format!("turn{}_{}", turn.turn_number, i),
|
||||
name: tc.name.clone(),
|
||||
arguments: tc.parameters.clone(),
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Assistant message declaring the tool calls (no text content)
|
||||
messages.push(ChatMessage::assistant_with_tool_calls(None, tool_calls));
|
||||
|
||||
// Individual tool result messages, truncated to limit context size.
|
||||
for (i, tc) in turn.tool_calls.iter().enumerate() {
|
||||
let call_id = format!("turn{}_{}", turn.turn_number, i);
|
||||
let content = if let Some(ref err) = tc.error {
|
||||
// .error already contains the full error text;
|
||||
// pass through without wrapping to avoid double-prefix.
|
||||
truncate_preview(err, 1000)
|
||||
} else if let Some(ref res) = tc.result {
|
||||
let raw = match res {
|
||||
serde_json::Value::String(s) => s.clone(),
|
||||
other => other.to_string(),
|
||||
};
|
||||
truncate_preview(&raw, 1000)
|
||||
} else {
|
||||
"OK".to_string()
|
||||
};
|
||||
messages.push(ChatMessage::tool_result(call_id, &tc.name, content));
|
||||
}
|
||||
}
|
||||
if let Some(ref response) = turn.response {
|
||||
messages.push(ChatMessage::assistant(response));
|
||||
}
|
||||
@@ -349,13 +396,16 @@ impl Thread {
|
||||
|
||||
/// Restore thread state from a checkpoint's messages.
|
||||
///
|
||||
/// Clears existing turns and rebuilds from message pairs.
|
||||
/// Messages should alternate: user, assistant, user, assistant...
|
||||
/// Clears existing turns and rebuilds from the message sequence.
|
||||
/// Handles the full message pattern including tool messages:
|
||||
/// `user → [assistant_with_tool_calls → tool_result*] → assistant`
|
||||
///
|
||||
/// Also supports the legacy pattern (user/assistant pairs only) for
|
||||
/// backward compatibility with old checkpoint data.
|
||||
pub fn restore_from_messages(&mut self, messages: Vec<ChatMessage>) {
|
||||
self.turns.clear();
|
||||
self.state = ThreadState::Idle;
|
||||
|
||||
// Messages alternate: user, assistant, user, assistant...
|
||||
let mut iter = messages.into_iter().peekable();
|
||||
let mut turn_number = 0;
|
||||
|
||||
@@ -363,18 +413,58 @@ impl Thread {
|
||||
if msg.role == crate::llm::Role::User {
|
||||
let mut turn = Turn::new(turn_number, &msg.content);
|
||||
|
||||
// Check if next is assistant response
|
||||
if let Some(next) = iter.peek()
|
||||
&& next.role == crate::llm::Role::Assistant
|
||||
{
|
||||
// iter.next() is guaranteed Some after a successful peek()
|
||||
if let Some(response) = iter.next() {
|
||||
turn.complete(&response.content);
|
||||
// Consume tool call sequences (assistant_with_tool_calls + tool_results).
|
||||
// A single turn may contain multiple rounds of tool calls, so we
|
||||
// track the cumulative base index into turn.tool_calls.
|
||||
while let Some(next) = iter.peek() {
|
||||
if next.role == crate::llm::Role::Assistant && next.tool_calls.is_some() {
|
||||
let call_base_idx = turn.tool_calls.len();
|
||||
|
||||
if let Some(assistant_msg) = iter.next()
|
||||
&& let Some(ref tcs) = assistant_msg.tool_calls
|
||||
{
|
||||
for tc in tcs {
|
||||
turn.record_tool_call(&tc.name, tc.arguments.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// Consume the corresponding tool_result messages,
|
||||
// indexing relative to this batch's base offset.
|
||||
let mut pos = 0;
|
||||
while let Some(tr) = iter.peek() {
|
||||
if tr.role != crate::llm::Role::Tool {
|
||||
break;
|
||||
}
|
||||
if let Some(tool_msg) = iter.next() {
|
||||
let idx = call_base_idx + pos;
|
||||
if idx < turn.tool_calls.len() {
|
||||
// Store as result — the error/success distinction
|
||||
// is for the live turn only; restored context just
|
||||
// needs the content the LLM originally saw.
|
||||
turn.tool_calls[idx].result =
|
||||
Some(serde_json::Value::String(tool_msg.content.clone()));
|
||||
}
|
||||
}
|
||||
pos += 1;
|
||||
}
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Check if next is the final assistant response for this turn
|
||||
let is_final_assistant = iter.peek().is_some_and(|n| {
|
||||
n.role == crate::llm::Role::Assistant && n.tool_calls.is_none()
|
||||
});
|
||||
if is_final_assistant && let Some(response) = iter.next() {
|
||||
turn.complete(&response.content);
|
||||
}
|
||||
|
||||
self.turns.push(turn);
|
||||
turn_number += 1;
|
||||
} else {
|
||||
// Skip non-user messages that aren't anchored to a turn
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -976,6 +1066,7 @@ mod tests {
|
||||
tool_call_id: "call_123".to_string(),
|
||||
context_messages: vec![ChatMessage::user("do it")],
|
||||
deferred_tool_calls: vec![],
|
||||
user_timezone: None,
|
||||
};
|
||||
|
||||
thread.await_approval(approval);
|
||||
@@ -1001,6 +1092,7 @@ mod tests {
|
||||
tool_call_id: "call_456".to_string(),
|
||||
context_messages: vec![],
|
||||
deferred_tool_calls: vec![],
|
||||
user_timezone: None,
|
||||
};
|
||||
|
||||
thread.await_approval(approval);
|
||||
@@ -1029,4 +1121,225 @@ mod tests {
|
||||
ThreadState::Processing
|
||||
);
|
||||
}
|
||||
|
||||
// Regression tests for #568: tool call history must survive hydration.
|
||||
|
||||
#[test]
|
||||
fn test_messages_includes_tool_calls() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
thread.start_turn("Search for X");
|
||||
{
|
||||
let turn = thread.turns.last_mut().unwrap();
|
||||
turn.record_tool_call("memory_search", serde_json::json!({"query": "X"}));
|
||||
turn.record_tool_result(serde_json::json!("Found X in doc.md"));
|
||||
}
|
||||
thread.complete_turn("I found X in doc.md.");
|
||||
|
||||
let messages = thread.messages();
|
||||
// user + assistant_with_tool_calls + tool_result + assistant = 4
|
||||
assert_eq!(messages.len(), 4);
|
||||
|
||||
assert_eq!(messages[0].role, crate::llm::Role::User);
|
||||
assert_eq!(messages[0].content, "Search for X");
|
||||
|
||||
assert_eq!(messages[1].role, crate::llm::Role::Assistant);
|
||||
assert!(messages[1].tool_calls.is_some());
|
||||
let tcs = messages[1].tool_calls.as_ref().unwrap();
|
||||
assert_eq!(tcs.len(), 1);
|
||||
assert_eq!(tcs[0].name, "memory_search");
|
||||
|
||||
assert_eq!(messages[2].role, crate::llm::Role::Tool);
|
||||
assert!(messages[2].content.contains("Found X"));
|
||||
|
||||
assert_eq!(messages[3].role, crate::llm::Role::Assistant);
|
||||
assert_eq!(messages[3].content, "I found X in doc.md.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_messages_multiple_tool_calls_per_turn() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
thread.start_turn("Do two things");
|
||||
{
|
||||
let turn = thread.turns.last_mut().unwrap();
|
||||
turn.record_tool_call("echo", serde_json::json!({"msg": "a"}));
|
||||
turn.record_tool_result(serde_json::json!("a"));
|
||||
turn.record_tool_call("time", serde_json::json!({}));
|
||||
turn.record_tool_error("timeout");
|
||||
}
|
||||
thread.complete_turn("Done.");
|
||||
|
||||
let messages = thread.messages();
|
||||
// user + assistant_with_calls(2) + tool_result + tool_result + assistant = 5
|
||||
assert_eq!(messages.len(), 5);
|
||||
|
||||
let tcs = messages[1].tool_calls.as_ref().unwrap();
|
||||
assert_eq!(tcs.len(), 2);
|
||||
|
||||
// First tool: success
|
||||
assert_eq!(messages[2].content, "a");
|
||||
// Second tool: error (passed through directly, no wrapping)
|
||||
assert!(messages[3].content.contains("timeout"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_restore_from_messages_with_tool_calls() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
// Build a message sequence with tool calls
|
||||
let tc = ToolCall {
|
||||
id: "call_0".to_string(),
|
||||
name: "search".to_string(),
|
||||
arguments: serde_json::json!({"q": "test"}),
|
||||
};
|
||||
let messages = vec![
|
||||
ChatMessage::user("Find test"),
|
||||
ChatMessage::assistant_with_tool_calls(None, vec![tc]),
|
||||
ChatMessage::tool_result("call_0", "search", "result: found"),
|
||||
ChatMessage::assistant("Found it."),
|
||||
];
|
||||
|
||||
thread.restore_from_messages(messages);
|
||||
|
||||
assert_eq!(thread.turns.len(), 1);
|
||||
let turn = &thread.turns[0];
|
||||
assert_eq!(turn.user_input, "Find test");
|
||||
assert_eq!(turn.tool_calls.len(), 1);
|
||||
assert_eq!(turn.tool_calls[0].name, "search");
|
||||
assert_eq!(
|
||||
turn.tool_calls[0].result,
|
||||
Some(serde_json::Value::String("result: found".to_string()))
|
||||
);
|
||||
assert_eq!(turn.response, Some("Found it.".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_restore_from_messages_with_tool_error() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
let tc = ToolCall {
|
||||
id: "call_0".to_string(),
|
||||
name: "http".to_string(),
|
||||
arguments: serde_json::json!({}),
|
||||
};
|
||||
let messages = vec![
|
||||
ChatMessage::user("Fetch URL"),
|
||||
ChatMessage::assistant_with_tool_calls(None, vec![tc]),
|
||||
ChatMessage::tool_result("call_0", "http", "Error: timeout"),
|
||||
ChatMessage::assistant("The request timed out."),
|
||||
];
|
||||
|
||||
thread.restore_from_messages(messages);
|
||||
|
||||
// restore_from_messages stores all tool content as result (not error),
|
||||
// because it can't reliably distinguish errors from results that happen
|
||||
// to start with "Error: ". The content is preserved for LLM context.
|
||||
let turn = &thread.turns[0];
|
||||
assert_eq!(
|
||||
turn.tool_calls[0].result,
|
||||
Some(serde_json::Value::String("Error: timeout".to_string()))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_messages_round_trip_with_tools() {
|
||||
// Build a thread with tool calls, get messages(), restore, get messages() again
|
||||
// The two message sequences should be equivalent.
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
thread.start_turn("Do search");
|
||||
{
|
||||
let turn = thread.turns.last_mut().unwrap();
|
||||
turn.record_tool_call("search", serde_json::json!({"q": "test"}));
|
||||
turn.record_tool_result(serde_json::json!("found"));
|
||||
}
|
||||
thread.complete_turn("Here are results.");
|
||||
|
||||
let messages_original = thread.messages();
|
||||
|
||||
// Restore into a new thread
|
||||
let mut thread2 = Thread::new(Uuid::new_v4());
|
||||
thread2.restore_from_messages(messages_original.clone());
|
||||
|
||||
let messages_restored = thread2.messages();
|
||||
|
||||
// Same number of messages
|
||||
assert_eq!(messages_original.len(), messages_restored.len());
|
||||
|
||||
// Same roles
|
||||
for (orig, rest) in messages_original.iter().zip(messages_restored.iter()) {
|
||||
assert_eq!(orig.role, rest.role);
|
||||
}
|
||||
|
||||
// Same final response
|
||||
assert_eq!(
|
||||
messages_original.last().unwrap().content,
|
||||
messages_restored.last().unwrap().content
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_restore_multi_stage_tool_calls() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
let tc1 = ToolCall {
|
||||
id: "call_a".to_string(),
|
||||
name: "search".to_string(),
|
||||
arguments: serde_json::json!({"q": "data"}),
|
||||
};
|
||||
let tc2 = ToolCall {
|
||||
id: "call_b".to_string(),
|
||||
name: "write".to_string(),
|
||||
arguments: serde_json::json!({"path": "out.txt"}),
|
||||
};
|
||||
let messages = vec![
|
||||
ChatMessage::user("Find and save"),
|
||||
ChatMessage::assistant_with_tool_calls(None, vec![tc1]),
|
||||
ChatMessage::tool_result("call_a", "search", "found data"),
|
||||
ChatMessage::assistant_with_tool_calls(None, vec![tc2]),
|
||||
ChatMessage::tool_result("call_b", "write", "written"),
|
||||
ChatMessage::assistant("Done, saved to out.txt"),
|
||||
];
|
||||
|
||||
thread.restore_from_messages(messages);
|
||||
|
||||
assert_eq!(thread.turns.len(), 1);
|
||||
let turn = &thread.turns[0];
|
||||
assert_eq!(turn.tool_calls.len(), 2);
|
||||
assert_eq!(turn.tool_calls[0].name, "search");
|
||||
assert_eq!(turn.tool_calls[1].name, "write");
|
||||
assert_eq!(
|
||||
turn.tool_calls[0].result,
|
||||
Some(serde_json::Value::String("found data".to_string()))
|
||||
);
|
||||
assert_eq!(
|
||||
turn.tool_calls[1].result,
|
||||
Some(serde_json::Value::String("written".to_string()))
|
||||
);
|
||||
assert_eq!(turn.response, Some("Done, saved to out.txt".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_messages_truncates_large_tool_results() {
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
|
||||
thread.start_turn("Read big file");
|
||||
{
|
||||
let turn = thread.turns.last_mut().unwrap();
|
||||
turn.record_tool_call("read_file", serde_json::json!({"path": "big.txt"}));
|
||||
let big_result = "x".repeat(2000);
|
||||
turn.record_tool_result(serde_json::json!(big_result));
|
||||
}
|
||||
thread.complete_turn("Here's the file content.");
|
||||
|
||||
let messages = thread.messages();
|
||||
let tool_result_content = &messages[2].content;
|
||||
assert!(
|
||||
tool_result_content.len() <= 1010,
|
||||
"Tool result should be truncated, got {} chars",
|
||||
tool_result_content.len()
|
||||
);
|
||||
assert!(tool_result_content.ends_with("..."));
|
||||
}
|
||||
}
|
||||
|
||||
+305
-60
@@ -20,7 +20,7 @@ use crate::channels::web::util::truncate_preview;
|
||||
use crate::channels::{IncomingMessage, StatusUpdate};
|
||||
use crate::context::JobContext;
|
||||
use crate::error::Error;
|
||||
use crate::llm::ChatMessage;
|
||||
use crate::llm::{ChatMessage, ToolCall};
|
||||
use crate::tools::redact_params;
|
||||
|
||||
impl Agent {
|
||||
@@ -66,16 +66,7 @@ impl Agent {
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
msg_count = db_messages.len();
|
||||
chat_messages = db_messages
|
||||
.iter()
|
||||
.filter_map(|m| match m.role.as_str() {
|
||||
"user" => Some(ChatMessage::user(&m.content)),
|
||||
"assistant" => Some(ChatMessage::assistant(&m.content)),
|
||||
// tool_calls rows are UI metadata (tool name + preview),
|
||||
// not part of the LLM conversation context.
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
chat_messages = rebuild_chat_messages_from_db(&db_messages);
|
||||
} else {
|
||||
msg_count = 0;
|
||||
}
|
||||
@@ -230,7 +221,7 @@ impl Agent {
|
||||
)
|
||||
.await;
|
||||
|
||||
let compactor = ContextCompactor::new(self.llm().clone(), self.safety().clone());
|
||||
let compactor = ContextCompactor::new(self.llm().clone());
|
||||
if let Err(e) = compactor
|
||||
.compact(thread, strategy, self.workspace().map(|w| w.as_ref()))
|
||||
.await
|
||||
@@ -340,10 +331,10 @@ impl Agent {
|
||||
};
|
||||
|
||||
thread.complete_turn(&response);
|
||||
let tool_calls = thread
|
||||
let (turn_number, tool_calls) = thread
|
||||
.turns
|
||||
.last()
|
||||
.map(|t| t.tool_calls.clone())
|
||||
.map(|t| (t.turn_number, t.tool_calls.clone()))
|
||||
.unwrap_or_default();
|
||||
let _ = self
|
||||
.channels
|
||||
@@ -355,7 +346,7 @@ impl Agent {
|
||||
.await;
|
||||
|
||||
// Persist tool calls then assistant response (user message already persisted at turn start)
|
||||
self.persist_tool_calls(thread_id, &message.user_id, &tool_calls)
|
||||
self.persist_tool_calls(thread_id, &message.user_id, turn_number, &tool_calls)
|
||||
.await;
|
||||
self.persist_assistant_response(thread_id, &message.user_id, &response)
|
||||
.await;
|
||||
@@ -464,6 +455,7 @@ impl Agent {
|
||||
&self,
|
||||
thread_id: Uuid,
|
||||
user_id: &str,
|
||||
turn_number: usize,
|
||||
tool_calls: &[crate::agent::session::TurnToolCall],
|
||||
) {
|
||||
if tool_calls.is_empty() {
|
||||
@@ -477,14 +469,24 @@ impl Agent {
|
||||
|
||||
let summaries: Vec<serde_json::Value> = tool_calls
|
||||
.iter()
|
||||
.map(|tc| {
|
||||
let mut obj = serde_json::json!({ "name": tc.name });
|
||||
.enumerate()
|
||||
.map(|(i, tc)| {
|
||||
let mut obj = serde_json::json!({
|
||||
"name": tc.name,
|
||||
"call_id": format!("turn{}_{}", turn_number, i),
|
||||
});
|
||||
if let Some(ref result) = tc.result {
|
||||
let preview = match result {
|
||||
serde_json::Value::String(s) => truncate_preview(s, 500),
|
||||
other => truncate_preview(&other.to_string(), 500),
|
||||
};
|
||||
obj["result_preview"] = serde_json::Value::String(preview);
|
||||
// Store full result (truncated to ~1000 chars) for LLM context rebuild
|
||||
let full_result = match result {
|
||||
serde_json::Value::String(s) => truncate_preview(s, 1000),
|
||||
other => truncate_preview(&other.to_string(), 1000),
|
||||
};
|
||||
obj["result"] = serde_json::Value::String(full_result);
|
||||
}
|
||||
if let Some(ref error) = tc.error {
|
||||
obj["error"] = serde_json::Value::String(truncate_preview(error, 200));
|
||||
@@ -627,7 +629,7 @@ impl Agent {
|
||||
crate::agent::context_monitor::CompactionStrategy::Summarize { keep_recent: 5 },
|
||||
);
|
||||
|
||||
let compactor = ContextCompactor::new(self.llm().clone(), self.safety().clone());
|
||||
let compactor = ContextCompactor::new(self.llm().clone());
|
||||
match compactor
|
||||
.compact(thread, strategy, self.workspace().map(|w| w.as_ref()))
|
||||
.await
|
||||
@@ -746,6 +748,16 @@ impl Agent {
|
||||
let mut job_ctx =
|
||||
JobContext::with_user(&message.user_id, "chat", "Interactive chat session");
|
||||
job_ctx.http_interceptor = self.deps.http_interceptor.clone();
|
||||
// Prefer a valid timezone from the approval message, fall back to the
|
||||
// resolved timezone stored when the approval was originally requested.
|
||||
let tz_candidate = message
|
||||
.timezone
|
||||
.as_deref()
|
||||
.filter(|tz| crate::timezone::parse_timezone(tz).is_some())
|
||||
.or(pending.user_timezone.as_deref());
|
||||
if let Some(tz) = tz_candidate {
|
||||
job_ctx.user_timezone = tz.to_string();
|
||||
}
|
||||
|
||||
let _ = self
|
||||
.channels
|
||||
@@ -797,19 +809,33 @@ impl Agent {
|
||||
let mut context_messages = pending.context_messages;
|
||||
let deferred_tool_calls = pending.deferred_tool_calls;
|
||||
|
||||
// Record result in thread
|
||||
// Sanitize tool result, then record the cleaned version in the
|
||||
// thread. Must happen before auth intercept check which may return early.
|
||||
let is_tool_error = tool_result.is_err();
|
||||
let result_content = match &tool_result {
|
||||
Ok(output) => {
|
||||
let sanitized = self
|
||||
.safety()
|
||||
.sanitize_tool_output(&pending.tool_name, output);
|
||||
self.safety().wrap_for_llm(
|
||||
&pending.tool_name,
|
||||
&sanitized.content,
|
||||
sanitized.was_modified,
|
||||
)
|
||||
}
|
||||
Err(e) => format!("Error: {}", e),
|
||||
};
|
||||
|
||||
// Record sanitized result in thread
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id)
|
||||
&& let Some(turn) = thread.last_turn_mut()
|
||||
{
|
||||
match &tool_result {
|
||||
Ok(output) => {
|
||||
turn.record_tool_result(serde_json::json!(output));
|
||||
}
|
||||
Err(e) => {
|
||||
turn.record_tool_error(e.to_string());
|
||||
}
|
||||
if is_tool_error {
|
||||
turn.record_tool_error(result_content.clone());
|
||||
} else {
|
||||
turn.record_tool_result(serde_json::json!(result_content));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -831,21 +857,6 @@ impl Agent {
|
||||
return Ok(SubmissionResult::response(instructions));
|
||||
}
|
||||
|
||||
// Add tool result to context
|
||||
let result_content = match tool_result {
|
||||
Ok(output) => {
|
||||
let sanitized = self
|
||||
.safety()
|
||||
.sanitize_tool_output(&pending.tool_name, &output);
|
||||
self.safety().wrap_for_llm(
|
||||
&pending.tool_name,
|
||||
&sanitized.content,
|
||||
sanitized.was_modified,
|
||||
)
|
||||
}
|
||||
Err(e) => format!("Error: {}", e),
|
||||
};
|
||||
|
||||
context_messages.push(ChatMessage::tool_result(
|
||||
&pending.tool_call_id,
|
||||
&pending.tool_name,
|
||||
@@ -1050,15 +1061,31 @@ impl Agent {
|
||||
.await;
|
||||
}
|
||||
|
||||
// Record in thread
|
||||
// Sanitize first, then record the cleaned version in thread.
|
||||
// Must happen before auth detection which may set deferred_auth.
|
||||
let is_deferred_error = deferred_result.is_err();
|
||||
let deferred_content = match &deferred_result {
|
||||
Ok(output) => {
|
||||
let sanitized = self.safety().sanitize_tool_output(&tc.name, output);
|
||||
self.safety().wrap_for_llm(
|
||||
&tc.name,
|
||||
&sanitized.content,
|
||||
sanitized.was_modified,
|
||||
)
|
||||
}
|
||||
Err(e) => format!("Error: {}", e),
|
||||
};
|
||||
|
||||
// Record sanitized result in thread
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id)
|
||||
&& let Some(turn) = thread.last_turn_mut()
|
||||
{
|
||||
match &deferred_result {
|
||||
Ok(output) => turn.record_tool_result(serde_json::json!(output)),
|
||||
Err(e) => turn.record_tool_error(e.to_string()),
|
||||
if is_deferred_error {
|
||||
turn.record_tool_error(deferred_content.clone());
|
||||
} else {
|
||||
turn.record_tool_result(serde_json::json!(deferred_content));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1080,18 +1107,6 @@ impl Agent {
|
||||
deferred_auth = Some(instructions);
|
||||
}
|
||||
|
||||
let deferred_content = match deferred_result {
|
||||
Ok(output) => {
|
||||
let sanitized = self.safety().sanitize_tool_output(&tc.name, &output);
|
||||
self.safety().wrap_for_llm(
|
||||
&tc.name,
|
||||
&sanitized.content,
|
||||
sanitized.was_modified,
|
||||
)
|
||||
}
|
||||
Err(e) => format!("Error: {}", e),
|
||||
};
|
||||
|
||||
context_messages.push(ChatMessage::tool_result(&tc.id, &tc.name, deferred_content));
|
||||
}
|
||||
|
||||
@@ -1111,6 +1126,8 @@ impl Agent {
|
||||
tool_call_id: tc.id.clone(),
|
||||
context_messages: context_messages.clone(),
|
||||
deferred_tool_calls: deferred_tool_calls[approval_idx + 1..].to_vec(),
|
||||
// Carry forward the resolved timezone from the original pending approval
|
||||
user_timezone: pending.user_timezone.clone(),
|
||||
};
|
||||
|
||||
let request_id = new_pending.request_id;
|
||||
@@ -1157,13 +1174,13 @@ impl Agent {
|
||||
match result {
|
||||
Ok(AgenticLoopResult::Response(response)) => {
|
||||
thread.complete_turn(&response);
|
||||
let tool_calls = thread
|
||||
let (turn_number, tool_calls) = thread
|
||||
.turns
|
||||
.last()
|
||||
.map(|t| t.tool_calls.clone())
|
||||
.map(|t| (t.turn_number, t.tool_calls.clone()))
|
||||
.unwrap_or_default();
|
||||
// User message already persisted at turn start; save tool calls then assistant response
|
||||
self.persist_tool_calls(thread_id, &message.user_id, &tool_calls)
|
||||
self.persist_tool_calls(thread_id, &message.user_id, turn_number, &tool_calls)
|
||||
.await;
|
||||
self.persist_assistant_response(thread_id, &message.user_id, &response)
|
||||
.await;
|
||||
@@ -1478,3 +1495,231 @@ impl Agent {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Rebuild full LLM-compatible `ChatMessage` sequence from DB messages.
|
||||
///
|
||||
/// Parses `role="tool_calls"` rows to reconstruct `assistant_with_tool_calls`
|
||||
/// and `tool_result` messages so that the LLM sees the complete tool execution
|
||||
/// history on thread hydration. Falls back gracefully for legacy rows that
|
||||
/// lack the enriched fields (`call_id`, `parameters`, `result`).
|
||||
fn rebuild_chat_messages_from_db(
|
||||
db_messages: &[crate::history::ConversationMessage],
|
||||
) -> Vec<ChatMessage> {
|
||||
let mut result = Vec::new();
|
||||
|
||||
for msg in db_messages {
|
||||
match msg.role.as_str() {
|
||||
"user" => result.push(ChatMessage::user(&msg.content)),
|
||||
"assistant" => result.push(ChatMessage::assistant(&msg.content)),
|
||||
"tool_calls" => {
|
||||
// Try to parse the enriched JSON and rebuild tool messages.
|
||||
if let Ok(calls) = serde_json::from_str::<Vec<serde_json::Value>>(&msg.content) {
|
||||
if calls.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Check if this is an enriched row (has call_id) or legacy
|
||||
let has_call_id = calls
|
||||
.first()
|
||||
.and_then(|c| c.get("call_id"))
|
||||
.and_then(|v| v.as_str())
|
||||
.is_some();
|
||||
|
||||
if has_call_id {
|
||||
// Build assistant_with_tool_calls + tool_result messages
|
||||
let tool_calls: Vec<ToolCall> = calls
|
||||
.iter()
|
||||
.map(|c| ToolCall {
|
||||
id: c["call_id"].as_str().unwrap_or("call_0").to_string(),
|
||||
name: c["name"].as_str().unwrap_or("unknown").to_string(),
|
||||
arguments: c
|
||||
.get("parameters")
|
||||
.cloned()
|
||||
.unwrap_or(serde_json::json!({})),
|
||||
})
|
||||
.collect();
|
||||
|
||||
// The assistant text for tool_calls is always None here;
|
||||
// the final assistant response comes as a separate
|
||||
// "assistant" row after this tool_calls row.
|
||||
result.push(ChatMessage::assistant_with_tool_calls(None, tool_calls));
|
||||
|
||||
// Emit tool_result messages for each call
|
||||
for c in &calls {
|
||||
let call_id = c["call_id"].as_str().unwrap_or("call_0").to_string();
|
||||
let name = c["name"].as_str().unwrap_or("unknown").to_string();
|
||||
let content = if let Some(err) = c.get("error").and_then(|v| v.as_str())
|
||||
{
|
||||
format!("Error: {}", err)
|
||||
} else if let Some(res) = c.get("result").and_then(|v| v.as_str()) {
|
||||
res.to_string()
|
||||
} else if let Some(preview) =
|
||||
c.get("result_preview").and_then(|v| v.as_str())
|
||||
{
|
||||
preview.to_string()
|
||||
} else {
|
||||
"OK".to_string()
|
||||
};
|
||||
result.push(ChatMessage::tool_result(call_id, name, content));
|
||||
}
|
||||
}
|
||||
// Legacy rows without call_id: skip (will appear as
|
||||
// simple user/assistant pairs, same as before this fix).
|
||||
}
|
||||
}
|
||||
_ => {} // Skip unknown roles
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_rebuild_chat_messages_user_assistant_only() {
|
||||
let messages = vec![
|
||||
make_db_msg("user", "Hello"),
|
||||
make_db_msg("assistant", "Hi there!"),
|
||||
];
|
||||
let result = rebuild_chat_messages_from_db(&messages);
|
||||
assert_eq!(result.len(), 2);
|
||||
assert_eq!(result[0].role, crate::llm::Role::User);
|
||||
assert_eq!(result[1].role, crate::llm::Role::Assistant);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rebuild_chat_messages_with_enriched_tool_calls() {
|
||||
let tool_json = serde_json::json!([
|
||||
{
|
||||
"name": "memory_search",
|
||||
"call_id": "call_0",
|
||||
"parameters": {"query": "test"},
|
||||
"result": "Found 3 results",
|
||||
"result_preview": "Found 3 re..."
|
||||
},
|
||||
{
|
||||
"name": "echo",
|
||||
"call_id": "call_1",
|
||||
"parameters": {"message": "hi"},
|
||||
"error": "timeout"
|
||||
}
|
||||
]);
|
||||
let messages = vec![
|
||||
make_db_msg("user", "Search for test"),
|
||||
make_db_msg("tool_calls", &tool_json.to_string()),
|
||||
make_db_msg("assistant", "I found some results."),
|
||||
];
|
||||
let result = rebuild_chat_messages_from_db(&messages);
|
||||
|
||||
// user + assistant_with_tool_calls + tool_result*2 + assistant
|
||||
assert_eq!(result.len(), 5);
|
||||
|
||||
// user
|
||||
assert_eq!(result[0].role, crate::llm::Role::User);
|
||||
|
||||
// assistant with tool_calls
|
||||
assert_eq!(result[1].role, crate::llm::Role::Assistant);
|
||||
assert!(result[1].tool_calls.is_some());
|
||||
let tcs = result[1].tool_calls.as_ref().unwrap();
|
||||
assert_eq!(tcs.len(), 2);
|
||||
assert_eq!(tcs[0].name, "memory_search");
|
||||
assert_eq!(tcs[0].id, "call_0");
|
||||
assert_eq!(tcs[1].name, "echo");
|
||||
|
||||
// tool results
|
||||
assert_eq!(result[2].role, crate::llm::Role::Tool);
|
||||
assert_eq!(result[2].tool_call_id, Some("call_0".to_string()));
|
||||
assert!(result[2].content.contains("Found 3 results"));
|
||||
|
||||
assert_eq!(result[3].role, crate::llm::Role::Tool);
|
||||
assert_eq!(result[3].tool_call_id, Some("call_1".to_string()));
|
||||
assert!(result[3].content.contains("Error: timeout"));
|
||||
|
||||
// final assistant
|
||||
assert_eq!(result[4].role, crate::llm::Role::Assistant);
|
||||
assert_eq!(result[4].content, "I found some results.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rebuild_chat_messages_legacy_tool_calls_skipped() {
|
||||
// Legacy format: no call_id field
|
||||
let tool_json = serde_json::json!([
|
||||
{"name": "echo", "result_preview": "hello"}
|
||||
]);
|
||||
let messages = vec![
|
||||
make_db_msg("user", "Hi"),
|
||||
make_db_msg("tool_calls", &tool_json.to_string()),
|
||||
make_db_msg("assistant", "Done"),
|
||||
];
|
||||
let result = rebuild_chat_messages_from_db(&messages);
|
||||
|
||||
// Legacy rows are skipped, only user + assistant
|
||||
assert_eq!(result.len(), 2);
|
||||
assert_eq!(result[0].role, crate::llm::Role::User);
|
||||
assert_eq!(result[1].role, crate::llm::Role::Assistant);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rebuild_chat_messages_empty() {
|
||||
let result = rebuild_chat_messages_from_db(&[]);
|
||||
assert!(result.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rebuild_chat_messages_malformed_tool_calls_json() {
|
||||
let messages = vec![
|
||||
make_db_msg("user", "Hi"),
|
||||
make_db_msg("tool_calls", "not valid json"),
|
||||
make_db_msg("assistant", "Done"),
|
||||
];
|
||||
let result = rebuild_chat_messages_from_db(&messages);
|
||||
// Malformed JSON is silently skipped
|
||||
assert_eq!(result.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rebuild_chat_messages_multi_turn_with_tools() {
|
||||
let tool_json_1 = serde_json::json!([
|
||||
{"name": "search", "call_id": "call_0", "parameters": {}, "result": "found it"}
|
||||
]);
|
||||
let tool_json_2 = serde_json::json!([
|
||||
{"name": "write", "call_id": "call_0", "parameters": {"path": "a.txt"}, "result": "ok"}
|
||||
]);
|
||||
let messages = vec![
|
||||
make_db_msg("user", "Find X"),
|
||||
make_db_msg("tool_calls", &tool_json_1.to_string()),
|
||||
make_db_msg("assistant", "Found X"),
|
||||
make_db_msg("user", "Write it"),
|
||||
make_db_msg("tool_calls", &tool_json_2.to_string()),
|
||||
make_db_msg("assistant", "Written"),
|
||||
];
|
||||
let result = rebuild_chat_messages_from_db(&messages);
|
||||
|
||||
// Turn 1: user + assistant_with_calls + tool_result + assistant = 4
|
||||
// Turn 2: user + assistant_with_calls + tool_result + assistant = 4
|
||||
assert_eq!(result.len(), 8);
|
||||
|
||||
// Verify turn boundaries
|
||||
assert_eq!(result[0].content, "Find X");
|
||||
assert!(result[1].tool_calls.is_some());
|
||||
assert_eq!(result[2].role, crate::llm::Role::Tool);
|
||||
assert_eq!(result[3].content, "Found X");
|
||||
|
||||
assert_eq!(result[4].content, "Write it");
|
||||
assert!(result[5].tool_calls.is_some());
|
||||
assert_eq!(result[6].role, crate::llm::Role::Tool);
|
||||
assert_eq!(result[7].content, "Written");
|
||||
}
|
||||
|
||||
fn make_db_msg(role: &str, content: &str) -> crate::history::ConversationMessage {
|
||||
crate::history::ConversationMessage {
|
||||
id: uuid::Uuid::new_v4(),
|
||||
role: role.to_string(),
|
||||
content: content.to_string(),
|
||||
created_at: chrono::Utc::now(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+1
-1
@@ -212,7 +212,7 @@ impl Worker {
|
||||
let job_ctx = self.context_manager().get_context(self.job_id).await?;
|
||||
|
||||
// Create reasoning engine
|
||||
let reasoning = Reasoning::new(self.llm().clone(), self.safety().clone());
|
||||
let reasoning = Reasoning::new(self.llm().clone());
|
||||
|
||||
// Build initial reasoning context (tool definitions refreshed each iteration in execution_loop)
|
||||
let mut reason_ctx = ReasoningContext::new().with_job(&job_ctx.description);
|
||||
|
||||
+228
-102
@@ -21,7 +21,7 @@ use crate::secrets::SecretsStore;
|
||||
use crate::skills::SkillRegistry;
|
||||
use crate::skills::catalog::SkillCatalog;
|
||||
use crate::tools::ToolRegistry;
|
||||
use crate::tools::mcp::McpSessionManager;
|
||||
use crate::tools::mcp::{McpProcessManager, McpSessionManager};
|
||||
use crate::tools::wasm::SharedCredentialRegistry;
|
||||
use crate::tools::wasm::WasmToolRuntime;
|
||||
use crate::workspace::{EmbeddingProvider, Workspace};
|
||||
@@ -41,6 +41,7 @@ pub struct AppComponents {
|
||||
pub workspace: Option<Arc<Workspace>>,
|
||||
pub extension_manager: Option<Arc<ExtensionManager>>,
|
||||
pub mcp_session_manager: Arc<McpSessionManager>,
|
||||
pub mcp_process_manager: Arc<McpProcessManager>,
|
||||
pub wasm_tool_runtime: Option<Arc<WasmToolRuntime>>,
|
||||
pub log_broadcaster: Arc<LogBroadcaster>,
|
||||
pub context_manager: Arc<ContextManager>,
|
||||
@@ -255,15 +256,18 @@ impl AppBuilder {
|
||||
self.libsql_db.take();
|
||||
}
|
||||
|
||||
// Re-resolve config with OS credentials
|
||||
if let Some(ref db) = self.db {
|
||||
let toml_path = self.toml_path.as_deref();
|
||||
if let Ok(refreshed) =
|
||||
Config::from_db_with_toml(db.as_ref(), "default", toml_path).await
|
||||
{
|
||||
self.config = refreshed;
|
||||
tracing::debug!("LlmConfig re-resolved after OS credential injection");
|
||||
}
|
||||
// Re-resolve only the LLM config with OS credentials.
|
||||
let store: Option<&(dyn crate::db::SettingsStore + Sync)> =
|
||||
self.db.as_ref().map(|db| db.as_ref() as _);
|
||||
let toml_path = self.toml_path.as_deref();
|
||||
if let Err(e) = self
|
||||
.config
|
||||
.re_resolve_llm(store, "default", toml_path)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
"Failed to re-resolve LLM config after OS credential injection: {e}"
|
||||
);
|
||||
}
|
||||
|
||||
return Ok(());
|
||||
@@ -308,18 +312,16 @@ impl AppBuilder {
|
||||
// Inject LLM API keys from encrypted storage
|
||||
crate::config::inject_llm_keys_from_secrets(secrets.as_ref(), "default").await;
|
||||
|
||||
// Re-resolve config with newly available keys
|
||||
if let Some(ref db) = self.db {
|
||||
let toml_path = self.toml_path.as_deref();
|
||||
match Config::from_db_with_toml(db.as_ref(), "default", toml_path).await {
|
||||
Ok(refreshed) => {
|
||||
self.config = refreshed;
|
||||
tracing::debug!("LlmConfig re-resolved after secret injection");
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to re-resolve config after secret injection: {}", e);
|
||||
}
|
||||
}
|
||||
// Re-resolve only the LLM config with newly available keys.
|
||||
let store: Option<&(dyn crate::db::SettingsStore + Sync)> =
|
||||
self.db.as_ref().map(|db| db.as_ref() as _);
|
||||
let toml_path = self.toml_path.as_deref();
|
||||
if let Err(e) = self
|
||||
.config
|
||||
.re_resolve_llm(store, "default", toml_path)
|
||||
.await
|
||||
{
|
||||
tracing::warn!("Failed to re-resolve LLM config after secret injection: {e}");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -332,7 +334,7 @@ impl AppBuilder {
|
||||
/// Delegates to `build_provider_chain` which applies all decorators
|
||||
/// (retry, smart routing, failover, circuit breaker, response cache).
|
||||
#[allow(clippy::type_complexity)]
|
||||
pub fn init_llm(
|
||||
pub async fn init_llm(
|
||||
&self,
|
||||
) -> Result<
|
||||
(
|
||||
@@ -343,7 +345,7 @@ impl AppBuilder {
|
||||
anyhow::Error,
|
||||
> {
|
||||
let (llm, cheap_llm, recording_handle) =
|
||||
crate::llm::build_provider_chain(&self.config.llm, self.session.clone())?;
|
||||
crate::llm::build_provider_chain(&self.config.llm, self.session.clone()).await?;
|
||||
Ok((llm, cheap_llm, recording_handle))
|
||||
}
|
||||
|
||||
@@ -398,16 +400,55 @@ impl AppBuilder {
|
||||
None
|
||||
};
|
||||
|
||||
// Register image/vision tools if we have a workspace and LLM API credentials
|
||||
if workspace.is_some() {
|
||||
let (api_base, api_key_opt) = if let Some(ref provider) = self.config.llm.provider {
|
||||
(
|
||||
provider.base_url.clone(),
|
||||
provider.api_key.as_ref().map(|s| {
|
||||
use secrecy::ExposeSecret;
|
||||
s.expose_secret().to_string()
|
||||
}),
|
||||
)
|
||||
} else {
|
||||
(
|
||||
self.config.llm.nearai.base_url.clone(),
|
||||
self.config.llm.nearai.api_key.as_ref().map(|s| {
|
||||
use secrecy::ExposeSecret;
|
||||
s.expose_secret().to_string()
|
||||
}),
|
||||
)
|
||||
};
|
||||
|
||||
if let Some(api_key) = api_key_opt {
|
||||
// Check for image generation models
|
||||
let model_name = self
|
||||
.config
|
||||
.llm
|
||||
.provider
|
||||
.as_ref()
|
||||
.map(|p| p.model.clone())
|
||||
.unwrap_or_else(|| self.config.llm.nearai.model.clone());
|
||||
let models = vec![model_name.clone()];
|
||||
let gen_model = crate::llm::image_models::suggest_image_model(&models)
|
||||
.unwrap_or("flux-1.1-pro")
|
||||
.to_string();
|
||||
tools.register_image_tools(api_base.clone(), api_key.clone(), gen_model, None);
|
||||
|
||||
// Check for vision models
|
||||
let vision_model = crate::llm::vision_models::suggest_vision_model(&models)
|
||||
.unwrap_or(&model_name)
|
||||
.to_string();
|
||||
tools.register_vision_tools(api_base, api_key, vision_model, None);
|
||||
}
|
||||
}
|
||||
|
||||
// Register builder tool if enabled
|
||||
if self.config.builder.enabled
|
||||
&& (self.config.agent.allow_local_tools || !self.config.sandbox.enabled)
|
||||
{
|
||||
tools
|
||||
.register_builder_tool(
|
||||
llm.clone(),
|
||||
safety.clone(),
|
||||
Some(self.config.builder.to_builder_config()),
|
||||
)
|
||||
.register_builder_tool(llm.clone(), Some(self.config.builder.to_builder_config()))
|
||||
.await;
|
||||
tracing::info!("Builder mode enabled");
|
||||
}
|
||||
@@ -423,6 +464,7 @@ impl AppBuilder {
|
||||
) -> Result<
|
||||
(
|
||||
Arc<McpSessionManager>,
|
||||
Arc<McpProcessManager>,
|
||||
Option<Arc<WasmToolRuntime>>,
|
||||
Option<Arc<ExtensionManager>>,
|
||||
Vec<crate::extensions::RegistryEntry>,
|
||||
@@ -430,10 +472,13 @@ impl AppBuilder {
|
||||
),
|
||||
anyhow::Error,
|
||||
> {
|
||||
use crate::tools::mcp::{McpClient, config::load_mcp_servers_from_db, is_authenticated};
|
||||
use crate::tools::mcp::{
|
||||
McpClient, McpTransport, config::load_mcp_servers_from_db, is_authenticated,
|
||||
};
|
||||
use crate::tools::wasm::{WasmToolLoader, load_dev_tools};
|
||||
|
||||
let mcp_session_manager = Arc::new(McpSessionManager::new());
|
||||
let mcp_process_manager = Arc::new(McpProcessManager::new());
|
||||
|
||||
// Create WASM tool runtime eagerly so extensions installed after startup
|
||||
// (e.g. via the web UI) can still be activated. The tools directory is only
|
||||
@@ -509,97 +554,175 @@ impl AppBuilder {
|
||||
let db = self.db.clone();
|
||||
let tools = Arc::clone(tools);
|
||||
let mcp_sm = Arc::clone(&mcp_session_manager);
|
||||
let pm = Arc::clone(&mcp_process_manager);
|
||||
async move {
|
||||
if let Some(ref secrets) = secrets_store {
|
||||
let servers_result = if let Some(ref d) = db {
|
||||
load_mcp_servers_from_db(d.as_ref(), "default").await
|
||||
} else {
|
||||
crate::tools::mcp::config::load_mcp_servers().await
|
||||
};
|
||||
match servers_result {
|
||||
Ok(servers) => {
|
||||
let enabled: Vec<_> = servers.enabled_servers().cloned().collect();
|
||||
if !enabled.is_empty() {
|
||||
tracing::info!(
|
||||
"Loading {} configured MCP server(s)...",
|
||||
enabled.len()
|
||||
);
|
||||
}
|
||||
let servers_result = if let Some(ref d) = db {
|
||||
load_mcp_servers_from_db(d.as_ref(), "default").await
|
||||
} else {
|
||||
crate::tools::mcp::config::load_mcp_servers().await
|
||||
};
|
||||
match servers_result {
|
||||
Ok(servers) => {
|
||||
let enabled: Vec<_> = servers.enabled_servers().cloned().collect();
|
||||
if !enabled.is_empty() {
|
||||
tracing::info!("Loading {} configured MCP server(s)...", enabled.len());
|
||||
}
|
||||
|
||||
let mut join_set = tokio::task::JoinSet::new();
|
||||
for server in enabled {
|
||||
let mcp_sm = Arc::clone(&mcp_sm);
|
||||
let secrets = Arc::clone(secrets);
|
||||
let tools = Arc::clone(&tools);
|
||||
let mut join_set = tokio::task::JoinSet::new();
|
||||
for server in enabled {
|
||||
let mcp_sm = Arc::clone(&mcp_sm);
|
||||
let secrets = secrets_store.clone();
|
||||
let tools = Arc::clone(&tools);
|
||||
let pm = Arc::clone(&pm);
|
||||
|
||||
join_set.spawn(async move {
|
||||
let server_name = server.name.clone();
|
||||
let has_tokens =
|
||||
is_authenticated(&server, &secrets, "default").await;
|
||||
join_set.spawn(async move {
|
||||
let server_name = server.name.clone();
|
||||
|
||||
let client = if has_tokens || server.requires_auth() {
|
||||
McpClient::new_authenticated(
|
||||
server, mcp_sm, secrets, "default",
|
||||
)
|
||||
} else {
|
||||
McpClient::new_with_name(&server_name, &server.url)
|
||||
};
|
||||
|
||||
match client.list_tools().await {
|
||||
Ok(mcp_tools) => {
|
||||
let tool_count = mcp_tools.len();
|
||||
match client.create_tools().await {
|
||||
Ok(tool_impls) => {
|
||||
for tool in tool_impls {
|
||||
tools.register(tool).await;
|
||||
}
|
||||
tracing::info!(
|
||||
"Loaded {} tools from MCP server '{}'",
|
||||
tool_count,
|
||||
server_name
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Failed to create tools from MCP server '{}': {}",
|
||||
server_name,
|
||||
e
|
||||
);
|
||||
}
|
||||
let client: McpClient = match server.effective_transport() {
|
||||
crate::tools::mcp::config::EffectiveTransport::Stdio {
|
||||
command,
|
||||
args,
|
||||
env,
|
||||
} => {
|
||||
match pm
|
||||
.spawn_stdio(
|
||||
&server_name,
|
||||
command,
|
||||
args.to_vec(),
|
||||
env.clone(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(transport) => McpClient::new_with_transport(
|
||||
&server_name,
|
||||
transport as Arc<dyn McpTransport>,
|
||||
None,
|
||||
secrets,
|
||||
"default",
|
||||
Some(server),
|
||||
),
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Failed to spawn stdio MCP server '{}': {}",
|
||||
server_name,
|
||||
e
|
||||
);
|
||||
return;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
let err_str = e.to_string();
|
||||
if err_str.contains("401")
|
||||
|| err_str.contains("authentication")
|
||||
{
|
||||
}
|
||||
#[cfg(unix)]
|
||||
crate::tools::mcp::config::EffectiveTransport::Unix {
|
||||
socket_path,
|
||||
} => {
|
||||
match crate::tools::mcp::unix_transport::UnixMcpTransport::connect(
|
||||
&server_name,
|
||||
socket_path,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(transport) => McpClient::new_with_transport(
|
||||
&server_name,
|
||||
Arc::new(transport) as Arc<dyn McpTransport>,
|
||||
None,
|
||||
secrets,
|
||||
"default",
|
||||
Some(server),
|
||||
),
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"MCP server '{}' requires authentication. \
|
||||
Run: ironclaw mcp auth {}",
|
||||
"Failed to connect to Unix MCP server '{}': {}",
|
||||
server_name,
|
||||
e
|
||||
);
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
#[cfg(not(unix))]
|
||||
crate::tools::mcp::config::EffectiveTransport::Unix { .. } => {
|
||||
tracing::warn!(
|
||||
"Unix socket transport is not supported on this platform (server '{}')",
|
||||
server_name
|
||||
);
|
||||
return;
|
||||
}
|
||||
crate::tools::mcp::config::EffectiveTransport::Http => {
|
||||
if let Some(ref secrets) = secrets {
|
||||
let has_tokens =
|
||||
is_authenticated(&server, secrets, "default")
|
||||
.await;
|
||||
|
||||
if has_tokens || server.requires_auth() {
|
||||
McpClient::new_authenticated(
|
||||
server,
|
||||
Arc::clone(&mcp_sm),
|
||||
Arc::clone(secrets),
|
||||
"default",
|
||||
)
|
||||
} else {
|
||||
McpClient::new_with_config(server)
|
||||
}
|
||||
} else {
|
||||
McpClient::new_with_config(server)
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
match client.list_tools().await {
|
||||
Ok(mcp_tools) => {
|
||||
let tool_count = mcp_tools.len();
|
||||
match client.create_tools().await {
|
||||
Ok(tool_impls) => {
|
||||
for tool in tool_impls {
|
||||
tools.register(tool).await;
|
||||
}
|
||||
tracing::info!(
|
||||
"Loaded {} tools from MCP server '{}'",
|
||||
tool_count,
|
||||
server_name
|
||||
);
|
||||
} else {
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Failed to connect to MCP server '{}': {}",
|
||||
"Failed to create tools from MCP server '{}': {}",
|
||||
server_name,
|
||||
e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
while let Some(result) = join_set.join_next().await {
|
||||
if let Err(e) = result {
|
||||
tracing::warn!("MCP server loading task panicked: {}", e);
|
||||
Err(e) => {
|
||||
let err_str = e.to_string();
|
||||
if err_str.contains("401")
|
||||
|| err_str.contains("authentication")
|
||||
{
|
||||
tracing::warn!(
|
||||
"MCP server '{}' requires authentication. \
|
||||
Run: ironclaw mcp auth {}",
|
||||
server_name,
|
||||
server_name
|
||||
);
|
||||
} else {
|
||||
tracing::warn!(
|
||||
"Failed to connect to MCP server '{}': {}",
|
||||
server_name,
|
||||
e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
while let Some(result) = join_set.join_next().await {
|
||||
if let Err(e) = result {
|
||||
tracing::warn!("MCP server loading task panicked: {}", e);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::debug!("No MCP servers configured ({})", e);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::debug!("No MCP servers configured ({})", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -670,6 +793,7 @@ impl AppBuilder {
|
||||
|
||||
Ok((
|
||||
mcp_session_manager,
|
||||
mcp_process_manager,
|
||||
wasm_tool_runtime,
|
||||
extension_manager,
|
||||
catalog_entries,
|
||||
@@ -696,7 +820,7 @@ impl AppBuilder {
|
||||
let (llm, cheap_llm, recording_handle) = if let Some(llm) = self.llm_override.take() {
|
||||
(llm, None, None)
|
||||
} else {
|
||||
self.init_llm()?
|
||||
self.init_llm().await?
|
||||
};
|
||||
let (safety, tools, embeddings, workspace) = self.init_tools(&llm).await?;
|
||||
|
||||
@@ -705,6 +829,7 @@ impl AppBuilder {
|
||||
|
||||
let (
|
||||
mcp_session_manager,
|
||||
mcp_process_manager,
|
||||
wasm_tool_runtime,
|
||||
extension_manager,
|
||||
catalog_entries,
|
||||
@@ -802,6 +927,7 @@ impl AppBuilder {
|
||||
workspace,
|
||||
extension_manager,
|
||||
mcp_session_manager,
|
||||
mcp_process_manager,
|
||||
wasm_tool_runtime,
|
||||
log_broadcaster: self.log_broadcaster,
|
||||
context_manager,
|
||||
|
||||
@@ -414,10 +414,103 @@ pub enum MigrationError {
|
||||
Io(String),
|
||||
}
|
||||
|
||||
// ── PID Lock ──────────────────────────────────────────────────────────────
|
||||
|
||||
/// Path to the PID lock file: `~/.ironclaw/ironclaw.pid`.
|
||||
pub fn pid_lock_path() -> PathBuf {
|
||||
ironclaw_base_dir().join("ironclaw.pid")
|
||||
}
|
||||
|
||||
/// A PID-based lock that prevents multiple IronClaw instances from running
|
||||
/// simultaneously.
|
||||
///
|
||||
/// Uses `fs4::try_lock_exclusive()` for atomic locking (no TOCTOU race),
|
||||
/// then writes the current PID into the locked file for diagnostics.
|
||||
/// The OS-level lock is held for the lifetime of this struct and
|
||||
/// automatically released on drop (along with the PID file cleanup).
|
||||
#[derive(Debug)]
|
||||
pub struct PidLock {
|
||||
path: PathBuf,
|
||||
/// Held open to maintain the OS-level exclusive lock.
|
||||
_file: std::fs::File,
|
||||
}
|
||||
|
||||
/// Errors from PID lock acquisition.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum PidLockError {
|
||||
#[error("Another IronClaw instance is already running (PID {pid})")]
|
||||
AlreadyRunning { pid: u32 },
|
||||
#[error("Failed to acquire PID lock: {0}")]
|
||||
Io(#[from] std::io::Error),
|
||||
}
|
||||
|
||||
impl PidLock {
|
||||
/// Try to acquire the PID lock.
|
||||
///
|
||||
/// Uses an exclusive file lock (`flock`/`LockFileEx`) so that two
|
||||
/// concurrent processes cannot both acquire the lock — no TOCTOU race.
|
||||
/// If the lock file exists but the holding process is gone (stale),
|
||||
/// the lock is reclaimed automatically by the OS.
|
||||
pub fn acquire() -> Result<Self, PidLockError> {
|
||||
Self::acquire_at(pid_lock_path())
|
||||
}
|
||||
|
||||
/// Acquire at a specific path (for testing).
|
||||
fn acquire_at(path: PathBuf) -> Result<Self, PidLockError> {
|
||||
use fs4::FileExt;
|
||||
use std::fs::OpenOptions;
|
||||
use std::io::Write;
|
||||
|
||||
// Ensure parent directory exists
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
|
||||
// Open (or create) the lock file
|
||||
let mut file = OpenOptions::new()
|
||||
.read(true)
|
||||
.write(true)
|
||||
.create(true)
|
||||
.truncate(false)
|
||||
.open(&path)?;
|
||||
|
||||
// Try non-blocking exclusive lock — if another process holds it,
|
||||
// this fails immediately instead of blocking.
|
||||
if let Err(e) = file.try_lock_exclusive() {
|
||||
if e.kind() == std::io::ErrorKind::WouldBlock {
|
||||
// Lock held by another process — read its PID for the error message
|
||||
let pid = std::fs::read_to_string(&path)
|
||||
.ok()
|
||||
.and_then(|s| s.trim().parse::<u32>().ok())
|
||||
.unwrap_or(0);
|
||||
return Err(PidLockError::AlreadyRunning { pid });
|
||||
}
|
||||
// Other errors (permissions, unsupported filesystem, etc.)
|
||||
return Err(PidLockError::Io(e));
|
||||
}
|
||||
|
||||
// We hold the exclusive lock — write our PID
|
||||
file.set_len(0)?; // truncate
|
||||
write!(file, "{}", std::process::id())?;
|
||||
|
||||
Ok(PidLock { path, _file: file })
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for PidLock {
|
||||
fn drop(&mut self) {
|
||||
// Remove the PID file; the OS-level lock is released when _file is dropped.
|
||||
let _ = std::fs::remove_file(&self.path);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::process::Command;
|
||||
use std::sync::Mutex;
|
||||
use std::thread;
|
||||
use std::time::{Duration, Instant};
|
||||
use tempfile::tempdir;
|
||||
|
||||
static ENV_MUTEX: Mutex<()> = Mutex::new(());
|
||||
@@ -986,4 +1079,162 @@ INJECTED="pwned"#;
|
||||
unsafe { std::env::remove_var("IRONCLAW_BASE_DIR") };
|
||||
}
|
||||
}
|
||||
|
||||
// ── PID Lock tests ───────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn test_pid_lock_acquire_and_drop() {
|
||||
let dir = tempdir().unwrap();
|
||||
let pid_path = dir.path().join("ironclaw.pid");
|
||||
|
||||
// Acquire lock
|
||||
let lock = PidLock::acquire_at(pid_path.clone()).unwrap();
|
||||
assert!(pid_path.exists());
|
||||
|
||||
// PID file should contain our PID
|
||||
let contents = std::fs::read_to_string(&pid_path).unwrap();
|
||||
assert_eq!(contents.trim().parse::<u32>().unwrap(), std::process::id());
|
||||
|
||||
// Drop should remove the file
|
||||
drop(lock);
|
||||
assert!(!pid_path.exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pid_lock_rejects_second_acquire() {
|
||||
let dir = tempdir().unwrap();
|
||||
let pid_path = dir.path().join("ironclaw.pid");
|
||||
|
||||
// First lock succeeds
|
||||
let _lock1 = PidLock::acquire_at(pid_path.clone()).unwrap();
|
||||
|
||||
// Second acquire on same file must fail (exclusive flock held)
|
||||
let result = PidLock::acquire_at(pid_path.clone());
|
||||
assert!(result.is_err());
|
||||
match result.unwrap_err() {
|
||||
PidLockError::AlreadyRunning { pid } => {
|
||||
assert_eq!(pid, std::process::id());
|
||||
}
|
||||
other => panic!("expected AlreadyRunning, got: {}", other),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pid_lock_reclaims_after_drop() {
|
||||
let dir = tempdir().unwrap();
|
||||
let pid_path = dir.path().join("ironclaw.pid");
|
||||
|
||||
// Acquire and release
|
||||
let lock = PidLock::acquire_at(pid_path.clone()).unwrap();
|
||||
drop(lock);
|
||||
|
||||
// Should succeed — OS lock was released on drop
|
||||
let lock2 = PidLock::acquire_at(pid_path).unwrap();
|
||||
drop(lock2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pid_lock_reclaims_stale_file_without_flock() {
|
||||
let dir = tempdir().unwrap();
|
||||
let pid_path = dir.path().join("ironclaw.pid");
|
||||
|
||||
// Write a stale PID file manually (no flock held)
|
||||
std::fs::write(&pid_path, "4294967294").unwrap();
|
||||
|
||||
// Should succeed because no OS lock is held on the file
|
||||
let lock = PidLock::acquire_at(pid_path.clone()).unwrap();
|
||||
let contents = std::fs::read_to_string(&pid_path).unwrap();
|
||||
assert_eq!(contents.trim().parse::<u32>().unwrap(), std::process::id());
|
||||
drop(lock);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pid_lock_handles_corrupt_pid_file() {
|
||||
let dir = tempdir().unwrap();
|
||||
let pid_path = dir.path().join("ironclaw.pid");
|
||||
|
||||
// Write garbage (no flock held)
|
||||
std::fs::write(&pid_path, "not-a-number").unwrap();
|
||||
|
||||
// Should succeed — no OS lock held, file is reclaimed
|
||||
let lock = PidLock::acquire_at(pid_path).unwrap();
|
||||
drop(lock);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pid_lock_creates_parent_dirs() {
|
||||
let dir = tempdir().unwrap();
|
||||
let pid_path = dir.path().join("nested").join("deep").join("ironclaw.pid");
|
||||
|
||||
let lock = PidLock::acquire_at(pid_path.clone()).unwrap();
|
||||
assert!(pid_path.exists());
|
||||
drop(lock);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pid_lock_child_helper_holds_lock() {
|
||||
if std::env::var("IRONCLAW_PID_LOCK_CHILD").ok().as_deref() != Some("1") {
|
||||
return;
|
||||
}
|
||||
|
||||
let pid_path = PathBuf::from(
|
||||
std::env::var("IRONCLAW_PID_LOCK_PATH").expect("IRONCLAW_PID_LOCK_PATH missing"),
|
||||
);
|
||||
let hold_ms = std::env::var("IRONCLAW_PID_LOCK_HOLD_MS")
|
||||
.ok()
|
||||
.and_then(|s| s.parse::<u64>().ok())
|
||||
.unwrap_or(3000);
|
||||
|
||||
let _lock = PidLock::acquire_at(pid_path).expect("child failed to acquire pid lock");
|
||||
thread::sleep(Duration::from_millis(hold_ms));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pid_lock_rejects_lock_held_by_other_process() {
|
||||
let dir = tempdir().unwrap();
|
||||
let pid_path = dir.path().join("ironclaw.pid");
|
||||
|
||||
let current_exe = std::env::current_exe().unwrap();
|
||||
let mut child = Command::new(current_exe)
|
||||
.args([
|
||||
"--exact",
|
||||
"bootstrap::tests::test_pid_lock_child_helper_holds_lock",
|
||||
"--nocapture",
|
||||
"--test-threads=1",
|
||||
])
|
||||
.env("IRONCLAW_PID_LOCK_CHILD", "1")
|
||||
.env("IRONCLAW_PID_LOCK_PATH", pid_path.display().to_string())
|
||||
.env("IRONCLAW_PID_LOCK_HOLD_MS", "3000")
|
||||
.spawn()
|
||||
.unwrap();
|
||||
|
||||
let started = Instant::now();
|
||||
while started.elapsed() < Duration::from_secs(2) {
|
||||
if pid_path.exists() {
|
||||
break;
|
||||
}
|
||||
if let Some(status) = child.try_wait().unwrap() {
|
||||
panic!("child exited before acquiring lock: {}", status);
|
||||
}
|
||||
thread::sleep(Duration::from_millis(20));
|
||||
}
|
||||
assert!(
|
||||
pid_path.exists(),
|
||||
"child did not create lock file in time: {}",
|
||||
pid_path.display()
|
||||
);
|
||||
|
||||
let result = PidLock::acquire_at(pid_path.clone());
|
||||
match result.unwrap_err() {
|
||||
PidLockError::AlreadyRunning { .. } => {}
|
||||
other => panic!("expected AlreadyRunning, got: {}", other),
|
||||
}
|
||||
|
||||
let status = child.wait().unwrap();
|
||||
assert!(status.success(), "child process failed: {}", status);
|
||||
|
||||
// After the child exits, lock should be released and reacquirable.
|
||||
let lock = PidLock::acquire_at(pid_path).unwrap();
|
||||
drop(lock);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -79,6 +79,8 @@ pub struct IncomingMessage {
|
||||
pub received_at: DateTime<Utc>,
|
||||
/// Channel-specific metadata.
|
||||
pub metadata: serde_json::Value,
|
||||
/// IANA timezone string from the client (e.g. "America/New_York").
|
||||
pub timezone: Option<String>,
|
||||
/// File or media attachments on this message.
|
||||
pub attachments: Vec<IncomingAttachment>,
|
||||
}
|
||||
@@ -99,6 +101,7 @@ impl IncomingMessage {
|
||||
thread_id: None,
|
||||
received_at: Utc::now(),
|
||||
metadata: serde_json::Value::Null,
|
||||
timezone: None,
|
||||
attachments: Vec::new(),
|
||||
}
|
||||
}
|
||||
@@ -121,6 +124,12 @@ impl IncomingMessage {
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the client timezone.
|
||||
pub fn with_timezone(mut self, tz: impl Into<String>) -> Self {
|
||||
self.timezone = Some(tz.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// Set attachments.
|
||||
pub fn with_attachments(mut self, attachments: Vec<IncomingAttachment>) -> Self {
|
||||
self.attachments = attachments;
|
||||
@@ -222,6 +231,13 @@ pub enum StatusUpdate {
|
||||
success: bool,
|
||||
message: String,
|
||||
},
|
||||
/// An image was generated by a tool.
|
||||
ImageGenerated {
|
||||
/// Base64 data URL of the generated image.
|
||||
data_url: String,
|
||||
/// Optional workspace path where the image was saved.
|
||||
path: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
impl StatusUpdate {
|
||||
@@ -454,4 +470,10 @@ mod tests {
|
||||
panic!("expected ToolCompleted variant");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_incoming_message_with_timezone() {
|
||||
let msg = IncomingMessage::new("test", "user1", "hello").with_timezone("America/New_York");
|
||||
assert_eq!(msg.timezone.as_deref(), Some("America/New_York"));
|
||||
}
|
||||
}
|
||||
|
||||
+126
-6
@@ -17,7 +17,9 @@ use tokio::sync::{RwLock, mpsc, oneshot};
|
||||
use tokio_stream::wrappers::ReceiverStream;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::channels::{Channel, IncomingMessage, MessageStream, OutgoingResponse};
|
||||
use crate::channels::{
|
||||
AttachmentKind, Channel, IncomingAttachment, IncomingMessage, MessageStream, OutgoingResponse,
|
||||
};
|
||||
use crate::config::HttpConfig;
|
||||
use crate::error::ChannelError;
|
||||
|
||||
@@ -46,8 +48,9 @@ struct RateLimitState {
|
||||
request_count: u32,
|
||||
}
|
||||
|
||||
/// Maximum JSON body size for webhook requests (64 KB).
|
||||
const MAX_BODY_BYTES: usize = 64 * 1024;
|
||||
/// Maximum JSON body size for webhook requests (15 MB, to support base64 image attachments
|
||||
/// with ~33% overhead from base64 encoding).
|
||||
const MAX_BODY_BYTES: usize = 15 * 1024 * 1024;
|
||||
|
||||
/// Maximum number of pending wait-for-response requests.
|
||||
const MAX_PENDING_RESPONSES: usize = 100;
|
||||
@@ -115,8 +118,34 @@ struct WebhookRequest {
|
||||
/// Whether to wait for a synchronous response.
|
||||
#[serde(default)]
|
||||
wait_for_response: bool,
|
||||
/// Optional file attachments (base64-encoded).
|
||||
#[serde(default)]
|
||||
attachments: Vec<AttachmentData>,
|
||||
}
|
||||
|
||||
/// A file attachment in a webhook request.
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct AttachmentData {
|
||||
/// MIME type (e.g. "image/png", "application/pdf").
|
||||
mime_type: String,
|
||||
/// Optional filename.
|
||||
#[serde(default)]
|
||||
filename: Option<String>,
|
||||
/// Base64-encoded file data.
|
||||
#[serde(default)]
|
||||
data_base64: Option<String>,
|
||||
/// URL to fetch the file from (not downloaded server-side for SSRF prevention).
|
||||
#[serde(default)]
|
||||
url: Option<String>,
|
||||
}
|
||||
|
||||
/// Maximum size per attachment (5 MB decoded).
|
||||
const MAX_ATTACHMENT_BYTES: usize = 5 * 1024 * 1024;
|
||||
/// Maximum total attachment size (10 MB decoded).
|
||||
const MAX_TOTAL_ATTACHMENT_BYTES: usize = 10 * 1024 * 1024;
|
||||
/// Maximum number of attachments per request.
|
||||
const MAX_ATTACHMENTS: usize = 5;
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct WebhookResponse {
|
||||
/// Message ID assigned to this request.
|
||||
@@ -211,15 +240,106 @@ async fn webhook_handler(
|
||||
);
|
||||
}
|
||||
|
||||
let msg = IncomingMessage::new("http", &state.user_id, &req.content).with_metadata(
|
||||
// Validate and decode attachments
|
||||
let attachments = if !req.attachments.is_empty() {
|
||||
if req.attachments.len() > MAX_ATTACHMENTS {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(WebhookResponse {
|
||||
message_id: Uuid::nil(),
|
||||
status: "error".to_string(),
|
||||
response: Some(format!("Too many attachments (max {})", MAX_ATTACHMENTS)),
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
let mut decoded_attachments = Vec::new();
|
||||
let mut total_bytes: usize = 0;
|
||||
for att in &req.attachments {
|
||||
if let Some(ref b64) = att.data_base64 {
|
||||
use base64::Engine;
|
||||
let data = match base64::engine::general_purpose::STANDARD.decode(b64) {
|
||||
Ok(d) => d,
|
||||
Err(_) => {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(WebhookResponse {
|
||||
message_id: Uuid::nil(),
|
||||
status: "error".to_string(),
|
||||
response: Some("Invalid base64 in attachment".to_string()),
|
||||
}),
|
||||
);
|
||||
}
|
||||
};
|
||||
if data.len() > MAX_ATTACHMENT_BYTES {
|
||||
return (
|
||||
StatusCode::PAYLOAD_TOO_LARGE,
|
||||
Json(WebhookResponse {
|
||||
message_id: Uuid::nil(),
|
||||
status: "error".to_string(),
|
||||
response: Some(format!(
|
||||
"Attachment too large (max {} bytes)",
|
||||
MAX_ATTACHMENT_BYTES
|
||||
)),
|
||||
}),
|
||||
);
|
||||
}
|
||||
total_bytes += data.len();
|
||||
if total_bytes > MAX_TOTAL_ATTACHMENT_BYTES {
|
||||
return (
|
||||
StatusCode::PAYLOAD_TOO_LARGE,
|
||||
Json(WebhookResponse {
|
||||
message_id: Uuid::nil(),
|
||||
status: "error".to_string(),
|
||||
response: Some("Total attachment size exceeds limit".to_string()),
|
||||
}),
|
||||
);
|
||||
}
|
||||
decoded_attachments.push(IncomingAttachment {
|
||||
id: Uuid::new_v4().to_string(),
|
||||
kind: AttachmentKind::from_mime_type(&att.mime_type),
|
||||
mime_type: att.mime_type.clone(),
|
||||
filename: att.filename.clone(),
|
||||
size_bytes: Some(data.len() as u64),
|
||||
source_url: None,
|
||||
storage_key: None,
|
||||
extracted_text: None,
|
||||
data,
|
||||
duration_secs: None,
|
||||
});
|
||||
} else if let Some(ref url) = att.url {
|
||||
// URL-only attachment: set source_url but don't download (SSRF prevention)
|
||||
decoded_attachments.push(IncomingAttachment {
|
||||
id: Uuid::new_v4().to_string(),
|
||||
kind: AttachmentKind::from_mime_type(&att.mime_type),
|
||||
mime_type: att.mime_type.clone(),
|
||||
filename: att.filename.clone(),
|
||||
size_bytes: None,
|
||||
source_url: Some(url.clone()),
|
||||
storage_key: None,
|
||||
extracted_text: None,
|
||||
data: Vec::new(),
|
||||
duration_secs: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
decoded_attachments
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
|
||||
let mut msg = IncomingMessage::new("http", &state.user_id, &req.content).with_metadata(
|
||||
serde_json::json!({
|
||||
"wait_for_response": req.wait_for_response,
|
||||
}),
|
||||
);
|
||||
|
||||
if !attachments.is_empty() {
|
||||
msg = msg.with_attachments(attachments);
|
||||
}
|
||||
|
||||
if let Some(thread_id) = &req.thread_id {
|
||||
let msg = msg.with_thread(thread_id);
|
||||
return process_message(state, msg, req.wait_for_response).await;
|
||||
msg = msg.with_thread(thread_id);
|
||||
}
|
||||
|
||||
process_message(state, msg, req.wait_for_response).await
|
||||
|
||||
+57
-9
@@ -18,7 +18,7 @@
|
||||
//! - `Esc` - Interrupt current operation
|
||||
|
||||
use std::borrow::Cow;
|
||||
use std::io::{self, Write};
|
||||
use std::io::{self, IsTerminal, Write};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
@@ -297,10 +297,15 @@ impl Channel for ReplChannel {
|
||||
let esc_interrupt_triggered_for_thread = Arc::new(AtomicBool::new(false));
|
||||
|
||||
std::thread::spawn(move || {
|
||||
let sys_tz = crate::timezone::detect_system_timezone().name().to_string();
|
||||
|
||||
// Single message mode: send it and return
|
||||
if let Some(msg) = single_message {
|
||||
let incoming = IncomingMessage::new("repl", "default", &msg);
|
||||
let incoming = IncomingMessage::new("repl", "default", &msg).with_timezone(&sys_tz);
|
||||
let _ = tx.blocking_send(incoming);
|
||||
// Ensure the agent exits after handling exactly one turn in -m mode,
|
||||
// even when other channels (gateway/http) are enabled.
|
||||
let _ = tx.blocking_send(IncomingMessage::new("repl", "default", "/quit"));
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -361,7 +366,8 @@ impl Channel for ReplChannel {
|
||||
"/quit" | "/exit" => {
|
||||
// Forward shutdown command so the agent loop exits even
|
||||
// when other channels (e.g. web gateway) are still active.
|
||||
let msg = IncomingMessage::new("repl", "default", "/quit");
|
||||
let msg = IncomingMessage::new("repl", "default", "/quit")
|
||||
.with_timezone(&sys_tz);
|
||||
let _ = tx.blocking_send(msg);
|
||||
break;
|
||||
}
|
||||
@@ -382,7 +388,8 @@ impl Channel for ReplChannel {
|
||||
_ => {}
|
||||
}
|
||||
|
||||
let msg = IncomingMessage::new("repl", "default", line);
|
||||
let msg =
|
||||
IncomingMessage::new("repl", "default", line).with_timezone(&sys_tz);
|
||||
if tx.blocking_send(msg).is_err() {
|
||||
break;
|
||||
}
|
||||
@@ -390,21 +397,29 @@ impl Channel for ReplChannel {
|
||||
Err(ReadlineError::Interrupted) => {
|
||||
if esc_interrupt_triggered_for_thread.swap(false, Ordering::Relaxed) {
|
||||
// Esc: interrupt current operation and keep REPL open.
|
||||
let msg = IncomingMessage::new("repl", "default", "/interrupt");
|
||||
let msg = IncomingMessage::new("repl", "default", "/interrupt")
|
||||
.with_timezone(&sys_tz);
|
||||
if tx.blocking_send(msg).is_err() {
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
// Ctrl+C (VINTR): request graceful shutdown.
|
||||
let msg = IncomingMessage::new("repl", "default", "/quit");
|
||||
let msg = IncomingMessage::new("repl", "default", "/quit")
|
||||
.with_timezone(&sys_tz);
|
||||
let _ = tx.blocking_send(msg);
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(ReadlineError::Eof) => {
|
||||
// Ctrl+D: send /quit so the agent loop runs graceful shutdown
|
||||
let msg = IncomingMessage::new("repl", "default", "/quit");
|
||||
let _ = tx.blocking_send(msg);
|
||||
// Ctrl+D in interactive mode: graceful shutdown.
|
||||
// In daemon mode (stdin = /dev/null, no TTY), EOF arrives
|
||||
// immediately — just drop the REPL thread silently so other
|
||||
// channels (gateway, telegram, …) keep running.
|
||||
if std::io::stdin().is_terminal() {
|
||||
let msg = IncomingMessage::new("repl", "default", "/quit")
|
||||
.with_timezone(&sys_tz);
|
||||
let _ = tx.blocking_send(msg);
|
||||
}
|
||||
break;
|
||||
}
|
||||
Err(e) => {
|
||||
@@ -585,6 +600,13 @@ impl Channel for ReplChannel {
|
||||
eprintln!("\x1b[31m {extension_name}: {message}\x1b[0m");
|
||||
}
|
||||
}
|
||||
StatusUpdate::ImageGenerated { path, .. } => {
|
||||
if let Some(ref p) = path {
|
||||
eprintln!("\x1b[36m [image] {p}\x1b[0m");
|
||||
} else {
|
||||
eprintln!("\x1b[36m [image generated]\x1b[0m");
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -614,3 +636,29 @@ impl Channel for ReplChannel {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use futures::StreamExt;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn single_message_mode_sends_message_then_quit() {
|
||||
let repl = ReplChannel::with_message("hi".to_string());
|
||||
let mut stream = repl.start().await.expect("repl start should succeed");
|
||||
|
||||
let first = stream.next().await.expect("first message missing");
|
||||
assert_eq!(first.channel, "repl");
|
||||
assert_eq!(first.content, "hi");
|
||||
|
||||
let second = stream.next().await.expect("quit message missing");
|
||||
assert_eq!(second.channel, "repl");
|
||||
assert_eq!(second.content, "/quit");
|
||||
|
||||
assert!(
|
||||
stream.next().await.is_none(),
|
||||
"stream should end after /quit"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2809,6 +2809,14 @@ fn status_to_wit(status: &StatusUpdate, metadata: &serde_json::Value) -> wit_cha
|
||||
),
|
||||
metadata_json,
|
||||
},
|
||||
StatusUpdate::ImageGenerated { path, .. } => wit_channel::StatusUpdate {
|
||||
status: wit_channel::StatusType::Status,
|
||||
message: match path {
|
||||
Some(p) => format!("[image] {}", p),
|
||||
None => "[image generated]".to_string(),
|
||||
},
|
||||
metadata_json,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -264,7 +264,7 @@ pub async fn routines_runs_handler(
|
||||
/// Convert a Routine to the trimmed RoutineInfo for list display.
|
||||
fn routine_to_info(r: &crate::agent::routine::Routine) -> RoutineInfo {
|
||||
let (trigger_type, trigger_summary) = match &r.trigger {
|
||||
crate::agent::routine::Trigger::Cron { schedule } => {
|
||||
crate::agent::routine::Trigger::Cron { schedule, .. } => {
|
||||
("cron".to_string(), format!("cron: {}", schedule))
|
||||
}
|
||||
crate::agent::routine::Trigger::Event {
|
||||
|
||||
@@ -386,6 +386,11 @@ impl Channel for GatewayChannel {
|
||||
success,
|
||||
message,
|
||||
},
|
||||
StatusUpdate::ImageGenerated { data_url, path } => SseEvent::ImageGenerated {
|
||||
data_url,
|
||||
path,
|
||||
thread_id,
|
||||
},
|
||||
};
|
||||
|
||||
self.state.sse.broadcast(event);
|
||||
|
||||
@@ -351,7 +351,7 @@ pub async fn start_server(
|
||||
.merge(statics)
|
||||
.merge(projects)
|
||||
.merge(protected)
|
||||
.layer(DefaultBodyLimit::max(1024 * 1024)) // 1 MB max request body
|
||||
.layer(DefaultBodyLimit::max(10 * 1024 * 1024)) // 10 MB max request body (image uploads)
|
||||
.layer(cors)
|
||||
.layer(SetResponseHeaderLayer::if_not_present(
|
||||
header::X_CONTENT_TYPE_OPTIONS,
|
||||
@@ -608,8 +608,59 @@ async fn oauth_callback_handler(
|
||||
|
||||
// --- Chat handlers ---
|
||||
|
||||
/// Convert web gateway `ImageData` to `IncomingAttachment` objects.
|
||||
pub(crate) fn images_to_attachments(
|
||||
images: &[ImageData],
|
||||
) -> Vec<crate::channels::IncomingAttachment> {
|
||||
use base64::Engine;
|
||||
images
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(i, img)| {
|
||||
if !img.media_type.starts_with("image/") {
|
||||
tracing::warn!(
|
||||
"Skipping image {i}: invalid media type '{}' (must start with 'image/')",
|
||||
img.media_type
|
||||
);
|
||||
return None;
|
||||
}
|
||||
let data = match base64::engine::general_purpose::STANDARD.decode(&img.data) {
|
||||
Ok(d) => d,
|
||||
Err(e) => {
|
||||
tracing::warn!("Skipping image {i}: invalid base64 data: {e}");
|
||||
return None;
|
||||
}
|
||||
};
|
||||
Some(crate::channels::IncomingAttachment {
|
||||
id: format!("web-image-{i}"),
|
||||
kind: crate::channels::AttachmentKind::Image,
|
||||
mime_type: img.media_type.clone(),
|
||||
filename: Some(format!("image-{i}.{}", mime_to_ext(&img.media_type))),
|
||||
size_bytes: Some(data.len() as u64),
|
||||
source_url: None,
|
||||
storage_key: None,
|
||||
extracted_text: None,
|
||||
data,
|
||||
duration_secs: None,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Map MIME type to file extension.
|
||||
fn mime_to_ext(mime: &str) -> &str {
|
||||
match mime {
|
||||
"image/png" => "png",
|
||||
"image/gif" => "gif",
|
||||
"image/webp" => "webp",
|
||||
"image/svg+xml" => "svg",
|
||||
_ => "jpg",
|
||||
}
|
||||
}
|
||||
|
||||
async fn chat_send_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
headers: axum::http::HeaderMap,
|
||||
Json(req): Json<SendMessageRequest>,
|
||||
) -> Result<(StatusCode, Json<SendMessageResponse>), (StatusCode, String)> {
|
||||
tracing::debug!(
|
||||
@@ -626,17 +677,32 @@ async fn chat_send_handler(
|
||||
}
|
||||
|
||||
let mut msg = IncomingMessage::new("gateway", &state.user_id, &req.content);
|
||||
// Prefer timezone from JSON body, fall back to X-Timezone header
|
||||
let tz = req
|
||||
.timezone
|
||||
.as_deref()
|
||||
.or_else(|| headers.get("X-Timezone").and_then(|v| v.to_str().ok()));
|
||||
if let Some(tz) = tz {
|
||||
msg = msg.with_timezone(tz);
|
||||
}
|
||||
|
||||
if let Some(ref thread_id) = req.thread_id {
|
||||
msg = msg.with_thread(thread_id);
|
||||
msg = msg.with_metadata(serde_json::json!({"thread_id": thread_id}));
|
||||
}
|
||||
|
||||
// Convert uploaded images to IncomingAttachments
|
||||
if !req.images.is_empty() {
|
||||
let attachments = images_to_attachments(&req.images);
|
||||
msg = msg.with_attachments(attachments);
|
||||
}
|
||||
|
||||
let msg_id = msg.id;
|
||||
tracing::debug!(
|
||||
"[chat_send_handler] Created message id={}, content={:?}",
|
||||
"[chat_send_handler] Created message id={}, content={:?}, images={}",
|
||||
msg_id,
|
||||
req.content
|
||||
req.content,
|
||||
req.images.len()
|
||||
);
|
||||
|
||||
let tx_guard = state.msg_tx.read().await;
|
||||
@@ -2115,7 +2181,7 @@ async fn routines_runs_handler(
|
||||
/// Convert a Routine to the trimmed RoutineInfo for list display.
|
||||
fn routine_to_info(r: &crate::agent::routine::Routine) -> RoutineInfo {
|
||||
let (trigger_type, trigger_summary) = match &r.trigger {
|
||||
crate::agent::routine::Trigger::Cron { schedule } => {
|
||||
crate::agent::routine::Trigger::Cron { schedule, .. } => {
|
||||
("cron".to_string(), format!("cron: {}", schedule))
|
||||
}
|
||||
crate::agent::routine::Trigger::Event {
|
||||
|
||||
@@ -142,6 +142,7 @@ impl SseManager {
|
||||
SseEvent::JobStatus { .. } => "job_status",
|
||||
SseEvent::JobResult { .. } => "job_result",
|
||||
SseEvent::Heartbeat => "heartbeat",
|
||||
SseEvent::ImageGenerated { .. } => "image_generated",
|
||||
SseEvent::ExtensionStatus { .. } => "extension_status",
|
||||
};
|
||||
Ok(Event::default().event(event_type).data(data))
|
||||
|
||||
+150
-32
@@ -18,6 +18,7 @@ let unreadThreads = new Map(); // thread_id -> unread count
|
||||
let _loadThreadsTimer = null;
|
||||
const JOB_EVENTS_CAP = 500;
|
||||
const MEMORY_SEARCH_QUERY_MAX_LENGTH = 100;
|
||||
let stagedImages = [];
|
||||
|
||||
// --- Slash Commands ---
|
||||
|
||||
@@ -181,6 +182,7 @@ function confirmRestart() {
|
||||
body: {
|
||||
content: '/restart',
|
||||
thread_id: currentThreadId,
|
||||
timezone: Intl.DateTimeFormat().resolvedOptions().timeZone,
|
||||
},
|
||||
})
|
||||
.then((response) => {
|
||||
@@ -223,23 +225,6 @@ function updateRestartButtonVisibility() {
|
||||
}
|
||||
}
|
||||
|
||||
function startGatewayStatusPolling() {
|
||||
fetchGatewayStatus();
|
||||
// Poll every 5 seconds
|
||||
setInterval(fetchGatewayStatus, 5000);
|
||||
}
|
||||
|
||||
function fetchGatewayStatus() {
|
||||
apiFetch('/api/gateway/status')
|
||||
.then((data) => {
|
||||
restartEnabled = data.restart_enabled || false;
|
||||
updateRestartButtonVisibility();
|
||||
})
|
||||
.catch((err) => {
|
||||
console.warn('[gateway status] Failed to fetch:', err);
|
||||
});
|
||||
}
|
||||
|
||||
// --- SSE ---
|
||||
|
||||
function connectSSE() {
|
||||
@@ -388,6 +373,12 @@ function connectSSE() {
|
||||
if (currentTab === 'extensions') loadExtensions();
|
||||
});
|
||||
|
||||
eventSource.addEventListener('image_generated', (e) => {
|
||||
const data = JSON.parse(e.data);
|
||||
if (!isCurrentThread(data.thread_id)) return;
|
||||
addGeneratedImage(data.data_url, data.path);
|
||||
});
|
||||
|
||||
eventSource.addEventListener('error', (e) => {
|
||||
if (e.data) {
|
||||
const data = JSON.parse(e.data);
|
||||
@@ -445,16 +436,23 @@ function sendMessage() {
|
||||
return;
|
||||
}
|
||||
const content = input.value.trim();
|
||||
if (!content) return;
|
||||
if (!content && stagedImages.length === 0) return;
|
||||
|
||||
addMessage('user', content);
|
||||
addMessage('user', content || '(images attached)');
|
||||
input.value = '';
|
||||
autoResizeTextarea(input);
|
||||
input.focus();
|
||||
|
||||
const body = { content, thread_id: currentThreadId || undefined, timezone: Intl.DateTimeFormat().resolvedOptions().timeZone };
|
||||
if (stagedImages.length > 0) {
|
||||
body.images = stagedImages.map(img => ({ media_type: img.media_type, data: img.data }));
|
||||
stagedImages = [];
|
||||
renderImagePreviews();
|
||||
}
|
||||
|
||||
apiFetch('/api/chat/send', {
|
||||
method: 'POST',
|
||||
body: { content, thread_id: currentThreadId || undefined },
|
||||
body: body,
|
||||
}).catch((err) => {
|
||||
addMessage('system', 'Failed to send: ' + err.message);
|
||||
});
|
||||
@@ -471,6 +469,104 @@ function enableChatInput() {
|
||||
if (btn) btn.disabled = false;
|
||||
}
|
||||
|
||||
// --- Image Upload ---
|
||||
|
||||
function renderImagePreviews() {
|
||||
const strip = document.getElementById('image-preview-strip');
|
||||
strip.innerHTML = '';
|
||||
stagedImages.forEach((img, idx) => {
|
||||
const container = document.createElement('div');
|
||||
container.className = 'image-preview-container';
|
||||
|
||||
const preview = document.createElement('img');
|
||||
preview.className = 'image-preview';
|
||||
preview.src = img.dataUrl;
|
||||
preview.alt = 'Attached image';
|
||||
|
||||
const removeBtn = document.createElement('button');
|
||||
removeBtn.className = 'image-preview-remove';
|
||||
removeBtn.textContent = '\u00d7';
|
||||
removeBtn.addEventListener('click', () => {
|
||||
stagedImages.splice(idx, 1);
|
||||
renderImagePreviews();
|
||||
});
|
||||
|
||||
container.appendChild(preview);
|
||||
container.appendChild(removeBtn);
|
||||
strip.appendChild(container);
|
||||
});
|
||||
}
|
||||
|
||||
const MAX_IMAGE_SIZE_BYTES = 5 * 1024 * 1024; // 5 MB per image
|
||||
const MAX_STAGED_IMAGES = 5;
|
||||
|
||||
function handleImageFiles(files) {
|
||||
Array.from(files).forEach(file => {
|
||||
if (!file.type.startsWith('image/')) return;
|
||||
if (file.size > MAX_IMAGE_SIZE_BYTES) {
|
||||
alert(`Image "${file.name}" exceeds 5 MB limit (${(file.size / 1024 / 1024).toFixed(1)} MB)`);
|
||||
return;
|
||||
}
|
||||
if (stagedImages.length >= MAX_STAGED_IMAGES) {
|
||||
alert(`Maximum ${MAX_STAGED_IMAGES} images allowed per message`);
|
||||
return;
|
||||
}
|
||||
const reader = new FileReader();
|
||||
reader.onload = function(e) {
|
||||
const dataUrl = e.target.result;
|
||||
const commaIdx = dataUrl.indexOf(',');
|
||||
const meta = dataUrl.substring(0, commaIdx); // e.g. "data:image/png;base64"
|
||||
const base64 = dataUrl.substring(commaIdx + 1);
|
||||
const mediaType = meta.replace('data:', '').replace(';base64', '');
|
||||
stagedImages.push({ media_type: mediaType, data: base64, dataUrl: dataUrl });
|
||||
renderImagePreviews();
|
||||
};
|
||||
reader.readAsDataURL(file);
|
||||
});
|
||||
}
|
||||
|
||||
document.getElementById('attach-btn').addEventListener('click', () => {
|
||||
document.getElementById('image-file-input').click();
|
||||
});
|
||||
|
||||
document.getElementById('image-file-input').addEventListener('change', (e) => {
|
||||
handleImageFiles(e.target.files);
|
||||
e.target.value = '';
|
||||
});
|
||||
|
||||
document.getElementById('chat-input').addEventListener('paste', (e) => {
|
||||
const items = (e.clipboardData || e.originalEvent.clipboardData).items;
|
||||
for (let i = 0; i < items.length; i++) {
|
||||
if (items[i].kind === 'file' && items[i].type.startsWith('image/')) {
|
||||
const file = items[i].getAsFile();
|
||||
if (file) handleImageFiles([file]);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
function addGeneratedImage(dataUrl, path) {
|
||||
const container = document.getElementById('chat-messages');
|
||||
const card = document.createElement('div');
|
||||
card.className = 'generated-image-card';
|
||||
|
||||
const img = document.createElement('img');
|
||||
img.className = 'generated-image';
|
||||
img.src = dataUrl;
|
||||
img.alt = 'Generated image';
|
||||
|
||||
card.appendChild(img);
|
||||
|
||||
if (path) {
|
||||
const pathLabel = document.createElement('div');
|
||||
pathLabel.className = 'generated-image-path';
|
||||
pathLabel.textContent = path;
|
||||
card.appendChild(pathLabel);
|
||||
}
|
||||
|
||||
container.appendChild(card);
|
||||
container.scrollTop = container.scrollHeight;
|
||||
}
|
||||
|
||||
// --- Slash Autocomplete ---
|
||||
|
||||
function showSlashAutocomplete(matches) {
|
||||
@@ -563,6 +659,13 @@ function sendApprovalAction(requestId, action) {
|
||||
|
||||
function renderMarkdown(text) {
|
||||
if (typeof marked !== 'undefined') {
|
||||
// Escape raw HTML error pages instead of rendering them as markup.
|
||||
// Only triggers when the text *starts with* a doctype or <html> tag
|
||||
// (after optional whitespace), so normal messages that mention HTML
|
||||
// tags in prose or code fences are not affected. See #263.
|
||||
if (/^\s*<!doctype\s/i.test(text) || /^\s*<html[\s>]/i.test(text)) {
|
||||
return escapeHtml(text);
|
||||
}
|
||||
let html = marked.parse(text);
|
||||
// Sanitize HTML output to prevent XSS from tool output or LLM responses.
|
||||
html = sanitizeRenderedHtml(html);
|
||||
@@ -1473,7 +1576,7 @@ chatInput.addEventListener('keydown', (e) => {
|
||||
}
|
||||
}
|
||||
|
||||
if (e.key === 'Enter' && !e.shiftKey) {
|
||||
if (e.key === 'Enter' && !e.shiftKey && !e.isComposing) {
|
||||
e.preventDefault();
|
||||
hideSlashAutocomplete();
|
||||
sendMessage();
|
||||
@@ -3408,6 +3511,10 @@ function shortModelName(model) {
|
||||
|
||||
function fetchGatewayStatus() {
|
||||
apiFetch('/api/gateway/status').then(function(data) {
|
||||
// Update restart button visibility
|
||||
restartEnabled = data.restart_enabled || false;
|
||||
updateRestartButtonVisibility();
|
||||
|
||||
var popover = document.getElementById('gateway-popover');
|
||||
var html = '';
|
||||
|
||||
@@ -3473,10 +3580,15 @@ let teeReportCache = null;
|
||||
let teeReportLoading = false;
|
||||
|
||||
function teeApiBase() {
|
||||
var parts = window.location.hostname.split('.');
|
||||
if (parts.length < 2) return null;
|
||||
var domain = parts.slice(1).join('.');
|
||||
return window.location.protocol + '//api.' + domain;
|
||||
var hostname = window.location.hostname;
|
||||
// Skip IP addresses (IPv4 and IPv6) and localhost
|
||||
if (hostname === "localhost" || /^(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$/.test(hostname) || hostname.indexOf(":") !== -1) {
|
||||
return null;
|
||||
}
|
||||
var parts = hostname.split(".");
|
||||
if (parts.length < 2) return null;
|
||||
var domain = parts.slice(1).join(".");
|
||||
return window.location.protocol + "//api." + domain;
|
||||
}
|
||||
|
||||
function teeInstanceName() {
|
||||
@@ -3487,13 +3599,19 @@ function checkTeeStatus() {
|
||||
var base = teeApiBase();
|
||||
if (!base) return;
|
||||
var name = teeInstanceName();
|
||||
fetch(base + '/instances/' + encodeURIComponent(name) + '/attestation').then(function(res) {
|
||||
if (!res.ok) throw new Error(res.status);
|
||||
return res.json();
|
||||
}).then(function(data) {
|
||||
teeInfo = data;
|
||||
document.getElementById('tee-shield').style.display = 'flex';
|
||||
}).catch(function() {});
|
||||
try {
|
||||
fetch(base + '/instances/' + encodeURIComponent(name) + '/attestation').then(function(res) {
|
||||
if (!res.ok) throw new Error(res.status);
|
||||
return res.json();
|
||||
}).then(function(data) {
|
||||
teeInfo = data;
|
||||
document.getElementById('tee-shield').style.display = 'flex';
|
||||
}).catch(function(err) {
|
||||
console.warn('Failed to fetch TEE attestation:', err);
|
||||
});
|
||||
} catch (e) {
|
||||
console.warn("Failed to check TEE status:", e);
|
||||
}
|
||||
}
|
||||
|
||||
function fetchTeeReport() {
|
||||
|
||||
@@ -99,7 +99,7 @@
|
||||
<span id="sse-status">Connected</span>
|
||||
<div class="gateway-popover" id="gateway-popover"></div>
|
||||
</div>
|
||||
<button class="restart-btn" id="restart-btn" onclick="triggerRestart()" title="Gracefully restart the process">
|
||||
<button class="restart-btn" id="restart-btn" onclick="triggerRestart()" title="Gracefully restart the process" style="display: none;">
|
||||
<svg id="restart-icon" width="13" height="13" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
|
||||
<path d="M23 4v6h-6"></path>
|
||||
<path d="M1 20v-6h6"></path>
|
||||
@@ -130,7 +130,10 @@
|
||||
<div class="chat-messages" id="chat-messages"></div>
|
||||
<div id="slash-autocomplete" class="slash-autocomplete" style="display:none"></div>
|
||||
<div class="chat-input">
|
||||
<div id="image-preview-strip" class="image-preview-strip"></div>
|
||||
<textarea id="chat-input" placeholder="Message or / for commands..." rows="1"></textarea>
|
||||
<input type="file" id="image-file-input" accept="image/*" multiple style="display:none">
|
||||
<button id="attach-btn" class="attach-btn" title="Attach images" aria-label="Attach images">📎</button>
|
||||
<button id="send-btn" onclick="sendMessage()">Send</button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
@@ -159,7 +159,7 @@ body {
|
||||
flex-shrink: 0;
|
||||
}
|
||||
|
||||
.tab-bar button:not(.status-logs-btn) {
|
||||
.tab-bar button:not(.status-logs-btn):not(.restart-btn) {
|
||||
padding: 10px 20px;
|
||||
background: none;
|
||||
border: none;
|
||||
@@ -171,11 +171,11 @@ body {
|
||||
transition: color 0.2s, border-color 0.2s;
|
||||
}
|
||||
|
||||
.tab-bar button:not(.status-logs-btn):hover {
|
||||
.tab-bar button:not(.status-logs-btn):not(.restart-btn):hover {
|
||||
color: var(--text);
|
||||
}
|
||||
|
||||
.tab-bar button:not(.status-logs-btn).active {
|
||||
.tab-bar button:not(.status-logs-btn):not(.restart-btn).active {
|
||||
color: var(--accent);
|
||||
border-bottom-color: var(--accent);
|
||||
}
|
||||
@@ -260,42 +260,42 @@ body {
|
||||
}
|
||||
|
||||
/* Restart Button */
|
||||
.restart-btn {
|
||||
.tab-bar .restart-btn {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 0.375rem;
|
||||
margin: 0.375rem;
|
||||
padding: 0.25rem 0.75rem;
|
||||
border-radius: 0.5rem;
|
||||
font-size: 0.8rem;
|
||||
border: 1px solid;
|
||||
border-color: #00d894;
|
||||
border: 1px solid #00d894;
|
||||
color: #00d894;
|
||||
background-color: transparent;
|
||||
cursor: pointer;
|
||||
transition: color 150ms, background-color 150ms, border-color 150ms;
|
||||
}
|
||||
|
||||
.restart-btn:hover:not(:disabled) {
|
||||
.tab-bar .restart-btn:hover:not(:disabled) {
|
||||
background-color: rgba(0, 216, 148, 0.1);
|
||||
}
|
||||
|
||||
.restart-btn:disabled {
|
||||
.tab-bar .restart-btn:disabled {
|
||||
border-color: #333;
|
||||
color: #666;
|
||||
cursor: not-allowed;
|
||||
}
|
||||
|
||||
.restart-btn:disabled:hover {
|
||||
.tab-bar .restart-btn:disabled:hover {
|
||||
background-color: transparent;
|
||||
}
|
||||
|
||||
.restart-btn svg {
|
||||
.tab-bar .restart-btn svg {
|
||||
flex-shrink: 0;
|
||||
width: 13px;
|
||||
height: 13px;
|
||||
}
|
||||
|
||||
.restart-btn svg.spinning {
|
||||
.tab-bar .restart-btn svg.spinning {
|
||||
animation: spin-icon 1s linear infinite;
|
||||
}
|
||||
|
||||
@@ -1272,6 +1272,7 @@ body {
|
||||
/* Chat input */
|
||||
.chat-input {
|
||||
display: flex;
|
||||
flex-wrap: wrap;
|
||||
padding: 12px 16px max(12px, env(safe-area-inset-bottom)) 16px;
|
||||
gap: 8px;
|
||||
background: var(--bg-secondary);
|
||||
@@ -3761,3 +3762,93 @@ mark {
|
||||
text-overflow: ellipsis;
|
||||
white-space: nowrap;
|
||||
}
|
||||
|
||||
/* Image Upload */
|
||||
.attach-btn {
|
||||
background: none;
|
||||
border: none;
|
||||
cursor: pointer;
|
||||
font-size: 1.2em;
|
||||
padding: 8px;
|
||||
align-self: flex-end;
|
||||
color: var(--text-secondary);
|
||||
transition: color 0.2s;
|
||||
min-height: 40px;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
}
|
||||
|
||||
.attach-btn:hover {
|
||||
color: var(--text);
|
||||
}
|
||||
|
||||
.image-preview-strip {
|
||||
display: flex;
|
||||
flex-direction: row;
|
||||
gap: 8px;
|
||||
padding: 4px;
|
||||
overflow-x: auto;
|
||||
min-height: 0;
|
||||
width: 100%;
|
||||
}
|
||||
|
||||
.image-preview-strip:empty {
|
||||
display: none;
|
||||
}
|
||||
|
||||
.image-preview-container {
|
||||
position: relative;
|
||||
display: inline-block;
|
||||
flex-shrink: 0;
|
||||
}
|
||||
|
||||
.image-preview {
|
||||
width: 60px;
|
||||
height: 60px;
|
||||
border-radius: 6px;
|
||||
object-fit: cover;
|
||||
display: block;
|
||||
}
|
||||
|
||||
.image-preview-remove {
|
||||
position: absolute;
|
||||
top: -6px;
|
||||
right: -6px;
|
||||
width: 18px;
|
||||
height: 18px;
|
||||
border-radius: 50%;
|
||||
background: var(--danger);
|
||||
color: #fff;
|
||||
border: none;
|
||||
font-size: 12px;
|
||||
line-height: 18px;
|
||||
text-align: center;
|
||||
cursor: pointer;
|
||||
padding: 0;
|
||||
}
|
||||
|
||||
.image-preview-remove:hover {
|
||||
background: #c33;
|
||||
}
|
||||
|
||||
/* Generated Image */
|
||||
.generated-image-card {
|
||||
max-width: 512px;
|
||||
margin: 8px 0;
|
||||
border-radius: 8px;
|
||||
overflow: hidden;
|
||||
border: 1px solid var(--border);
|
||||
}
|
||||
|
||||
.generated-image {
|
||||
max-width: 100%;
|
||||
display: block;
|
||||
}
|
||||
|
||||
.generated-image-path {
|
||||
font-size: 12px;
|
||||
color: var(--text-secondary);
|
||||
padding: 4px 8px;
|
||||
background: var(--bg-secondary);
|
||||
}
|
||||
|
||||
@@ -5,10 +5,23 @@ use uuid::Uuid;
|
||||
|
||||
// --- Chat ---
|
||||
|
||||
/// Base64-encoded image data sent from the web frontend.
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct ImageData {
|
||||
/// MIME type (e.g., "image/png", "image/jpeg").
|
||||
pub media_type: String,
|
||||
/// Base64-encoded image data (without data: URL prefix).
|
||||
pub data: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct SendMessageRequest {
|
||||
pub content: String,
|
||||
pub thread_id: Option<String>,
|
||||
pub timezone: Option<String>,
|
||||
/// Optional images attached to the message.
|
||||
#[serde(default)]
|
||||
pub images: Vec<ImageData>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
@@ -219,6 +232,16 @@ pub enum SseEvent {
|
||||
session_id: Option<String>,
|
||||
},
|
||||
|
||||
/// An image was generated by a tool.
|
||||
#[serde(rename = "image_generated")]
|
||||
ImageGenerated {
|
||||
data_url: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
path: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
thread_id: Option<String>,
|
||||
},
|
||||
|
||||
/// Extension activation status change (WASM channels).
|
||||
#[serde(rename = "extension_status")]
|
||||
ExtensionStatus {
|
||||
@@ -613,6 +636,10 @@ pub enum WsClientMessage {
|
||||
Message {
|
||||
content: String,
|
||||
thread_id: Option<String>,
|
||||
timezone: Option<String>,
|
||||
/// Optional images attached to the message.
|
||||
#[serde(default)]
|
||||
images: Vec<ImageData>,
|
||||
},
|
||||
/// Approve or deny a pending tool execution.
|
||||
#[serde(rename = "approval")]
|
||||
@@ -679,6 +706,7 @@ impl WsServerMessage {
|
||||
SseEvent::JobToolResult { .. } => "job_tool_result",
|
||||
SseEvent::JobStatus { .. } => "job_status",
|
||||
SseEvent::JobResult { .. } => "job_result",
|
||||
SseEvent::ImageGenerated { .. } => "image_generated",
|
||||
SseEvent::ExtensionStatus { .. } => "extension_status",
|
||||
};
|
||||
let data = serde_json::to_value(event).unwrap_or(serde_json::Value::Null);
|
||||
@@ -798,7 +826,9 @@ mod tests {
|
||||
let json = r#"{"type":"message","content":"hello","thread_id":"t1"}"#;
|
||||
let msg: WsClientMessage = serde_json::from_str(json).unwrap();
|
||||
match msg {
|
||||
WsClientMessage::Message { content, thread_id } => {
|
||||
WsClientMessage::Message {
|
||||
content, thread_id, ..
|
||||
} => {
|
||||
assert_eq!(content, "hello");
|
||||
assert_eq!(thread_id.as_deref(), Some("t1"));
|
||||
}
|
||||
@@ -811,7 +841,9 @@ mod tests {
|
||||
let json = r#"{"type":"message","content":"hi"}"#;
|
||||
let msg: WsClientMessage = serde_json::from_str(json).unwrap();
|
||||
match msg {
|
||||
WsClientMessage::Message { content, thread_id } => {
|
||||
WsClientMessage::Message {
|
||||
content, thread_id, ..
|
||||
} => {
|
||||
assert_eq!(content, "hi");
|
||||
assert!(thread_id.is_none());
|
||||
}
|
||||
|
||||
@@ -3,6 +3,10 @@
|
||||
use crate::channels::web::types::{ToolCallInfo, TurnInfo};
|
||||
|
||||
/// Truncate a string to at most `max_bytes` bytes at a char boundary, appending "...".
|
||||
///
|
||||
/// If the input is wrapped in `<tool_output …>…</tool_output>` and truncation
|
||||
/// removes the closing tag, the tag is re-appended so downstream XML parsers
|
||||
/// never see an unclosed element.
|
||||
pub fn truncate_preview(s: &str, max_bytes: usize) -> String {
|
||||
if s.len() <= max_bytes {
|
||||
return s.to_string();
|
||||
@@ -12,7 +16,14 @@ pub fn truncate_preview(s: &str, max_bytes: usize) -> String {
|
||||
while end > 0 && !s.is_char_boundary(end) {
|
||||
end -= 1;
|
||||
}
|
||||
format!("{}...", &s[..end])
|
||||
let mut result = format!("{}...", &s[..end]);
|
||||
|
||||
// Re-close <tool_output> if truncation cut through the closing tag.
|
||||
if s.starts_with("<tool_output") && !result.ends_with("</tool_output>") {
|
||||
result.push_str("\n</tool_output>");
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
/// Build TurnInfo pairs from flat DB messages (user/tool_calls/assistant triples).
|
||||
@@ -162,6 +173,33 @@ mod tests {
|
||||
assert_eq!(truncate_preview("hello", 0), "...");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_preview_closes_tool_output_tag() {
|
||||
let s = "<tool_output name=\"search\" sanitized=\"true\">\nSome very long content here\n</tool_output>";
|
||||
// Truncate so it cuts before the closing tag
|
||||
let result = truncate_preview(s, 60);
|
||||
assert!(result.ends_with("</tool_output>"));
|
||||
assert!(result.contains("..."));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_preview_no_extra_close_when_intact() {
|
||||
let s = "<tool_output name=\"echo\" sanitized=\"false\">\nshort\n</tool_output>";
|
||||
// The string is short enough not to be truncated
|
||||
let result = truncate_preview(s, 500);
|
||||
assert_eq!(result, s);
|
||||
// Should not have a duplicate closing tag
|
||||
assert_eq!(result.matches("</tool_output>").count(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_preview_non_xml_unaffected() {
|
||||
let s = "Just a plain long string that gets truncated";
|
||||
let result = truncate_preview(s, 10);
|
||||
assert_eq!(result, "Just a pla...");
|
||||
assert!(!result.contains("</tool_output>"));
|
||||
}
|
||||
|
||||
// ---- build_turns_from_db_messages tests ----
|
||||
|
||||
fn make_msg(role: &str, content: &str, offset_ms: i64) -> crate::history::ConversationMessage {
|
||||
|
||||
+19
-1
@@ -156,12 +156,26 @@ async fn handle_client_message(
|
||||
direct_tx: &mpsc::Sender<WsServerMessage>,
|
||||
) {
|
||||
match msg {
|
||||
WsClientMessage::Message { content, thread_id } => {
|
||||
WsClientMessage::Message {
|
||||
content,
|
||||
thread_id,
|
||||
timezone,
|
||||
images,
|
||||
} => {
|
||||
let mut incoming = IncomingMessage::new("gateway", user_id, &content);
|
||||
if let Some(ref tz) = timezone {
|
||||
incoming = incoming.with_timezone(tz);
|
||||
}
|
||||
if let Some(ref tid) = thread_id {
|
||||
incoming = incoming.with_thread(tid);
|
||||
}
|
||||
|
||||
// Convert uploaded images to IncomingAttachments
|
||||
if !images.is_empty() {
|
||||
let attachments = crate::channels::web::server::images_to_attachments(&images);
|
||||
incoming = incoming.with_attachments(attachments);
|
||||
}
|
||||
|
||||
let tx_guard = state.msg_tx.read().await;
|
||||
if let Some(ref tx) = *tx_guard {
|
||||
if tx.send(incoming).await.is_err() {
|
||||
@@ -349,6 +363,8 @@ mod tests {
|
||||
WsClientMessage::Message {
|
||||
content: "hello agent".to_string(),
|
||||
thread_id: Some("t1".to_string()),
|
||||
timezone: None,
|
||||
images: Vec::new(),
|
||||
},
|
||||
&state,
|
||||
"user1",
|
||||
@@ -373,6 +389,8 @@ mod tests {
|
||||
WsClientMessage::Message {
|
||||
content: "hello".to_string(),
|
||||
thread_id: None,
|
||||
timezone: None,
|
||||
images: Vec::new(),
|
||||
},
|
||||
&state,
|
||||
"user1",
|
||||
|
||||
+243
-121
@@ -2,52 +2,79 @@
|
||||
//!
|
||||
//! Commands for adding, removing, authenticating, and testing MCP servers.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::io::Write;
|
||||
use std::sync::Arc;
|
||||
|
||||
use clap::Subcommand;
|
||||
use clap::{Args, Subcommand};
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::db::Database;
|
||||
#[cfg(feature = "postgres")]
|
||||
use crate::secrets::PostgresSecretsStore;
|
||||
use crate::secrets::{SecretsCrypto, SecretsStore};
|
||||
use crate::tools::mcp::{
|
||||
McpClient, McpServerConfig, McpSessionManager, OAuthConfig,
|
||||
auth::{authorize_mcp_server, is_authenticated},
|
||||
config::{self, McpServersFile},
|
||||
config::{self, EffectiveTransport, McpServersFile},
|
||||
};
|
||||
|
||||
/// Arguments for the `mcp add` subcommand.
|
||||
#[derive(Args, Debug, Clone)]
|
||||
pub struct McpAddArgs {
|
||||
/// Server name (e.g., "notion", "github")
|
||||
pub name: String,
|
||||
|
||||
/// Server URL (e.g., "https://mcp.notion.com") -- required for http transport
|
||||
pub url: Option<String>,
|
||||
|
||||
/// Transport type: http (default), stdio, unix
|
||||
#[arg(long, default_value = "http")]
|
||||
pub transport: String,
|
||||
|
||||
/// Command to run (stdio transport)
|
||||
#[arg(long)]
|
||||
pub command: Option<String>,
|
||||
|
||||
/// Command arguments (stdio transport, can be repeated)
|
||||
#[arg(long = "arg", num_args = 1..)]
|
||||
pub cmd_args: Vec<String>,
|
||||
|
||||
/// Environment variables (stdio transport, KEY=VALUE format, can be repeated)
|
||||
#[arg(long = "env", value_parser = parse_env_var)]
|
||||
pub env: Vec<(String, String)>,
|
||||
|
||||
/// Unix socket path (unix transport)
|
||||
#[arg(long)]
|
||||
pub socket: Option<String>,
|
||||
|
||||
/// Custom HTTP headers (KEY:VALUE format, can be repeated)
|
||||
#[arg(long = "header", value_parser = parse_header)]
|
||||
pub headers: Vec<(String, String)>,
|
||||
|
||||
/// OAuth client ID (if authentication is required)
|
||||
#[arg(long)]
|
||||
pub client_id: Option<String>,
|
||||
|
||||
/// OAuth authorization URL (optional, can be discovered)
|
||||
#[arg(long)]
|
||||
pub auth_url: Option<String>,
|
||||
|
||||
/// OAuth token URL (optional, can be discovered)
|
||||
#[arg(long)]
|
||||
pub token_url: Option<String>,
|
||||
|
||||
/// Scopes to request (comma-separated)
|
||||
#[arg(long)]
|
||||
pub scopes: Option<String>,
|
||||
|
||||
/// Server description
|
||||
#[arg(long)]
|
||||
pub description: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Subcommand, Debug, Clone)]
|
||||
pub enum McpCommand {
|
||||
/// Add an MCP server
|
||||
Add {
|
||||
/// Server name (e.g., "notion", "github")
|
||||
name: String,
|
||||
|
||||
/// Server URL (e.g., "https://mcp.notion.com")
|
||||
url: String,
|
||||
|
||||
/// OAuth client ID (if authentication is required)
|
||||
#[arg(long)]
|
||||
client_id: Option<String>,
|
||||
|
||||
/// OAuth authorization URL (optional, can be discovered)
|
||||
#[arg(long)]
|
||||
auth_url: Option<String>,
|
||||
|
||||
/// OAuth token URL (optional, can be discovered)
|
||||
#[arg(long)]
|
||||
token_url: Option<String>,
|
||||
|
||||
/// Scopes to request (comma-separated)
|
||||
#[arg(long)]
|
||||
scopes: Option<String>,
|
||||
|
||||
/// Server description
|
||||
#[arg(long)]
|
||||
description: Option<String>,
|
||||
},
|
||||
Add(Box<McpAddArgs>),
|
||||
|
||||
/// Remove an MCP server
|
||||
Remove {
|
||||
@@ -97,29 +124,24 @@ pub enum McpCommand {
|
||||
},
|
||||
}
|
||||
|
||||
fn parse_header(s: &str) -> Result<(String, String), String> {
|
||||
let pos = s
|
||||
.find(':')
|
||||
.ok_or_else(|| format!("invalid header format '{}', expected KEY:VALUE", s))?;
|
||||
Ok((s[..pos].trim().to_string(), s[pos + 1..].trim().to_string()))
|
||||
}
|
||||
|
||||
fn parse_env_var(s: &str) -> Result<(String, String), String> {
|
||||
let pos = s
|
||||
.find('=')
|
||||
.ok_or_else(|| format!("invalid env var format '{}', expected KEY=VALUE", s))?;
|
||||
Ok((s[..pos].to_string(), s[pos + 1..].to_string()))
|
||||
}
|
||||
|
||||
/// Run an MCP command.
|
||||
pub async fn run_mcp_command(cmd: McpCommand) -> anyhow::Result<()> {
|
||||
match cmd {
|
||||
McpCommand::Add {
|
||||
name,
|
||||
url,
|
||||
client_id,
|
||||
auth_url,
|
||||
token_url,
|
||||
scopes,
|
||||
description,
|
||||
} => {
|
||||
add_server(
|
||||
name,
|
||||
url,
|
||||
client_id,
|
||||
auth_url,
|
||||
token_url,
|
||||
scopes,
|
||||
description,
|
||||
)
|
||||
.await
|
||||
}
|
||||
McpCommand::Add(args) => add_server(*args).await,
|
||||
McpCommand::Remove { name } => remove_server(name).await,
|
||||
McpCommand::List { verbose } => list_servers(verbose).await,
|
||||
McpCommand::Auth { name, user } => auth_server(name, user).await,
|
||||
@@ -133,16 +155,58 @@ pub async fn run_mcp_command(cmd: McpCommand) -> anyhow::Result<()> {
|
||||
}
|
||||
|
||||
/// Add a new MCP server.
|
||||
async fn add_server(
|
||||
name: String,
|
||||
url: String,
|
||||
client_id: Option<String>,
|
||||
auth_url: Option<String>,
|
||||
token_url: Option<String>,
|
||||
scopes: Option<String>,
|
||||
description: Option<String>,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut config = McpServerConfig::new(&name, &url);
|
||||
async fn add_server(args: McpAddArgs) -> anyhow::Result<()> {
|
||||
let McpAddArgs {
|
||||
name,
|
||||
url,
|
||||
transport,
|
||||
command,
|
||||
cmd_args,
|
||||
env,
|
||||
socket,
|
||||
headers,
|
||||
client_id,
|
||||
auth_url,
|
||||
token_url,
|
||||
scopes,
|
||||
description,
|
||||
} = args;
|
||||
|
||||
let transport_lower = transport.to_lowercase();
|
||||
|
||||
let mut config = match transport_lower.as_str() {
|
||||
"stdio" => {
|
||||
let cmd = command
|
||||
.clone()
|
||||
.ok_or_else(|| anyhow::anyhow!("--command is required for stdio transport"))?;
|
||||
let env_map: HashMap<String, String> = env.into_iter().collect();
|
||||
McpServerConfig::new_stdio(&name, &cmd, cmd_args.clone(), env_map)
|
||||
}
|
||||
"unix" => {
|
||||
let socket_path = socket
|
||||
.clone()
|
||||
.ok_or_else(|| anyhow::anyhow!("--socket is required for unix transport"))?;
|
||||
McpServerConfig::new_unix(&name, &socket_path)
|
||||
}
|
||||
"http" => {
|
||||
let url_val = url
|
||||
.as_deref()
|
||||
.ok_or_else(|| anyhow::anyhow!("URL is required for http transport"))?;
|
||||
McpServerConfig::new(&name, url_val)
|
||||
}
|
||||
other => {
|
||||
anyhow::bail!(
|
||||
"Unknown transport type '{}'. Supported: http, stdio, unix",
|
||||
other
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
// Apply headers if any
|
||||
if !headers.is_empty() {
|
||||
let headers_map: HashMap<String, String> = headers.into_iter().collect();
|
||||
config = config.with_headers(headers_map);
|
||||
}
|
||||
|
||||
if let Some(desc) = description {
|
||||
config = config.with_description(desc);
|
||||
@@ -151,8 +215,12 @@ async fn add_server(
|
||||
// Track if auth is required
|
||||
let requires_auth = client_id.is_some();
|
||||
|
||||
// Set up OAuth if client_id is provided
|
||||
// Set up OAuth if client_id is provided (HTTP transport only)
|
||||
if let Some(client_id) = client_id {
|
||||
if transport_lower != "http" {
|
||||
anyhow::bail!("OAuth authentication is only supported with http transport");
|
||||
}
|
||||
|
||||
let mut oauth = OAuthConfig::new(client_id);
|
||||
|
||||
if let (Some(auth), Some(token)) = (auth_url, token_url) {
|
||||
@@ -181,7 +249,24 @@ async fn add_server(
|
||||
|
||||
println!();
|
||||
println!(" ✓ Added MCP server '{}'", name);
|
||||
println!(" URL: {}", url);
|
||||
|
||||
match transport_lower.as_str() {
|
||||
"stdio" => {
|
||||
println!(
|
||||
" Transport: stdio (command: {})",
|
||||
command.as_deref().unwrap_or("")
|
||||
);
|
||||
}
|
||||
"unix" => {
|
||||
println!(
|
||||
" Transport: unix (socket: {})",
|
||||
socket.as_deref().unwrap_or("")
|
||||
);
|
||||
}
|
||||
_ => {
|
||||
println!(" URL: {}", url.as_deref().unwrap_or(""));
|
||||
}
|
||||
}
|
||||
|
||||
if requires_auth {
|
||||
println!();
|
||||
@@ -236,9 +321,40 @@ async fn list_servers(verbose: bool) -> anyhow::Result<()> {
|
||||
""
|
||||
};
|
||||
|
||||
let effective = server.effective_transport();
|
||||
|
||||
let transport_label = match &effective {
|
||||
EffectiveTransport::Http => "http".to_string(),
|
||||
EffectiveTransport::Stdio { command, .. } => {
|
||||
format!("stdio ({})", command)
|
||||
}
|
||||
EffectiveTransport::Unix { socket_path } => {
|
||||
format!("unix ({})", socket_path)
|
||||
}
|
||||
};
|
||||
|
||||
if verbose {
|
||||
println!(" {} {}{}", status, server.name, auth_status);
|
||||
println!(" URL: {}", server.url);
|
||||
println!(" Transport: {}", transport_label);
|
||||
match &effective {
|
||||
EffectiveTransport::Http => {
|
||||
println!(" URL: {}", server.url);
|
||||
}
|
||||
EffectiveTransport::Stdio { command, args, env } => {
|
||||
println!(" Command: {}", command);
|
||||
if !args.is_empty() {
|
||||
println!(" Args: {}", args.join(", "));
|
||||
}
|
||||
if !env.is_empty() {
|
||||
// Only print env var names, not values (may contain secrets).
|
||||
let env_keys: Vec<&str> = env.keys().map(|k| k.as_str()).collect();
|
||||
println!(" Env: {}", env_keys.join(", "));
|
||||
}
|
||||
}
|
||||
EffectiveTransport::Unix { socket_path } => {
|
||||
println!(" Socket: {}", socket_path);
|
||||
}
|
||||
}
|
||||
if let Some(ref desc) = server.description {
|
||||
println!(" Description: {}", desc);
|
||||
}
|
||||
@@ -248,11 +364,27 @@ async fn list_servers(verbose: bool) -> anyhow::Result<()> {
|
||||
println!(" Scopes: {}", oauth.scopes.join(", "));
|
||||
}
|
||||
}
|
||||
if !server.headers.is_empty() {
|
||||
let header_keys: Vec<&String> = server.headers.keys().collect();
|
||||
println!(
|
||||
" Headers: {}",
|
||||
header_keys
|
||||
.iter()
|
||||
.map(|k| k.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ")
|
||||
);
|
||||
}
|
||||
println!();
|
||||
} else {
|
||||
let display = match &effective {
|
||||
EffectiveTransport::Http => server.url.clone(),
|
||||
EffectiveTransport::Stdio { command, .. } => command.to_string(),
|
||||
EffectiveTransport::Unix { socket_path } => socket_path.to_string(),
|
||||
};
|
||||
println!(
|
||||
" {} {} - {}{}",
|
||||
status, server.name, server.url, auth_status
|
||||
" {} {} - {} [{}]{}",
|
||||
status, server.name, display, transport_label, auth_status
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -374,7 +506,7 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
return Ok(());
|
||||
} else {
|
||||
// No OAuth and no tokens - try unauthenticated
|
||||
McpClient::new_with_name(&server.name, &server.url)
|
||||
McpClient::new_with_config(server.clone())
|
||||
};
|
||||
|
||||
// Test connection
|
||||
@@ -504,61 +636,9 @@ async fn get_secrets_store() -> anyhow::Result<Arc<dyn SecretsStore + Send + Syn
|
||||
)
|
||||
})?;
|
||||
|
||||
let crypto = SecretsCrypto::new(master_key.clone())?;
|
||||
let crypto = Arc::new(SecretsCrypto::new(master_key.clone())?);
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
{
|
||||
let store = crate::history::Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
Ok(Arc::new(PostgresSecretsStore::new(
|
||||
store.pool(),
|
||||
Arc::new(crypto),
|
||||
)))
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "libsql", not(feature = "postgres")))]
|
||||
{
|
||||
use crate::db::Database as _;
|
||||
use crate::db::libsql::LibSqlBackend;
|
||||
use secrecy::ExposeSecret as _;
|
||||
|
||||
let default_path = crate::config::default_libsql_path();
|
||||
let db_path = config
|
||||
.database
|
||||
.libsql_path
|
||||
.as_deref()
|
||||
.unwrap_or(&default_path);
|
||||
|
||||
let backend = if let Some(ref url) = config.database.libsql_url {
|
||||
let token = config.database.libsql_auth_token.as_ref().ok_or_else(|| {
|
||||
anyhow::anyhow!("LIBSQL_AUTH_TOKEN is required when LIBSQL_URL is set")
|
||||
})?;
|
||||
LibSqlBackend::new_remote_replica(db_path, url, token.expose_secret())
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?
|
||||
} else {
|
||||
LibSqlBackend::new_local(db_path)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?
|
||||
};
|
||||
backend
|
||||
.run_migrations()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
|
||||
Ok(Arc::new(crate::secrets::LibSqlSecretsStore::new(
|
||||
backend.shared_db(),
|
||||
Arc::new(crypto),
|
||||
)))
|
||||
}
|
||||
|
||||
#[cfg(not(any(feature = "postgres", feature = "libsql")))]
|
||||
{
|
||||
let _ = crypto;
|
||||
anyhow::bail!(
|
||||
"No database backend available for secrets. Enable 'postgres' or 'libsql' feature."
|
||||
);
|
||||
}
|
||||
Ok(crate::db::create_secrets_store(&config.database, crypto).await?)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -579,4 +659,46 @@ mod tests {
|
||||
|
||||
TestCli::command().debug_assert();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_header_valid() {
|
||||
let result = parse_header("Authorization: Bearer token123").unwrap();
|
||||
assert_eq!(result.0, "Authorization");
|
||||
assert_eq!(result.1, "Bearer token123");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_header_no_spaces() {
|
||||
let result = parse_header("X-Api-Key:abc123").unwrap();
|
||||
assert_eq!(result.0, "X-Api-Key");
|
||||
assert_eq!(result.1, "abc123");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_header_invalid() {
|
||||
let result = parse_header("no-colon-here");
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("invalid header format"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_env_var_valid() {
|
||||
let result = parse_env_var("NODE_ENV=production").unwrap();
|
||||
assert_eq!(result.0, "NODE_ENV");
|
||||
assert_eq!(result.1, "production");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_env_var_with_equals_in_value() {
|
||||
let result = parse_env_var("KEY=value=with=equals").unwrap();
|
||||
assert_eq!(result.0, "KEY");
|
||||
assert_eq!(result.1, "value=with=equals");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_env_var_invalid() {
|
||||
let result = parse_env_var("no-equals-here");
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("invalid env var format"));
|
||||
}
|
||||
}
|
||||
|
||||
+1
-1
@@ -132,7 +132,7 @@ pub enum Command {
|
||||
about = "Manage MCP servers",
|
||||
long_about = "Add, auth, list, or test MCP servers.\nExample: ironclaw mcp add notion https://mcp.notion.com"
|
||||
)]
|
||||
Mcp(McpCommand),
|
||||
Mcp(Box<McpCommand>),
|
||||
|
||||
/// Query and manage workspace memory
|
||||
#[command(
|
||||
|
||||
+123
-1
@@ -8,9 +8,35 @@ use std::path::PathBuf;
|
||||
use crate::bootstrap::ironclaw_base_dir;
|
||||
use crate::settings::Settings;
|
||||
|
||||
/// Load settings from JSON and TOML config files, matching the runtime
|
||||
/// priority: TOML overlay > settings.json > defaults.
|
||||
///
|
||||
/// This mirrors the loading chain in `Config::from_env_with_toml()` but
|
||||
/// without resolving the full `Config` (which requires async + secrets).
|
||||
fn load_settings() -> Settings {
|
||||
load_settings_from(&Settings::default_path(), &Settings::default_toml_path())
|
||||
}
|
||||
|
||||
/// Inner implementation with injectable paths (testable).
|
||||
fn load_settings_from(json_path: &std::path::Path, toml_path: &std::path::Path) -> Settings {
|
||||
let mut settings = Settings::load_from(json_path);
|
||||
|
||||
match Settings::load_toml(toml_path) {
|
||||
Ok(Some(toml_settings)) => {
|
||||
settings.merge_from(&toml_settings);
|
||||
}
|
||||
Ok(None) => {} // File not found — fine for default path
|
||||
Err(e) => {
|
||||
eprintln!("Warning: failed to parse {}: {}", toml_path.display(), e);
|
||||
}
|
||||
}
|
||||
|
||||
settings
|
||||
}
|
||||
|
||||
/// Run the status command, printing system health info.
|
||||
pub async fn run_status_command() -> anyhow::Result<()> {
|
||||
let settings = Settings::default();
|
||||
let settings = load_settings();
|
||||
|
||||
println!("IronClaw Status");
|
||||
println!("===============\n");
|
||||
@@ -209,3 +235,99 @@ fn default_tools_dir() -> PathBuf {
|
||||
fn default_channels_dir() -> PathBuf {
|
||||
ironclaw_base_dir().join("channels")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::load_settings_from;
|
||||
|
||||
/// Regression test for #354: load_settings_from must read config.toml.
|
||||
#[test]
|
||||
fn reads_toml_heartbeat_enabled() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let json_path = dir.path().join("settings.json");
|
||||
let toml_path = dir.path().join("config.toml");
|
||||
|
||||
// No JSON file — only TOML
|
||||
std::fs::write(
|
||||
&toml_path,
|
||||
"[heartbeat]\nenabled = true\ninterval_secs = 600",
|
||||
)
|
||||
.expect("write toml");
|
||||
|
||||
let settings = load_settings_from(&json_path, &toml_path);
|
||||
assert!(settings.heartbeat.enabled);
|
||||
assert_eq!(settings.heartbeat.interval_secs, 600);
|
||||
}
|
||||
|
||||
/// Without any config files, defaults are returned.
|
||||
#[test]
|
||||
fn defaults_without_config_files() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let settings = load_settings_from(
|
||||
&dir.path().join("nonexistent.json"),
|
||||
&dir.path().join("nonexistent.toml"),
|
||||
);
|
||||
assert!(!settings.heartbeat.enabled);
|
||||
}
|
||||
|
||||
/// settings.json is respected.
|
||||
#[test]
|
||||
fn reads_json_heartbeat_enabled() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let json_path = dir.path().join("settings.json");
|
||||
let toml_path = dir.path().join("nonexistent.toml");
|
||||
|
||||
std::fs::write(
|
||||
&json_path,
|
||||
r#"{"heartbeat":{"enabled":true,"interval_secs":900}}"#,
|
||||
)
|
||||
.expect("write json");
|
||||
|
||||
let settings = load_settings_from(&json_path, &toml_path);
|
||||
assert!(settings.heartbeat.enabled);
|
||||
assert_eq!(settings.heartbeat.interval_secs, 900);
|
||||
}
|
||||
|
||||
/// TOML overlay wins over JSON settings.
|
||||
#[test]
|
||||
fn toml_overlay_wins_over_json() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let json_path = dir.path().join("settings.json");
|
||||
let toml_path = dir.path().join("config.toml");
|
||||
|
||||
std::fs::write(
|
||||
&json_path,
|
||||
r#"{"heartbeat":{"enabled":false,"interval_secs":100}}"#,
|
||||
)
|
||||
.expect("write json");
|
||||
std::fs::write(
|
||||
&toml_path,
|
||||
"[heartbeat]\nenabled = true\ninterval_secs = 200",
|
||||
)
|
||||
.expect("write toml");
|
||||
|
||||
let settings = load_settings_from(&json_path, &toml_path);
|
||||
assert!(settings.heartbeat.enabled);
|
||||
assert_eq!(settings.heartbeat.interval_secs, 200);
|
||||
}
|
||||
|
||||
/// Invalid TOML is warned but doesn't crash; falls back to JSON/defaults.
|
||||
#[test]
|
||||
fn invalid_toml_falls_back_gracefully() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let json_path = dir.path().join("settings.json");
|
||||
let toml_path = dir.path().join("config.toml");
|
||||
|
||||
std::fs::write(
|
||||
&json_path,
|
||||
r#"{"heartbeat":{"enabled":true,"interval_secs":500}}"#,
|
||||
)
|
||||
.expect("write json");
|
||||
std::fs::write(&toml_path, "this is not valid toml [[[").expect("write bad toml");
|
||||
|
||||
let settings = load_settings_from(&json_path, &toml_path);
|
||||
// Should fall back to JSON values, not crash
|
||||
assert!(settings.heartbeat.enabled);
|
||||
assert_eq!(settings.heartbeat.interval_secs, 500);
|
||||
}
|
||||
}
|
||||
|
||||
+2
-56
@@ -11,10 +11,6 @@ use tokio::fs;
|
||||
|
||||
use crate::bootstrap::ironclaw_base_dir;
|
||||
use crate::config::Config;
|
||||
#[allow(unused_imports)]
|
||||
use crate::db::Database;
|
||||
#[cfg(feature = "postgres")]
|
||||
use crate::secrets::PostgresSecretsStore;
|
||||
use crate::secrets::{CreateSecretParams, SecretsCrypto, SecretsStore};
|
||||
use crate::tools::wasm::{CapabilitiesFile, compute_binary_hash};
|
||||
|
||||
@@ -563,59 +559,9 @@ async fn init_secrets_store() -> anyhow::Result<Arc<dyn SecretsStore + Send + Sy
|
||||
)
|
||||
})?;
|
||||
|
||||
let crypto = SecretsCrypto::new(master_key.clone())?;
|
||||
let crypto = Arc::new(SecretsCrypto::new(master_key.clone())?);
|
||||
|
||||
let store: Arc<dyn SecretsStore + Send + Sync> = {
|
||||
#[cfg(feature = "postgres")]
|
||||
{
|
||||
let store = crate::history::Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
Arc::new(PostgresSecretsStore::new(store.pool(), Arc::new(crypto)))
|
||||
}
|
||||
#[cfg(all(feature = "libsql", not(feature = "postgres")))]
|
||||
{
|
||||
use crate::db::Database as _;
|
||||
use crate::db::libsql::LibSqlBackend;
|
||||
use secrecy::ExposeSecret as _;
|
||||
|
||||
let default_path = crate::config::default_libsql_path();
|
||||
let db_path = config
|
||||
.database
|
||||
.libsql_path
|
||||
.as_deref()
|
||||
.unwrap_or(&default_path);
|
||||
|
||||
let backend = if let Some(ref url) = config.database.libsql_url {
|
||||
let token = config.database.libsql_auth_token.as_ref().ok_or_else(|| {
|
||||
anyhow::anyhow!("LIBSQL_AUTH_TOKEN is required when LIBSQL_URL is set")
|
||||
})?;
|
||||
LibSqlBackend::new_remote_replica(db_path, url, token.expose_secret())
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?
|
||||
} else {
|
||||
LibSqlBackend::new_local(db_path)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?
|
||||
};
|
||||
backend
|
||||
.run_migrations()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
|
||||
Arc::new(crate::secrets::LibSqlSecretsStore::new(
|
||||
backend.shared_db(),
|
||||
Arc::new(crypto),
|
||||
))
|
||||
}
|
||||
#[cfg(not(any(feature = "postgres", feature = "libsql")))]
|
||||
{
|
||||
let _ = crypto;
|
||||
anyhow::bail!(
|
||||
"No database backend available for secrets. Enable 'postgres' or 'libsql' feature."
|
||||
);
|
||||
}
|
||||
};
|
||||
Ok(store)
|
||||
Ok(crate::db::create_secrets_store(&config.database, crypto).await?)
|
||||
}
|
||||
|
||||
/// Configure authentication for a tool.
|
||||
|
||||
@@ -27,6 +27,8 @@ pub struct AgentConfig {
|
||||
pub max_tool_iterations: usize,
|
||||
/// When true, skip tool approval checks entirely. For benchmarks/CI.
|
||||
pub auto_approve_tools: bool,
|
||||
/// Default timezone for new sessions (IANA name, e.g. "America/New_York").
|
||||
pub default_timezone: String,
|
||||
}
|
||||
|
||||
impl AgentConfig {
|
||||
@@ -47,6 +49,7 @@ impl AgentConfig {
|
||||
max_actions_per_hour: None,
|
||||
max_tool_iterations: 10,
|
||||
auto_approve_tools: true,
|
||||
default_timezone: "UTC".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -89,6 +92,40 @@ impl AgentConfig {
|
||||
"AGENT_AUTO_APPROVE_TOOLS",
|
||||
settings.agent.auto_approve_tools,
|
||||
)?,
|
||||
default_timezone: {
|
||||
let tz: String = parse_optional_env(
|
||||
"DEFAULT_TIMEZONE",
|
||||
settings.agent.default_timezone.clone(),
|
||||
)?;
|
||||
if crate::timezone::parse_timezone(&tz).is_none() {
|
||||
return Err(ConfigError::InvalidValue {
|
||||
key: "DEFAULT_TIMEZONE".into(),
|
||||
message: format!("invalid IANA timezone: '{tz}'"),
|
||||
});
|
||||
}
|
||||
tz
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_default_timezone_rejects_invalid() {
|
||||
let mut settings = Settings::default();
|
||||
settings.agent.default_timezone = "Fake/Zone".to_string();
|
||||
|
||||
let result = AgentConfig::resolve(&settings);
|
||||
assert!(result.is_err(), "invalid IANA timezone should be rejected");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_default_timezone_accepts_valid() {
|
||||
let settings = Settings::default(); // default is "UTC"
|
||||
let config = AgentConfig::resolve(&settings).expect("resolve");
|
||||
assert_eq!(config.default_timezone, "UTC");
|
||||
}
|
||||
}
|
||||
|
||||
+102
-1
@@ -1,4 +1,4 @@
|
||||
use crate::config::helpers::{optional_env, parse_bool_env, parse_optional_env};
|
||||
use crate::config::helpers::{optional_env, parse_bool_env, parse_option_env, parse_optional_env};
|
||||
use crate::error::ConfigError;
|
||||
use crate::settings::Settings;
|
||||
|
||||
@@ -13,6 +13,12 @@ pub struct HeartbeatConfig {
|
||||
pub notify_channel: Option<String>,
|
||||
/// User ID to notify on heartbeat findings.
|
||||
pub notify_user: Option<String>,
|
||||
/// Hour (0-23) when quiet hours start.
|
||||
pub quiet_hours_start: Option<u32>,
|
||||
/// Hour (0-23) when quiet hours end.
|
||||
pub quiet_hours_end: Option<u32>,
|
||||
/// Timezone for quiet hours evaluation (IANA name).
|
||||
pub timezone: Option<String>,
|
||||
}
|
||||
|
||||
impl Default for HeartbeatConfig {
|
||||
@@ -22,6 +28,9 @@ impl Default for HeartbeatConfig {
|
||||
interval_secs: 1800, // 30 minutes
|
||||
notify_channel: None,
|
||||
notify_user: None,
|
||||
quiet_hours_start: None,
|
||||
quiet_hours_end: None,
|
||||
timezone: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -38,6 +47,98 @@ impl HeartbeatConfig {
|
||||
.or_else(|| settings.heartbeat.notify_channel.clone()),
|
||||
notify_user: optional_env("HEARTBEAT_NOTIFY_USER")?
|
||||
.or_else(|| settings.heartbeat.notify_user.clone()),
|
||||
quiet_hours_start: parse_option_env::<u32>("HEARTBEAT_QUIET_START")?
|
||||
.or(settings.heartbeat.quiet_hours_start)
|
||||
.map(|h| {
|
||||
if h > 23 {
|
||||
return Err(ConfigError::InvalidValue {
|
||||
key: "HEARTBEAT_QUIET_START".into(),
|
||||
message: "must be 0-23".into(),
|
||||
});
|
||||
}
|
||||
Ok(h)
|
||||
})
|
||||
.transpose()?,
|
||||
quiet_hours_end: parse_option_env::<u32>("HEARTBEAT_QUIET_END")?
|
||||
.or(settings.heartbeat.quiet_hours_end)
|
||||
.map(|h| {
|
||||
if h > 23 {
|
||||
return Err(ConfigError::InvalidValue {
|
||||
key: "HEARTBEAT_QUIET_END".into(),
|
||||
message: "must be 0-23".into(),
|
||||
});
|
||||
}
|
||||
Ok(h)
|
||||
})
|
||||
.transpose()?,
|
||||
timezone: {
|
||||
let tz = optional_env("HEARTBEAT_TIMEZONE")?
|
||||
.or_else(|| settings.heartbeat.timezone.clone());
|
||||
if let Some(ref tz_str) = tz
|
||||
&& crate::timezone::parse_timezone(tz_str).is_none()
|
||||
{
|
||||
return Err(ConfigError::InvalidValue {
|
||||
key: "HEARTBEAT_TIMEZONE".into(),
|
||||
message: format!("invalid IANA timezone: '{tz_str}'"),
|
||||
});
|
||||
}
|
||||
tz
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_quiet_hours_settings_fallback() {
|
||||
// When env vars are not set, settings values should be used
|
||||
let mut settings = Settings::default();
|
||||
settings.heartbeat.quiet_hours_start = Some(22);
|
||||
settings.heartbeat.quiet_hours_end = Some(6);
|
||||
|
||||
let config = HeartbeatConfig::resolve(&settings).expect("resolve");
|
||||
assert_eq!(config.quiet_hours_start, Some(22));
|
||||
assert_eq!(config.quiet_hours_end, Some(6));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_quiet_hours_rejects_invalid_hour() {
|
||||
let mut settings = Settings::default();
|
||||
settings.heartbeat.quiet_hours_start = Some(24);
|
||||
|
||||
let result = HeartbeatConfig::resolve(&settings);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_quiet_hours_accepts_boundary_values() {
|
||||
let mut settings = Settings::default();
|
||||
settings.heartbeat.quiet_hours_start = Some(0);
|
||||
settings.heartbeat.quiet_hours_end = Some(23);
|
||||
|
||||
let config = HeartbeatConfig::resolve(&settings).expect("resolve");
|
||||
assert_eq!(config.quiet_hours_start, Some(0));
|
||||
assert_eq!(config.quiet_hours_end, Some(23));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_heartbeat_timezone_rejects_invalid() {
|
||||
let mut settings = Settings::default();
|
||||
settings.heartbeat.timezone = Some("Fake/Zone".to_string());
|
||||
|
||||
let result = HeartbeatConfig::resolve(&settings);
|
||||
assert!(result.is_err(), "invalid IANA timezone should be rejected");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_heartbeat_timezone_accepts_valid() {
|
||||
let mut settings = Settings::default();
|
||||
settings.heartbeat.timezone = Some("America/New_York".to_string());
|
||||
|
||||
let config = HeartbeatConfig::resolve(&settings).expect("resolve");
|
||||
assert_eq!(config.timezone.as_deref(), Some("America/New_York"));
|
||||
}
|
||||
}
|
||||
|
||||
+96
-4
@@ -86,6 +86,19 @@ pub struct RegistryProviderConfig {
|
||||
pub oauth_token: Option<SecretString>,
|
||||
}
|
||||
|
||||
/// Configuration for AWS Bedrock (native Converse API).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BedrockConfig {
|
||||
/// AWS region (e.g. "us-east-1").
|
||||
pub region: String,
|
||||
/// Bedrock model ID (e.g. "anthropic.claude-opus-4-6-v1").
|
||||
pub model: String,
|
||||
/// Cross-region inference prefix: "us", "eu", "apac", "global", or None.
|
||||
pub cross_region: Option<String>,
|
||||
/// AWS named profile (for SSO / assume-role workflows).
|
||||
pub profile: Option<String>,
|
||||
}
|
||||
|
||||
/// LLM provider configuration.
|
||||
///
|
||||
/// NearAI remains the default backend with its own config struct (session auth).
|
||||
@@ -101,8 +114,14 @@ pub struct LlmConfig {
|
||||
/// NEAR AI config (always populated, also used for embeddings).
|
||||
pub nearai: NearAiConfig,
|
||||
/// Resolved provider config for registry-based providers.
|
||||
/// `None` when backend is "nearai".
|
||||
/// `None` when backend is "nearai" or "bedrock".
|
||||
pub provider: Option<RegistryProviderConfig>,
|
||||
/// AWS Bedrock config (populated when backend=bedrock, requires --features bedrock).
|
||||
pub bedrock: Option<BedrockConfig>,
|
||||
/// HTTP request timeout in seconds for LLM API calls.
|
||||
/// Default: 120. Increase for local LLMs (Ollama, vLLM, LM Studio) that
|
||||
/// need more time for prompt evaluation on consumer hardware.
|
||||
pub request_timeout_secs: u64,
|
||||
}
|
||||
|
||||
/// NEAR AI configuration.
|
||||
@@ -165,6 +184,8 @@ impl LlmConfig {
|
||||
smart_routing_cascade: false,
|
||||
},
|
||||
provider: None,
|
||||
bedrock: None,
|
||||
request_timeout_secs: 120,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -195,8 +216,10 @@ impl LlmConfig {
|
||||
let backend_lower = backend.to_lowercase();
|
||||
let is_nearai =
|
||||
backend_lower == "nearai" || backend_lower == "near_ai" || backend_lower == "near";
|
||||
let is_bedrock =
|
||||
backend_lower == "bedrock" || backend_lower == "aws_bedrock" || backend_lower == "aws";
|
||||
|
||||
if !is_nearai && registry.find(&backend_lower).is_none() {
|
||||
if !is_nearai && !is_bedrock && registry.find(&backend_lower).is_none() {
|
||||
tracing::warn!(
|
||||
"Unknown LLM backend '{}'. Will attempt as openai_compatible fallback.",
|
||||
backend
|
||||
@@ -243,8 +266,8 @@ impl LlmConfig {
|
||||
smart_routing_cascade: parse_optional_env("SMART_ROUTING_CASCADE", true)?,
|
||||
};
|
||||
|
||||
// Resolve registry provider config (for non-NearAI backends)
|
||||
let provider = if is_nearai {
|
||||
// Resolve registry provider config (for non-NearAI, non-Bedrock backends)
|
||||
let provider = if is_nearai || is_bedrock {
|
||||
None
|
||||
} else {
|
||||
Some(Self::resolve_registry_provider(
|
||||
@@ -254,9 +277,50 @@ impl LlmConfig {
|
||||
)?)
|
||||
};
|
||||
|
||||
let bedrock = if is_bedrock {
|
||||
let explicit_region =
|
||||
optional_env("BEDROCK_REGION")?.or_else(|| settings.bedrock_region.clone());
|
||||
if explicit_region.is_none() {
|
||||
tracing::info!("BEDROCK_REGION not set, defaulting to us-east-1");
|
||||
}
|
||||
let region = explicit_region.unwrap_or_else(|| "us-east-1".to_string());
|
||||
let model = optional_env("BEDROCK_MODEL")?
|
||||
.or_else(|| settings.selected_model.clone())
|
||||
.ok_or_else(|| ConfigError::MissingRequired {
|
||||
key: "BEDROCK_MODEL".to_string(),
|
||||
hint: "Set BEDROCK_MODEL when LLM_BACKEND=bedrock".to_string(),
|
||||
})?;
|
||||
let cross_region = optional_env("BEDROCK_CROSS_REGION")?
|
||||
.or_else(|| settings.bedrock_cross_region.clone());
|
||||
if let Some(ref cr) = cross_region
|
||||
&& !matches!(cr.as_str(), "us" | "eu" | "apac" | "global")
|
||||
{
|
||||
return Err(ConfigError::InvalidValue {
|
||||
key: "BEDROCK_CROSS_REGION".to_string(),
|
||||
message: format!(
|
||||
"'{}' is not valid, expected one of: us, eu, apac, global",
|
||||
cr
|
||||
),
|
||||
});
|
||||
}
|
||||
let profile = optional_env("AWS_PROFILE")?.or_else(|| settings.bedrock_profile.clone());
|
||||
Some(BedrockConfig {
|
||||
region,
|
||||
model,
|
||||
cross_region,
|
||||
profile,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let request_timeout_secs = parse_optional_env("LLM_REQUEST_TIMEOUT_SECS", 120)?;
|
||||
|
||||
Ok(Self {
|
||||
backend: if is_nearai {
|
||||
"nearai".to_string()
|
||||
} else if is_bedrock {
|
||||
"bedrock".to_string()
|
||||
} else if let Some(ref p) = provider {
|
||||
p.provider_id.clone()
|
||||
} else {
|
||||
@@ -265,6 +329,8 @@ impl LlmConfig {
|
||||
session,
|
||||
nearai,
|
||||
provider,
|
||||
bedrock,
|
||||
request_timeout_secs,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1016,4 +1082,30 @@ mod tests {
|
||||
assert_eq!(parsed, variant, "round-trip failed for {s}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_request_timeout_defaults_to_120() {
|
||||
let _guard = ENV_MUTEX.lock().expect("env mutex poisoned");
|
||||
// SAFETY: Under ENV_MUTEX.
|
||||
unsafe {
|
||||
std::env::remove_var("LLM_REQUEST_TIMEOUT_SECS");
|
||||
}
|
||||
let config = LlmConfig::resolve(&Settings::default()).expect("resolve");
|
||||
assert_eq!(config.request_timeout_secs, 120);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_request_timeout_configurable() {
|
||||
let _guard = ENV_MUTEX.lock().expect("env mutex poisoned");
|
||||
// SAFETY: Under ENV_MUTEX.
|
||||
unsafe {
|
||||
std::env::set_var("LLM_REQUEST_TIMEOUT_SECS", "300");
|
||||
}
|
||||
let config = LlmConfig::resolve(&Settings::default()).expect("resolve");
|
||||
assert_eq!(config.request_timeout_secs, 300);
|
||||
// SAFETY: Cleanup
|
||||
unsafe {
|
||||
std::env::remove_var("LLM_REQUEST_TIMEOUT_SECS");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+29
-1
@@ -37,7 +37,9 @@ pub use self::database::{DatabaseBackend, DatabaseConfig, SslMode, default_libsq
|
||||
pub use self::embeddings::EmbeddingsConfig;
|
||||
pub use self::heartbeat::HeartbeatConfig;
|
||||
pub use self::hygiene::HygieneConfig;
|
||||
pub use self::llm::{CacheRetention, LlmConfig, NearAiConfig, RegistryProviderConfig};
|
||||
pub use self::llm::{
|
||||
BedrockConfig, CacheRetention, LlmConfig, NearAiConfig, RegistryProviderConfig,
|
||||
};
|
||||
pub use self::routines::RoutineConfig;
|
||||
pub use self::safety::SafetyConfig;
|
||||
pub use self::sandbox::{ClaudeCodeConfig, SandboxModeConfig};
|
||||
@@ -257,6 +259,32 @@ impl Config {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Re-resolve only the LLM config after credential injection.
|
||||
///
|
||||
/// Called by `AppBuilder::init_secrets()` after injecting API keys into
|
||||
/// the env overlay. Only rebuilds `self.llm` — all other config fields
|
||||
/// are unaffected, preserving values from the initial config load (or
|
||||
/// from `Config::for_testing()` in test mode).
|
||||
pub async fn re_resolve_llm(
|
||||
&mut self,
|
||||
store: Option<&(dyn crate::db::SettingsStore + Sync)>,
|
||||
user_id: &str,
|
||||
toml_path: Option<&std::path::Path>,
|
||||
) -> Result<(), ConfigError> {
|
||||
let settings = if let Some(store) = store {
|
||||
let mut s = match store.get_all_settings(user_id).await {
|
||||
Ok(map) => Settings::from_db_map(&map),
|
||||
Err(_) => Settings::default(),
|
||||
};
|
||||
Self::apply_toml_overlay(&mut s, toml_path)?;
|
||||
s
|
||||
} else {
|
||||
Settings::default()
|
||||
};
|
||||
self.llm = LlmConfig::resolve(&settings)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Build config from settings (shared by from_env and from_db).
|
||||
async fn build(settings: &Settings) -> Result<Self, ConfigError> {
|
||||
Ok(Self {
|
||||
|
||||
@@ -164,6 +164,8 @@ pub struct JobContext {
|
||||
/// previous results by ID via `$tool_call_id` parameter syntax.
|
||||
#[serde(skip)]
|
||||
pub tool_output_stash: Arc<tokio::sync::RwLock<HashMap<String, String>>>,
|
||||
/// User's preferred timezone (IANA name, e.g. "America/New_York"). Defaults to "UTC".
|
||||
pub user_timezone: String,
|
||||
}
|
||||
|
||||
impl JobContext {
|
||||
@@ -203,9 +205,16 @@ impl JobContext {
|
||||
http_interceptor: None,
|
||||
metadata: serde_json::Value::Null,
|
||||
tool_output_stash: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
|
||||
user_timezone: "UTC".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the user timezone on this context.
|
||||
pub fn with_timezone(mut self, tz: impl Into<String>) -> Self {
|
||||
self.user_timezone = tz.into();
|
||||
self
|
||||
}
|
||||
|
||||
/// Transition to a new state.
|
||||
pub fn transition_to(
|
||||
&mut self,
|
||||
|
||||
@@ -121,6 +121,9 @@ impl JobStore for LibSqlBackend {
|
||||
tool_output_stash: std::sync::Arc::new(tokio::sync::RwLock::new(
|
||||
std::collections::HashMap::new(),
|
||||
)),
|
||||
// TODO(#661): persist user_timezone in agent_jobs table so
|
||||
// background/routine jobs retain the session's timezone context.
|
||||
user_timezone: "UTC".to_string(),
|
||||
}))
|
||||
}
|
||||
None => Ok(None),
|
||||
|
||||
+44
-99
@@ -170,16 +170,16 @@ pub(crate) fn parse_timestamp(s: &str) -> Result<DateTime<Utc>, String> {
|
||||
// Naive with fractional seconds (legacy or SQLite datetime() output)
|
||||
if let Ok(ndt) = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f") {
|
||||
tracing::warn!(
|
||||
timestamp = s,
|
||||
"parsing naive timestamp without timezone; assuming UTC — consider re-running migrations"
|
||||
timestamp = %s,
|
||||
"parsed naive timestamp without timezone; assuming UTC for backward compatibility"
|
||||
);
|
||||
return Ok(ndt.and_utc());
|
||||
}
|
||||
// Naive without fractional seconds (legacy format)
|
||||
if let Ok(ndt) = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S") {
|
||||
tracing::warn!(
|
||||
timestamp = s,
|
||||
"parsing naive timestamp without timezone; assuming UTC — consider re-running migrations"
|
||||
timestamp = %s,
|
||||
"parsed naive timestamp without timezone; assuming UTC for backward compatibility"
|
||||
);
|
||||
return Ok(ndt.and_utc());
|
||||
}
|
||||
@@ -410,8 +410,47 @@ pub(crate) fn row_to_routine_run_libsql(row: &libsql::Row) -> Result<RoutineRun,
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use chrono::{TimeZone, Utc};
|
||||
|
||||
use crate::db::Database;
|
||||
use crate::db::libsql::LibSqlBackend;
|
||||
use crate::db::libsql::{LibSqlBackend, parse_timestamp};
|
||||
|
||||
#[test]
|
||||
fn test_parse_timestamp_accepts_rfc3339_and_legacy_naive_formats() {
|
||||
let expected = Utc.with_ymd_and_hms(2026, 3, 7, 12, 34, 56).unwrap();
|
||||
|
||||
let with_millis = parse_timestamp("2026-03-07T12:34:56.789Z").unwrap();
|
||||
assert_eq!(with_millis, expected + chrono::Duration::milliseconds(789));
|
||||
|
||||
let naive_with_millis = parse_timestamp("2026-03-07 12:34:56.789").unwrap();
|
||||
assert_eq!(
|
||||
naive_with_millis,
|
||||
expected + chrono::Duration::milliseconds(789)
|
||||
);
|
||||
|
||||
let naive_without_millis = parse_timestamp("2026-03-07 12:34:56").unwrap();
|
||||
assert_eq!(naive_without_millis, expected);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_libsql_now_format_is_rfc3339_and_parseable() {
|
||||
let backend = LibSqlBackend::new_memory().await.unwrap();
|
||||
backend.run_migrations().await.unwrap();
|
||||
|
||||
let conn = backend.connect().await.unwrap();
|
||||
let mut rows = conn
|
||||
.query("SELECT strftime('%Y-%m-%dT%H:%M:%fZ', 'now')", ())
|
||||
.await
|
||||
.unwrap();
|
||||
let row = rows.next().await.unwrap().unwrap();
|
||||
let ts: String = row.get(0).unwrap();
|
||||
|
||||
let parsed = parse_timestamp(&ts).unwrap();
|
||||
assert_eq!(
|
||||
ts,
|
||||
parsed.to_rfc3339_opts(chrono::SecondsFormat::Millis, true)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_wal_mode_after_migrations() {
|
||||
@@ -518,98 +557,4 @@ mod tests {
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_timestamp_rfc3339() {
|
||||
use super::parse_timestamp;
|
||||
|
||||
// Standard RFC 3339 with Z suffix
|
||||
let dt = parse_timestamp("2024-01-15T10:30:00.123Z").unwrap();
|
||||
assert_eq!(
|
||||
dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
|
||||
"2024-01-15T10:30:00.123Z"
|
||||
);
|
||||
|
||||
// RFC 3339 with +00:00 offset
|
||||
let dt = parse_timestamp("2024-01-15T10:30:00.000+00:00").unwrap();
|
||||
assert_eq!(
|
||||
dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
|
||||
"2024-01-15T10:30:00.000Z"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_timestamp_naive_fallback() {
|
||||
use super::parse_timestamp;
|
||||
|
||||
// Naive with fractional seconds (legacy datetime('now') output)
|
||||
let dt = parse_timestamp("2024-01-15 10:30:00.123").unwrap();
|
||||
assert_eq!(
|
||||
dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
|
||||
"2024-01-15T10:30:00.123Z"
|
||||
);
|
||||
|
||||
// Naive without fractional seconds
|
||||
let dt = parse_timestamp("2024-01-15 10:30:00").unwrap();
|
||||
assert_eq!(
|
||||
dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
|
||||
"2024-01-15T10:30:00.000Z"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_timestamp_invalid() {
|
||||
use super::parse_timestamp;
|
||||
|
||||
assert!(parse_timestamp("not-a-timestamp").is_err());
|
||||
assert!(parse_timestamp("").is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_default_timestamps_are_rfc3339() {
|
||||
// Verify that DEFAULT column values produce RFC 3339 timestamps
|
||||
// after the migration change from datetime('now') to strftime.
|
||||
// Use file-based DB because in-memory doesn't share schema across connections.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db_path = dir.path().join("test_ts.db");
|
||||
let backend = LibSqlBackend::new_local(&db_path).await.unwrap();
|
||||
backend.run_migrations().await.unwrap();
|
||||
|
||||
let conn = backend.connect().await.unwrap();
|
||||
let id = uuid::Uuid::new_v4().to_string();
|
||||
conn.execute(
|
||||
"INSERT INTO conversations (id, channel, user_id) VALUES (?1, ?2, ?3)",
|
||||
libsql::params![id.clone(), "test", "user1"],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut rows = conn
|
||||
.query(
|
||||
"SELECT started_at, last_activity FROM conversations WHERE id = ?1",
|
||||
libsql::params![id],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let row = rows.next().await.unwrap().unwrap();
|
||||
let started_at: String = row.get(0).unwrap();
|
||||
let last_activity: String = row.get(1).unwrap();
|
||||
|
||||
// Must end with 'Z' (RFC 3339 UTC) and contain 'T' separator
|
||||
assert!(
|
||||
started_at.ends_with('Z') && started_at.contains('T'),
|
||||
"started_at should be RFC 3339, got: {started_at}"
|
||||
);
|
||||
assert!(
|
||||
last_activity.ends_with('Z') && last_activity.contains('T'),
|
||||
"last_activity should be RFC 3339, got: {last_activity}"
|
||||
);
|
||||
|
||||
// Must be parseable by the RFC 3339 parser directly (not just naive fallback)
|
||||
use chrono::DateTime;
|
||||
assert!(
|
||||
DateTime::parse_from_rfc3339(&started_at).is_ok(),
|
||||
"started_at not valid RFC 3339: {started_at}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
+101
@@ -91,6 +91,64 @@ pub async fn connect_from_config(
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a secrets store from database and secrets configuration.
|
||||
///
|
||||
/// This is the shared factory for CLI commands and other call sites that need
|
||||
/// a `SecretsStore` without going through the full `AppBuilder`. Mirrors the
|
||||
/// pattern of [`connect_from_config`] but returns a secrets-specific store.
|
||||
pub async fn create_secrets_store(
|
||||
config: &crate::config::DatabaseConfig,
|
||||
crypto: Arc<crate::secrets::SecretsCrypto>,
|
||||
) -> Result<Arc<dyn crate::secrets::SecretsStore + Send + Sync>, DatabaseError> {
|
||||
match config.backend {
|
||||
#[cfg(feature = "libsql")]
|
||||
crate::config::DatabaseBackend::LibSql => {
|
||||
use secrecy::ExposeSecret as _;
|
||||
|
||||
let default_path = crate::config::default_libsql_path();
|
||||
let db_path = config.libsql_path.as_deref().unwrap_or(&default_path);
|
||||
|
||||
let backend = if let Some(ref url) = config.libsql_url {
|
||||
let token = config.libsql_auth_token.as_ref().ok_or_else(|| {
|
||||
DatabaseError::Pool(
|
||||
"LIBSQL_AUTH_TOKEN required when LIBSQL_URL is set".to_string(),
|
||||
)
|
||||
})?;
|
||||
libsql::LibSqlBackend::new_remote_replica(db_path, url, token.expose_secret())
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Pool(e.to_string()))?
|
||||
} else {
|
||||
libsql::LibSqlBackend::new_local(db_path)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Pool(e.to_string()))?
|
||||
};
|
||||
backend.run_migrations().await?;
|
||||
|
||||
Ok(Arc::new(crate::secrets::LibSqlSecretsStore::new(
|
||||
backend.shared_db(),
|
||||
crypto,
|
||||
)))
|
||||
}
|
||||
#[cfg(feature = "postgres")]
|
||||
_ => {
|
||||
let pg = postgres::PgBackend::new(config)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Pool(e.to_string()))?;
|
||||
pg.run_migrations().await?;
|
||||
|
||||
Ok(Arc::new(crate::secrets::PostgresSecretsStore::new(
|
||||
pg.pool(),
|
||||
crypto,
|
||||
)))
|
||||
}
|
||||
#[cfg(not(feature = "postgres"))]
|
||||
_ => Err(DatabaseError::Pool(
|
||||
"No database backend available for secrets. Enable 'postgres' or 'libsql' feature."
|
||||
.to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== Sub-traits ====================
|
||||
//
|
||||
// Each sub-trait groups related persistence methods. The `Database` supertrait
|
||||
@@ -435,3 +493,46 @@ pub trait Database:
|
||||
/// Run schema migrations for this backend.
|
||||
async fn run_migrations(&self) -> Result<(), DatabaseError>;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Regression test: `create_secrets_store` selects the correct backend at
|
||||
/// runtime based on `DatabaseConfig`, not at compile time. Previously the
|
||||
/// CLI duplicated this logic with compile-time `#[cfg]` gates that always
|
||||
/// chose postgres when both features were enabled (PR #209).
|
||||
#[cfg(feature = "libsql")]
|
||||
#[tokio::test]
|
||||
async fn test_create_secrets_store_libsql_backend() {
|
||||
use secrecy::SecretString;
|
||||
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let db_path = tmp.path().join("test.db");
|
||||
|
||||
let config = crate::config::DatabaseConfig {
|
||||
backend: crate::config::DatabaseBackend::LibSql,
|
||||
libsql_path: Some(db_path),
|
||||
libsql_url: None,
|
||||
libsql_auth_token: None,
|
||||
url: SecretString::from("unused://libsql".to_string()),
|
||||
pool_size: 1,
|
||||
ssl_mode: crate::config::SslMode::default(),
|
||||
};
|
||||
|
||||
let master_key = SecretString::from("a]".repeat(16));
|
||||
let crypto = Arc::new(crate::secrets::SecretsCrypto::new(master_key).unwrap());
|
||||
|
||||
let store = create_secrets_store(&config, crypto).await;
|
||||
assert!(
|
||||
store.is_ok(),
|
||||
"create_secrets_store should succeed for libsql backend"
|
||||
);
|
||||
|
||||
// Verify basic operation works
|
||||
let store = store.unwrap();
|
||||
let exists = store.exists("test_user", "nonexistent_secret").await;
|
||||
assert!(exists.is_ok());
|
||||
assert!(!exists.unwrap());
|
||||
}
|
||||
}
|
||||
|
||||
+123
-227
@@ -1,13 +1,10 @@
|
||||
//! Success evaluation for jobs.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::context::{ActionRecord, JobContext};
|
||||
use crate::error::EvaluationError;
|
||||
use crate::llm::LlmProvider;
|
||||
|
||||
/// Result of evaluating job success.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
@@ -64,233 +61,132 @@ pub trait SuccessEvaluator: Send + Sync {
|
||||
) -> Result<EvaluationResult, EvaluationError>;
|
||||
}
|
||||
|
||||
/// Rule-based success evaluator.
|
||||
pub struct RuleBasedEvaluator {
|
||||
/// Minimum success rate for actions.
|
||||
min_action_success_rate: f64,
|
||||
/// Maximum allowed failures.
|
||||
max_failures: u32,
|
||||
}
|
||||
|
||||
impl RuleBasedEvaluator {
|
||||
/// Create a new rule-based evaluator.
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
min_action_success_rate: 0.8,
|
||||
max_failures: 3,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set minimum action success rate.
|
||||
#[allow(dead_code)] // Public API for configuring evaluation threshold
|
||||
pub fn with_min_success_rate(mut self, rate: f64) -> Self {
|
||||
self.min_action_success_rate = rate;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set maximum failures.
|
||||
#[allow(dead_code)] // Public API for configuring failure tolerance
|
||||
pub fn with_max_failures(mut self, max: u32) -> Self {
|
||||
self.max_failures = max;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for RuleBasedEvaluator {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl SuccessEvaluator for RuleBasedEvaluator {
|
||||
async fn evaluate(
|
||||
&self,
|
||||
job: &JobContext,
|
||||
actions: &[ActionRecord],
|
||||
_output: Option<&str>,
|
||||
) -> Result<EvaluationResult, EvaluationError> {
|
||||
let mut issues = Vec::new();
|
||||
|
||||
// Check if there were any actions
|
||||
if actions.is_empty() {
|
||||
return Ok(EvaluationResult::failure(
|
||||
"No actions were taken",
|
||||
vec!["No actions recorded".to_string()],
|
||||
));
|
||||
}
|
||||
|
||||
// Calculate action success rate
|
||||
let successful = actions.iter().filter(|a| a.success).count();
|
||||
let total = actions.len();
|
||||
let success_rate = successful as f64 / total as f64;
|
||||
|
||||
if success_rate < self.min_action_success_rate {
|
||||
issues.push(format!(
|
||||
"Action success rate {:.1}% below threshold {:.1}%",
|
||||
success_rate * 100.0,
|
||||
self.min_action_success_rate * 100.0
|
||||
));
|
||||
}
|
||||
|
||||
// Count failures
|
||||
let failures = actions.iter().filter(|a| !a.success).count() as u32;
|
||||
if failures > self.max_failures {
|
||||
issues.push(format!(
|
||||
"Too many failures: {} (max {})",
|
||||
failures, self.max_failures
|
||||
));
|
||||
}
|
||||
|
||||
// Check for critical errors
|
||||
for action in actions.iter().filter(|a| !a.success) {
|
||||
if let Some(ref error) = action.error
|
||||
&& (error.to_lowercase().contains("critical")
|
||||
|| error.to_lowercase().contains("fatal"))
|
||||
{
|
||||
issues.push(format!("Critical error in {}: {}", action.tool_name, error));
|
||||
}
|
||||
}
|
||||
|
||||
// Check job state
|
||||
if job.state != crate::context::JobState::Completed
|
||||
&& job.state != crate::context::JobState::Submitted
|
||||
{
|
||||
issues.push(format!("Job not in completed state: {:?}", job.state));
|
||||
}
|
||||
|
||||
// Calculate quality score
|
||||
let quality_score = if issues.is_empty() {
|
||||
let base_score = (success_rate * 80.0) as u32;
|
||||
let completion_bonus = if job.state == crate::context::JobState::Completed {
|
||||
20
|
||||
} else {
|
||||
0
|
||||
};
|
||||
(base_score + completion_bonus).min(100)
|
||||
} else {
|
||||
((success_rate * 50.0) as u32).min(50)
|
||||
};
|
||||
|
||||
if issues.is_empty() {
|
||||
Ok(EvaluationResult::success(
|
||||
format!(
|
||||
"Job completed successfully with {}/{} actions succeeding ({:.1}%)",
|
||||
successful,
|
||||
total,
|
||||
success_rate * 100.0
|
||||
),
|
||||
quality_score,
|
||||
))
|
||||
} else {
|
||||
Ok(EvaluationResult {
|
||||
success: false,
|
||||
confidence: 0.85,
|
||||
reasoning: format!("Job had {} issues", issues.len()),
|
||||
issues,
|
||||
suggestions: vec![
|
||||
"Review failed actions for common patterns".to_string(),
|
||||
"Consider adjusting retry logic".to_string(),
|
||||
],
|
||||
quality_score,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// LLM-based success evaluator for more nuanced evaluation.
|
||||
pub struct LlmEvaluator {
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
}
|
||||
|
||||
impl LlmEvaluator {
|
||||
/// Create a new LLM-based evaluator.
|
||||
#[allow(dead_code)] // Public API for LLM-based evaluation
|
||||
pub fn new(llm: Arc<dyn LlmProvider>) -> Self {
|
||||
Self { llm }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl SuccessEvaluator for LlmEvaluator {
|
||||
async fn evaluate(
|
||||
&self,
|
||||
job: &JobContext,
|
||||
actions: &[ActionRecord],
|
||||
output: Option<&str>,
|
||||
) -> Result<EvaluationResult, EvaluationError> {
|
||||
// Build evaluation prompt
|
||||
let actions_summary: Vec<String> = actions
|
||||
.iter()
|
||||
.map(|a| {
|
||||
format!(
|
||||
"- {}: {} ({})",
|
||||
a.tool_name,
|
||||
if a.success { "success" } else { "failed" },
|
||||
a.error.as_deref().unwrap_or("ok")
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
|
||||
let prompt = format!(
|
||||
r#"Evaluate if this job was completed successfully.
|
||||
|
||||
Job: {}
|
||||
Description: {}
|
||||
State: {:?}
|
||||
|
||||
Actions taken:
|
||||
{}
|
||||
|
||||
{}
|
||||
|
||||
Respond in JSON format:
|
||||
{{
|
||||
"success": true/false,
|
||||
"confidence": 0.0-1.0,
|
||||
"reasoning": "...",
|
||||
"issues": ["..."],
|
||||
"suggestions": ["..."],
|
||||
"quality_score": 0-100
|
||||
}}"#,
|
||||
job.title,
|
||||
job.description,
|
||||
job.state,
|
||||
actions_summary.join("\n"),
|
||||
output
|
||||
.map(|o| format!("Output:\n{}", o))
|
||||
.unwrap_or_default()
|
||||
);
|
||||
|
||||
let request =
|
||||
crate::llm::CompletionRequest::new(vec![crate::llm::ChatMessage::user(prompt)])
|
||||
.with_max_tokens(1024)
|
||||
.with_temperature(0.1);
|
||||
|
||||
let response = self
|
||||
.llm
|
||||
.complete(request)
|
||||
.await
|
||||
.map_err(|e| EvaluationError::Failed {
|
||||
job_id: job.job_id,
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
// Parse the response
|
||||
let result: EvaluationResult =
|
||||
serde_json::from_str(&response.content).map_err(|e| EvaluationError::Failed {
|
||||
job_id: job.job_id,
|
||||
reason: format!("Failed to parse LLM evaluation: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::context::JobContext;
|
||||
use crate::context::{ActionRecord, JobContext};
|
||||
use crate::error::EvaluationError;
|
||||
|
||||
/// Rule-based success evaluator (test-only; no production callers).
|
||||
struct RuleBasedEvaluator {
|
||||
min_action_success_rate: f64,
|
||||
max_failures: u32,
|
||||
}
|
||||
|
||||
impl RuleBasedEvaluator {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
min_action_success_rate: 0.8,
|
||||
max_failures: 3,
|
||||
}
|
||||
}
|
||||
|
||||
fn with_min_success_rate(mut self, rate: f64) -> Self {
|
||||
self.min_action_success_rate = rate;
|
||||
self
|
||||
}
|
||||
|
||||
fn with_max_failures(mut self, max: u32) -> Self {
|
||||
self.max_failures = max;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for RuleBasedEvaluator {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl SuccessEvaluator for RuleBasedEvaluator {
|
||||
async fn evaluate(
|
||||
&self,
|
||||
job: &JobContext,
|
||||
actions: &[ActionRecord],
|
||||
_output: Option<&str>,
|
||||
) -> Result<EvaluationResult, EvaluationError> {
|
||||
let mut issues = Vec::new();
|
||||
|
||||
if actions.is_empty() {
|
||||
return Ok(EvaluationResult::failure(
|
||||
"No actions were taken",
|
||||
vec!["No actions recorded".to_string()],
|
||||
));
|
||||
}
|
||||
|
||||
let successful = actions.iter().filter(|a| a.success).count();
|
||||
let total = actions.len();
|
||||
let success_rate = successful as f64 / total as f64;
|
||||
|
||||
if success_rate < self.min_action_success_rate {
|
||||
issues.push(format!(
|
||||
"Action success rate {:.1}% below threshold {:.1}%",
|
||||
success_rate * 100.0,
|
||||
self.min_action_success_rate * 100.0
|
||||
));
|
||||
}
|
||||
|
||||
let failures = actions.iter().filter(|a| !a.success).count() as u32;
|
||||
if failures > self.max_failures {
|
||||
issues.push(format!(
|
||||
"Too many failures: {} (max {})",
|
||||
failures, self.max_failures
|
||||
));
|
||||
}
|
||||
|
||||
for action in actions.iter().filter(|a| !a.success) {
|
||||
if let Some(ref error) = action.error
|
||||
&& (error.to_lowercase().contains("critical")
|
||||
|| error.to_lowercase().contains("fatal"))
|
||||
{
|
||||
issues.push(format!("Critical error in {}: {}", action.tool_name, error));
|
||||
}
|
||||
}
|
||||
|
||||
if job.state != crate::context::JobState::Completed
|
||||
&& job.state != crate::context::JobState::Submitted
|
||||
{
|
||||
issues.push(format!("Job not in completed state: {:?}", job.state));
|
||||
}
|
||||
|
||||
let quality_score = if issues.is_empty() {
|
||||
let base_score = (success_rate * 80.0) as u32;
|
||||
let completion_bonus = if job.state == crate::context::JobState::Completed {
|
||||
20
|
||||
} else {
|
||||
0
|
||||
};
|
||||
(base_score + completion_bonus).min(100)
|
||||
} else {
|
||||
((success_rate * 50.0) as u32).min(50)
|
||||
};
|
||||
|
||||
if issues.is_empty() {
|
||||
Ok(EvaluationResult::success(
|
||||
format!(
|
||||
"Job completed successfully with {}/{} actions succeeding ({:.1}%)",
|
||||
successful,
|
||||
total,
|
||||
success_rate * 100.0
|
||||
),
|
||||
quality_score,
|
||||
))
|
||||
} else {
|
||||
Ok(EvaluationResult {
|
||||
success: false,
|
||||
confidence: 0.85,
|
||||
reasoning: format!("Job had {} issues", issues.len()),
|
||||
issues,
|
||||
suggestions: vec![
|
||||
"Review failed actions for common patterns".to_string(),
|
||||
"Consider adjusting retry logic".to_string(),
|
||||
],
|
||||
quality_score,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_rule_based_evaluator_success() {
|
||||
|
||||
@@ -1405,38 +1405,6 @@ impl ExtensionManager {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[allow(dead_code)] // Used by upcoming hot-activation flow
|
||||
async fn install_bundled_channel_from_artifacts(
|
||||
&self,
|
||||
name: &str,
|
||||
) -> Result<InstallResult, ExtensionError> {
|
||||
// Check if already installed
|
||||
let channel_wasm = self.wasm_channels_dir.join(format!("{}.wasm", name));
|
||||
if channel_wasm.exists() {
|
||||
return Err(ExtensionError::AlreadyInstalled(name.to_string()));
|
||||
}
|
||||
|
||||
crate::channels::wasm::install_bundled_channel(name, &self.wasm_channels_dir, false)
|
||||
.await
|
||||
.map_err(ExtensionError::InstallFailed)?;
|
||||
|
||||
tracing::info!(
|
||||
"Installed bundled channel '{}' to {}",
|
||||
name,
|
||||
self.wasm_channels_dir.display()
|
||||
);
|
||||
|
||||
Ok(InstallResult {
|
||||
name: name.to_string(),
|
||||
kind: ExtensionKind::WasmChannel,
|
||||
message: format!(
|
||||
"Channel '{}' installed. \
|
||||
Run tool_auth('{}') to configure authentication, then activate.",
|
||||
name, name,
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
/// Install a WASM extension from local build artifacts (WasmBuildable source).
|
||||
///
|
||||
/// Resolves the build directory (relative to `CARGO_MANIFEST_DIR` or absolute),
|
||||
@@ -1628,6 +1596,7 @@ impl ExtensionManager {
|
||||
&metadata.scopes_supported,
|
||||
Some(&pkce),
|
||||
&std::collections::HashMap::new(),
|
||||
None,
|
||||
);
|
||||
|
||||
// Store pending auth for later callback handling
|
||||
@@ -2508,7 +2477,7 @@ impl ExtensionManager {
|
||||
&self.user_id,
|
||||
)
|
||||
} else {
|
||||
McpClient::new_with_name(&server.name, &server.url)
|
||||
McpClient::new_with_config(server.clone())
|
||||
};
|
||||
|
||||
// Try to list and create tools
|
||||
|
||||
@@ -241,6 +241,9 @@ impl Store {
|
||||
tool_output_stash: std::sync::Arc::new(tokio::sync::RwLock::new(
|
||||
std::collections::HashMap::new(),
|
||||
)),
|
||||
// TODO(#661): persist user_timezone in agent_jobs table so
|
||||
// background/routine jobs retain the session's timezone context.
|
||||
user_timezone: "UTC".to_string(),
|
||||
}))
|
||||
}
|
||||
None => Ok(None),
|
||||
|
||||
@@ -66,6 +66,7 @@ pub mod service;
|
||||
pub mod settings;
|
||||
pub mod setup;
|
||||
pub mod skills;
|
||||
pub mod timezone;
|
||||
pub mod tools;
|
||||
pub mod tracing_fmt;
|
||||
pub mod transcription;
|
||||
|
||||
+1148
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,95 @@
|
||||
//! Image generation model detection utilities.
|
||||
|
||||
/// Known image generation model families.
|
||||
const IMAGE_GEN_PATTERNS: &[&str] = &[
|
||||
"flux",
|
||||
"dall-e",
|
||||
"dalle",
|
||||
"stable-diffusion",
|
||||
"sdxl",
|
||||
"imagen",
|
||||
"midjourney",
|
||||
"ideogram",
|
||||
"playground",
|
||||
];
|
||||
|
||||
/// Check if a model name indicates an image generation model.
|
||||
pub fn is_image_generation_model(model: &str) -> bool {
|
||||
let lower = model.to_lowercase();
|
||||
IMAGE_GEN_PATTERNS.iter().any(|p| lower.contains(p))
|
||||
}
|
||||
|
||||
/// Suggest the best image generation model from a list of available models.
|
||||
///
|
||||
/// Priority: FLUX > DALL-E > Stable Diffusion > others.
|
||||
pub fn suggest_image_model(models: &[String]) -> Option<&str> {
|
||||
let priorities: &[&str] = &[
|
||||
"flux",
|
||||
"dall-e",
|
||||
"dalle",
|
||||
"stable-diffusion",
|
||||
"sdxl",
|
||||
"imagen",
|
||||
];
|
||||
for priority in priorities {
|
||||
if let Some(model) = models.iter().find(|m| m.to_lowercase().contains(priority)) {
|
||||
return Some(model);
|
||||
}
|
||||
}
|
||||
// Fall back to any image gen model
|
||||
models.iter().find_map(|m| {
|
||||
if is_image_generation_model(m) {
|
||||
Some(m.as_str())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn detects_flux_models() {
|
||||
assert!(is_image_generation_model(
|
||||
"black-forest-labs/FLUX.1-schnell"
|
||||
));
|
||||
assert!(is_image_generation_model("flux-pro"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_dalle_models() {
|
||||
assert!(is_image_generation_model("dall-e-3"));
|
||||
assert!(is_image_generation_model("dalle-3"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_non_image_models() {
|
||||
assert!(!is_image_generation_model("gpt-4o"));
|
||||
assert!(!is_image_generation_model("claude-3-sonnet"));
|
||||
assert!(!is_image_generation_model("llama-3.1-70b"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn suggests_flux_first() {
|
||||
let models = vec![
|
||||
"gpt-4o".to_string(),
|
||||
"dall-e-3".to_string(),
|
||||
"flux-pro".to_string(),
|
||||
];
|
||||
assert_eq!(suggest_image_model(&models), Some("flux-pro"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn suggests_dalle_without_flux() {
|
||||
let models = vec!["gpt-4o".to_string(), "dall-e-3".to_string()];
|
||||
assert_eq!(suggest_image_model(&models), Some("dall-e-3"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_none_when_no_image_models() {
|
||||
let models = vec!["gpt-4o".to_string(), "claude-3-sonnet".to_string()];
|
||||
assert_eq!(suggest_image_model(&models), None);
|
||||
}
|
||||
}
|
||||
+64
-7
@@ -6,8 +6,11 @@
|
||||
//! - **Anthropic**: Direct API access with your own key
|
||||
//! - **Ollama**: Local model inference
|
||||
//! - **OpenAI-compatible**: Any endpoint that speaks the OpenAI API
|
||||
//! - **AWS Bedrock**: Native Converse API via aws-sdk-bedrockruntime
|
||||
|
||||
mod anthropic_oauth;
|
||||
#[cfg(feature = "bedrock")]
|
||||
mod bedrock;
|
||||
pub mod circuit_breaker;
|
||||
pub mod costs;
|
||||
pub mod failover;
|
||||
@@ -22,6 +25,9 @@ mod rig_adapter;
|
||||
pub mod session;
|
||||
pub mod smart_routing;
|
||||
|
||||
pub mod image_models;
|
||||
pub mod vision_models;
|
||||
|
||||
pub use circuit_breaker::{CircuitBreakerConfig, CircuitBreakerProvider};
|
||||
pub use failover::{CooldownConfig, FailoverProvider};
|
||||
pub use nearai_chat::{ModelInfo, NearAiChatProvider};
|
||||
@@ -54,12 +60,29 @@ use crate::error::LlmError;
|
||||
///
|
||||
/// - NearAI backend: Uses session manager for authentication
|
||||
/// - Registry providers: Looked up by protocol and constructed generically
|
||||
pub fn create_llm_provider(
|
||||
pub async fn create_llm_provider(
|
||||
config: &LlmConfig,
|
||||
session: Arc<SessionManager>,
|
||||
) -> Result<Arc<dyn LlmProvider>, LlmError> {
|
||||
let timeout = config.request_timeout_secs;
|
||||
|
||||
if config.backend == "nearai" || config.backend == "near_ai" || config.backend == "near" {
|
||||
return create_llm_provider_with_config(&config.nearai, session);
|
||||
return create_llm_provider_with_config(&config.nearai, session, timeout);
|
||||
}
|
||||
|
||||
// Bedrock uses a native AWS SDK, not the rig-core registry
|
||||
if config.backend == "bedrock" {
|
||||
#[cfg(feature = "bedrock")]
|
||||
{
|
||||
return create_bedrock_provider(config).await;
|
||||
}
|
||||
#[cfg(not(feature = "bedrock"))]
|
||||
{
|
||||
return Err(LlmError::RequestFailed {
|
||||
provider: "bedrock".to_string(),
|
||||
reason: "Bedrock support not compiled. Rebuild with --features bedrock".to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
let reg_config = config
|
||||
@@ -79,6 +102,7 @@ pub fn create_llm_provider(
|
||||
pub fn create_llm_provider_with_config(
|
||||
config: &NearAiConfig,
|
||||
session: Arc<SessionManager>,
|
||||
request_timeout_secs: u64,
|
||||
) -> Result<Arc<dyn LlmProvider>, LlmError> {
|
||||
let auth_mode = if config.api_key.is_some() {
|
||||
"API key"
|
||||
@@ -89,9 +113,14 @@ pub fn create_llm_provider_with_config(
|
||||
model = %config.model,
|
||||
base_url = %config.base_url,
|
||||
auth = auth_mode,
|
||||
timeout_secs = request_timeout_secs,
|
||||
"Using NEAR AI (Chat Completions API)"
|
||||
);
|
||||
Ok(Arc::new(NearAiChatProvider::new(config.clone(), session)?))
|
||||
Ok(Arc::new(NearAiChatProvider::new_with_timeout(
|
||||
config.clone(),
|
||||
session,
|
||||
request_timeout_secs,
|
||||
)?))
|
||||
}
|
||||
|
||||
/// Create a provider from a registry-resolved config.
|
||||
@@ -109,6 +138,24 @@ fn create_registry_provider(
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "bedrock")]
|
||||
async fn create_bedrock_provider(config: &LlmConfig) -> Result<Arc<dyn LlmProvider>, LlmError> {
|
||||
let br = config
|
||||
.bedrock
|
||||
.as_ref()
|
||||
.ok_or_else(|| LlmError::AuthFailed {
|
||||
provider: "bedrock".to_string(),
|
||||
})?;
|
||||
|
||||
let provider = bedrock::BedrockProvider::new(br).await?;
|
||||
tracing::info!(
|
||||
"Using AWS Bedrock (Converse API, region: {}, model: {})",
|
||||
br.region,
|
||||
provider.active_model_name(),
|
||||
);
|
||||
Ok(Arc::new(provider))
|
||||
}
|
||||
|
||||
fn create_openai_compat_from_registry(
|
||||
config: &RegistryProviderConfig,
|
||||
) -> Result<Arc<dyn LlmProvider>, LlmError> {
|
||||
@@ -333,7 +380,7 @@ pub fn create_cheap_llm_provider(
|
||||
/// This is the single source of truth for provider chain construction,
|
||||
/// called by both `main.rs` and `app.rs`.
|
||||
#[allow(clippy::type_complexity)]
|
||||
pub fn build_provider_chain(
|
||||
pub async fn build_provider_chain(
|
||||
config: &LlmConfig,
|
||||
session: Arc<SessionManager>,
|
||||
) -> Result<
|
||||
@@ -344,7 +391,7 @@ pub fn build_provider_chain(
|
||||
),
|
||||
LlmError,
|
||||
> {
|
||||
let llm = create_llm_provider(config, session.clone())?;
|
||||
let llm = create_llm_provider(config, session.clone()).await?;
|
||||
tracing::info!("LLM provider initialized: {}", llm.model_name());
|
||||
|
||||
// 1. Retry
|
||||
@@ -365,7 +412,11 @@ pub fn build_provider_chain(
|
||||
let llm: Arc<dyn LlmProvider> = if let Some(ref cheap_model) = config.nearai.cheap_model {
|
||||
let mut cheap_config = config.nearai.clone();
|
||||
cheap_config.model = cheap_model.clone();
|
||||
let cheap = create_llm_provider_with_config(&cheap_config, session.clone())?;
|
||||
let cheap = create_llm_provider_with_config(
|
||||
&cheap_config,
|
||||
session.clone(),
|
||||
config.request_timeout_secs,
|
||||
)?;
|
||||
let cheap: Arc<dyn LlmProvider> = if retry_config.max_retries > 0 {
|
||||
Arc::new(RetryProvider::new(cheap, retry_config.clone()))
|
||||
} else {
|
||||
@@ -397,7 +448,11 @@ pub fn build_provider_chain(
|
||||
}
|
||||
let mut fallback_config = config.nearai.clone();
|
||||
fallback_config.model = fallback_model.clone();
|
||||
let fallback = create_llm_provider_with_config(&fallback_config, session.clone())?;
|
||||
let fallback = create_llm_provider_with_config(
|
||||
&fallback_config,
|
||||
session.clone(),
|
||||
config.request_timeout_secs,
|
||||
)?;
|
||||
tracing::info!(
|
||||
primary = %llm.model_name(),
|
||||
fallback = %fallback.model_name(),
|
||||
@@ -503,6 +558,8 @@ mod tests {
|
||||
session: SessionConfig::default(),
|
||||
nearai: test_nearai_config(),
|
||||
provider: None,
|
||||
bedrock: None,
|
||||
request_timeout_secs: 120,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+15
-4
@@ -58,17 +58,28 @@ impl NearAiChatProvider {
|
||||
/// By default this enables tool-message flattening for compatibility with
|
||||
/// providers that reject `role: "tool"` messages.
|
||||
pub fn new(config: NearAiConfig, session: Arc<SessionManager>) -> Result<Self, LlmError> {
|
||||
Self::new_with_flatten(config, session, true)
|
||||
Self::new_with_options(config, session, true, 120)
|
||||
}
|
||||
|
||||
/// Create a chat completions provider with configurable tool-message flattening.
|
||||
pub fn new_with_flatten(
|
||||
/// Create a new provider with a custom request timeout.
|
||||
pub fn new_with_timeout(
|
||||
config: NearAiConfig,
|
||||
session: Arc<SessionManager>,
|
||||
request_timeout_secs: u64,
|
||||
) -> Result<Self, LlmError> {
|
||||
Self::new_with_options(config, session, true, request_timeout_secs)
|
||||
}
|
||||
|
||||
/// Create a chat completions provider with configurable tool-message flattening
|
||||
/// and request timeout.
|
||||
pub fn new_with_options(
|
||||
config: NearAiConfig,
|
||||
session: Arc<SessionManager>,
|
||||
flatten_tool_messages: bool,
|
||||
request_timeout_secs: u64,
|
||||
) -> Result<Self, LlmError> {
|
||||
let client = Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(120))
|
||||
.timeout(std::time::Duration::from_secs(request_timeout_secs))
|
||||
.build()
|
||||
.map_err(|e| LlmError::RequestFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
|
||||
+2
-12
@@ -11,7 +11,6 @@ use crate::llm::{
|
||||
ChatMessage, CompletionRequest, LlmProvider, Role, ToolCall, ToolCompletionRequest,
|
||||
ToolDefinition,
|
||||
};
|
||||
use crate::safety::SafetyLayer;
|
||||
|
||||
/// Token the agent returns when it has nothing to say (e.g. in group chats).
|
||||
/// The dispatcher should check for this and suppress the message.
|
||||
@@ -343,8 +342,6 @@ pub struct RespondOutput {
|
||||
/// Reasoning engine for the agent.
|
||||
pub struct Reasoning {
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
#[allow(dead_code)] // Will be used for sanitizing tool outputs
|
||||
safety: Arc<SafetyLayer>,
|
||||
/// Optional workspace for loading identity/system prompts.
|
||||
workspace_system_prompt: Option<String>,
|
||||
/// Optional skill context block to inject into system prompt.
|
||||
@@ -362,10 +359,9 @@ pub struct Reasoning {
|
||||
|
||||
impl Reasoning {
|
||||
/// Create a new reasoning engine.
|
||||
pub fn new(llm: Arc<dyn LlmProvider>, safety: Arc<SafetyLayer>) -> Self {
|
||||
pub fn new(llm: Arc<dyn LlmProvider>) -> Self {
|
||||
Self {
|
||||
llm,
|
||||
safety,
|
||||
workspace_system_prompt: None,
|
||||
skill_context: None,
|
||||
channel: None,
|
||||
@@ -2117,15 +2113,9 @@ That's my plan."#;
|
||||
// ---- System prompt building tests (issue #565) ----
|
||||
|
||||
fn make_test_reasoning() -> Reasoning {
|
||||
use crate::config::SafetyConfig;
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::testing::StubLlm;
|
||||
let llm = Arc::new(StubLlm::new("test"));
|
||||
let safety = Arc::new(SafetyLayer::new(&SafetyConfig {
|
||||
max_output_length: 100_000,
|
||||
injection_check_enabled: false,
|
||||
}));
|
||||
Reasoning::new(llm, safety)
|
||||
Reasoning::new(llm)
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
+2
-1
@@ -200,9 +200,10 @@ impl SessionManager {
|
||||
|
||||
let status = response.status();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
let preview = crate::agent::truncate_for_preview(&body, 200);
|
||||
Err(LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("Validation failed: HTTP {}: {}", status, body),
|
||||
reason: format!("Validation failed: HTTP {status}: {preview}"),
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,104 @@
|
||||
//! Vision model detection utilities.
|
||||
|
||||
/// Known vision-capable model families.
|
||||
const VISION_PATTERNS: &[&str] = &[
|
||||
"claude-3",
|
||||
"claude-4",
|
||||
"gpt-4o",
|
||||
"gpt-4-turbo",
|
||||
"gpt-4-vision",
|
||||
"gemini-pro-vision",
|
||||
"gemini-1.5",
|
||||
"gemini-2",
|
||||
"llava",
|
||||
"cogvlm",
|
||||
"internvl",
|
||||
"qwen-vl",
|
||||
"qwen2-vl",
|
||||
"pixtral",
|
||||
];
|
||||
|
||||
/// Check if a model name indicates vision capabilities.
|
||||
pub fn is_vision_model(model: &str) -> bool {
|
||||
let lower = model.to_lowercase();
|
||||
VISION_PATTERNS.iter().any(|p| lower.contains(p))
|
||||
}
|
||||
|
||||
/// Suggest the best vision model from a list of available models.
|
||||
///
|
||||
/// Priority: Claude > GPT-4 > Gemini > others.
|
||||
pub fn suggest_vision_model(models: &[String]) -> Option<&str> {
|
||||
let priorities: &[&str] = &[
|
||||
"claude-3",
|
||||
"claude-4",
|
||||
"gpt-4o",
|
||||
"gpt-4-turbo",
|
||||
"gpt-4-vision",
|
||||
"gemini",
|
||||
"llava",
|
||||
"pixtral",
|
||||
];
|
||||
for priority in priorities {
|
||||
if let Some(model) = models.iter().find(|m| m.to_lowercase().contains(priority)) {
|
||||
return Some(model);
|
||||
}
|
||||
}
|
||||
models.iter().find_map(|m| {
|
||||
if is_vision_model(m) {
|
||||
Some(m.as_str())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn detects_claude_vision() {
|
||||
assert!(is_vision_model("claude-3-5-sonnet-20241022"));
|
||||
assert!(is_vision_model("claude-3-opus"));
|
||||
assert!(is_vision_model("claude-4-sonnet"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_gpt4_vision() {
|
||||
assert!(is_vision_model("gpt-4o"));
|
||||
assert!(is_vision_model("gpt-4-turbo"));
|
||||
assert!(is_vision_model("gpt-4-vision-preview"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_other_vision_models() {
|
||||
assert!(is_vision_model("gemini-1.5-pro"));
|
||||
assert!(is_vision_model("llava-v1.6"));
|
||||
assert!(is_vision_model("pixtral-12b"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_non_vision_models() {
|
||||
assert!(!is_vision_model("gpt-3.5-turbo"));
|
||||
assert!(!is_vision_model("llama-3.1-70b"));
|
||||
assert!(!is_vision_model("mistral-7b"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn suggests_claude_first() {
|
||||
let models = vec![
|
||||
"gpt-4o".to_string(),
|
||||
"claude-3-5-sonnet-20241022".to_string(),
|
||||
];
|
||||
assert_eq!(
|
||||
suggest_vision_model(&models),
|
||||
Some("claude-3-5-sonnet-20241022")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_none_when_no_vision_models() {
|
||||
let models = vec!["gpt-3.5-turbo".to_string(), "llama-3.1-70b".to_string()];
|
||||
assert_eq!(suggest_vision_model(&models), None);
|
||||
}
|
||||
}
|
||||
+27
-7
@@ -75,7 +75,7 @@ async fn async_main() -> anyhow::Result<()> {
|
||||
}
|
||||
Some(Command::Mcp(mcp_cmd)) => {
|
||||
init_cli_tracing();
|
||||
return run_mcp_command(mcp_cmd.clone()).await;
|
||||
return run_mcp_command(*mcp_cmd.clone()).await;
|
||||
}
|
||||
Some(Command::Memory(mem_cmd)) => {
|
||||
init_cli_tracing();
|
||||
@@ -145,6 +145,24 @@ async fn async_main() -> anyhow::Result<()> {
|
||||
}
|
||||
}
|
||||
|
||||
// ── PID lock (prevent multiple instances) ────────────────────────
|
||||
let _pid_lock = match ironclaw::bootstrap::PidLock::acquire() {
|
||||
Ok(lock) => Some(lock),
|
||||
Err(ironclaw::bootstrap::PidLockError::AlreadyRunning { pid }) => {
|
||||
anyhow::bail!(
|
||||
"Another IronClaw instance is already running (PID {}). \
|
||||
If this is incorrect, remove the stale PID file: {}",
|
||||
pid,
|
||||
ironclaw::bootstrap::pid_lock_path().display()
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("Warning: Could not acquire PID lock: {}", e);
|
||||
eprintln!("Continuing without PID lock protection.");
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
// ── Agent startup ──────────────────────────────────────────────────
|
||||
|
||||
// Enhanced first-run detection
|
||||
@@ -166,13 +184,12 @@ async fn async_main() -> anyhow::Result<()> {
|
||||
let config = match Config::from_env_with_toml(toml_path).await {
|
||||
Ok(c) => c,
|
||||
Err(ironclaw::error::ConfigError::MissingRequired { key, hint }) => {
|
||||
eprintln!("Configuration error: Missing required setting '{}'", key);
|
||||
eprintln!(" {}", hint);
|
||||
eprintln!();
|
||||
eprintln!(
|
||||
"Run 'ironclaw onboard' to configure, or set the required environment variables."
|
||||
anyhow::bail!(
|
||||
"Configuration error: Missing required setting '{}'. {}. \
|
||||
Run 'ironclaw onboard' to configure, or set the required environment variables.",
|
||||
key,
|
||||
hint
|
||||
);
|
||||
std::process::exit(1);
|
||||
}
|
||||
Err(e) => return Err(e.into()),
|
||||
};
|
||||
@@ -706,6 +723,9 @@ async fn async_main() -> anyhow::Result<()> {
|
||||
|
||||
// ── Shutdown ────────────────────────────────────────────────────────
|
||||
|
||||
// Shut down all stdio MCP server child processes.
|
||||
components.mcp_process_manager.shutdown_all().await;
|
||||
|
||||
// Flush LLM trace recording if enabled
|
||||
if let Some(ref recorder) = components.recording_handle
|
||||
&& let Err(e) = recorder.flush().await
|
||||
|
||||
+66
-1
@@ -42,8 +42,12 @@ pub struct Settings {
|
||||
#[serde(default)]
|
||||
pub secrets_master_key_source: KeySource,
|
||||
|
||||
/// Generated master key hex (env var mode only, written to .env by wizard).
|
||||
#[serde(default, skip_serializing)]
|
||||
pub secrets_master_key_hex: Option<String>,
|
||||
|
||||
// === Step 3: Inference Provider ===
|
||||
/// LLM backend: "nearai", "anthropic", "openai", "ollama", "openai_compatible".
|
||||
/// LLM backend: "nearai", "anthropic", "openai", "ollama", "openai_compatible", "tinfoil", "bedrock".
|
||||
#[serde(default)]
|
||||
pub llm_backend: Option<String>,
|
||||
|
||||
@@ -55,6 +59,18 @@ pub struct Settings {
|
||||
#[serde(default)]
|
||||
pub openai_compatible_base_url: Option<String>,
|
||||
|
||||
/// Bedrock region (when llm_backend = "bedrock").
|
||||
#[serde(default)]
|
||||
pub bedrock_region: Option<String>,
|
||||
|
||||
/// Bedrock cross-region inference prefix (when llm_backend = "bedrock").
|
||||
#[serde(default)]
|
||||
pub bedrock_cross_region: Option<String>,
|
||||
|
||||
/// AWS profile name for Bedrock (when llm_backend = "bedrock").
|
||||
#[serde(default)]
|
||||
pub bedrock_profile: Option<String>,
|
||||
|
||||
// === Step 4: Model Selection ===
|
||||
/// Currently selected model.
|
||||
#[serde(default)]
|
||||
@@ -291,6 +307,18 @@ pub struct HeartbeatSettings {
|
||||
/// User ID to notify on heartbeat findings.
|
||||
#[serde(default)]
|
||||
pub notify_user: Option<String>,
|
||||
|
||||
/// Hour (0-23) when quiet hours start (heartbeat skipped).
|
||||
#[serde(default)]
|
||||
pub quiet_hours_start: Option<u32>,
|
||||
|
||||
/// Hour (0-23) when quiet hours end (heartbeat resumes).
|
||||
#[serde(default)]
|
||||
pub quiet_hours_end: Option<u32>,
|
||||
|
||||
/// Timezone for quiet hours evaluation (IANA name, e.g. "America/New_York").
|
||||
#[serde(default)]
|
||||
pub timezone: Option<String>,
|
||||
}
|
||||
|
||||
fn default_heartbeat_interval() -> u64 {
|
||||
@@ -304,6 +332,9 @@ impl Default for HeartbeatSettings {
|
||||
interval_secs: default_heartbeat_interval(),
|
||||
notify_channel: None,
|
||||
notify_user: None,
|
||||
quiet_hours_start: None,
|
||||
quiet_hours_end: None,
|
||||
timezone: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -351,6 +382,10 @@ pub struct AgentSettings {
|
||||
/// When true, skip tool approval checks entirely. For benchmarks/CI.
|
||||
#[serde(default)]
|
||||
pub auto_approve_tools: bool,
|
||||
|
||||
/// Default timezone for new sessions (IANA name, e.g. "America/New_York").
|
||||
#[serde(default = "default_timezone")]
|
||||
pub default_timezone: String,
|
||||
}
|
||||
|
||||
fn default_agent_name() -> String {
|
||||
@@ -385,6 +420,10 @@ fn default_max_tool_iterations() -> usize {
|
||||
50
|
||||
}
|
||||
|
||||
fn default_timezone() -> String {
|
||||
"UTC".to_string()
|
||||
}
|
||||
|
||||
fn default_true() -> bool {
|
||||
true
|
||||
}
|
||||
@@ -402,6 +441,7 @@ impl Default for AgentSettings {
|
||||
session_idle_timeout_secs: default_session_idle_timeout(),
|
||||
max_tool_iterations: default_max_tool_iterations(),
|
||||
auto_approve_tools: false,
|
||||
default_timezone: default_timezone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1174,6 +1214,31 @@ mod tests {
|
||||
assert_eq!(loaded.heartbeat.interval_secs, 900);
|
||||
}
|
||||
|
||||
/// Regression test: /model command must persist selected_model to TOML config.
|
||||
/// Prior to the fix, `set_model()` only changed the in-memory provider and the
|
||||
/// choice was lost on restart.
|
||||
#[test]
|
||||
fn toml_selected_model_update_persists() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("config.toml");
|
||||
|
||||
// Start with a config that has a different model.
|
||||
let settings = Settings {
|
||||
selected_model: Some("old-model".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
settings.save_toml(&path).unwrap();
|
||||
|
||||
// Simulate what persist_selected_model does: load, update, save.
|
||||
let mut loaded = Settings::load_toml(&path).unwrap().unwrap();
|
||||
loaded.selected_model = Some("new-model".to_string());
|
||||
loaded.save_toml(&path).unwrap();
|
||||
|
||||
// Verify the change survived a reload.
|
||||
let reloaded = Settings::load_toml(&path).unwrap().unwrap();
|
||||
assert_eq!(reloaded.selected_model, Some("new-model".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn toml_missing_file_returns_none() {
|
||||
let result = Settings::load_toml(std::path::Path::new("/tmp/nonexistent_config.toml"));
|
||||
|
||||
+2
-1
@@ -174,6 +174,7 @@ env-var mode or skipped secrets.
|
||||
| Ollama | None | - | - |
|
||||
| OpenRouter¹ | API key | `llm_compatible_api_key` | `LLM_API_KEY` |
|
||||
| OpenAI-compatible¹ | Optional API key | `llm_compatible_api_key` | `LLM_API_KEY` |
|
||||
| AWS Bedrock | AWS credentials (IAM, SSO, instance roles) | - | - |
|
||||
|
||||
¹ OpenRouter and OpenAI-compatible share the same secret name and env var because
|
||||
OpenRouter is stored as `llm_backend = "openai_compatible"` under the hood.
|
||||
@@ -479,7 +480,7 @@ pub struct Settings {
|
||||
pub secrets_master_key_source: KeySource, // Keychain | Env | None
|
||||
|
||||
// Step 3: Inference
|
||||
pub llm_backend: Option<String>, // "nearai" | "anthropic" | "openai" | "ollama" | "openai_compatible"
|
||||
pub llm_backend: Option<String>, // "nearai" | "anthropic" | "openai" | "ollama" | "openai_compatible" | "bedrock"
|
||||
pub ollama_base_url: Option<String>,
|
||||
pub openai_compatible_base_url: Option<String>,
|
||||
|
||||
|
||||
+330
-43
@@ -769,13 +769,28 @@ impl SetupWizard {
|
||||
print_success("Master key generated and stored in OS keychain");
|
||||
}
|
||||
1 => {
|
||||
// Env var mode
|
||||
print_info("Generate a key and add it to your environment:");
|
||||
// Env var mode — generate key, init crypto, and persist to .env
|
||||
let key_hex = crate::secrets::keychain::generate_master_key_hex();
|
||||
|
||||
// Initialize crypto so subsequent wizard steps (channel setup,
|
||||
// API key storage) can encrypt secrets immediately.
|
||||
self.secrets_crypto = Some(Arc::new(
|
||||
SecretsCrypto::new(SecretString::from(key_hex.clone()))
|
||||
.map_err(|e| SetupError::Config(e.to_string()))?,
|
||||
));
|
||||
|
||||
// Make visible to optional_env() for any subsequent config resolution.
|
||||
crate::config::inject_single_var("SECRETS_MASTER_KEY", &key_hex);
|
||||
|
||||
// Store hex for write_bootstrap_env to persist to ~/.ironclaw/.env.
|
||||
self.settings.secrets_master_key_hex = Some(key_hex.clone());
|
||||
|
||||
println!();
|
||||
println!(" export SECRETS_MASTER_KEY={}", key_hex);
|
||||
print_info("Master key generated and will be saved to ~/.ironclaw/.env");
|
||||
println!();
|
||||
print_info("Add this to your shell profile or .env file.");
|
||||
println!(" SECRETS_MASTER_KEY={}", key_hex);
|
||||
println!();
|
||||
print_info("You can also copy this to another .env file or CI secrets.");
|
||||
|
||||
self.settings.secrets_master_key_source = KeySource::Env;
|
||||
print_success("Configured for environment variable");
|
||||
@@ -812,9 +827,16 @@ impl SetupWizard {
|
||||
print_info(&format!("Current provider: {}", display));
|
||||
println!();
|
||||
|
||||
let is_known = current == "nearai" || registry.is_known(¤t);
|
||||
let is_known =
|
||||
current == "nearai" || current == "bedrock" || registry.is_known(¤t);
|
||||
|
||||
if is_known && confirm("Keep current provider?", true).map_err(SetupError::Io)? {
|
||||
if current == "bedrock" {
|
||||
// Keeping the existing Bedrock config — no need to re-run
|
||||
// the full setup flow (region, auth, cross-region).
|
||||
print_info("Keeping existing AWS Bedrock configuration.");
|
||||
return Ok(());
|
||||
}
|
||||
return self.run_provider_setup(¤t, ®istry).await;
|
||||
}
|
||||
|
||||
@@ -829,10 +851,10 @@ impl SetupWizard {
|
||||
print_info("Select your inference provider:");
|
||||
println!();
|
||||
|
||||
// Build menu: NearAI first, then all registry providers with setup hints
|
||||
// Build menu: NearAI first, then all registry providers with setup hints, then Bedrock
|
||||
let selectable = registry.selectable();
|
||||
let mut options: Vec<String> = Vec::with_capacity(1 + selectable.len());
|
||||
let mut provider_ids: Vec<String> = Vec::with_capacity(1 + selectable.len());
|
||||
let mut options: Vec<String> = Vec::with_capacity(2 + selectable.len());
|
||||
let mut provider_ids: Vec<String> = Vec::with_capacity(2 + selectable.len());
|
||||
|
||||
options.push("NEAR AI - multi-model access via NEAR account".to_string());
|
||||
provider_ids.push("nearai".to_string());
|
||||
@@ -850,11 +872,19 @@ impl SetupWizard {
|
||||
provider_ids.push(def.id.clone());
|
||||
}
|
||||
|
||||
// Bedrock is a special case (native AWS SDK, not registry-based)
|
||||
options.push("AWS Bedrock - Claude & other models via AWS (IAM, SSO)".to_string());
|
||||
provider_ids.push("bedrock".to_string());
|
||||
|
||||
let option_refs: Vec<&str> = options.iter().map(|s| s.as_str()).collect();
|
||||
let choice = select_one("Provider:", &option_refs).map_err(SetupError::Io)?;
|
||||
let selected_id = &provider_ids[choice];
|
||||
|
||||
self.run_provider_setup(selected_id, ®istry).await?;
|
||||
if selected_id == "bedrock" {
|
||||
self.setup_bedrock().await?;
|
||||
} else {
|
||||
self.run_provider_setup(selected_id, ®istry).await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -1022,10 +1052,11 @@ impl SetupWizard {
|
||||
|
||||
/// Anthropic OAuth setup: extract token from `claude login` credentials.
|
||||
async fn setup_anthropic_oauth(&mut self) -> Result<(), SetupError> {
|
||||
self.settings.llm_backend = Some("anthropic".to_string());
|
||||
if self.settings.selected_model.is_some() {
|
||||
// Clear model only when switching providers (old model may be invalid)
|
||||
if self.settings.llm_backend.as_deref() != Some("anthropic") {
|
||||
self.settings.selected_model = None;
|
||||
}
|
||||
self.settings.llm_backend = Some("anthropic".to_string());
|
||||
|
||||
// Try to extract existing OAuth token from Claude Code credentials
|
||||
if let Some(token) = crate::config::ClaudeCodeConfig::extract_oauth_token() {
|
||||
@@ -1119,10 +1150,11 @@ impl SetupWizard {
|
||||
other => other,
|
||||
});
|
||||
|
||||
self.settings.llm_backend = Some(backend.to_string());
|
||||
if self.settings.selected_model.is_some() {
|
||||
// Clear model only when switching providers (old model may be invalid)
|
||||
if self.settings.llm_backend.as_deref() != Some(backend) {
|
||||
self.settings.selected_model = None;
|
||||
}
|
||||
self.settings.llm_backend = Some(backend.to_string());
|
||||
|
||||
// Check env var first
|
||||
if let Ok(existing) = std::env::var(env_var) {
|
||||
@@ -1181,10 +1213,11 @@ impl SetupWizard {
|
||||
&mut self,
|
||||
def: &crate::llm::ProviderDefinition,
|
||||
) -> Result<(), SetupError> {
|
||||
self.settings.llm_backend = Some(def.id.clone());
|
||||
if self.settings.selected_model.is_some() {
|
||||
// Clear model only when switching providers (old model may be invalid)
|
||||
if self.settings.llm_backend.as_deref() != Some(&def.id) {
|
||||
self.settings.selected_model = None;
|
||||
}
|
||||
self.settings.llm_backend = Some(def.id.clone());
|
||||
|
||||
let default_url = self
|
||||
.settings
|
||||
@@ -1212,6 +1245,95 @@ impl SetupWizard {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// AWS Bedrock provider setup: region, auth, and cross-region config.
|
||||
async fn setup_bedrock(&mut self) -> Result<(), SetupError> {
|
||||
if self.settings.llm_backend.as_deref() != Some("bedrock") {
|
||||
self.settings.selected_model = None;
|
||||
}
|
||||
self.settings.llm_backend = Some("bedrock".to_string());
|
||||
|
||||
// Region
|
||||
let default_region = self
|
||||
.settings
|
||||
.bedrock_region
|
||||
.as_deref()
|
||||
.unwrap_or("us-east-1");
|
||||
|
||||
let region_input =
|
||||
optional_input("AWS region", Some(&format!("default: {}", default_region)))
|
||||
.map_err(SetupError::Io)?;
|
||||
|
||||
let region = region_input.unwrap_or_else(|| default_region.to_string());
|
||||
self.settings.bedrock_region = Some(region.clone());
|
||||
|
||||
// Auth method
|
||||
print_info("Select authentication method:");
|
||||
println!();
|
||||
let auth_options = &[
|
||||
"AWS default credentials (env vars, ~/.aws/credentials, IAM roles)",
|
||||
"AWS named profile (SSO / assume-role)",
|
||||
];
|
||||
let auth_choice = select_one("Auth:", auth_options).map_err(SetupError::Io)?;
|
||||
|
||||
match auth_choice {
|
||||
0 => {
|
||||
// Default AWS credentials — clear any stale named profile
|
||||
self.settings.bedrock_profile = None;
|
||||
print_info(
|
||||
"Using default AWS credential chain (env vars, ~/.aws/credentials, IAM roles).",
|
||||
);
|
||||
}
|
||||
1 => {
|
||||
// Named profile
|
||||
let profile =
|
||||
input("AWS profile name (from ~/.aws/config)").map_err(SetupError::Io)?;
|
||||
if profile.trim().is_empty() {
|
||||
// Empty input clears any previously configured profile
|
||||
self.settings.bedrock_profile = None;
|
||||
print_info("AWS profile cleared; using default AWS credential chain instead.");
|
||||
} else {
|
||||
self.settings.bedrock_profile = Some(profile.clone());
|
||||
print_success(&format!("AWS profile '{}' saved", profile));
|
||||
}
|
||||
}
|
||||
_ => return Err(SetupError::Config("Invalid auth selection".to_string())),
|
||||
}
|
||||
|
||||
self.setup_bedrock_cross_region()
|
||||
}
|
||||
|
||||
/// Bedrock cross-region inference prefix selection (sub-step of setup_bedrock).
|
||||
fn setup_bedrock_cross_region(&mut self) -> Result<(), SetupError> {
|
||||
print_info("Cross-region inference routes requests across AWS regions for capacity:");
|
||||
println!();
|
||||
let cross_options = &[
|
||||
"us - route within US regions (recommended for us-east-1)",
|
||||
"global - route to any AWS region worldwide",
|
||||
"eu - route within European regions",
|
||||
"apac - route within Asia-Pacific regions",
|
||||
"none - single-region only (no cross-region routing)",
|
||||
];
|
||||
let cross_choice = select_one("Cross-region:", cross_options).map_err(SetupError::Io)?;
|
||||
|
||||
let cross_region = match cross_choice {
|
||||
0 => Some("us".to_string()),
|
||||
1 => Some("global".to_string()),
|
||||
2 => Some("eu".to_string()),
|
||||
3 => Some("apac".to_string()),
|
||||
4 => None,
|
||||
_ => None,
|
||||
};
|
||||
self.settings.bedrock_cross_region = cross_region;
|
||||
|
||||
let region = self
|
||||
.settings
|
||||
.bedrock_region
|
||||
.as_deref()
|
||||
.unwrap_or("us-east-1");
|
||||
print_success(&format!("AWS Bedrock configured (region: {})", region));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Generic OpenAI-compatible setup: base URL + optional API key.
|
||||
async fn setup_openai_compatible_generic(
|
||||
&mut self,
|
||||
@@ -1219,10 +1341,11 @@ impl SetupWizard {
|
||||
secret_name: &str,
|
||||
display_name: &str,
|
||||
) -> Result<(), SetupError> {
|
||||
self.settings.llm_backend = Some(backend_id.to_string());
|
||||
if self.settings.selected_model.is_some() {
|
||||
// Clear model only when switching providers (old model may be invalid)
|
||||
if self.settings.llm_backend.as_deref() != Some(backend_id) {
|
||||
self.settings.selected_model = None;
|
||||
}
|
||||
self.settings.llm_backend = Some(backend_id.to_string());
|
||||
|
||||
let existing_url = self
|
||||
.settings
|
||||
@@ -1393,6 +1516,14 @@ impl SetupWizard {
|
||||
self.settings.selected_model = Some(model_id.clone());
|
||||
print_success(&format!("Selected {}", model_id));
|
||||
}
|
||||
} else if backend == "bedrock" {
|
||||
let model_id = input("Bedrock model ID (e.g., anthropic.claude-opus-4-6-v1)")
|
||||
.map_err(SetupError::Io)?;
|
||||
if model_id.is_empty() {
|
||||
return Err(SetupError::Config("Model ID is required".to_string()));
|
||||
}
|
||||
self.settings.selected_model = Some(model_id.clone());
|
||||
print_success(&format!("Selected {}", model_id));
|
||||
} else {
|
||||
// Unknown provider, manual entry
|
||||
let model_id = input("Model name (e.g., meta-llama/Llama-3-8b-chat-hf)")
|
||||
@@ -1476,9 +1607,11 @@ impl SetupWizard {
|
||||
smart_routing_cascade: true,
|
||||
},
|
||||
provider: None,
|
||||
bedrock: None,
|
||||
request_timeout_secs: 120,
|
||||
};
|
||||
|
||||
match create_llm_provider(&config, session) {
|
||||
match create_llm_provider(&config, session).await {
|
||||
Ok(provider) => match provider.list_models().await {
|
||||
Ok(models) => models,
|
||||
Err(e) => {
|
||||
@@ -1592,48 +1725,50 @@ impl SetupWizard {
|
||||
};
|
||||
|
||||
// Create backend-appropriate secrets store.
|
||||
// Respect the user's selected backend when both features are compiled,
|
||||
// so we don't accidentally use a postgres pool from DATABASE_URL when
|
||||
// libsql was chosen (or vice versa).
|
||||
// Use runtime dispatch based on the user's selected backend.
|
||||
// Default to whichever backend is compiled in. When only libsql is
|
||||
// available, we must not default to "postgres" or we'd skip store creation.
|
||||
let default_backend = {
|
||||
#[cfg(feature = "postgres")]
|
||||
{
|
||||
"postgres"
|
||||
}
|
||||
#[cfg(not(feature = "postgres"))]
|
||||
{
|
||||
"libsql"
|
||||
}
|
||||
};
|
||||
let selected_backend = self
|
||||
.settings
|
||||
.database_backend
|
||||
.as_deref()
|
||||
.unwrap_or("postgres");
|
||||
.unwrap_or(default_backend);
|
||||
|
||||
#[cfg(all(feature = "libsql", feature = "postgres"))]
|
||||
{
|
||||
if selected_backend == "libsql" {
|
||||
match selected_backend {
|
||||
#[cfg(feature = "libsql")]
|
||||
"libsql" | "turso" | "sqlite" => {
|
||||
if let Some(store) = self.create_libsql_secrets_store(&crypto)? {
|
||||
return Ok(SecretsContext::from_store(store, "default"));
|
||||
}
|
||||
// Fallback to postgres if libsql store creation returned None
|
||||
#[cfg(feature = "postgres")]
|
||||
if let Some(store) = self.create_postgres_secrets_store(&crypto).await? {
|
||||
return Ok(SecretsContext::from_store(store, "default"));
|
||||
}
|
||||
} else {
|
||||
}
|
||||
#[cfg(feature = "postgres")]
|
||||
_ => {
|
||||
if let Some(store) = self.create_postgres_secrets_store(&crypto).await? {
|
||||
return Ok(SecretsContext::from_store(store, "default"));
|
||||
}
|
||||
// Fallback to libsql if postgres store creation returned None
|
||||
#[cfg(feature = "libsql")]
|
||||
if let Some(store) = self.create_libsql_secrets_store(&crypto)? {
|
||||
return Ok(SecretsContext::from_store(store, "default"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "postgres", not(feature = "libsql")))]
|
||||
{
|
||||
let _ = selected_backend;
|
||||
if let Some(store) = self.create_postgres_secrets_store(&crypto).await? {
|
||||
return Ok(SecretsContext::from_store(store, "default"));
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "libsql", not(feature = "postgres")))]
|
||||
{
|
||||
let _ = selected_backend;
|
||||
if let Some(store) = self.create_libsql_secrets_store(&crypto)? {
|
||||
return Ok(SecretsContext::from_store(store, "default"));
|
||||
}
|
||||
#[cfg(not(feature = "postgres"))]
|
||||
_ => {}
|
||||
}
|
||||
|
||||
Err(SetupError::Config(
|
||||
@@ -2293,12 +2428,29 @@ impl SetupWizard {
|
||||
if let Some(ref url) = self.settings.ollama_base_url {
|
||||
env_vars.push(("OLLAMA_BASE_URL".to_string(), url.clone()));
|
||||
}
|
||||
if let Some(ref region) = self.settings.bedrock_region {
|
||||
env_vars.push(("BEDROCK_REGION".to_string(), region.clone()));
|
||||
}
|
||||
if self.settings.llm_backend.as_deref() == Some("bedrock") {
|
||||
if let Some(ref model) = self.settings.selected_model {
|
||||
env_vars.push(("BEDROCK_MODEL".to_string(), model.clone()));
|
||||
}
|
||||
if let Some(ref cross) = self.settings.bedrock_cross_region {
|
||||
env_vars.push(("BEDROCK_CROSS_REGION".to_string(), cross.clone()));
|
||||
}
|
||||
if let Some(ref profile) = self.settings.bedrock_profile {
|
||||
env_vars.push(("AWS_PROFILE".to_string(), profile.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
// Model name: same chicken-and-egg — Config::from_env() resolves the
|
||||
// model before the DB is connected, so we must persist it to .env.
|
||||
// Write the backend-specific env var so the correct resolution path
|
||||
// picks it up (looked up from the provider registry).
|
||||
if let Some(ref model) = self.settings.selected_model {
|
||||
// Bedrock model is already written above as BEDROCK_MODEL, skip here.
|
||||
if self.settings.llm_backend.as_deref() != Some("bedrock")
|
||||
&& let Some(ref model) = self.settings.selected_model
|
||||
{
|
||||
let backend_str = self.settings.llm_backend.as_deref().unwrap_or("nearai");
|
||||
let model_env = registry.model_env_var(backend_str);
|
||||
env_vars.push((model_env.to_string(), model.clone()));
|
||||
@@ -2324,6 +2476,12 @@ impl SetupWizard {
|
||||
env_vars.push(("NEARAI_API_KEY".to_string(), api_key));
|
||||
}
|
||||
|
||||
// Secrets master key (env var mode): write to .env so it's available
|
||||
// on next startup before the DB is connected.
|
||||
if let Some(ref key_hex) = self.settings.secrets_master_key_hex {
|
||||
env_vars.push(("SECRETS_MASTER_KEY".to_string(), key_hex.clone()));
|
||||
}
|
||||
|
||||
// Always write ONBOARD_COMPLETED so that check_onboard_needed()
|
||||
// (which runs before the DB is connected) knows to skip re-onboarding.
|
||||
if self.settings.onboard_completed {
|
||||
@@ -2577,6 +2735,7 @@ impl SetupWizard {
|
||||
"openai" => "OpenAI",
|
||||
"ollama" => "Ollama",
|
||||
"openai_compatible" => "OpenAI-compatible",
|
||||
"bedrock" => "AWS Bedrock",
|
||||
other => other,
|
||||
};
|
||||
println!(" Provider: {}", display);
|
||||
@@ -3499,6 +3658,108 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
/// Regression test for #600: re-running provider setup for the same backend
|
||||
/// must NOT clear selected_model. Only switching to a different backend should.
|
||||
#[test]
|
||||
fn test_same_provider_preserves_selected_model() {
|
||||
let mut wizard = SetupWizard::new();
|
||||
wizard.settings.llm_backend = Some("ollama".to_string());
|
||||
wizard.settings.selected_model = Some("llama3".to_string());
|
||||
|
||||
// Simulate re-entering the same provider -- model should survive
|
||||
// (This is the check that each setup_* function now performs)
|
||||
if wizard.settings.llm_backend.as_deref() != Some("ollama") {
|
||||
wizard.settings.selected_model = None;
|
||||
}
|
||||
wizard.settings.llm_backend = Some("ollama".to_string());
|
||||
|
||||
assert_eq!(
|
||||
wizard.settings.selected_model.as_deref(),
|
||||
Some("llama3"),
|
||||
"model should be preserved when re-selecting the same provider"
|
||||
);
|
||||
}
|
||||
|
||||
/// Regression test for #600: switching to a different provider must clear
|
||||
/// selected_model since the old model may not be valid for the new backend.
|
||||
#[test]
|
||||
fn test_different_provider_clears_selected_model() {
|
||||
let mut wizard = SetupWizard::new();
|
||||
wizard.settings.llm_backend = Some("ollama".to_string());
|
||||
wizard.settings.selected_model = Some("llama3".to_string());
|
||||
|
||||
// Simulate switching to a different provider -- model should be cleared
|
||||
if wizard.settings.llm_backend.as_deref() != Some("openai") {
|
||||
wizard.settings.selected_model = None;
|
||||
}
|
||||
wizard.settings.llm_backend = Some("openai".to_string());
|
||||
|
||||
assert!(
|
||||
wizard.settings.selected_model.is_none(),
|
||||
"model should be cleared when switching providers"
|
||||
);
|
||||
}
|
||||
|
||||
/// Regression: Bedrock setup_bedrock() should preserve selected_model
|
||||
/// when re-entering the same provider (matches pattern from #600).
|
||||
#[test]
|
||||
fn test_bedrock_same_provider_preserves_model() {
|
||||
let mut wizard = SetupWizard::new();
|
||||
wizard.settings.llm_backend = Some("bedrock".to_string());
|
||||
wizard.settings.selected_model = Some("anthropic.claude-opus-4-6-v1".to_string());
|
||||
|
||||
// Simulate the conditional clearing logic from setup_bedrock()
|
||||
if wizard.settings.llm_backend.as_deref() != Some("bedrock") {
|
||||
wizard.settings.selected_model = None;
|
||||
}
|
||||
wizard.settings.llm_backend = Some("bedrock".to_string());
|
||||
|
||||
assert_eq!(
|
||||
wizard.settings.selected_model.as_deref(),
|
||||
Some("anthropic.claude-opus-4-6-v1"),
|
||||
"bedrock model should be preserved when re-selecting bedrock"
|
||||
);
|
||||
}
|
||||
|
||||
/// Regression: switching from another provider to bedrock must clear
|
||||
/// selected_model, and choosing "default credentials" must clear
|
||||
/// bedrock_profile.
|
||||
#[test]
|
||||
fn test_bedrock_clears_stale_profile_on_default_creds() {
|
||||
let mut wizard = SetupWizard::new();
|
||||
wizard.settings.llm_backend = Some("bedrock".to_string());
|
||||
wizard.settings.bedrock_profile = Some("old-sso-profile".to_string());
|
||||
|
||||
// Simulate auth_choice == 0 (default credentials) clearing the profile
|
||||
wizard.settings.bedrock_profile = None;
|
||||
|
||||
assert!(
|
||||
wizard.settings.bedrock_profile.is_none(),
|
||||
"bedrock_profile should be cleared when selecting default credentials"
|
||||
);
|
||||
}
|
||||
|
||||
/// Regression: empty profile input in named-profile auth should clear
|
||||
/// any previously configured profile instead of leaving it stale.
|
||||
#[test]
|
||||
fn test_bedrock_empty_profile_clears_existing() {
|
||||
let mut wizard = SetupWizard::new();
|
||||
wizard.settings.bedrock_profile = Some("old-profile".to_string());
|
||||
|
||||
// Simulate auth_choice == 1 with empty input
|
||||
let profile = "".to_string();
|
||||
if profile.trim().is_empty() {
|
||||
wizard.settings.bedrock_profile = None;
|
||||
} else {
|
||||
wizard.settings.bedrock_profile = Some(profile);
|
||||
}
|
||||
|
||||
assert!(
|
||||
wizard.settings.bedrock_profile.is_none(),
|
||||
"empty profile input should clear existing bedrock_profile"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_run_provider_setup_no_setup_hint() {
|
||||
// A provider with setup: None should not error. It should set the
|
||||
@@ -3536,4 +3797,30 @@ mod tests {
|
||||
"backend should be set even without setup hint"
|
||||
);
|
||||
}
|
||||
|
||||
/// Regression test for #666: env-var security option must initialize
|
||||
/// secrets_crypto so subsequent steps can encrypt API keys.
|
||||
#[test]
|
||||
fn test_env_var_security_initializes_crypto() {
|
||||
use crate::secrets::SecretsCrypto;
|
||||
use secrecy::SecretString;
|
||||
|
||||
// Simulate what option 1 in step_security() does after the fix:
|
||||
let key_hex = crate::secrets::keychain::generate_master_key_hex();
|
||||
|
||||
// The fix: create SecretsCrypto from the generated key.
|
||||
// Before the fix, this was skipped, leaving secrets_crypto = None.
|
||||
let crypto = SecretsCrypto::new(SecretString::from(key_hex.clone()));
|
||||
assert!(
|
||||
crypto.is_ok(),
|
||||
"generated key hex must produce valid SecretsCrypto"
|
||||
);
|
||||
|
||||
// Verify the key is stored for bootstrap env persistence.
|
||||
let settings = Settings {
|
||||
secrets_master_key_hex: Some(key_hex),
|
||||
..Settings::default()
|
||||
};
|
||||
assert!(settings.secrets_master_key_hex.is_some());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -142,6 +142,7 @@ mod tests {
|
||||
content_hash: "sha256:000".to_string(),
|
||||
compiled_patterns: vec![],
|
||||
lowercased_keywords: vec![],
|
||||
lowercased_exclude_keywords: vec![],
|
||||
lowercased_tags: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
@@ -98,6 +98,10 @@ pub struct ActivationCriteria {
|
||||
/// Capped at `MAX_KEYWORDS_PER_SKILL` during loading.
|
||||
#[serde(default)]
|
||||
pub keywords: Vec<String>,
|
||||
/// Keywords that veto this skill — if any match, score is 0 regardless of
|
||||
/// keyword/pattern matches. Prevents cross-skill interference.
|
||||
#[serde(default)]
|
||||
pub exclude_keywords: Vec<String>,
|
||||
/// Regex patterns for more complex matching.
|
||||
/// Capped at `MAX_PATTERNS_PER_SKILL` during loading.
|
||||
#[serde(default)]
|
||||
@@ -118,6 +122,9 @@ impl ActivationCriteria {
|
||||
pub fn enforce_limits(&mut self) {
|
||||
self.keywords.retain(|k| k.len() >= MIN_KEYWORD_TAG_LENGTH);
|
||||
self.keywords.truncate(MAX_KEYWORDS_PER_SKILL);
|
||||
self.exclude_keywords
|
||||
.retain(|k| k.len() >= MIN_KEYWORD_TAG_LENGTH);
|
||||
self.exclude_keywords.truncate(MAX_KEYWORDS_PER_SKILL);
|
||||
self.patterns.truncate(MAX_PATTERNS_PER_SKILL);
|
||||
self.tags.retain(|t| t.len() >= MIN_KEYWORD_TAG_LENGTH);
|
||||
self.tags.truncate(MAX_TAGS_PER_SKILL);
|
||||
@@ -199,6 +206,9 @@ pub struct LoadedSkill {
|
||||
/// Pre-computed lowercased keywords for scoring (avoids per-message allocation).
|
||||
/// Derived from `manifest.activation.keywords` at load time — do not mutate independently.
|
||||
pub lowercased_keywords: Vec<String>,
|
||||
/// Pre-computed lowercased exclude keywords for veto scoring.
|
||||
/// Derived from `manifest.activation.exclude_keywords` at load time.
|
||||
pub lowercased_exclude_keywords: Vec<String>,
|
||||
/// Pre-computed lowercased tags for scoring (avoids per-message allocation).
|
||||
/// Derived from `manifest.activation.tags` at load time — do not mutate independently.
|
||||
pub lowercased_tags: Vec<String>,
|
||||
@@ -513,6 +523,7 @@ metadata:
|
||||
content_hash: "sha256:000".to_string(),
|
||||
compiled_patterns: vec![],
|
||||
lowercased_keywords: vec![],
|
||||
lowercased_exclude_keywords: vec![],
|
||||
lowercased_tags: vec![],
|
||||
};
|
||||
assert_eq!(skill.name(), "test");
|
||||
|
||||
+8
-12
@@ -24,6 +24,10 @@ use crate::skills::{
|
||||
/// Prevents resource exhaustion from a directory with thousands of entries.
|
||||
const MAX_DISCOVERED_SKILLS: usize = 100;
|
||||
|
||||
fn to_lowercase_vec(items: &[String]) -> Vec<String> {
|
||||
items.iter().map(|s| s.to_lowercase()).collect()
|
||||
}
|
||||
|
||||
/// Error type for skill registry operations.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum SkillRegistryError {
|
||||
@@ -582,18 +586,9 @@ async fn load_and_validate_skill(
|
||||
let compiled_patterns = LoadedSkill::compile_patterns(&manifest.activation.patterns);
|
||||
|
||||
// Pre-compute lowercased keywords and tags for efficient scoring
|
||||
let lowercased_keywords = manifest
|
||||
.activation
|
||||
.keywords
|
||||
.iter()
|
||||
.map(|k| k.to_lowercase())
|
||||
.collect();
|
||||
let lowercased_tags = manifest
|
||||
.activation
|
||||
.tags
|
||||
.iter()
|
||||
.map(|t| t.to_lowercase())
|
||||
.collect();
|
||||
let lowercased_keywords = to_lowercase_vec(&manifest.activation.keywords);
|
||||
let lowercased_exclude_keywords = to_lowercase_vec(&manifest.activation.exclude_keywords);
|
||||
let lowercased_tags = to_lowercase_vec(&manifest.activation.tags);
|
||||
|
||||
let name = manifest.name.clone();
|
||||
let skill = LoadedSkill {
|
||||
@@ -604,6 +599,7 @@ async fn load_and_validate_skill(
|
||||
content_hash,
|
||||
compiled_patterns,
|
||||
lowercased_keywords,
|
||||
lowercased_exclude_keywords,
|
||||
lowercased_tags,
|
||||
};
|
||||
|
||||
|
||||
@@ -99,6 +99,15 @@ pub fn prefilter_skills<'a>(
|
||||
|
||||
/// Score a skill against a user message.
|
||||
fn score_skill(skill: &LoadedSkill, message_lower: &str, message_original: &str) -> u32 {
|
||||
// Exclusion veto: if any exclude_keyword is present in the message, score 0
|
||||
if skill
|
||||
.lowercased_exclude_keywords
|
||||
.iter()
|
||||
.any(|excl| message_lower.contains(excl.as_str()))
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
|
||||
let mut score: u32 = 0;
|
||||
|
||||
// Keyword scoring with cap to prevent gaming via keyword stuffing
|
||||
@@ -158,6 +167,7 @@ mod tests {
|
||||
description: format!("{} skill", name),
|
||||
activation: ActivationCriteria {
|
||||
keywords: kw_vec,
|
||||
exclude_keywords: vec![],
|
||||
patterns: pattern_strings,
|
||||
tags: tag_vec,
|
||||
max_context_tokens: 1000,
|
||||
@@ -170,6 +180,7 @@ mod tests {
|
||||
content_hash: "sha256:000".to_string(),
|
||||
compiled_patterns: compiled,
|
||||
lowercased_keywords,
|
||||
lowercased_exclude_keywords: vec![],
|
||||
lowercased_tags,
|
||||
}
|
||||
}
|
||||
@@ -368,4 +379,111 @@ mod tests {
|
||||
let result = prefilter_skills("test", &skills, 5, 1);
|
||||
assert_eq!(result.len(), 1);
|
||||
}
|
||||
|
||||
fn make_skill_with_excludes(
|
||||
name: &str,
|
||||
keywords: &[&str],
|
||||
exclude_keywords: &[&str],
|
||||
tags: &[&str],
|
||||
patterns: &[&str],
|
||||
) -> LoadedSkill {
|
||||
let mut skill = make_skill(name, keywords, tags, patterns);
|
||||
let excl_vec: Vec<String> = exclude_keywords.iter().map(|s| s.to_string()).collect();
|
||||
skill.lowercased_exclude_keywords = excl_vec.iter().map(|k| k.to_lowercase()).collect();
|
||||
skill.manifest.activation.exclude_keywords = excl_vec;
|
||||
skill
|
||||
}
|
||||
|
||||
// --- exclude_keywords tests ---
|
||||
|
||||
#[test]
|
||||
fn test_exclude_keyword_vetos_match() {
|
||||
// Skill matches on "write" but exclude_keywords: ["route"] — message contains "route"
|
||||
// so the skill should score 0 and be excluded.
|
||||
let skills = vec![make_skill_with_excludes(
|
||||
"writer",
|
||||
&["write"],
|
||||
&["route"],
|
||||
&[],
|
||||
&[],
|
||||
)];
|
||||
let result = prefilter_skills(
|
||||
"route this write request to another agent",
|
||||
&skills,
|
||||
3,
|
||||
MAX_SKILL_CONTEXT_TOKENS,
|
||||
);
|
||||
assert!(
|
||||
result.is_empty(),
|
||||
"skill with matching exclude_keyword should score 0"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exclude_keyword_absent_does_not_block() {
|
||||
// Same skill, message does NOT contain the exclude keyword — should activate normally.
|
||||
let skills = vec![make_skill_with_excludes(
|
||||
"writer",
|
||||
&["write"],
|
||||
&["route"],
|
||||
&[],
|
||||
&[],
|
||||
)];
|
||||
let result = prefilter_skills(
|
||||
"help me write an email",
|
||||
&skills,
|
||||
3,
|
||||
MAX_SKILL_CONTEXT_TOKENS,
|
||||
);
|
||||
assert_eq!(
|
||||
result.len(),
|
||||
1,
|
||||
"skill should activate when no exclude_keyword is present"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exclude_keyword_veto_wins_over_positive_match() {
|
||||
// Both a keyword match AND an exclude_keyword match are present.
|
||||
// The veto must win regardless of how high the positive score is.
|
||||
let skills = vec![make_skill_with_excludes(
|
||||
"writer",
|
||||
&["write", "draft", "compose"],
|
||||
&["redirect"],
|
||||
&[],
|
||||
&[],
|
||||
)];
|
||||
let result = prefilter_skills(
|
||||
"write and draft and compose — but redirect this somewhere else",
|
||||
&skills,
|
||||
3,
|
||||
MAX_SKILL_CONTEXT_TOKENS,
|
||||
);
|
||||
assert!(
|
||||
result.is_empty(),
|
||||
"exclude_keyword veto must win even when multiple positive keywords match"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exclude_keyword_case_insensitive() {
|
||||
// exclude_keywords are pre-lowercased; the veto must fire regardless of case in the message.
|
||||
let skills = vec![make_skill_with_excludes(
|
||||
"writer",
|
||||
&["write"],
|
||||
&["Route"],
|
||||
&[],
|
||||
&[],
|
||||
)];
|
||||
let result = prefilter_skills(
|
||||
"please ROUTE this write request",
|
||||
&skills,
|
||||
3,
|
||||
MAX_SKILL_CONTEXT_TOKENS,
|
||||
);
|
||||
assert!(
|
||||
result.is_empty(),
|
||||
"exclude_keyword veto should be case-insensitive"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1009,6 +1009,7 @@ mod tests {
|
||||
enabled: true,
|
||||
trigger: Trigger::Cron {
|
||||
schedule: "0 * * * *".to_string(),
|
||||
timezone: None,
|
||||
},
|
||||
action: RoutineAction::Lightweight {
|
||||
prompt: "Check status".to_string(),
|
||||
|
||||
+110
@@ -0,0 +1,110 @@
|
||||
//! Timezone resolution and utilities.
|
||||
|
||||
use chrono::{DateTime, NaiveDate, Utc};
|
||||
use chrono_tz::Tz;
|
||||
|
||||
/// Resolve the effective timezone from a priority chain.
|
||||
///
|
||||
/// Priority: client_tz > user_setting > config_default > UTC
|
||||
pub fn resolve_timezone(
|
||||
client_tz: Option<&str>,
|
||||
user_setting: Option<&str>,
|
||||
config_default: &str,
|
||||
) -> Tz {
|
||||
// Try each in priority order, skipping invalid values
|
||||
for candidate in [client_tz, user_setting, Some(config_default)] {
|
||||
if let Some(tz) = candidate.and_then(parse_timezone) {
|
||||
return tz;
|
||||
}
|
||||
}
|
||||
Tz::UTC
|
||||
}
|
||||
|
||||
/// Parse a timezone string (IANA name) into a `Tz`.
|
||||
pub fn parse_timezone(s: &str) -> Option<Tz> {
|
||||
s.parse::<Tz>().ok()
|
||||
}
|
||||
|
||||
/// Get today's date in the given timezone.
|
||||
pub fn today_in_tz(tz: Tz) -> NaiveDate {
|
||||
Utc::now().with_timezone(&tz).date_naive()
|
||||
}
|
||||
|
||||
/// Get the current time in the given timezone.
|
||||
pub fn now_in_tz(tz: Tz) -> DateTime<Tz> {
|
||||
Utc::now().with_timezone(&tz)
|
||||
}
|
||||
|
||||
/// Detect the system's timezone, falling back to UTC.
|
||||
pub fn detect_system_timezone() -> Tz {
|
||||
iana_time_zone::get_timezone()
|
||||
.ok()
|
||||
.and_then(|s| parse_timezone(&s))
|
||||
.unwrap_or(Tz::UTC)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use chrono::Datelike;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_resolve_client_wins() {
|
||||
let tz = resolve_timezone(Some("America/New_York"), Some("Europe/London"), "UTC");
|
||||
assert_eq!(tz, chrono_tz::America::New_York);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_user_setting_fallback() {
|
||||
let tz = resolve_timezone(None, Some("Europe/London"), "UTC");
|
||||
assert_eq!(tz, chrono_tz::Europe::London);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_config_fallback() {
|
||||
let tz = resolve_timezone(None, None, "Asia/Tokyo");
|
||||
assert_eq!(tz, chrono_tz::Asia::Tokyo);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_all_none_utc() {
|
||||
let tz = resolve_timezone(None, None, "UTC");
|
||||
assert_eq!(tz, Tz::UTC);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_invalid_client_skipped() {
|
||||
let tz = resolve_timezone(Some("Fake/Zone"), Some("Europe/London"), "UTC");
|
||||
assert_eq!(tz, chrono_tz::Europe::London);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_valid() {
|
||||
assert_eq!(
|
||||
parse_timezone("America/Chicago"),
|
||||
Some(chrono_tz::America::Chicago)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_invalid() {
|
||||
assert_eq!(parse_timezone("Fake/Zone"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_system_tz() {
|
||||
// Should always return a valid Tz (at minimum UTC)
|
||||
let tz = detect_system_timezone();
|
||||
let _ = now_in_tz(tz); // Should not panic
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_today_in_tz_returns_valid_date() {
|
||||
let date = today_in_tz(Tz::UTC);
|
||||
// Verify it returns a valid date (year, month, day are all positive)
|
||||
assert!(date.year() > 0);
|
||||
assert!((1..=12).contains(&date.month()));
|
||||
assert!((1..=31).contains(&date.day()));
|
||||
}
|
||||
}
|
||||
@@ -43,7 +43,6 @@ use crate::error::ToolError as AgentToolError;
|
||||
use crate::llm::{
|
||||
ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult, ToolDefinition,
|
||||
};
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::tools::ToolRegistry;
|
||||
use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput};
|
||||
|
||||
@@ -251,29 +250,18 @@ pub trait SoftwareBuilder: Send + Sync {
|
||||
pub struct LlmSoftwareBuilder {
|
||||
config: BuilderConfig,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
tools: Arc<ToolRegistry>,
|
||||
}
|
||||
|
||||
impl LlmSoftwareBuilder {
|
||||
/// Create a new LLM-based software builder.
|
||||
pub fn new(
|
||||
config: BuilderConfig,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
tools: Arc<ToolRegistry>,
|
||||
) -> Self {
|
||||
pub fn new(config: BuilderConfig, llm: Arc<dyn LlmProvider>, tools: Arc<ToolRegistry>) -> Self {
|
||||
// Ensure build directory exists
|
||||
if let Err(e) = std::fs::create_dir_all(&config.build_dir) {
|
||||
tracing::warn!("Failed to create build directory: {}", e);
|
||||
}
|
||||
|
||||
Self {
|
||||
config,
|
||||
llm,
|
||||
safety,
|
||||
tools,
|
||||
}
|
||||
Self { config, llm, tools }
|
||||
}
|
||||
|
||||
/// Get the build tools available for the build loop.
|
||||
@@ -521,7 +509,7 @@ Create alongside the .wasm file to grant capabilities:
|
||||
let mut iteration = 0;
|
||||
|
||||
// Create reasoning engine
|
||||
let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone());
|
||||
let reasoning = Reasoning::new(self.llm.clone());
|
||||
|
||||
// Build initial context
|
||||
let tool_defs = self.get_build_tools().await;
|
||||
@@ -822,7 +810,7 @@ Create alongside the .wasm file to grant capabilities:
|
||||
impl SoftwareBuilder for LlmSoftwareBuilder {
|
||||
async fn analyze(&self, description: &str) -> Result<BuildRequirement, AgentToolError> {
|
||||
// Use LLM to parse the description
|
||||
let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone());
|
||||
let reasoning = Reasoning::new(self.llm.clone());
|
||||
|
||||
let prompt = format!(
|
||||
r#"Analyze this software requirement and extract structured information.
|
||||
|
||||
@@ -0,0 +1,250 @@
|
||||
//! Image analysis tool using vision-capable LLM models.
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use base64::Engine;
|
||||
use secrecy::{ExposeSecret, SecretString};
|
||||
|
||||
use crate::context::JobContext;
|
||||
use crate::tools::builtin::path_utils::validate_path;
|
||||
use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput};
|
||||
|
||||
/// Tool for analyzing images using a vision-capable model.
|
||||
pub struct ImageAnalyzeTool {
|
||||
/// API base URL.
|
||||
api_base_url: String,
|
||||
/// Bearer token for API auth.
|
||||
api_key: SecretString,
|
||||
/// Vision-capable model name.
|
||||
model: String,
|
||||
/// HTTP client.
|
||||
client: reqwest::Client,
|
||||
/// Optional base directory for resolving relative image paths.
|
||||
base_dir: Option<PathBuf>,
|
||||
}
|
||||
|
||||
impl ImageAnalyzeTool {
|
||||
/// Create a new image analysis tool.
|
||||
pub fn new(
|
||||
api_base_url: String,
|
||||
api_key: String,
|
||||
model: String,
|
||||
base_dir: Option<PathBuf>,
|
||||
) -> Self {
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(120))
|
||||
.build()
|
||||
.unwrap_or_default();
|
||||
Self {
|
||||
api_base_url,
|
||||
api_key: SecretString::from(api_key),
|
||||
model,
|
||||
client,
|
||||
base_dir,
|
||||
}
|
||||
}
|
||||
|
||||
/// Read binary image bytes from filesystem.
|
||||
///
|
||||
/// Validates the path against the base directory sandbox to prevent
|
||||
/// path traversal attacks, then reads the file bytes.
|
||||
async fn read_image_bytes(&self, image_path: &str) -> Result<Vec<u8>, ToolError> {
|
||||
let resolved = validate_path(image_path, self.base_dir.as_deref())?;
|
||||
|
||||
tokio::fs::read(&resolved)
|
||||
.await
|
||||
.map_err(|e| ToolError::ExecutionFailed(format!("Failed to read image file: {e}")))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ImageAnalyzeTool {
|
||||
fn name(&self) -> &str {
|
||||
"image_analyze"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Analyze an image using a vision-capable AI model. Provide a workspace path to the image and an optional analysis question."
|
||||
}
|
||||
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"image_path": {
|
||||
"type": "string",
|
||||
"description": "Path to the image file in the workspace (e.g., 'images/photo.jpg')"
|
||||
},
|
||||
"question": {
|
||||
"type": "string",
|
||||
"description": "Specific question to answer about the image. Defaults to general analysis.",
|
||||
"default": "Describe this image in detail."
|
||||
}
|
||||
},
|
||||
"required": ["image_path"]
|
||||
})
|
||||
}
|
||||
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
}
|
||||
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
params: serde_json::Value,
|
||||
_ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
let image_path = params
|
||||
.get("image_path")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| {
|
||||
ToolError::InvalidParameters("Missing required 'image_path' parameter".to_string())
|
||||
})?;
|
||||
|
||||
let question = params
|
||||
.get("question")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("Describe this image in detail.");
|
||||
|
||||
// Read binary image bytes directly from filesystem
|
||||
let image_bytes = self.read_image_bytes(image_path).await?;
|
||||
if image_bytes.is_empty() {
|
||||
return Err(ToolError::ExecutionFailed(
|
||||
"Image file is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let media_type = super::media_type_from_path(image_path);
|
||||
let b64 = base64::engine::general_purpose::STANDARD.encode(&image_bytes);
|
||||
let data_url = format!("data:{media_type};base64,{b64}");
|
||||
|
||||
// Call vision model via chat completions API
|
||||
let url = format!(
|
||||
"{}/v1/chat/completions",
|
||||
self.api_base_url.trim_end_matches('/')
|
||||
);
|
||||
|
||||
let request_body = serde_json::json!({
|
||||
"model": &self.model,
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": question
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": data_url
|
||||
}
|
||||
}
|
||||
]
|
||||
}],
|
||||
"max_tokens": 2048
|
||||
});
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.bearer_auth(self.api_key.expose_secret())
|
||||
.json(&request_body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| ToolError::ExecutionFailed(format!("Vision API request failed: {e}")))?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
return Err(ToolError::ExecutionFailed(format!(
|
||||
"Vision API returned {status}: {body}"
|
||||
)));
|
||||
}
|
||||
|
||||
let resp: serde_json::Value = response.json().await.map_err(|e| {
|
||||
ToolError::ExecutionFailed(format!("Failed to parse vision API response: {e}"))
|
||||
})?;
|
||||
|
||||
let analysis = resp
|
||||
.pointer("/choices/0/message/content")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("No analysis available.");
|
||||
|
||||
Ok(ToolOutput::text(analysis, start.elapsed()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::super::media_type_from_path;
|
||||
use super::*;
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[test]
|
||||
fn test_media_type_detection() {
|
||||
assert_eq!(media_type_from_path("photo.png"), "image/png");
|
||||
assert_eq!(media_type_from_path("photo.jpg"), "image/jpeg");
|
||||
assert_eq!(media_type_from_path("photo.jpeg"), "image/jpeg");
|
||||
assert_eq!(media_type_from_path("photo.gif"), "image/gif");
|
||||
assert_eq!(media_type_from_path("photo.webp"), "image/webp");
|
||||
assert_eq!(media_type_from_path("photo.bmp"), "image/bmp");
|
||||
assert_eq!(media_type_from_path("photo.svg"), "image/svg+xml");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_requires_approval_returns_unless_auto_approved() {
|
||||
let tool = ImageAnalyzeTool::new(
|
||||
"https://api.example.com".to_string(),
|
||||
"test-key".to_string(),
|
||||
"gpt-4o".to_string(),
|
||||
None,
|
||||
);
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({})),
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_image_bytes_rejects_path_traversal() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let tool = ImageAnalyzeTool::new(
|
||||
"https://api.example.com".to_string(),
|
||||
"test-key".to_string(),
|
||||
"gpt-4o".to_string(),
|
||||
Some(dir.path().to_path_buf()),
|
||||
);
|
||||
|
||||
let result = tool.read_image_bytes("../../etc/passwd").await;
|
||||
assert!(
|
||||
result.is_err(),
|
||||
"Should reject path traversal, got: {:?}",
|
||||
result
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_image_bytes_rejects_absolute_path_outside_sandbox() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let tool = ImageAnalyzeTool::new(
|
||||
"https://api.example.com".to_string(),
|
||||
"test-key".to_string(),
|
||||
"gpt-4o".to_string(),
|
||||
Some(dir.path().to_path_buf()),
|
||||
);
|
||||
|
||||
let result = tool.read_image_bytes("/etc/passwd").await;
|
||||
assert!(
|
||||
result.is_err(),
|
||||
"Should reject absolute path outside sandbox, got: {:?}",
|
||||
result
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,322 @@
|
||||
//! Image editing tool using cloud API.
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use secrecy::{ExposeSecret, SecretString};
|
||||
|
||||
use crate::context::JobContext;
|
||||
use crate::tools::builtin::path_utils::validate_path;
|
||||
use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput};
|
||||
|
||||
/// Tool for editing images using an AI image editing API.
|
||||
pub struct ImageEditTool {
|
||||
/// API base URL.
|
||||
api_base_url: String,
|
||||
/// Bearer token for API auth.
|
||||
api_key: SecretString,
|
||||
/// Model to use.
|
||||
model: String,
|
||||
/// HTTP client.
|
||||
client: reqwest::Client,
|
||||
/// Optional base directory for resolving relative image paths.
|
||||
base_dir: Option<PathBuf>,
|
||||
}
|
||||
|
||||
impl ImageEditTool {
|
||||
/// Create a new image edit tool.
|
||||
pub fn new(
|
||||
api_base_url: String,
|
||||
api_key: String,
|
||||
model: String,
|
||||
base_dir: Option<PathBuf>,
|
||||
) -> Self {
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(180))
|
||||
.build()
|
||||
.unwrap_or_default();
|
||||
Self {
|
||||
api_base_url,
|
||||
api_key: SecretString::from(api_key),
|
||||
model,
|
||||
client,
|
||||
base_dir,
|
||||
}
|
||||
}
|
||||
|
||||
/// Read binary image bytes from filesystem.
|
||||
///
|
||||
/// Validates the path against the base directory sandbox to prevent
|
||||
/// path traversal attacks, then reads the file bytes.
|
||||
async fn read_image_bytes(&self, image_path: &str) -> Result<Vec<u8>, ToolError> {
|
||||
let resolved = validate_path(image_path, self.base_dir.as_deref())?;
|
||||
|
||||
tokio::fs::read(&resolved)
|
||||
.await
|
||||
.map_err(|e| ToolError::ExecutionFailed(format!("Failed to read image file: {e}")))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ImageEditTool {
|
||||
fn name(&self) -> &str {
|
||||
"image_edit"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Edit an existing image using an AI model. Provide the workspace path to the source image and a text prompt describing the desired edits."
|
||||
}
|
||||
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"prompt": {
|
||||
"type": "string",
|
||||
"description": "Text description of the edits to apply to the image",
|
||||
"maxLength": 4000
|
||||
},
|
||||
"image_path": {
|
||||
"type": "string",
|
||||
"description": "Path to the source image in the workspace (e.g., 'images/photo.jpg')"
|
||||
}
|
||||
},
|
||||
"required": ["prompt", "image_path"]
|
||||
})
|
||||
}
|
||||
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
}
|
||||
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
params: serde_json::Value,
|
||||
_ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
let prompt = params
|
||||
.get("prompt")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| {
|
||||
ToolError::InvalidParameters("Missing required 'prompt' parameter".to_string())
|
||||
})?;
|
||||
|
||||
let image_path = params
|
||||
.get("image_path")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| {
|
||||
ToolError::InvalidParameters("Missing required 'image_path' parameter".to_string())
|
||||
})?;
|
||||
|
||||
if prompt.len() > 4000 {
|
||||
return Err(ToolError::InvalidParameters(
|
||||
"Prompt exceeds 4000 character limit".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// Read binary image bytes directly from filesystem
|
||||
let image_bytes = self.read_image_bytes(image_path).await?;
|
||||
if image_bytes.is_empty() {
|
||||
return Err(ToolError::ExecutionFailed(
|
||||
"Source image file is empty".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let media_type = super::media_type_from_path(image_path);
|
||||
|
||||
// Use multipart form for image edit API
|
||||
let url = format!(
|
||||
"{}/v1/images/edits",
|
||||
self.api_base_url.trim_end_matches('/')
|
||||
);
|
||||
|
||||
let form = reqwest::multipart::Form::new()
|
||||
.text("model", self.model.clone())
|
||||
.text("prompt", prompt.to_string())
|
||||
.text("response_format", "b64_json")
|
||||
.part(
|
||||
"image",
|
||||
reqwest::multipart::Part::bytes(image_bytes)
|
||||
.mime_str(&media_type)
|
||||
.map_err(|e| ToolError::ExecutionFailed(format!("Invalid media type: {e}")))?
|
||||
.file_name("image"),
|
||||
);
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.bearer_auth(self.api_key.expose_secret())
|
||||
.multipart(form)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| ToolError::ExecutionFailed(format!("Image edit request failed: {e}")))?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
|
||||
// Fall back to generation if edits endpoint not available
|
||||
if status.as_u16() == 404 {
|
||||
tracing::warn!(
|
||||
"Image edit endpoint returned 404, falling back to generation API. \
|
||||
Note: the source image will NOT be used — a new image will be generated from the prompt alone."
|
||||
);
|
||||
return self.fallback_generate(prompt, start).await;
|
||||
}
|
||||
|
||||
return Err(ToolError::ExecutionFailed(format!(
|
||||
"Image edit API returned {status}: {body}"
|
||||
)));
|
||||
}
|
||||
|
||||
let resp: serde_json::Value = response.json().await.map_err(|e| {
|
||||
ToolError::ExecutionFailed(format!("Failed to parse image edit response: {e}"))
|
||||
})?;
|
||||
|
||||
let edited_data = resp
|
||||
.pointer("/data/0/b64_json")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| {
|
||||
ToolError::ExecutionFailed("No image data in edit response".to_string())
|
||||
})?;
|
||||
|
||||
let sentinel = serde_json::json!({
|
||||
"type": "image_generated",
|
||||
"data": format!("data:image/png;base64,{}", edited_data),
|
||||
"media_type": "image/png",
|
||||
"prompt": prompt,
|
||||
"source_path": image_path
|
||||
});
|
||||
|
||||
Ok(ToolOutput::text(sentinel.to_string(), start.elapsed()))
|
||||
}
|
||||
}
|
||||
|
||||
impl ImageEditTool {
|
||||
/// Fallback: generate a new image from the prompt when the edit endpoint is unavailable.
|
||||
///
|
||||
/// The source image is NOT used — this generates a completely new image.
|
||||
/// The response includes a `note` field warning the user.
|
||||
async fn fallback_generate(
|
||||
&self,
|
||||
prompt: &str,
|
||||
start: std::time::Instant,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let url = format!(
|
||||
"{}/v1/images/generations",
|
||||
self.api_base_url.trim_end_matches('/')
|
||||
);
|
||||
|
||||
let request_body = serde_json::json!({
|
||||
"model": &self.model,
|
||||
"prompt": prompt,
|
||||
"size": "1024x1024",
|
||||
"response_format": "b64_json",
|
||||
"n": 1
|
||||
});
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.bearer_auth(self.api_key.expose_secret())
|
||||
.json(&request_body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| {
|
||||
ToolError::ExecutionFailed(format!("Fallback image generation failed: {e}"))
|
||||
})?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
return Err(ToolError::ExecutionFailed(format!(
|
||||
"Fallback generation API returned {status}: {body}"
|
||||
)));
|
||||
}
|
||||
|
||||
let resp: serde_json::Value = response.json().await.map_err(|e| {
|
||||
ToolError::ExecutionFailed(format!("Failed to parse fallback response: {e}"))
|
||||
})?;
|
||||
|
||||
let image_data = resp
|
||||
.pointer("/data/0/b64_json")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| {
|
||||
ToolError::ExecutionFailed("No image data in fallback response".to_string())
|
||||
})?;
|
||||
|
||||
let sentinel = serde_json::json!({
|
||||
"type": "image_generated",
|
||||
"data": format!("data:image/png;base64,{}", image_data),
|
||||
"media_type": "image/png",
|
||||
"prompt": prompt,
|
||||
"note": "Generated new image (edit endpoint unavailable — source image was NOT used)"
|
||||
});
|
||||
|
||||
Ok(ToolOutput::text(sentinel.to_string(), start.elapsed()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[test]
|
||||
fn test_tool_metadata() {
|
||||
let tool = ImageEditTool::new(
|
||||
"https://api.example.com".to_string(),
|
||||
"test-key".to_string(),
|
||||
"flux-1".to_string(),
|
||||
None,
|
||||
);
|
||||
assert_eq!(tool.name(), "image_edit");
|
||||
assert!(!tool.requires_sanitization());
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({})),
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_image_bytes_rejects_path_traversal() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let tool = ImageEditTool::new(
|
||||
"https://api.example.com".to_string(),
|
||||
"test-key".to_string(),
|
||||
"flux-1".to_string(),
|
||||
Some(dir.path().to_path_buf()),
|
||||
);
|
||||
|
||||
let result = tool.read_image_bytes("../../etc/passwd").await;
|
||||
assert!(
|
||||
result.is_err(),
|
||||
"Should reject path traversal, got: {:?}",
|
||||
result
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_image_bytes_rejects_absolute_path_outside_sandbox() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let tool = ImageEditTool::new(
|
||||
"https://api.example.com".to_string(),
|
||||
"test-key".to_string(),
|
||||
"flux-1".to_string(),
|
||||
Some(dir.path().to_path_buf()),
|
||||
);
|
||||
|
||||
let result = tool.read_image_bytes("/etc/passwd").await;
|
||||
assert!(
|
||||
result.is_err(),
|
||||
"Should reject absolute path outside sandbox, got: {:?}",
|
||||
result
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,251 @@
|
||||
//! Image generation tool using cloud API.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use secrecy::{ExposeSecret, SecretString};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::context::JobContext;
|
||||
use crate::tools::tool::ApprovalRequirement;
|
||||
use crate::tools::{Tool, ToolError, ToolOutput};
|
||||
|
||||
/// Tool for generating images using FLUX or compatible image generation APIs.
|
||||
pub struct ImageGenerateTool {
|
||||
/// API base URL (e.g., "https://cloud-api.near.ai").
|
||||
api_base_url: String,
|
||||
/// Bearer token for API auth.
|
||||
api_key: SecretString,
|
||||
/// Model to use (e.g., "black-forest-labs/FLUX.1-schnell").
|
||||
model: String,
|
||||
/// HTTP client.
|
||||
client: reqwest::Client,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct ImageGenRequest {
|
||||
model: String,
|
||||
prompt: String,
|
||||
size: String,
|
||||
response_format: String,
|
||||
n: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ImageGenResponse {
|
||||
data: Vec<ImageGenData>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[allow(dead_code)]
|
||||
struct ImageGenData {
|
||||
b64_json: Option<String>,
|
||||
url: Option<String>,
|
||||
}
|
||||
|
||||
impl ImageGenerateTool {
|
||||
/// Create a new image generation tool.
|
||||
pub fn new(api_base_url: String, api_key: String, model: String) -> Self {
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(180))
|
||||
.build()
|
||||
.unwrap_or_default();
|
||||
Self {
|
||||
api_base_url,
|
||||
api_key: SecretString::from(api_key),
|
||||
model,
|
||||
client,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ImageGenerateTool {
|
||||
fn name(&self) -> &str {
|
||||
"image_generate"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Generate an image from a text prompt using an AI image generation model (e.g., FLUX). Returns the generated image data."
|
||||
}
|
||||
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"prompt": {
|
||||
"type": "string",
|
||||
"description": "Text description of the image to generate (max 4000 chars)",
|
||||
"maxLength": 4000
|
||||
},
|
||||
"size": {
|
||||
"type": "string",
|
||||
"description": "Image dimensions",
|
||||
"enum": ["1024x1024", "1792x1024", "1024x1792"],
|
||||
"default": "1024x1024"
|
||||
}
|
||||
},
|
||||
"required": ["prompt"]
|
||||
})
|
||||
}
|
||||
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
}
|
||||
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
params: serde_json::Value,
|
||||
_ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
let prompt = params
|
||||
.get("prompt")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| {
|
||||
ToolError::InvalidParameters("Missing required 'prompt' parameter".to_string())
|
||||
})?;
|
||||
|
||||
if prompt.len() > 4000 {
|
||||
return Err(ToolError::InvalidParameters(
|
||||
"Prompt exceeds 4000 character limit".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let size = params
|
||||
.get("size")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("1024x1024");
|
||||
|
||||
// Validate size
|
||||
if !["1024x1024", "1792x1024", "1024x1792"].contains(&size) {
|
||||
return Err(ToolError::InvalidParameters(format!(
|
||||
"Invalid size '{}'. Must be 1024x1024, 1792x1024, or 1024x1792",
|
||||
size
|
||||
)));
|
||||
}
|
||||
|
||||
let url = format!(
|
||||
"{}/v1/images/generations",
|
||||
self.api_base_url.trim_end_matches('/')
|
||||
);
|
||||
|
||||
let request_body = ImageGenRequest {
|
||||
model: self.model.clone(),
|
||||
prompt: prompt.to_string(),
|
||||
size: size.to_string(),
|
||||
response_format: "b64_json".to_string(),
|
||||
n: 1,
|
||||
};
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.bearer_auth(self.api_key.expose_secret())
|
||||
.json(&request_body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| {
|
||||
ToolError::ExecutionFailed(format!("Image generation request failed: {e}"))
|
||||
})?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
return Err(ToolError::ExecutionFailed(format!(
|
||||
"Image generation API returned {status}: {body}"
|
||||
)));
|
||||
}
|
||||
|
||||
let gen_response: ImageGenResponse = response.json().await.map_err(|e| {
|
||||
ToolError::ExecutionFailed(format!("Failed to parse image generation response: {e}"))
|
||||
})?;
|
||||
|
||||
let image_data = gen_response
|
||||
.data
|
||||
.first()
|
||||
.and_then(|d| d.b64_json.as_deref())
|
||||
.ok_or_else(|| ToolError::ExecutionFailed("No image data in response".to_string()))?;
|
||||
|
||||
// Return sentinel JSON for image display
|
||||
let sentinel = serde_json::json!({
|
||||
"type": "image_generated",
|
||||
"data": format!("data:image/png;base64,{}", image_data),
|
||||
"media_type": "image/png",
|
||||
"prompt": prompt,
|
||||
"size": size
|
||||
});
|
||||
|
||||
Ok(ToolOutput::text(sentinel.to_string(), start.elapsed()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_tool_metadata() {
|
||||
let tool = ImageGenerateTool::new(
|
||||
"https://api.example.com".to_string(),
|
||||
"test-key".to_string(),
|
||||
"flux-1".to_string(),
|
||||
);
|
||||
assert_eq!(tool.name(), "image_generate");
|
||||
assert_eq!(
|
||||
tool.requires_approval(&serde_json::json!({})),
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
);
|
||||
|
||||
let schema = tool.parameters_schema();
|
||||
assert!(schema["properties"]["prompt"].is_object());
|
||||
assert!(schema["properties"]["size"].is_object());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_missing_prompt() {
|
||||
let tool = ImageGenerateTool::new(
|
||||
"https://api.example.com".to_string(),
|
||||
"test-key".to_string(),
|
||||
"flux-1".to_string(),
|
||||
);
|
||||
let ctx = JobContext::default();
|
||||
let result = tool.execute(serde_json::json!({}), &ctx).await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_invalid_size() {
|
||||
let tool = ImageGenerateTool::new(
|
||||
"https://api.example.com".to_string(),
|
||||
"test-key".to_string(),
|
||||
"flux-1".to_string(),
|
||||
);
|
||||
let ctx = JobContext::default();
|
||||
let result = tool
|
||||
.execute(
|
||||
serde_json::json!({"prompt": "a cat", "size": "999x999"}),
|
||||
&ctx,
|
||||
)
|
||||
.await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_prompt_too_long() {
|
||||
let tool = ImageGenerateTool::new(
|
||||
"https://api.example.com".to_string(),
|
||||
"test-key".to_string(),
|
||||
"flux-1".to_string(),
|
||||
);
|
||||
let ctx = JobContext::default();
|
||||
let long_prompt = "x".repeat(4001);
|
||||
let result = tool
|
||||
.execute(serde_json::json!({"prompt": long_prompt}), &ctx)
|
||||
.await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
}
|
||||
@@ -1416,6 +1416,185 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_job_params() {
|
||||
let manager = Arc::new(ContextManager::new(5));
|
||||
let tool = CreateJobTool::new(manager);
|
||||
let ctx = JobContext::default();
|
||||
|
||||
let missing_title = tool
|
||||
.execute(serde_json::json!({ "description": "A test job" }), &ctx)
|
||||
.await;
|
||||
assert!(missing_title.is_err());
|
||||
assert!(
|
||||
missing_title
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("missing 'title' parameter")
|
||||
);
|
||||
|
||||
let missing_description = tool
|
||||
.execute(serde_json::json!({ "title": "Test Job" }), &ctx)
|
||||
.await;
|
||||
assert!(missing_description.is_err());
|
||||
assert!(
|
||||
missing_description
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("missing 'description' parameter")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_list_jobs_formatting() {
|
||||
let manager = Arc::new(ContextManager::new(10));
|
||||
let pending_id = manager
|
||||
.create_job_for_user("default", "Pending Job", "Todo")
|
||||
.await
|
||||
.unwrap();
|
||||
let completed_id = manager
|
||||
.create_job_for_user("default", "Completed Job", "Done")
|
||||
.await
|
||||
.unwrap();
|
||||
let failed_id = manager
|
||||
.create_job_for_user("default", "Failed Job", "Oops")
|
||||
.await
|
||||
.unwrap();
|
||||
manager
|
||||
.create_job_for_user("other-user", "Other User Job", "Ignore")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
manager
|
||||
.update_context(completed_id, |ctx| {
|
||||
ctx.transition_to(JobState::InProgress, None)?;
|
||||
ctx.transition_to(JobState::Completed, Some("done".to_string()))
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
manager
|
||||
.update_context(failed_id, |ctx| {
|
||||
ctx.transition_to(JobState::InProgress, None)?;
|
||||
ctx.transition_to(JobState::Failed, Some("boom".to_string()))
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
let tool = ListJobsTool::new(Arc::clone(&manager));
|
||||
let ctx = JobContext::default();
|
||||
let result = tool.execute(serde_json::json!({}), &ctx).await.unwrap();
|
||||
|
||||
let jobs = result.result.get("jobs").unwrap().as_array().unwrap();
|
||||
assert_eq!(jobs.len(), 3);
|
||||
assert!(jobs.iter().any(|job| {
|
||||
job.get("job_id").and_then(|v| v.as_str()) == Some(&pending_id.to_string())
|
||||
&& job.get("status").and_then(|v| v.as_str()) == Some("Pending")
|
||||
}));
|
||||
assert!(jobs.iter().any(|job| {
|
||||
job.get("job_id").and_then(|v| v.as_str()) == Some(&completed_id.to_string())
|
||||
&& job.get("status").and_then(|v| v.as_str()) == Some("Completed")
|
||||
}));
|
||||
assert!(jobs.iter().any(|job| {
|
||||
job.get("job_id").and_then(|v| v.as_str()) == Some(&failed_id.to_string())
|
||||
&& job.get("status").and_then(|v| v.as_str()) == Some("Failed")
|
||||
}));
|
||||
|
||||
let summary = result.result.get("summary").unwrap();
|
||||
assert_eq!(summary.get("total").and_then(|v| v.as_u64()), Some(3));
|
||||
assert_eq!(summary.get("pending").and_then(|v| v.as_u64()), Some(1));
|
||||
assert_eq!(summary.get("completed").and_then(|v| v.as_u64()), Some(1));
|
||||
assert_eq!(summary.get("failed").and_then(|v| v.as_u64()), Some(1));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_job_status_transitions() {
|
||||
let manager = Arc::new(ContextManager::new(5));
|
||||
let job_id = manager
|
||||
.create_job_for_user("default", "Transition Job", "Track me")
|
||||
.await
|
||||
.unwrap();
|
||||
manager
|
||||
.update_context(job_id, |ctx| {
|
||||
ctx.transition_to(JobState::InProgress, Some("started".to_string()))?;
|
||||
ctx.transition_to(JobState::Completed, Some("finished".to_string()))
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
let tool = JobStatusTool::new(Arc::clone(&manager));
|
||||
let ctx = JobContext::default();
|
||||
let result = tool
|
||||
.execute(serde_json::json!({ "job_id": job_id.to_string() }), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
result.result.get("status").and_then(|v| v.as_str()),
|
||||
Some("Completed")
|
||||
);
|
||||
assert!(result.result.get("started_at").unwrap().is_string());
|
||||
assert!(result.result.get("completed_at").unwrap().is_string());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cancel_job_running() {
|
||||
let manager = Arc::new(ContextManager::new(5));
|
||||
let job_id = manager
|
||||
.create_job_for_user("default", "Running Job", "In progress")
|
||||
.await
|
||||
.unwrap();
|
||||
manager
|
||||
.update_context(job_id, |ctx| ctx.transition_to(JobState::InProgress, None))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
let tool = CancelJobTool::new(Arc::clone(&manager));
|
||||
let ctx = JobContext::default();
|
||||
let result = tool
|
||||
.execute(serde_json::json!({ "job_id": job_id.to_string() }), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
result.result.get("status").and_then(|v| v.as_str()),
|
||||
Some("cancelled")
|
||||
);
|
||||
let updated = manager.get_context(job_id).await.unwrap();
|
||||
assert_eq!(updated.state, JobState::Cancelled);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cancel_job_completed() {
|
||||
let manager = Arc::new(ContextManager::new(5));
|
||||
let job_id = manager
|
||||
.create_job_for_user("default", "Completed Job", "Already done")
|
||||
.await
|
||||
.unwrap();
|
||||
manager
|
||||
.update_context(job_id, |ctx| {
|
||||
ctx.transition_to(JobState::InProgress, None)?;
|
||||
ctx.transition_to(JobState::Completed, Some("done".to_string()))
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
let tool = CancelJobTool::new(Arc::clone(&manager));
|
||||
let ctx = JobContext::default();
|
||||
let result = tool
|
||||
.execute(serde_json::json!({ "job_id": job_id.to_string() }), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let error = result.result.get("error").and_then(|v| v.as_str()).unwrap();
|
||||
assert!(error.contains("Cannot cancel job"));
|
||||
assert!(error.contains("completed"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_project_dir_auto() {
|
||||
let project_id = Uuid::new_v4();
|
||||
|
||||
@@ -172,7 +172,7 @@ impl Tool for MemoryWriteTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
params: serde_json::Value,
|
||||
_ctx: &JobContext,
|
||||
ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
@@ -239,11 +239,12 @@ impl Tool for MemoryWriteTool {
|
||||
paths::MEMORY.to_string()
|
||||
}
|
||||
"daily_log" => {
|
||||
let tz = crate::timezone::parse_timezone(&ctx.user_timezone)
|
||||
.unwrap_or(chrono_tz::Tz::UTC);
|
||||
self.workspace
|
||||
.append_daily_log(content)
|
||||
.append_daily_log_tz(content, tz)
|
||||
.await
|
||||
.map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?;
|
||||
format!("daily/{}.md", chrono::Utc::now().format("%Y-%m-%d"))
|
||||
.map_err(|e| ToolError::ExecutionFailed(format!("Write failed: {}", e)))?
|
||||
}
|
||||
"heartbeat" => {
|
||||
if append {
|
||||
|
||||
+184
-50
@@ -105,42 +105,47 @@ impl Tool for MessageTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
params: serde_json::Value,
|
||||
_ctx: &JobContext,
|
||||
ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
let content = require_str(¶ms, "content")?;
|
||||
|
||||
// Get channel: use param or fall back to default
|
||||
let channel = if let Some(c) = params.get("channel").and_then(|v| v.as_str()) {
|
||||
c.to_string()
|
||||
} else {
|
||||
self.default_channel
|
||||
// Get channel: use param → conversation default → job metadata → None (broadcast all)
|
||||
let channel: Option<String> =
|
||||
if let Some(c) = params.get("channel").and_then(|v| v.as_str()) {
|
||||
Some(c.to_string())
|
||||
} else if let Some(c) = self
|
||||
.default_channel
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.clone()
|
||||
.ok_or_else(|| {
|
||||
ToolError::ExecutionFailed(
|
||||
"No channel specified and no active conversation. Provide channel parameter."
|
||||
.to_string(),
|
||||
)
|
||||
})?
|
||||
};
|
||||
{
|
||||
Some(c)
|
||||
} else {
|
||||
ctx.metadata
|
||||
.get("notify_channel")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|c| c.to_string())
|
||||
};
|
||||
|
||||
// Get target: use param or fall back to default
|
||||
// Get target: use param → conversation default → job metadata
|
||||
let target = if let Some(t) = params.get("target").and_then(|v| v.as_str()) {
|
||||
t.to_string()
|
||||
} else if let Some(t) = self
|
||||
.default_target
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.clone()
|
||||
{
|
||||
t
|
||||
} else if let Some(t) = ctx.metadata.get("notify_user").and_then(|v| v.as_str()) {
|
||||
t.to_string()
|
||||
} else {
|
||||
self.default_target
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.clone()
|
||||
.ok_or_else(|| {
|
||||
ToolError::ExecutionFailed(
|
||||
"No target specified and no active conversation. Provide target parameter."
|
||||
.to_string(),
|
||||
)
|
||||
})?
|
||||
return Err(ToolError::ExecutionFailed(
|
||||
"No target specified and no active conversation. Provide target parameter."
|
||||
.to_string(),
|
||||
));
|
||||
};
|
||||
|
||||
let attachments: Vec<String> = match params.get("attachments") {
|
||||
@@ -181,36 +186,79 @@ impl Tool for MessageTool {
|
||||
response = response.with_attachments(attachments);
|
||||
}
|
||||
|
||||
match self
|
||||
.channel_manager
|
||||
.broadcast(&channel, &target, response)
|
||||
.await
|
||||
{
|
||||
Ok(()) => {
|
||||
tracing::info!(
|
||||
message_sent = true,
|
||||
channel = %channel,
|
||||
target = %target,
|
||||
attachments = attachment_count,
|
||||
"Message sent via message tool"
|
||||
);
|
||||
let msg = format!("Sent message to {}:{}", channel, target);
|
||||
Ok(ToolOutput::text(msg, start.elapsed()))
|
||||
if let Some(ref channel) = channel {
|
||||
// Send to a specific channel
|
||||
match self
|
||||
.channel_manager
|
||||
.broadcast(channel, &target, response)
|
||||
.await
|
||||
{
|
||||
Ok(()) => {
|
||||
tracing::info!(
|
||||
message_sent = true,
|
||||
channel = %channel,
|
||||
target = %target,
|
||||
attachments = attachment_count,
|
||||
"Message sent via message tool"
|
||||
);
|
||||
let msg = format!("Sent message to {}:{}", channel, target);
|
||||
Ok(ToolOutput::text(msg, start.elapsed()))
|
||||
}
|
||||
Err(e) => {
|
||||
let available = self.channel_manager.channel_names().await.join(", ");
|
||||
let err_msg = if available.is_empty() {
|
||||
format!(
|
||||
"Failed to send to {}:{}: {}. No channels connected.",
|
||||
channel, target, e
|
||||
)
|
||||
} else {
|
||||
format!(
|
||||
"Failed to send to {}:{}. Available channels: {}. Error: {}",
|
||||
channel, target, available, e
|
||||
)
|
||||
};
|
||||
Err(ToolError::ExecutionFailed(err_msg))
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
let available = self.channel_manager.channel_names().await.join(", ");
|
||||
let err_msg = if available.is_empty() {
|
||||
format!(
|
||||
"Failed to send to {}:{}: {}. No channels connected.",
|
||||
channel, target, e
|
||||
)
|
||||
} else {
|
||||
// No channel specified — broadcast to all channels (routine with notify.channel = None)
|
||||
let results = self.channel_manager.broadcast_all(&target, response).await;
|
||||
let mut succeeded = Vec::new();
|
||||
let mut failed: Vec<&str> = Vec::new();
|
||||
for (ch, result) in &results {
|
||||
match result {
|
||||
Ok(()) => succeeded.push(ch.as_str()),
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
channel = %ch,
|
||||
target = %target,
|
||||
"broadcast_all: channel failed: {}", e
|
||||
);
|
||||
failed.push(ch.as_str());
|
||||
}
|
||||
}
|
||||
}
|
||||
if succeeded.is_empty() {
|
||||
let err_msg = if failed.is_empty() {
|
||||
"No channels connected.".to_string()
|
||||
} else {
|
||||
format!(
|
||||
"Failed to send to {}:{}. Available channels: {}. Error: {}",
|
||||
channel, target, available, e
|
||||
)
|
||||
format!("All channels failed: {}", failed.join(", "))
|
||||
};
|
||||
Err(ToolError::ExecutionFailed(err_msg))
|
||||
} else {
|
||||
tracing::info!(
|
||||
message_sent = true,
|
||||
channels = ?succeeded,
|
||||
target = %target,
|
||||
attachments = attachment_count,
|
||||
"Message broadcast via message tool"
|
||||
);
|
||||
let msg = format!(
|
||||
"Broadcast message to {} (target: {})",
|
||||
succeeded.join(", "),
|
||||
target
|
||||
);
|
||||
Ok(ToolOutput::text(msg, start.elapsed()))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -576,4 +624,90 @@ mod tests {
|
||||
ApprovalRequirement::Never,
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn message_tool_falls_back_to_job_metadata() {
|
||||
// Regression: when no conversation context is set (e.g. routine full-job),
|
||||
// the message tool should fall back to notify_channel/notify_user from
|
||||
// JobContext metadata instead of returning "No target specified".
|
||||
let tool = MessageTool::new(Arc::new(ChannelManager::new()));
|
||||
|
||||
let mut ctx = crate::context::JobContext::new("routine-job", "price alert");
|
||||
ctx.metadata = serde_json::json!({
|
||||
"notify_channel": "telegram",
|
||||
"notify_user": "123456789",
|
||||
});
|
||||
|
||||
// No set_context called — simulates a routine full-job worker
|
||||
let result = tool
|
||||
.execute(serde_json::json!({"content": "NEAR price is $5"}), &ctx)
|
||||
.await;
|
||||
|
||||
// Should fail at channel broadcast (no real channel), NOT at
|
||||
// "No target specified and no active conversation"
|
||||
assert!(result.is_err());
|
||||
let err = result.unwrap_err().to_string();
|
||||
assert!(
|
||||
!err.contains("No target specified"),
|
||||
"Should not get 'No target specified' when metadata has notify_user, got: {}",
|
||||
err
|
||||
);
|
||||
assert!(
|
||||
!err.contains("No channel specified"),
|
||||
"Should not get 'No channel specified' when metadata has notify_channel, got: {}",
|
||||
err
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn message_tool_no_metadata_still_errors() {
|
||||
// When neither conversation context nor metadata is set, should still
|
||||
// return a clear error (target resolution fails).
|
||||
let tool = MessageTool::new(Arc::new(ChannelManager::new()));
|
||||
let ctx = crate::context::JobContext::new("orphan-job", "no notify config");
|
||||
|
||||
let result = tool
|
||||
.execute(serde_json::json!({"content": "hello"}), &ctx)
|
||||
.await;
|
||||
|
||||
assert!(result.is_err());
|
||||
let err = result.unwrap_err().to_string();
|
||||
assert!(
|
||||
err.contains("No target specified"),
|
||||
"Expected 'No target specified' error, got: {}",
|
||||
err
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn message_tool_broadcasts_all_when_no_channel() {
|
||||
// Regression: when notify.channel is None but notify_user is set,
|
||||
// the message tool should attempt broadcast_all instead of erroring
|
||||
// with "No channel specified".
|
||||
let tool = MessageTool::new(Arc::new(ChannelManager::new()));
|
||||
|
||||
let mut ctx = crate::context::JobContext::new("routine-job", "price alert");
|
||||
ctx.metadata = serde_json::json!({
|
||||
"notify_user": "123456789",
|
||||
});
|
||||
|
||||
let result = tool
|
||||
.execute(serde_json::json!({"content": "NEAR price is $5"}), &ctx)
|
||||
.await;
|
||||
|
||||
// Should fail because no channels are registered (empty ChannelManager),
|
||||
// NOT because "No channel specified".
|
||||
assert!(result.is_err());
|
||||
let err = result.unwrap_err().to_string();
|
||||
assert!(
|
||||
!err.contains("No channel specified"),
|
||||
"Should not get 'No channel specified' when broadcasting, got: {}",
|
||||
err
|
||||
);
|
||||
assert!(
|
||||
err.contains("No channels connected") || err.contains("All channels failed"),
|
||||
"Expected channel delivery error, got: {}",
|
||||
err
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -40,5 +40,21 @@ pub use shell::ShellTool;
|
||||
pub use skill_tools::{SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool};
|
||||
pub use time::TimeTool;
|
||||
mod html_converter;
|
||||
pub mod image_analyze;
|
||||
pub mod image_edit;
|
||||
pub mod image_gen;
|
||||
|
||||
pub use html_converter::convert_html_to_markdown;
|
||||
pub use image_analyze::ImageAnalyzeTool;
|
||||
pub use image_edit::ImageEditTool;
|
||||
pub use image_gen::ImageGenerateTool;
|
||||
|
||||
/// Detect image media type from file extension via `mime_guess`.
|
||||
/// Falls back to `image/jpeg` for unrecognized or non-image extensions.
|
||||
pub(crate) fn media_type_from_path(path: &str) -> String {
|
||||
mime_guess::from_path(path)
|
||||
.first_raw()
|
||||
.filter(|m| m.starts_with("image/"))
|
||||
.unwrap_or("image/jpeg")
|
||||
.to_string()
|
||||
}
|
||||
|
||||
@@ -107,6 +107,10 @@ impl Tool for RoutineCreateTool {
|
||||
"notify_user": {
|
||||
"type": "string",
|
||||
"description": "User/target to notify (e.g. username, chat ID). Defaults to 'default'."
|
||||
},
|
||||
"timezone": {
|
||||
"type": "string",
|
||||
"description": "IANA timezone for cron schedule evaluation (e.g. 'America/New_York'). Defaults to UTC."
|
||||
}
|
||||
},
|
||||
"required": ["name", "trigger_type", "prompt"]
|
||||
@@ -143,12 +147,26 @@ impl Tool for RoutineCreateTool {
|
||||
"cron trigger requires 'schedule'".to_string(),
|
||||
)
|
||||
})?;
|
||||
let timezone = params
|
||||
.get("timezone")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|tz| {
|
||||
crate::timezone::parse_timezone(tz)
|
||||
.map(|_| tz.to_string())
|
||||
.ok_or_else(|| {
|
||||
ToolError::InvalidParameters(format!(
|
||||
"invalid IANA timezone: '{tz}'"
|
||||
))
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
// Validate cron expression
|
||||
next_cron_fire(schedule).map_err(|e| {
|
||||
next_cron_fire(schedule, timezone.as_deref()).map_err(|e| {
|
||||
ToolError::InvalidParameters(format!("invalid cron schedule: {e}"))
|
||||
})?;
|
||||
Trigger::Cron {
|
||||
schedule: schedule.to_string(),
|
||||
timezone,
|
||||
}
|
||||
}
|
||||
"event" => {
|
||||
@@ -228,8 +246,12 @@ impl Tool for RoutineCreateTool {
|
||||
.unwrap_or(300);
|
||||
|
||||
// Compute next fire time for cron
|
||||
let next_fire = if let Trigger::Cron { ref schedule } = trigger {
|
||||
next_cron_fire(schedule).unwrap_or(None)
|
||||
let next_fire = if let Trigger::Cron {
|
||||
ref schedule,
|
||||
ref timezone,
|
||||
} = trigger
|
||||
{
|
||||
next_cron_fire(schedule, timezone.as_deref()).unwrap_or(None)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
@@ -412,6 +434,10 @@ impl Tool for RoutineUpdateTool {
|
||||
"type": "string",
|
||||
"description": "New cron schedule (for cron triggers)"
|
||||
},
|
||||
"timezone": {
|
||||
"type": "string",
|
||||
"description": "IANA timezone for cron schedule (e.g. 'America/New_York'). Only valid for cron triggers."
|
||||
},
|
||||
"description": {
|
||||
"type": "string",
|
||||
"description": "New description"
|
||||
@@ -453,15 +479,47 @@ impl Tool for RoutineUpdateTool {
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(schedule) = params.get("schedule").and_then(|v| v.as_str()) {
|
||||
// Validate
|
||||
next_cron_fire(schedule)
|
||||
.map_err(|e| ToolError::InvalidParameters(format!("invalid cron schedule: {e}")))?;
|
||||
// Validate timezone param if provided
|
||||
let new_timezone = params
|
||||
.get("timezone")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|tz| {
|
||||
crate::timezone::parse_timezone(tz)
|
||||
.map(|_| tz.to_string())
|
||||
.ok_or_else(|| {
|
||||
ToolError::InvalidParameters(format!("invalid IANA timezone: '{tz}'"))
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
|
||||
routine.trigger = Trigger::Cron {
|
||||
schedule: schedule.to_string(),
|
||||
let new_schedule = params.get("schedule").and_then(|v| v.as_str());
|
||||
|
||||
if new_schedule.is_some() || new_timezone.is_some() {
|
||||
// Extract existing cron fields (cloned to avoid borrow conflict)
|
||||
let existing_cron = match &routine.trigger {
|
||||
Trigger::Cron { schedule, timezone } => Some((schedule.clone(), timezone.clone())),
|
||||
_ => None,
|
||||
};
|
||||
routine.next_fire_at = next_cron_fire(schedule).unwrap_or(None);
|
||||
|
||||
if let Some((old_schedule, old_tz)) = existing_cron {
|
||||
let effective_schedule = new_schedule.unwrap_or(&old_schedule);
|
||||
let effective_tz = new_timezone.or(old_tz);
|
||||
// Validate
|
||||
next_cron_fire(effective_schedule, effective_tz.as_deref()).map_err(|e| {
|
||||
ToolError::InvalidParameters(format!("invalid cron schedule: {e}"))
|
||||
})?;
|
||||
|
||||
routine.trigger = Trigger::Cron {
|
||||
schedule: effective_schedule.to_string(),
|
||||
timezone: effective_tz.clone(),
|
||||
};
|
||||
routine.next_fire_at =
|
||||
next_cron_fire(effective_schedule, effective_tz.as_deref()).unwrap_or(None);
|
||||
} else {
|
||||
return Err(ToolError::InvalidParameters(
|
||||
"Cannot update schedule or timezone on a non-cron routine.".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
self.store
|
||||
|
||||
+472
-49
@@ -1,7 +1,8 @@
|
||||
//! Time utility tool.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use chrono::{DateTime, LocalResult, NaiveDate, NaiveDateTime, TimeZone, Utc};
|
||||
use chrono_tz::Tz;
|
||||
|
||||
use crate::context::JobContext;
|
||||
use crate::tools::tool::{Tool, ToolError, ToolOutput, require_str};
|
||||
@@ -16,7 +17,7 @@ impl Tool for TimeTool {
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Get current time, convert timezones, or calculate time differences."
|
||||
"Get current time, parse or format timestamps, convert timezones, or calculate time differences."
|
||||
}
|
||||
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
@@ -25,20 +26,40 @@ impl Tool for TimeTool {
|
||||
"properties": {
|
||||
"operation": {
|
||||
"type": "string",
|
||||
"enum": ["now", "parse", "format", "diff"],
|
||||
"enum": ["now", "parse", "convert", "format", "diff"],
|
||||
"description": "The time operation to perform"
|
||||
},
|
||||
"input": {
|
||||
"type": "string",
|
||||
"description": "Input timestamp. Accepts RFC 3339, or a naive timestamp when timezone/from_timezone is provided."
|
||||
},
|
||||
"timestamp": {
|
||||
"type": "string",
|
||||
"description": "ISO 8601 timestamp (for parse/format/diff operations)"
|
||||
"description": "Alias for input (kept for backward compatibility)."
|
||||
},
|
||||
"timezone": {
|
||||
"type": "string",
|
||||
"description": "IANA timezone name (e.g. 'America/New_York'). Used by now/format, and can also interpret naive timestamps."
|
||||
},
|
||||
"from_timezone": {
|
||||
"type": "string",
|
||||
"description": "Source IANA timezone for naive input timestamps during convert/format/diff."
|
||||
},
|
||||
"to_timezone": {
|
||||
"type": "string",
|
||||
"description": "Target IANA timezone for convert."
|
||||
},
|
||||
"format": {
|
||||
"type": "string",
|
||||
"description": "Output format string (for format operation)"
|
||||
"description": "strftime format string for format (kept for backward compatibility)."
|
||||
},
|
||||
"format_string": {
|
||||
"type": "string",
|
||||
"description": "strftime format string for format."
|
||||
},
|
||||
"timestamp2": {
|
||||
"type": "string",
|
||||
"description": "Second timestamp (for diff operation)"
|
||||
"description": "Second timestamp for diff."
|
||||
}
|
||||
},
|
||||
"required": ["operation"]
|
||||
@@ -48,55 +69,18 @@ impl Tool for TimeTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
params: serde_json::Value,
|
||||
_ctx: &JobContext,
|
||||
ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
let operation = require_str(¶ms, "operation")?;
|
||||
|
||||
let result = match operation {
|
||||
"now" => {
|
||||
let now = Utc::now();
|
||||
serde_json::json!({
|
||||
"iso": now.to_rfc3339(),
|
||||
"unix": now.timestamp(),
|
||||
"unix_millis": now.timestamp_millis()
|
||||
})
|
||||
}
|
||||
"parse" => {
|
||||
let timestamp = require_str(¶ms, "timestamp")?;
|
||||
|
||||
let dt: DateTime<Utc> = timestamp.parse().map_err(|e| {
|
||||
ToolError::InvalidParameters(format!("invalid timestamp: {}", e))
|
||||
})?;
|
||||
|
||||
serde_json::json!({
|
||||
"iso": dt.to_rfc3339(),
|
||||
"unix": dt.timestamp(),
|
||||
"unix_millis": dt.timestamp_millis()
|
||||
})
|
||||
}
|
||||
"diff" => {
|
||||
let ts1 = require_str(¶ms, "timestamp")?;
|
||||
|
||||
let ts2 = require_str(¶ms, "timestamp2")?;
|
||||
|
||||
let dt1: DateTime<Utc> = ts1.parse().map_err(|e| {
|
||||
ToolError::InvalidParameters(format!("invalid timestamp: {}", e))
|
||||
})?;
|
||||
let dt2: DateTime<Utc> = ts2.parse().map_err(|e| {
|
||||
ToolError::InvalidParameters(format!("invalid timestamp2: {}", e))
|
||||
})?;
|
||||
|
||||
let diff = dt2.signed_duration_since(dt1);
|
||||
|
||||
serde_json::json!({
|
||||
"seconds": diff.num_seconds(),
|
||||
"minutes": diff.num_minutes(),
|
||||
"hours": diff.num_hours(),
|
||||
"days": diff.num_days()
|
||||
})
|
||||
}
|
||||
"now" => execute_now(¶ms, ctx)?,
|
||||
"parse" => execute_parse(¶ms, ctx)?,
|
||||
"convert" => execute_convert(¶ms, ctx)?,
|
||||
"format" => execute_format(¶ms, ctx)?,
|
||||
"diff" => execute_diff(¶ms, ctx)?,
|
||||
_ => {
|
||||
return Err(ToolError::InvalidParameters(format!(
|
||||
"unknown operation: {}",
|
||||
@@ -112,3 +96,442 @@ impl Tool for TimeTool {
|
||||
false // Internal tool, no external data
|
||||
}
|
||||
}
|
||||
|
||||
fn execute_now(
|
||||
params: &serde_json::Value,
|
||||
ctx: &JobContext,
|
||||
) -> Result<serde_json::Value, ToolError> {
|
||||
let now = Utc::now();
|
||||
let mut result = serde_json::json!({
|
||||
"iso": now.to_rfc3339(),
|
||||
"utc_iso": now.to_rfc3339(),
|
||||
"unix": now.timestamp(),
|
||||
"unix_millis": now.timestamp_millis()
|
||||
});
|
||||
|
||||
if let Some((tz, tz_name)) = resolve_timezone_for_output(params, ctx)? {
|
||||
let local = now.with_timezone(&tz);
|
||||
result["local_iso"] = serde_json::Value::String(local.to_rfc3339());
|
||||
result["timezone"] = serde_json::Value::String(tz_name);
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
fn execute_parse(
|
||||
params: &serde_json::Value,
|
||||
ctx: &JobContext,
|
||||
) -> Result<serde_json::Value, ToolError> {
|
||||
let input = require_input(params)?;
|
||||
let parse_tz = resolve_parse_timezone(params, ctx)?;
|
||||
let dt = parse_timestamp(input, parse_tz.as_ref())?;
|
||||
|
||||
Ok(serde_json::json!({
|
||||
"iso": dt.to_rfc3339(),
|
||||
"unix": dt.timestamp(),
|
||||
"unix_millis": dt.timestamp_millis()
|
||||
}))
|
||||
}
|
||||
|
||||
fn execute_convert(
|
||||
params: &serde_json::Value,
|
||||
ctx: &JobContext,
|
||||
) -> Result<serde_json::Value, ToolError> {
|
||||
let input = require_input(params)?;
|
||||
let source_tz = optional_timezone(params, &["from_timezone", "timezone"])?;
|
||||
let dt = parse_timestamp(input, source_tz.as_ref())?;
|
||||
|
||||
let target_name = params
|
||||
.get("to_timezone")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| {
|
||||
ToolError::InvalidParameters("convert operation requires 'to_timezone'".to_string())
|
||||
})?;
|
||||
let target_tz = parse_timezone(target_name)?;
|
||||
let converted = dt.with_timezone(&target_tz);
|
||||
|
||||
let mut result = serde_json::json!({
|
||||
"input": input,
|
||||
"utc_iso": dt.to_rfc3339(),
|
||||
"output": converted.to_rfc3339(),
|
||||
"timezone": target_tz.to_string()
|
||||
});
|
||||
|
||||
if let Some((ctx_tz, ctx_tz_name)) = context_timezone(ctx)? {
|
||||
result["context_timezone"] = serde_json::Value::String(ctx_tz_name);
|
||||
result["context_iso"] = serde_json::Value::String(dt.with_timezone(&ctx_tz).to_rfc3339());
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
fn execute_format(
|
||||
params: &serde_json::Value,
|
||||
ctx: &JobContext,
|
||||
) -> Result<serde_json::Value, ToolError> {
|
||||
let input = require_input(params)?;
|
||||
let output_tz = resolve_timezone_for_output(params, ctx)?;
|
||||
let source_tz = optional_timezone(params, &["from_timezone"])?
|
||||
.or_else(|| output_tz.as_ref().map(|(tz, _)| *tz));
|
||||
let dt = parse_timestamp(input, source_tz.as_ref())?;
|
||||
let format_string = params
|
||||
.get("format_string")
|
||||
.and_then(|v| v.as_str())
|
||||
.or_else(|| params.get("format").and_then(|v| v.as_str()))
|
||||
.unwrap_or("%Y-%m-%d %H:%M:%S %Z");
|
||||
|
||||
let mut result = if let Some((tz, tz_name)) = output_tz {
|
||||
serde_json::json!({
|
||||
"formatted": dt.with_timezone(&tz).format(format_string).to_string(),
|
||||
"timezone": tz_name
|
||||
})
|
||||
} else {
|
||||
serde_json::json!({
|
||||
"formatted": dt.format(format_string).to_string()
|
||||
})
|
||||
};
|
||||
|
||||
result["utc_iso"] = serde_json::Value::String(dt.to_rfc3339());
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
fn execute_diff(
|
||||
params: &serde_json::Value,
|
||||
ctx: &JobContext,
|
||||
) -> Result<serde_json::Value, ToolError> {
|
||||
let parse_tz = resolve_parse_timezone(params, ctx)?;
|
||||
let ts1 = require_input(params)?;
|
||||
let ts2 = params
|
||||
.get("timestamp2")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| {
|
||||
ToolError::InvalidParameters("diff operation requires 'timestamp2'".to_string())
|
||||
})?;
|
||||
|
||||
let dt1 = parse_timestamp(ts1, parse_tz.as_ref())?;
|
||||
let dt2 = parse_timestamp(ts2, parse_tz.as_ref())?;
|
||||
let diff = dt2.signed_duration_since(dt1);
|
||||
|
||||
Ok(serde_json::json!({
|
||||
"seconds": diff.num_seconds(),
|
||||
"minutes": diff.num_minutes(),
|
||||
"hours": diff.num_hours(),
|
||||
"days": diff.num_days()
|
||||
}))
|
||||
}
|
||||
|
||||
fn require_input(params: &serde_json::Value) -> Result<&str, ToolError> {
|
||||
params
|
||||
.get("input")
|
||||
.and_then(|v| v.as_str())
|
||||
.or_else(|| params.get("timestamp").and_then(|v| v.as_str()))
|
||||
.ok_or_else(|| {
|
||||
ToolError::InvalidParameters(
|
||||
"missing 'input' (or legacy 'timestamp') parameter".to_string(),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn resolve_parse_timezone(
|
||||
params: &serde_json::Value,
|
||||
ctx: &JobContext,
|
||||
) -> Result<Option<Tz>, ToolError> {
|
||||
if let Some(tz) = optional_timezone(params, &["from_timezone", "timezone"])? {
|
||||
return Ok(Some(tz));
|
||||
}
|
||||
|
||||
Ok(context_timezone(ctx)?.map(|(tz, _)| tz))
|
||||
}
|
||||
|
||||
fn resolve_timezone_for_output(
|
||||
params: &serde_json::Value,
|
||||
ctx: &JobContext,
|
||||
) -> Result<Option<(Tz, String)>, ToolError> {
|
||||
if let Some(name) = params.get("timezone").and_then(|v| v.as_str()) {
|
||||
let tz = parse_timezone(name)?;
|
||||
return Ok(Some((tz, tz.to_string())));
|
||||
}
|
||||
|
||||
context_timezone(ctx)
|
||||
}
|
||||
|
||||
/// Resolve the user's timezone from the JobContext.
|
||||
///
|
||||
/// Uses `ctx.user_timezone` (set from main's timezone resolution) as the
|
||||
/// primary source. Falls back to metadata fields for backward compatibility.
|
||||
fn context_timezone(ctx: &JobContext) -> Result<Option<(Tz, String)>, ToolError> {
|
||||
// Primary: use the dedicated user_timezone field from JobContext
|
||||
if ctx.user_timezone != "UTC"
|
||||
&& !ctx.user_timezone.is_empty()
|
||||
&& let Some(tz) = crate::timezone::parse_timezone(&ctx.user_timezone)
|
||||
{
|
||||
return Ok(Some((tz, tz.to_string())));
|
||||
}
|
||||
|
||||
// Fallback: check metadata for backward compatibility
|
||||
let tz_name = ctx
|
||||
.metadata
|
||||
.get("user_timezone")
|
||||
.and_then(|v| v.as_str())
|
||||
.or_else(|| ctx.metadata.get("timezone").and_then(|v| v.as_str()));
|
||||
|
||||
match tz_name {
|
||||
Some(name) => {
|
||||
let tz = parse_timezone(name)?;
|
||||
Ok(Some((tz, tz.to_string())))
|
||||
}
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
fn optional_timezone(params: &serde_json::Value, keys: &[&str]) -> Result<Option<Tz>, ToolError> {
|
||||
for key in keys {
|
||||
if let Some(value) = params.get(*key).and_then(|v| v.as_str()) {
|
||||
return parse_timezone(value).map(Some);
|
||||
}
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn parse_timezone(value: &str) -> Result<Tz, ToolError> {
|
||||
value.parse::<Tz>().map_err(|_| {
|
||||
ToolError::InvalidParameters(format!(
|
||||
"Unknown timezone '{}'. Use IANA names like 'America/New_York' or 'Europe/London'.",
|
||||
value
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_timestamp(input: &str, fallback_tz: Option<&Tz>) -> Result<DateTime<Utc>, ToolError> {
|
||||
if let Ok(dt) = DateTime::parse_from_rfc3339(input) {
|
||||
return Ok(dt.with_timezone(&Utc));
|
||||
}
|
||||
|
||||
if let Some(naive) = parse_naive_datetime(input) {
|
||||
return localize_naive_datetime(naive, fallback_tz, input);
|
||||
}
|
||||
|
||||
Err(ToolError::InvalidParameters(format!(
|
||||
"invalid timestamp '{}': expected RFC 3339 or a naive timestamp with timezone/from_timezone",
|
||||
input
|
||||
)))
|
||||
}
|
||||
|
||||
fn parse_naive_datetime(input: &str) -> Option<NaiveDateTime> {
|
||||
const DATETIME_FORMATS: &[&str] = &[
|
||||
"%Y-%m-%d %H:%M:%S%.f",
|
||||
"%Y-%m-%dT%H:%M:%S%.f",
|
||||
"%Y-%m-%d %H:%M",
|
||||
"%Y-%m-%dT%H:%M",
|
||||
];
|
||||
const DATE_FORMATS: &[&str] = &["%Y-%m-%d"];
|
||||
|
||||
for format in DATETIME_FORMATS {
|
||||
if let Ok(value) = NaiveDateTime::parse_from_str(input, format) {
|
||||
return Some(value);
|
||||
}
|
||||
}
|
||||
|
||||
for format in DATE_FORMATS {
|
||||
if let Ok(date) = NaiveDate::parse_from_str(input, format) {
|
||||
return date.and_hms_opt(0, 0, 0);
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
fn localize_naive_datetime(
|
||||
naive: NaiveDateTime,
|
||||
fallback_tz: Option<&Tz>,
|
||||
original_input: &str,
|
||||
) -> Result<DateTime<Utc>, ToolError> {
|
||||
let tz = fallback_tz.ok_or_else(|| {
|
||||
ToolError::InvalidParameters(format!(
|
||||
"timestamp '{}' has no UTC offset; provide 'timezone' or 'from_timezone'",
|
||||
original_input
|
||||
))
|
||||
})?;
|
||||
|
||||
match tz.from_local_datetime(&naive) {
|
||||
LocalResult::Single(dt) => Ok(dt.with_timezone(&Utc)),
|
||||
LocalResult::Ambiguous(_, _) => Err(ToolError::InvalidParameters(format!(
|
||||
"timestamp '{}' is ambiguous in timezone '{}'; include an explicit UTC offset instead",
|
||||
original_input, tz
|
||||
))),
|
||||
LocalResult::None => Err(ToolError::InvalidParameters(format!(
|
||||
"timestamp '{}' does not exist in timezone '{}'",
|
||||
original_input, tz
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_now_accepts_explicit_timezone() {
|
||||
let tool = TimeTool;
|
||||
let ctx = JobContext::with_user("test", "chat", "test");
|
||||
|
||||
let output = tool
|
||||
.execute(
|
||||
serde_json::json!({
|
||||
"operation": "now",
|
||||
"timezone": "America/New_York"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.expect("execute");
|
||||
|
||||
assert_eq!(output.result["timezone"].as_str(), Some("America/New_York"));
|
||||
assert!(
|
||||
output.result.get("utc_iso").is_some(),
|
||||
"should have utc_iso"
|
||||
);
|
||||
assert!(
|
||||
output.result.get("local_iso").is_some(),
|
||||
"should have local_iso"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_now_includes_local_time_when_user_timezone_set() {
|
||||
let tool = TimeTool;
|
||||
let mut ctx = JobContext::with_user("test", "chat", "test");
|
||||
ctx.user_timezone = "America/New_York".to_string();
|
||||
|
||||
let output = tool
|
||||
.execute(serde_json::json!({"operation": "now"}), &ctx)
|
||||
.await
|
||||
.expect("execute");
|
||||
assert!(
|
||||
output.result.get("local_iso").is_some(),
|
||||
"should have local_iso"
|
||||
);
|
||||
assert_eq!(
|
||||
output.result["timezone"].as_str(),
|
||||
Some("America/New_York"),
|
||||
"should report timezone"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_now_uses_context_metadata_timezone_fallback() {
|
||||
let tool = TimeTool;
|
||||
let mut ctx = JobContext::with_user("test", "chat", "test");
|
||||
ctx.metadata = serde_json::json!({
|
||||
"user_timezone": "America/Los_Angeles"
|
||||
});
|
||||
|
||||
let output = tool
|
||||
.execute(serde_json::json!({"operation": "now"}), &ctx)
|
||||
.await
|
||||
.expect("execute");
|
||||
|
||||
assert_eq!(
|
||||
output.result["timezone"].as_str(),
|
||||
Some("America/Los_Angeles")
|
||||
);
|
||||
assert!(
|
||||
output.result.get("local_iso").is_some(),
|
||||
"should have local_iso"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_now_returns_utc_by_default() {
|
||||
let tool = TimeTool;
|
||||
let ctx = JobContext::with_user("test", "chat", "test");
|
||||
// Default user_timezone is "UTC" -- context_timezone skips UTC so no
|
||||
// local_iso is added, but iso and utc_iso are always present.
|
||||
let output = tool
|
||||
.execute(serde_json::json!({"operation": "now"}), &ctx)
|
||||
.await
|
||||
.expect("execute");
|
||||
assert!(output.result.get("iso").is_some(), "should have iso");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_convert_across_dst_boundary() {
|
||||
let tool = TimeTool;
|
||||
let ctx = JobContext::with_user("test", "chat", "test");
|
||||
|
||||
let output = tool
|
||||
.execute(
|
||||
serde_json::json!({
|
||||
"operation": "convert",
|
||||
"input": "2026-03-08T07:30:00Z",
|
||||
"to_timezone": "America/New_York"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.expect("execute");
|
||||
|
||||
assert_eq!(output.result["timezone"].as_str(), Some("America/New_York"));
|
||||
assert_eq!(
|
||||
output.result["output"].as_str(),
|
||||
Some("2026-03-08T03:30:00-04:00")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_format_with_timezone() {
|
||||
let tool = TimeTool;
|
||||
let ctx = JobContext::with_user("test", "chat", "test");
|
||||
|
||||
let output = tool
|
||||
.execute(
|
||||
serde_json::json!({
|
||||
"operation": "format",
|
||||
"input": "2026-03-08T07:30:00Z",
|
||||
"timezone": "America/New_York",
|
||||
"format_string": "%Y-%m-%d %H:%M:%S %Z"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.expect("execute");
|
||||
|
||||
assert_eq!(output.result["timezone"].as_str(), Some("America/New_York"));
|
||||
assert_eq!(
|
||||
output.result["formatted"].as_str(),
|
||||
Some("2026-03-08 03:30:00 EDT")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_invalid_timezone_returns_clear_error() {
|
||||
let tool = TimeTool;
|
||||
let ctx = JobContext::with_user("test", "chat", "test");
|
||||
|
||||
let err = tool
|
||||
.execute(
|
||||
serde_json::json!({
|
||||
"operation": "now",
|
||||
"timezone": "Mars/Olympus"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.expect_err("expected invalid timezone error");
|
||||
|
||||
match err {
|
||||
ToolError::InvalidParameters(message) => {
|
||||
assert!(message.contains("Unknown timezone 'Mars/Olympus'"));
|
||||
}
|
||||
other => panic!("unexpected error: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_naive_timestamp_with_timezone() {
|
||||
let dt = parse_timestamp("2026-03-08 03:30:00", Some(&chrono_tz::America::New_York))
|
||||
.expect("parse timestamp");
|
||||
|
||||
assert_eq!(dt.to_rfc3339(), "2026-03-08T07:30:00+00:00");
|
||||
}
|
||||
}
|
||||
|
||||
+570
-23
@@ -4,6 +4,7 @@
|
||||
//! See: https://spec.modelcontextprotocol.io/specification/2025-03-26/basic/authorization/
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::net::IpAddr;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
@@ -199,23 +200,285 @@ impl PkceChallenge {
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Well-known URI construction (RFC 8414 / RFC 9728)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Build a well-known URI according to RFC 8414 / RFC 9728.
|
||||
///
|
||||
/// The path component of the base URL is placed *after* the well-known suffix:
|
||||
/// ```text
|
||||
/// https://example.com/path + oauth-authorization-server
|
||||
/// -> https://example.com/.well-known/oauth-authorization-server/path
|
||||
/// ```
|
||||
pub fn build_well_known_uri(base_url: &str, suffix: &str) -> Result<String, AuthError> {
|
||||
let parsed = reqwest::Url::parse(base_url)
|
||||
.map_err(|e| AuthError::DiscoveryFailed(format!("Invalid URL: {}", e)))?;
|
||||
let origin = parsed.origin().ascii_serialization();
|
||||
let path = parsed.path().trim_end_matches('/');
|
||||
Ok(format!("{}/.well-known/{}{}", origin, suffix, path))
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// RFC 8707 resource parameter
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Compute the canonical resource URI for RFC 8707.
|
||||
///
|
||||
/// Strips fragments and trailing slashes from the server URL.
|
||||
pub fn canonical_resource_uri(server_url: &str) -> String {
|
||||
match reqwest::Url::parse(server_url) {
|
||||
Ok(mut parsed) => {
|
||||
parsed.set_fragment(None);
|
||||
let s = parsed.to_string();
|
||||
s.trim_end_matches('/').to_string()
|
||||
}
|
||||
Err(_) => server_url.trim_end_matches('/').to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// SSRF protection
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Check if an IP address is dangerous (loopback, link-local, private, etc.)
|
||||
fn is_dangerous_ip(ip: IpAddr) -> bool {
|
||||
match ip {
|
||||
IpAddr::V4(v4) => {
|
||||
v4.is_loopback()
|
||||
|| v4.is_private()
|
||||
|| v4.is_link_local()
|
||||
|| v4.is_broadcast()
|
||||
|| v4.is_unspecified()
|
||||
|| (v4.octets()[0] == 169 && v4.octets()[1] == 254) // link-local
|
||||
|| (v4.octets()[0] == 100 && (v4.octets()[1] & 0xC0) == 64) // CGNAT 100.64/10
|
||||
}
|
||||
IpAddr::V6(v6) => {
|
||||
let segs = v6.segments();
|
||||
v6.is_loopback()
|
||||
|| v6.is_unspecified()
|
||||
// Link-local (fe80::/10)
|
||||
|| (segs[0] & 0xffc0) == 0xfe80
|
||||
// Site-local / deprecated (fec0::/10)
|
||||
|| (segs[0] & 0xffc0) == 0xfec0
|
||||
// Unique local (fc00::/7)
|
||||
|| (segs[0] & 0xfe00) == 0xfc00
|
||||
// Documentation (2001:db8::/32)
|
||||
|| (segs[0] == 0x2001 && segs[1] == 0x0db8)
|
||||
// Check for IPv4-mapped IPv6 (::ffff:x.x.x.x)
|
||||
|| v6
|
||||
.to_ipv4_mapped()
|
||||
.is_some_and(|v4| is_dangerous_ip(IpAddr::V4(v4)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate that a URL is safe for server-side requests (SSRF protection).
|
||||
async fn validate_url_safe(url: &str) -> Result<(), AuthError> {
|
||||
let parsed = reqwest::Url::parse(url)
|
||||
.map_err(|e| AuthError::DiscoveryFailed(format!("Invalid URL: {}", e)))?;
|
||||
|
||||
// Must be HTTPS. HTTP is only allowed for localhost/loopback (dev scenarios).
|
||||
let scheme = parsed.scheme();
|
||||
if scheme != "https" && scheme != "http" {
|
||||
return Err(AuthError::DiscoveryFailed(format!(
|
||||
"Unsupported scheme: {}",
|
||||
scheme
|
||||
)));
|
||||
}
|
||||
if scheme == "http" {
|
||||
let host = parsed.host_str().unwrap_or("");
|
||||
let is_localhost =
|
||||
host == "localhost" || host == "127.0.0.1" || host == "::1" || host == "[::1]";
|
||||
if !is_localhost {
|
||||
return Err(AuthError::DiscoveryFailed(format!(
|
||||
"HTTP is only allowed for localhost; use HTTPS for '{}'",
|
||||
host
|
||||
)));
|
||||
}
|
||||
// Localhost HTTP is allowed for dev — skip SSRF checks since we've
|
||||
// already validated the host is localhost/loopback.
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let host = parsed
|
||||
.host_str()
|
||||
.ok_or_else(|| AuthError::DiscoveryFailed("URL has no host".to_string()))?;
|
||||
|
||||
// For IP literals, parse directly and check.
|
||||
if let Ok(ip) = host.parse::<IpAddr>()
|
||||
&& is_dangerous_ip(ip)
|
||||
{
|
||||
return Err(AuthError::DiscoveryFailed(format!(
|
||||
"URL points to a restricted IP address: {}",
|
||||
host
|
||||
)));
|
||||
}
|
||||
|
||||
// For hostnames, resolve DNS and check each resolved address.
|
||||
// This prevents DNS-based SSRF where a hostname resolves to an internal IP
|
||||
// (e.g., 169.254.169.254 for cloud metadata endpoints).
|
||||
if host.parse::<IpAddr>().is_err() {
|
||||
let addr = format!("{}:{}", host, parsed.port_or_known_default().unwrap_or(443));
|
||||
match tokio::net::lookup_host(&addr).await {
|
||||
Ok(addrs) => {
|
||||
for socket_addr in addrs {
|
||||
if is_dangerous_ip(socket_addr.ip()) {
|
||||
return Err(AuthError::DiscoveryFailed(format!(
|
||||
"URL hostname '{}' resolves to restricted IP address: {}",
|
||||
host,
|
||||
socket_addr.ip()
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
// DNS failure = fail closed (do not allow the request)
|
||||
return Err(AuthError::DiscoveryFailed(format!(
|
||||
"DNS resolution failed for '{}': {}",
|
||||
host, e
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Multi-strategy OAuth discovery helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Parse the resource_metadata URL from a WWW-Authenticate header value.
|
||||
fn parse_resource_metadata_url(www_authenticate: &str) -> Option<String> {
|
||||
// Try comma-separated parameters first
|
||||
for part in www_authenticate.split(',') {
|
||||
let part = part.trim();
|
||||
if let Some(rest) = part.strip_prefix("resource_metadata=\"") {
|
||||
return rest.strip_suffix('"').map(|s| s.to_string());
|
||||
}
|
||||
if let Some(rest) = part.strip_prefix("resource_metadata=") {
|
||||
let val = rest.trim_matches('"');
|
||||
return Some(val.to_string());
|
||||
}
|
||||
}
|
||||
// Also try whitespace-separated tokens (e.g. Bearer resource_metadata="url")
|
||||
for part in www_authenticate.split_whitespace() {
|
||||
if let Some(rest) = part.strip_prefix("resource_metadata=\"") {
|
||||
return rest
|
||||
.trim_end_matches(',')
|
||||
.strip_suffix('"')
|
||||
.map(|s| s.to_string());
|
||||
}
|
||||
if let Some(rest) = part.strip_prefix("resource_metadata=") {
|
||||
let val = rest.trim_matches('"').trim_end_matches(',');
|
||||
return Some(val.to_string());
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Fetch protected resource metadata from a URL.
|
||||
async fn fetch_resource_metadata(url: &str) -> Result<ProtectedResourceMetadata, AuthError> {
|
||||
validate_url_safe(url).await?;
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(10))
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.build()
|
||||
.map_err(|e| AuthError::Http(e.to_string()))?;
|
||||
|
||||
let response = client
|
||||
.get(url)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AuthError::DiscoveryFailed(e.to_string()))?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
return Err(AuthError::DiscoveryFailed(format!(
|
||||
"HTTP {}",
|
||||
response.status()
|
||||
)));
|
||||
}
|
||||
|
||||
response
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| AuthError::DiscoveryFailed(format!("Invalid metadata: {}", e)))
|
||||
}
|
||||
|
||||
/// Try to discover OAuth metadata via 401 challenge response.
|
||||
async fn discover_via_401(server_url: &str) -> Result<AuthorizationServerMetadata, AuthError> {
|
||||
validate_url_safe(server_url).await?;
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(10))
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.build()
|
||||
.map_err(|e| AuthError::Http(e.to_string()))?;
|
||||
|
||||
let response = client
|
||||
.post(server_url)
|
||||
.header("Content-Type", "application/json")
|
||||
.body("{}")
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AuthError::DiscoveryFailed(e.to_string()))?;
|
||||
|
||||
if response.status().as_u16() != 401 {
|
||||
return Err(AuthError::DiscoveryFailed(format!(
|
||||
"Expected 401, got {}",
|
||||
response.status()
|
||||
)));
|
||||
}
|
||||
|
||||
let www_auth = response
|
||||
.headers()
|
||||
.get("WWW-Authenticate")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.ok_or_else(|| {
|
||||
AuthError::DiscoveryFailed("No WWW-Authenticate header in 401 response".to_string())
|
||||
})?;
|
||||
|
||||
let resource_metadata_url = parse_resource_metadata_url(www_auth).ok_or_else(|| {
|
||||
AuthError::DiscoveryFailed(
|
||||
"No resource_metadata URL in WWW-Authenticate header".to_string(),
|
||||
)
|
||||
})?;
|
||||
|
||||
let resource_meta = fetch_resource_metadata(&resource_metadata_url).await?;
|
||||
try_discover_from_auth_servers(&resource_meta).await
|
||||
}
|
||||
|
||||
/// Try to discover auth server metadata from resource metadata's authorization_servers list.
|
||||
async fn try_discover_from_auth_servers(
|
||||
resource_meta: &ProtectedResourceMetadata,
|
||||
) -> Result<AuthorizationServerMetadata, AuthError> {
|
||||
let auth_server_url = resource_meta
|
||||
.authorization_servers
|
||||
.first()
|
||||
.ok_or_else(|| AuthError::DiscoveryFailed("No authorization servers listed".to_string()))?;
|
||||
|
||||
discover_authorization_server(auth_server_url).await
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Discovery functions
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Discover protected resource metadata from an MCP server.
|
||||
pub async fn discover_protected_resource(
|
||||
server_url: &str,
|
||||
) -> Result<ProtectedResourceMetadata, AuthError> {
|
||||
validate_url_safe(server_url).await?;
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(10))
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.build()
|
||||
.map_err(|e| AuthError::Http(e.to_string()))?;
|
||||
|
||||
// Parse the server URL to extract the origin (scheme + host + port)
|
||||
// The .well-known endpoints are always at the root of the origin, not under any path
|
||||
let parsed = reqwest::Url::parse(server_url)
|
||||
.map_err(|e| AuthError::DiscoveryFailed(format!("Invalid server URL: {}", e)))?;
|
||||
let origin = parsed.origin().ascii_serialization();
|
||||
|
||||
// Try the well-known endpoint at the origin root
|
||||
let well_known_url = format!("{}/.well-known/oauth-protected-resource", origin);
|
||||
let well_known_url = build_well_known_uri(server_url, "oauth-protected-resource")?;
|
||||
|
||||
let response = client
|
||||
.get(&well_known_url)
|
||||
@@ -237,13 +500,15 @@ pub async fn discover_protected_resource(
|
||||
pub async fn discover_authorization_server(
|
||||
auth_server_url: &str,
|
||||
) -> Result<AuthorizationServerMetadata, AuthError> {
|
||||
validate_url_safe(auth_server_url).await?;
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(10))
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.build()
|
||||
.map_err(|e| AuthError::Http(e.to_string()))?;
|
||||
|
||||
let base_url = auth_server_url.trim_end_matches('/');
|
||||
let well_known_url = format!("{}/.well-known/oauth-authorization-server", base_url);
|
||||
let well_known_url = build_well_known_uri(auth_server_url, "oauth-authorization-server")?;
|
||||
|
||||
let response = client
|
||||
.get(&well_known_url)
|
||||
@@ -298,20 +563,27 @@ pub async fn discover_oauth_endpoints(
|
||||
/// Discover full OAuth metadata including DCR support.
|
||||
///
|
||||
/// Returns authorization server metadata which includes registration_endpoint if DCR is supported.
|
||||
/// Uses a 3-strategy discovery chain:
|
||||
/// 1. **401-based**: POST to MCP server, parse WWW-Authenticate header for resource_metadata URL
|
||||
/// 2. **RFC 9728**: Discover protected resource metadata, then authorization server from it
|
||||
/// 3. **Direct**: Treat MCP server as its own auth server
|
||||
pub async fn discover_full_oauth_metadata(
|
||||
server_url: &str,
|
||||
) -> Result<AuthorizationServerMetadata, AuthError> {
|
||||
// Try to discover from the server
|
||||
let resource_meta = discover_protected_resource(server_url).await?;
|
||||
// Strategy 1: 401-based discovery
|
||||
if let Ok(meta) = discover_via_401(server_url).await {
|
||||
return Ok(meta);
|
||||
}
|
||||
|
||||
// Get the first authorization server
|
||||
let auth_server_url = resource_meta
|
||||
.authorization_servers
|
||||
.first()
|
||||
.ok_or_else(|| AuthError::DiscoveryFailed("No authorization servers listed".to_string()))?;
|
||||
// Strategy 2: RFC 9728 protected resource discovery
|
||||
if let Ok(resource_meta) = discover_protected_resource(server_url).await
|
||||
&& let Ok(meta) = try_discover_from_auth_servers(&resource_meta).await
|
||||
{
|
||||
return Ok(meta);
|
||||
}
|
||||
|
||||
// Discover the authorization server metadata
|
||||
discover_authorization_server(auth_server_url).await
|
||||
// Strategy 3: Direct - treat MCP server as its own auth server
|
||||
discover_authorization_server(server_url).await
|
||||
}
|
||||
|
||||
/// Perform Dynamic Client Registration with an authorization server.
|
||||
@@ -321,8 +593,11 @@ pub async fn register_client(
|
||||
registration_endpoint: &str,
|
||||
redirect_uri: &str,
|
||||
) -> Result<ClientRegistrationResponse, AuthError> {
|
||||
validate_url_safe(registration_endpoint).await?;
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(30))
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.build()
|
||||
.map_err(|e| AuthError::Http(e.to_string()))?;
|
||||
|
||||
@@ -417,7 +692,7 @@ pub async fn authorize_mcp_server(
|
||||
|
||||
println!(" Registering client dynamically...");
|
||||
let registration = register_client(®istration_endpoint, &redirect_uri).await?;
|
||||
println!(" ✓ Client registered: {}", registration.client_id);
|
||||
println!(" Client registered: {}", registration.client_id);
|
||||
|
||||
(
|
||||
registration.client_id,
|
||||
@@ -436,6 +711,15 @@ pub async fn authorize_mcp_server(
|
||||
None
|
||||
};
|
||||
|
||||
// Compute canonical resource URI for RFC 8707
|
||||
let resource = canonical_resource_uri(&server_config.url);
|
||||
|
||||
// Validate the discovered authorization URL to prevent a malicious MCP server
|
||||
// from redirecting the user to a phishing page or non-HTTPS endpoint.
|
||||
validate_url_safe(&authorization_url)
|
||||
.await
|
||||
.map_err(|e| AuthError::DiscoveryFailed(format!("Unsafe authorization endpoint: {}", e)))?;
|
||||
|
||||
// Build authorization URL
|
||||
let auth_url = build_authorization_url(
|
||||
&authorization_url,
|
||||
@@ -444,6 +728,7 @@ pub async fn authorize_mcp_server(
|
||||
&scopes,
|
||||
pkce.as_ref(),
|
||||
&extra_params,
|
||||
Some(&resource),
|
||||
);
|
||||
|
||||
// Open browser
|
||||
@@ -462,9 +747,15 @@ pub async fn authorize_mcp_server(
|
||||
println!(" Exchanging code for token...");
|
||||
|
||||
// Exchange code for token
|
||||
let token =
|
||||
exchange_code_for_token(&token_url, &client_id, &code, &redirect_uri, pkce.as_ref())
|
||||
.await?;
|
||||
let token = exchange_code_for_token(
|
||||
&token_url,
|
||||
&client_id,
|
||||
&code,
|
||||
&redirect_uri,
|
||||
pkce.as_ref(),
|
||||
Some(&resource),
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Store the tokens
|
||||
store_tokens(secrets, user_id, server_config, &token).await?;
|
||||
@@ -493,6 +784,7 @@ pub fn build_authorization_url(
|
||||
scopes: &[String],
|
||||
pkce: Option<&PkceChallenge>,
|
||||
extra_params: &HashMap<String, String>,
|
||||
resource: Option<&str>,
|
||||
) -> String {
|
||||
let mut url = format!(
|
||||
"{}?client_id={}&response_type=code&redirect_uri={}",
|
||||
@@ -523,6 +815,10 @@ pub fn build_authorization_url(
|
||||
));
|
||||
}
|
||||
|
||||
if let Some(resource) = resource {
|
||||
url.push_str(&format!("&resource={}", urlencoding::encode(resource)));
|
||||
}
|
||||
|
||||
url
|
||||
}
|
||||
|
||||
@@ -553,9 +849,13 @@ pub async fn exchange_code_for_token(
|
||||
code: &str,
|
||||
redirect_uri: &str,
|
||||
pkce: Option<&PkceChallenge>,
|
||||
resource: Option<&str>,
|
||||
) -> Result<AccessToken, AuthError> {
|
||||
validate_url_safe(token_url).await?;
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(30))
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.build()
|
||||
.map_err(|e| AuthError::Http(e.to_string()))?;
|
||||
|
||||
@@ -570,6 +870,10 @@ pub async fn exchange_code_for_token(
|
||||
params.push(("code_verifier", pkce.verifier.clone()));
|
||||
}
|
||||
|
||||
if let Some(resource) = resource {
|
||||
params.push(("resource", resource.to_string()));
|
||||
}
|
||||
|
||||
let response = client
|
||||
.post(token_url)
|
||||
.form(¶ms)
|
||||
@@ -738,15 +1042,22 @@ pub async fn refresh_access_token(
|
||||
auth_meta.token_endpoint
|
||||
};
|
||||
|
||||
validate_url_safe(&token_url).await?;
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(30))
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.build()
|
||||
.map_err(|e| AuthError::Http(e.to_string()))?;
|
||||
|
||||
// Compute canonical resource URI for RFC 8707
|
||||
let resource = canonical_resource_uri(&server_config.url);
|
||||
|
||||
let params = vec![
|
||||
("grant_type", "refresh_token".to_string()),
|
||||
("refresh_token", refresh_token.expose().to_string()),
|
||||
("client_id", client_id),
|
||||
("resource", resource),
|
||||
];
|
||||
|
||||
let response = client
|
||||
@@ -815,6 +1126,7 @@ mod tests {
|
||||
&["read".to_string(), "write".to_string()],
|
||||
None,
|
||||
&HashMap::new(),
|
||||
None,
|
||||
);
|
||||
|
||||
assert!(url.starts_with("https://auth.example.com/authorize?"));
|
||||
@@ -834,6 +1146,7 @@ mod tests {
|
||||
&[],
|
||||
Some(&pkce),
|
||||
&HashMap::new(),
|
||||
None,
|
||||
);
|
||||
|
||||
assert!(url.contains(&format!("code_challenge={}", pkce.challenge)));
|
||||
@@ -853,6 +1166,7 @@ mod tests {
|
||||
&[],
|
||||
None,
|
||||
&extra,
|
||||
None,
|
||||
);
|
||||
|
||||
assert!(url.contains("owner=user"));
|
||||
@@ -880,6 +1194,7 @@ mod tests {
|
||||
&[],
|
||||
None,
|
||||
&HashMap::new(),
|
||||
None,
|
||||
);
|
||||
|
||||
// With no scopes, the URL must not contain a scope parameter at all.
|
||||
@@ -895,6 +1210,7 @@ mod tests {
|
||||
&[],
|
||||
None,
|
||||
&HashMap::new(),
|
||||
None,
|
||||
);
|
||||
|
||||
// Spaces and ampersands in client_id must be percent-encoded.
|
||||
@@ -1164,4 +1480,235 @@ mod tests {
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// --- New tests for well-known URI construction ---
|
||||
|
||||
#[test]
|
||||
fn test_build_well_known_uri_no_path() {
|
||||
let uri =
|
||||
build_well_known_uri("https://example.com", "oauth-authorization-server").unwrap();
|
||||
assert_eq!(
|
||||
uri,
|
||||
"https://example.com/.well-known/oauth-authorization-server"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_well_known_uri_with_path() {
|
||||
let uri =
|
||||
build_well_known_uri("https://example.com/path", "oauth-authorization-server").unwrap();
|
||||
assert_eq!(
|
||||
uri,
|
||||
"https://example.com/.well-known/oauth-authorization-server/path"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_well_known_uri_with_trailing_slash() {
|
||||
let uri =
|
||||
build_well_known_uri("https://example.com/path/", "oauth-protected-resource").unwrap();
|
||||
assert_eq!(
|
||||
uri,
|
||||
"https://example.com/.well-known/oauth-protected-resource/path"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_well_known_uri_root_trailing_slash() {
|
||||
let uri =
|
||||
build_well_known_uri("https://example.com/", "oauth-authorization-server").unwrap();
|
||||
assert_eq!(
|
||||
uri,
|
||||
"https://example.com/.well-known/oauth-authorization-server"
|
||||
);
|
||||
}
|
||||
|
||||
// --- New tests for canonical_resource_uri ---
|
||||
|
||||
#[test]
|
||||
fn test_canonical_resource_uri_strips_fragment() {
|
||||
assert_eq!(
|
||||
canonical_resource_uri("https://mcp.example.com/v1#section"),
|
||||
"https://mcp.example.com/v1"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_canonical_resource_uri_strips_trailing_slash() {
|
||||
assert_eq!(
|
||||
canonical_resource_uri("https://mcp.example.com/v1/"),
|
||||
"https://mcp.example.com/v1"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_canonical_resource_uri_no_changes_needed() {
|
||||
assert_eq!(
|
||||
canonical_resource_uri("https://mcp.example.com/v1"),
|
||||
"https://mcp.example.com/v1"
|
||||
);
|
||||
}
|
||||
|
||||
// --- New tests for SSRF protection ---
|
||||
|
||||
#[test]
|
||||
fn test_is_dangerous_ip_loopback_v4() {
|
||||
assert!(is_dangerous_ip("127.0.0.1".parse().unwrap()));
|
||||
assert!(is_dangerous_ip("127.0.0.2".parse().unwrap()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_dangerous_ip_private_v4() {
|
||||
assert!(is_dangerous_ip("10.0.0.1".parse().unwrap()));
|
||||
assert!(is_dangerous_ip("172.16.0.1".parse().unwrap()));
|
||||
assert!(is_dangerous_ip("192.168.1.1".parse().unwrap()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_dangerous_ip_link_local_v4() {
|
||||
assert!(is_dangerous_ip("169.254.169.254".parse().unwrap()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_dangerous_ip_cgnat() {
|
||||
assert!(is_dangerous_ip("100.64.0.1".parse().unwrap()));
|
||||
assert!(is_dangerous_ip("100.127.255.254".parse().unwrap()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_dangerous_ip_safe_v4() {
|
||||
assert!(!is_dangerous_ip("8.8.8.8".parse().unwrap()));
|
||||
assert!(!is_dangerous_ip("1.1.1.1".parse().unwrap()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_dangerous_ip_ipv4_mapped_v6_loopback() {
|
||||
// ::ffff:127.0.0.1 must be blocked
|
||||
let ip: IpAddr = "::ffff:127.0.0.1".parse().unwrap();
|
||||
assert!(is_dangerous_ip(ip));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_dangerous_ip_ipv4_mapped_v6_link_local() {
|
||||
// ::ffff:169.254.169.254 must be blocked
|
||||
let ip: IpAddr = "::ffff:169.254.169.254".parse().unwrap();
|
||||
assert!(is_dangerous_ip(ip));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_dangerous_ip_unspecified() {
|
||||
assert!(is_dangerous_ip("0.0.0.0".parse().unwrap()));
|
||||
assert!(is_dangerous_ip("::".parse().unwrap()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_dangerous_ip_v6_loopback() {
|
||||
assert!(is_dangerous_ip("::1".parse().unwrap()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_validate_url_safe_https() {
|
||||
assert!(validate_url_safe("https://example.com/path").await.is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_validate_url_safe_http_localhost_allowed() {
|
||||
// HTTP is only allowed for localhost dev scenarios
|
||||
assert!(validate_url_safe("http://localhost/path").await.is_ok());
|
||||
assert!(
|
||||
validate_url_safe("http://localhost:8080/path")
|
||||
.await
|
||||
.is_ok()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_validate_url_safe_http_non_localhost_rejected() {
|
||||
// HTTP to non-localhost hosts must be rejected (plaintext credential risk)
|
||||
assert!(validate_url_safe("http://example.com/path").await.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_validate_url_safe_bad_scheme() {
|
||||
assert!(validate_url_safe("ftp://example.com/path").await.is_err());
|
||||
assert!(validate_url_safe("file:///etc/passwd").await.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_validate_url_safe_private_ip() {
|
||||
// 127.0.0.1 over HTTP is allowed (localhost dev scenario)
|
||||
assert!(validate_url_safe("http://127.0.0.1/path").await.is_ok());
|
||||
// Private/link-local IPs over HTTPS are blocked (SSRF protection)
|
||||
assert!(validate_url_safe("https://10.0.0.1/path").await.is_err());
|
||||
assert!(
|
||||
validate_url_safe("https://169.254.169.254/latest/meta-data")
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
// Private IPs over HTTP (non-localhost) are blocked
|
||||
assert!(validate_url_safe("http://10.0.0.1/path").await.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_validate_url_safe_public_ip() {
|
||||
assert!(validate_url_safe("https://8.8.8.8/dns").await.is_ok());
|
||||
}
|
||||
|
||||
// --- New tests for parse_resource_metadata_url ---
|
||||
|
||||
#[test]
|
||||
fn test_parse_resource_metadata_url_bearer() {
|
||||
let header = r#"Bearer resource_metadata="https://res.example.com/.well-known/oauth-protected-resource""#;
|
||||
let url = parse_resource_metadata_url(header);
|
||||
assert_eq!(
|
||||
url.as_deref(),
|
||||
Some("https://res.example.com/.well-known/oauth-protected-resource")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_resource_metadata_url_with_other_params() {
|
||||
let header = r#"Bearer realm="example", resource_metadata="https://res.example.com/meta""#;
|
||||
let url = parse_resource_metadata_url(header);
|
||||
assert_eq!(url.as_deref(), Some("https://res.example.com/meta"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_resource_metadata_url_missing() {
|
||||
let header = r#"Bearer realm="example""#;
|
||||
let url = parse_resource_metadata_url(header);
|
||||
assert!(url.is_none());
|
||||
}
|
||||
|
||||
// --- New tests for resource parameter in authorization URL ---
|
||||
|
||||
#[test]
|
||||
fn test_build_authorization_url_with_resource() {
|
||||
let url = build_authorization_url(
|
||||
"https://auth.example.com/authorize",
|
||||
"client-123",
|
||||
"http://localhost:9876/callback",
|
||||
&[],
|
||||
None,
|
||||
&HashMap::new(),
|
||||
Some("https://mcp.example.com/v1"),
|
||||
);
|
||||
|
||||
assert!(url.contains("resource=https%3A%2F%2Fmcp.example.com%2Fv1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_authorization_url_without_resource() {
|
||||
let url = build_authorization_url(
|
||||
"https://auth.example.com/authorize",
|
||||
"client-123",
|
||||
"http://localhost:9876/callback",
|
||||
&[],
|
||||
None,
|
||||
&HashMap::new(),
|
||||
None,
|
||||
);
|
||||
|
||||
assert!(!url.contains("resource="));
|
||||
}
|
||||
}
|
||||
|
||||
+253
-187
@@ -1,11 +1,11 @@
|
||||
//! MCP client for connecting to MCP servers.
|
||||
//!
|
||||
//! Supports both local (unauthenticated) and hosted (OAuth-authenticated) servers.
|
||||
//! Uses the Streamable HTTP transport with session management.
|
||||
//! Uses pluggable transports (HTTP, stdio, Unix) via the `McpTransport` trait.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use tokio::sync::RwLock;
|
||||
@@ -14,27 +14,29 @@ use crate::context::JobContext;
|
||||
use crate::secrets::SecretsStore;
|
||||
use crate::tools::mcp::auth::refresh_access_token;
|
||||
use crate::tools::mcp::config::McpServerConfig;
|
||||
use crate::tools::mcp::http_transport::HttpMcpTransport;
|
||||
use crate::tools::mcp::protocol::{
|
||||
CallToolResult, InitializeResult, ListToolsResult, McpRequest, McpResponse, McpTool,
|
||||
};
|
||||
use crate::tools::mcp::session::McpSessionManager;
|
||||
use crate::tools::mcp::transport::McpTransport;
|
||||
use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput};
|
||||
|
||||
/// MCP client for communicating with MCP servers.
|
||||
///
|
||||
/// Supports two modes:
|
||||
/// - Simple: Just a URL, no auth or session management (for local/test servers)
|
||||
/// - Authenticated: Full OAuth support with session management (for hosted servers)
|
||||
/// Supports multiple transport types:
|
||||
/// - HTTP: For remote MCP servers (created via `new`, `new_with_name`, `new_authenticated`)
|
||||
/// - Stdio/Unix: Via `new_with_transport` with a custom `McpTransport` implementation
|
||||
pub struct McpClient {
|
||||
/// Server URL (for HTTP transport).
|
||||
/// Transport for sending requests.
|
||||
transport: Arc<dyn McpTransport>,
|
||||
|
||||
/// Server URL (kept for accessor compatibility).
|
||||
server_url: String,
|
||||
|
||||
/// Server name (for logging and session management).
|
||||
server_name: String,
|
||||
|
||||
/// HTTP client.
|
||||
http_client: reqwest::Client,
|
||||
|
||||
/// Request ID counter.
|
||||
next_id: AtomicU64,
|
||||
|
||||
@@ -52,6 +54,9 @@ pub struct McpClient {
|
||||
|
||||
/// Server configuration (for token secret name lookup).
|
||||
server_config: Option<McpServerConfig>,
|
||||
|
||||
/// Custom headers to include in every request.
|
||||
custom_headers: HashMap<String, String>,
|
||||
}
|
||||
|
||||
impl McpClient {
|
||||
@@ -59,22 +64,21 @@ impl McpClient {
|
||||
///
|
||||
/// Use this for local development servers or servers that don't require auth.
|
||||
pub fn new(server_url: impl Into<String>) -> Self {
|
||||
let url = server_url.into();
|
||||
let url: String = server_url.into();
|
||||
let name = extract_server_name(&url);
|
||||
let transport = Arc::new(HttpMcpTransport::new(url.clone(), name.clone()));
|
||||
|
||||
Self {
|
||||
transport,
|
||||
server_url: url,
|
||||
server_name: name,
|
||||
http_client: reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(30))
|
||||
.build()
|
||||
.expect("Failed to create HTTP client"),
|
||||
next_id: AtomicU64::new(1),
|
||||
tools_cache: RwLock::new(None),
|
||||
session_manager: None,
|
||||
secrets: None,
|
||||
user_id: "default".to_string(),
|
||||
server_config: None,
|
||||
custom_headers: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -82,19 +86,52 @@ impl McpClient {
|
||||
///
|
||||
/// Use this when you have a configured server name but no authentication.
|
||||
pub fn new_with_name(server_name: impl Into<String>, server_url: impl Into<String>) -> Self {
|
||||
let name: String = server_name.into();
|
||||
let url: String = server_url.into();
|
||||
let transport = Arc::new(HttpMcpTransport::new(url.clone(), name.clone()));
|
||||
|
||||
Self {
|
||||
server_url: server_url.into(),
|
||||
server_name: server_name.into(),
|
||||
http_client: reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(30))
|
||||
.build()
|
||||
.expect("Failed to create HTTP client"),
|
||||
transport,
|
||||
server_url: url,
|
||||
server_name: name,
|
||||
next_id: AtomicU64::new(1),
|
||||
tools_cache: RwLock::new(None),
|
||||
session_manager: None,
|
||||
secrets: None,
|
||||
user_id: "default".to_string(),
|
||||
server_config: None,
|
||||
custom_headers: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a new simple MCP client from an HTTP server configuration (no authentication).
|
||||
///
|
||||
/// Use this when you have an `McpServerConfig` with custom headers but no OAuth.
|
||||
/// The config must use HTTP transport (the default); for stdio/UDS use `new_with_transport`.
|
||||
pub fn new_with_config(config: McpServerConfig) -> Self {
|
||||
assert!(
|
||||
matches!(
|
||||
config.effective_transport(),
|
||||
crate::tools::mcp::config::EffectiveTransport::Http
|
||||
),
|
||||
"new_with_config only supports HTTP transport; use new_with_transport for stdio/UDS"
|
||||
);
|
||||
let transport = Arc::new(HttpMcpTransport::new(
|
||||
config.url.clone(),
|
||||
config.name.clone(),
|
||||
));
|
||||
|
||||
Self {
|
||||
transport,
|
||||
server_url: config.url.clone(),
|
||||
server_name: config.name.clone(),
|
||||
next_id: AtomicU64::new(1),
|
||||
tools_cache: RwLock::new(None),
|
||||
session_manager: None,
|
||||
secrets: None,
|
||||
user_id: "default".to_string(),
|
||||
custom_headers: config.headers.clone(),
|
||||
server_config: Some(config),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -107,19 +144,59 @@ impl McpClient {
|
||||
secrets: Arc<dyn SecretsStore + Send + Sync>,
|
||||
user_id: impl Into<String>,
|
||||
) -> Self {
|
||||
let transport = Arc::new(
|
||||
HttpMcpTransport::new(config.url.clone(), config.name.clone())
|
||||
.with_session_manager(session_manager.clone()),
|
||||
);
|
||||
|
||||
let custom_headers = config.headers.clone();
|
||||
|
||||
Self {
|
||||
transport,
|
||||
server_url: config.url.clone(),
|
||||
server_name: config.name.clone(),
|
||||
http_client: reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(30))
|
||||
.build()
|
||||
.expect("Failed to create HTTP client"),
|
||||
next_id: AtomicU64::new(1),
|
||||
tools_cache: RwLock::new(None),
|
||||
session_manager: Some(session_manager),
|
||||
secrets: Some(secrets),
|
||||
user_id: user_id.into(),
|
||||
server_config: Some(config),
|
||||
custom_headers,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a new MCP client with a custom transport.
|
||||
///
|
||||
/// Use this for stdio, UDS, or other non-HTTP transports.
|
||||
pub fn new_with_transport(
|
||||
server_name: impl Into<String>,
|
||||
transport: Arc<dyn McpTransport>,
|
||||
session_manager: Option<Arc<McpSessionManager>>,
|
||||
secrets: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
||||
user_id: impl Into<String>,
|
||||
server_config: Option<McpServerConfig>,
|
||||
) -> Self {
|
||||
let name: String = server_name.into();
|
||||
let url = server_config
|
||||
.as_ref()
|
||||
.map(|c| c.url.clone())
|
||||
.unwrap_or_default();
|
||||
let custom_headers = server_config
|
||||
.as_ref()
|
||||
.map(|c| c.headers.clone())
|
||||
.unwrap_or_default();
|
||||
|
||||
Self {
|
||||
transport,
|
||||
server_url: url,
|
||||
server_name: name,
|
||||
next_id: AtomicU64::new(1),
|
||||
tools_cache: RwLock::new(None),
|
||||
session_manager,
|
||||
secrets,
|
||||
user_id: user_id.into(),
|
||||
server_config,
|
||||
custom_headers,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -139,19 +216,13 @@ impl McpClient {
|
||||
}
|
||||
|
||||
/// Get the access token for this server (if authenticated).
|
||||
///
|
||||
/// Returns the stored token regardless of whether OAuth was pre-configured
|
||||
/// or obtained via Dynamic Client Registration.
|
||||
async fn get_access_token(&self) -> Result<Option<String>, ToolError> {
|
||||
let Some(ref secrets) = self.secrets else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let Some(ref config) = self.server_config else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
// Try to get stored token (from either pre-configured OAuth or DCR)
|
||||
match secrets
|
||||
.get_decrypted(&self.user_id, &config.token_secret_name())
|
||||
.await
|
||||
@@ -165,46 +236,41 @@ impl McpClient {
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the headers map for a request (auth, session-id, custom headers).
|
||||
async fn build_request_headers(&self) -> Result<HashMap<String, String>, ToolError> {
|
||||
let mut headers = self.custom_headers.clone();
|
||||
if let Some(token) = self.get_access_token().await? {
|
||||
headers.insert("Authorization".to_string(), format!("Bearer {}", token));
|
||||
}
|
||||
if let Some(ref session_manager) = self.session_manager
|
||||
&& let Some(session_id) = session_manager.get_session_id(&self.server_name).await
|
||||
{
|
||||
headers.insert("Mcp-Session-Id".to_string(), session_id);
|
||||
}
|
||||
Ok(headers)
|
||||
}
|
||||
|
||||
/// Send a request to the MCP server with auth and session headers.
|
||||
/// Automatically attempts token refresh on 401 errors.
|
||||
/// Automatically attempts token refresh on 401 errors (HTTP transports only).
|
||||
async fn send_request(&self, request: McpRequest) -> Result<McpResponse, ToolError> {
|
||||
// Try up to 2 times: first attempt, then retry after token refresh
|
||||
// For non-HTTP transports, just send directly without retry logic
|
||||
if !self.transport.supports_http_features() {
|
||||
let headers = self.build_request_headers().await?;
|
||||
return self.transport.send(&request, &headers).await;
|
||||
}
|
||||
|
||||
// HTTP transport: try up to 2 times (first attempt, then retry after token refresh)
|
||||
for attempt in 0..2 {
|
||||
// Request both JSON and SSE as per MCP spec
|
||||
let mut req_builder = self
|
||||
.http_client
|
||||
.post(&self.server_url)
|
||||
.header("Accept", "application/json, text/event-stream")
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&request);
|
||||
let headers = self.build_request_headers().await?;
|
||||
let result = self.transport.send(&request, &headers).await;
|
||||
|
||||
// Add Authorization header if we have a token
|
||||
if let Some(token) = self.get_access_token().await? {
|
||||
req_builder = req_builder.header("Authorization", format!("Bearer {}", token));
|
||||
}
|
||||
|
||||
// Add Mcp-Session-Id header if we have a session
|
||||
if let Some(ref session_manager) = self.session_manager
|
||||
&& let Some(session_id) = session_manager.get_session_id(&self.server_name).await
|
||||
{
|
||||
req_builder = req_builder.header("Mcp-Session-Id", session_id);
|
||||
}
|
||||
|
||||
let response = req_builder.send().await.map_err(|e| {
|
||||
let mut chain = format!("MCP request failed: {}", e);
|
||||
let mut source = std::error::Error::source(&e);
|
||||
while let Some(cause) = source {
|
||||
chain.push_str(&format!(" -> {}", cause));
|
||||
source = cause.source();
|
||||
}
|
||||
ToolError::ExternalService(chain)
|
||||
})?;
|
||||
|
||||
// Check for 401 Unauthorized - try to refresh token on first attempt
|
||||
if response.status() == reqwest::StatusCode::UNAUTHORIZED {
|
||||
if attempt == 0 {
|
||||
// Try to refresh the token
|
||||
if let Some(ref secrets) = self.secrets
|
||||
match result {
|
||||
Ok(response) => return Ok(response),
|
||||
Err(ToolError::ExternalService(ref msg))
|
||||
if msg.contains("401") || msg.contains("Unauthorized") =>
|
||||
{
|
||||
if attempt == 0
|
||||
&& let Some(ref secrets) = self.secrets
|
||||
&& let Some(ref config) = self.server_config
|
||||
{
|
||||
tracing::debug!(
|
||||
@@ -214,7 +280,6 @@ impl McpClient {
|
||||
match refresh_access_token(config, secrets, &self.user_id).await {
|
||||
Ok(_) => {
|
||||
tracing::info!("MCP token refreshed for '{}'", self.server_name);
|
||||
// Continue to next iteration to retry with new token
|
||||
continue;
|
||||
}
|
||||
Err(e) => {
|
||||
@@ -223,108 +288,30 @@ impl McpClient {
|
||||
self.server_name,
|
||||
e
|
||||
);
|
||||
// Fall through to return auth error
|
||||
}
|
||||
}
|
||||
}
|
||||
return Err(ToolError::ExternalService(format!(
|
||||
"MCP server '{}' requires authentication. Run: ironclaw mcp auth {}",
|
||||
self.server_name, self.server_name
|
||||
)));
|
||||
}
|
||||
return Err(ToolError::ExternalService(format!(
|
||||
"MCP server '{}' requires authentication. Run: ironclaw mcp auth {}",
|
||||
self.server_name, self.server_name
|
||||
)));
|
||||
Err(e) => return Err(e),
|
||||
}
|
||||
|
||||
// Success path - return the parsed response
|
||||
return self.parse_response(response).await;
|
||||
}
|
||||
|
||||
// Should not reach here, but just in case
|
||||
Err(ToolError::ExternalService(
|
||||
"MCP request failed after retry".to_string(),
|
||||
))
|
||||
}
|
||||
|
||||
/// Parse the HTTP response into an MCP response.
|
||||
async fn parse_response(&self, response: reqwest::Response) -> Result<McpResponse, ToolError> {
|
||||
// Extract session ID from response header
|
||||
if let Some(ref session_manager) = self.session_manager
|
||||
&& let Some(session_id) = response
|
||||
.headers()
|
||||
.get("Mcp-Session-Id")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
{
|
||||
session_manager
|
||||
.update_session_id(&self.server_name, Some(session_id.to_string()))
|
||||
.await;
|
||||
}
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
return Err(ToolError::ExternalService(format!(
|
||||
"MCP server returned status: {} - {}",
|
||||
status, body
|
||||
)));
|
||||
}
|
||||
|
||||
// Check content type to handle SSE vs JSON responses
|
||||
let content_type = response
|
||||
.headers()
|
||||
.get("content-type")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
|
||||
if content_type.contains("text/event-stream") {
|
||||
// SSE response - read chunks until we get a complete JSON message
|
||||
use futures::StreamExt;
|
||||
|
||||
let mut stream = response.bytes_stream();
|
||||
let mut buffer = String::new();
|
||||
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = chunk.map_err(|e| {
|
||||
ToolError::ExternalService(format!("Failed to read SSE chunk: {}", e))
|
||||
})?;
|
||||
|
||||
buffer.push_str(&String::from_utf8_lossy(&chunk));
|
||||
|
||||
// Look for complete SSE data lines
|
||||
for line in buffer.lines() {
|
||||
if let Some(json_str) = line.strip_prefix("data: ") {
|
||||
// Try to parse - if valid JSON, we're done
|
||||
if let Ok(response) = serde_json::from_str::<McpResponse>(json_str) {
|
||||
return Ok(response);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Err(ToolError::ExternalService(format!(
|
||||
"No valid data in SSE response: {}",
|
||||
buffer
|
||||
)))
|
||||
} else {
|
||||
// JSON response
|
||||
response.json().await.map_err(|e| {
|
||||
ToolError::ExternalService(format!("Failed to parse MCP response: {}", e))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Initialize the connection to the MCP server.
|
||||
///
|
||||
/// This should be called once per session to establish capabilities.
|
||||
pub async fn initialize(&self) -> Result<InitializeResult, ToolError> {
|
||||
// Check if already initialized
|
||||
if let Some(ref session_manager) = self.session_manager
|
||||
&& session_manager.is_initialized(&self.server_name).await
|
||||
{
|
||||
// Return cached/default capabilities
|
||||
return Ok(InitializeResult::default());
|
||||
}
|
||||
|
||||
// Ensure we have a session
|
||||
if let Some(ref session_manager) = self.session_manager {
|
||||
session_manager
|
||||
.get_or_create(&self.server_name, &self.server_url)
|
||||
@@ -352,14 +339,11 @@ impl McpClient {
|
||||
})
|
||||
})?;
|
||||
|
||||
// Mark session as initialized
|
||||
if let Some(ref session_manager) = self.session_manager {
|
||||
session_manager.mark_initialized(&self.server_name).await;
|
||||
}
|
||||
|
||||
// Send initialized notification
|
||||
let notification = McpRequest::initialized_notification();
|
||||
// Fire and forget - notifications don't have responses
|
||||
let _ = self.send_request(notification).await;
|
||||
|
||||
Ok(result)
|
||||
@@ -367,12 +351,9 @@ impl McpClient {
|
||||
|
||||
/// List available tools from the MCP server.
|
||||
pub async fn list_tools(&self) -> Result<Vec<McpTool>, ToolError> {
|
||||
// Check cache first
|
||||
if let Some(tools) = self.tools_cache.read().await.as_ref() {
|
||||
return Ok(tools.clone());
|
||||
}
|
||||
|
||||
// Ensure initialized for authenticated sessions
|
||||
if self.session_manager.is_some() {
|
||||
self.initialize().await?;
|
||||
}
|
||||
@@ -395,9 +376,7 @@ impl McpClient {
|
||||
.map_err(|e| ToolError::ExternalService(format!("Invalid tools list: {}", e)))
|
||||
})?;
|
||||
|
||||
// Cache the tools
|
||||
*self.tools_cache.write().await = Some(result.tools.clone());
|
||||
|
||||
Ok(result.tools)
|
||||
}
|
||||
|
||||
@@ -407,7 +386,6 @@ impl McpClient {
|
||||
name: &str,
|
||||
arguments: serde_json::Value,
|
||||
) -> Result<CallToolResult, ToolError> {
|
||||
// Ensure initialized for authenticated sessions
|
||||
if self.session_manager.is_some() {
|
||||
self.initialize().await?;
|
||||
}
|
||||
@@ -440,7 +418,6 @@ impl McpClient {
|
||||
pub async fn create_tools(&self) -> Result<Vec<Arc<dyn Tool>>, ToolError> {
|
||||
let mcp_tools = self.list_tools().await?;
|
||||
let client = Arc::new(self.clone());
|
||||
|
||||
Ok(mcp_tools
|
||||
.into_iter()
|
||||
.map(|t| {
|
||||
@@ -465,15 +442,16 @@ impl McpClient {
|
||||
impl Clone for McpClient {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
transport: self.transport.clone(),
|
||||
server_url: self.server_url.clone(),
|
||||
server_name: self.server_name.clone(),
|
||||
http_client: self.http_client.clone(),
|
||||
next_id: AtomicU64::new(self.next_id.load(Ordering::SeqCst)),
|
||||
tools_cache: RwLock::new(None),
|
||||
session_manager: self.session_manager.clone(),
|
||||
secrets: self.secrets.clone(),
|
||||
user_id: self.user_id.clone(),
|
||||
server_config: self.server_config.clone(),
|
||||
custom_headers: self.custom_headers.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -490,7 +468,6 @@ fn extract_server_name(url: &str) -> String {
|
||||
/// Wrapper that implements Tool for an MCP tool.
|
||||
struct McpToolWrapper {
|
||||
tool: McpTool,
|
||||
/// Prefixed name (server_name_tool_name) for unique identification.
|
||||
prefixed_name: String,
|
||||
client: Arc<McpClient>,
|
||||
}
|
||||
@@ -500,11 +477,9 @@ impl Tool for McpToolWrapper {
|
||||
fn name(&self) -> &str {
|
||||
&self.prefixed_name
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
&self.tool.description
|
||||
}
|
||||
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
self.tool.input_schema.clone()
|
||||
}
|
||||
@@ -515,31 +490,24 @@ impl Tool for McpToolWrapper {
|
||||
_ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
// Use the original tool name (without prefix) for the actual call
|
||||
let result = self.client.call_tool(&self.tool.name, params).await?;
|
||||
|
||||
// Convert content blocks to a single result
|
||||
let content: String = result
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|block| block.as_text())
|
||||
.filter_map(|b| b.as_text())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
if result.is_error {
|
||||
return Err(ToolError::ExecutionFailed(content));
|
||||
}
|
||||
|
||||
Ok(ToolOutput::text(content, start.elapsed()))
|
||||
}
|
||||
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
true // MCP tools are external, always sanitize
|
||||
true
|
||||
}
|
||||
|
||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
||||
// Delegate to the MCP protocol type's own requires_approval() bool method
|
||||
if self.tool.requires_approval() {
|
||||
ApprovalRequirement::UnlessAutoApproved
|
||||
} else {
|
||||
@@ -548,6 +516,9 @@ impl Tool for McpToolWrapper {
|
||||
}
|
||||
}
|
||||
|
||||
/// Sanitize an HTTP error response body for safe display.
|
||||
///
|
||||
/// Detects full HTML error pages (containing `<html` or `<!DOCTYPE`) and
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -556,7 +527,7 @@ mod tests {
|
||||
fn test_mcp_request_list_tools() {
|
||||
let req = McpRequest::list_tools(1);
|
||||
assert_eq!(req.method, "tools/list");
|
||||
assert_eq!(req.id, 1);
|
||||
assert_eq!(req.id, Some(1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -660,42 +631,61 @@ mod tests {
|
||||
#[test]
|
||||
fn test_clone_preserves_fields() {
|
||||
let client = McpClient::new_with_name("cloned-server", "http://localhost:5555");
|
||||
// Bump the request ID a few times
|
||||
client.next_request_id();
|
||||
client.next_request_id();
|
||||
|
||||
let cloned = client.clone();
|
||||
assert_eq!(cloned.server_url(), "http://localhost:5555");
|
||||
assert_eq!(cloned.server_name(), "cloned-server");
|
||||
assert_eq!(cloned.user_id, "default");
|
||||
// The atomic counter value is copied
|
||||
assert_eq!(cloned.next_id.load(Ordering::SeqCst), 3);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_clone_resets_tools_cache() {
|
||||
let client = McpClient::new("http://localhost:5555");
|
||||
// The clone implementation resets tools_cache to None
|
||||
let cloned = client.clone();
|
||||
let cache = cloned.tools_cache.read().await;
|
||||
assert!(cache.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_new_with_config_carries_custom_headers() {
|
||||
let mut headers = HashMap::new();
|
||||
headers.insert("X-API-Key".to_string(), "secret".to_string());
|
||||
headers.insert("X-Custom".to_string(), "value".to_string());
|
||||
|
||||
let config = McpServerConfig::new("test", "http://localhost:8080").with_headers(headers);
|
||||
let client = McpClient::new_with_config(config.clone());
|
||||
|
||||
assert_eq!(client.server_name(), "test");
|
||||
assert_eq!(client.server_url(), "http://localhost:8080");
|
||||
assert_eq!(client.custom_headers.len(), 2);
|
||||
assert_eq!(client.custom_headers.get("X-API-Key").unwrap(), "secret");
|
||||
assert!(client.server_config.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_new_with_config_no_headers() {
|
||||
let config = McpServerConfig::new("bare", "http://localhost:9090");
|
||||
let client = McpClient::new_with_config(config);
|
||||
|
||||
assert_eq!(client.server_name(), "bare");
|
||||
assert!(client.custom_headers.is_empty());
|
||||
assert!(client.secrets.is_none());
|
||||
assert!(client.session_manager.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_next_request_id_monotonically_increasing() {
|
||||
let client = McpClient::new("http://localhost:1234");
|
||||
let id1 = client.next_request_id();
|
||||
let id2 = client.next_request_id();
|
||||
let id3 = client.next_request_id();
|
||||
assert_eq!(id1, 1);
|
||||
assert_eq!(id2, 2);
|
||||
assert_eq!(id3, 3);
|
||||
assert_eq!(client.next_request_id(), 1);
|
||||
assert_eq!(client.next_request_id(), 2);
|
||||
assert_eq!(client.next_request_id(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mcp_tool_requires_approval_destructive() {
|
||||
use crate::tools::mcp::protocol::{McpTool, McpToolAnnotations};
|
||||
|
||||
let tool = McpTool {
|
||||
name: "delete_all".to_string(),
|
||||
description: "Deletes everything".to_string(),
|
||||
@@ -713,7 +703,6 @@ mod tests {
|
||||
#[test]
|
||||
fn test_mcp_tool_no_approval_when_not_destructive() {
|
||||
use crate::tools::mcp::protocol::{McpTool, McpToolAnnotations};
|
||||
|
||||
let tool = McpTool {
|
||||
name: "read_data".to_string(),
|
||||
description: "Reads data".to_string(),
|
||||
@@ -731,7 +720,6 @@ mod tests {
|
||||
#[test]
|
||||
fn test_mcp_tool_no_approval_when_no_annotations() {
|
||||
use crate::tools::mcp::protocol::McpTool;
|
||||
|
||||
let tool = McpTool {
|
||||
name: "simple_tool".to_string(),
|
||||
description: "A simple tool".to_string(),
|
||||
@@ -740,4 +728,82 @@ mod tests {
|
||||
};
|
||||
assert!(!tool.requires_approval());
|
||||
}
|
||||
|
||||
/// Mock transport for testing transport abstraction behavior.
|
||||
struct MockTransport {
|
||||
supports_http: bool,
|
||||
responses: std::sync::Mutex<Vec<McpResponse>>,
|
||||
recorded_headers: std::sync::Mutex<Vec<HashMap<String, String>>>,
|
||||
}
|
||||
|
||||
impl MockTransport {
|
||||
fn new(supports_http: bool, responses: Vec<McpResponse>) -> Self {
|
||||
Self {
|
||||
supports_http,
|
||||
responses: std::sync::Mutex::new(responses),
|
||||
recorded_headers: std::sync::Mutex::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
fn recorded_headers(&self) -> Vec<HashMap<String, String>> {
|
||||
self.recorded_headers.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl McpTransport for MockTransport {
|
||||
async fn send(
|
||||
&self,
|
||||
_request: &McpRequest,
|
||||
headers: &HashMap<String, String>,
|
||||
) -> Result<McpResponse, ToolError> {
|
||||
self.recorded_headers.lock().unwrap().push(headers.clone());
|
||||
let mut responses = self.responses.lock().unwrap();
|
||||
if responses.is_empty() {
|
||||
return Err(ToolError::ExternalService(
|
||||
"No more mock responses".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(responses.remove(0))
|
||||
}
|
||||
async fn shutdown(&self) -> Result<(), ToolError> {
|
||||
Ok(())
|
||||
}
|
||||
fn supports_http_features(&self) -> bool {
|
||||
self.supports_http
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_non_http_transport_skips_401_retry() {
|
||||
let response = McpResponse {
|
||||
jsonrpc: "2.0".to_string(),
|
||||
id: Some(1),
|
||||
result: Some(serde_json::json!({"tools": []})),
|
||||
error: None,
|
||||
};
|
||||
let transport = Arc::new(MockTransport::new(false, vec![response]));
|
||||
let client = McpClient::new_with_transport(
|
||||
"test-stdio",
|
||||
transport.clone(),
|
||||
None,
|
||||
None,
|
||||
"default",
|
||||
None,
|
||||
);
|
||||
let result = client.list_tools().await;
|
||||
assert!(result.is_ok());
|
||||
assert_eq!(result.unwrap().len(), 0);
|
||||
let headers = transport.recorded_headers();
|
||||
assert_eq!(headers.len(), 1);
|
||||
assert!(!headers[0].contains_key("Authorization"));
|
||||
assert!(!headers[0].contains_key("Mcp-Session-Id"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_transport_supports_http_features_accessor() {
|
||||
let http_transport = HttpMcpTransport::new("http://localhost:8080", "test");
|
||||
assert!(http_transport.supports_http_features());
|
||||
let mock_non_http = MockTransport::new(false, vec![]);
|
||||
assert!(!mock_non_http.supports_http_features());
|
||||
}
|
||||
}
|
||||
|
||||
+382
-12
@@ -12,6 +12,24 @@ use tokio::fs;
|
||||
use crate::bootstrap::ironclaw_base_dir;
|
||||
use crate::tools::tool::ToolError;
|
||||
|
||||
/// Transport configuration for an MCP server.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "transport", rename_all = "lowercase")]
|
||||
pub enum McpTransportConfig {
|
||||
/// HTTP/HTTPS transport (uses the `url` field on McpServerConfig).
|
||||
Http,
|
||||
/// Stdio transport — spawns a child process.
|
||||
Stdio {
|
||||
command: String,
|
||||
#[serde(default)]
|
||||
args: Vec<String>,
|
||||
#[serde(default)]
|
||||
env: HashMap<String, String>,
|
||||
},
|
||||
/// Unix domain socket transport.
|
||||
Unix { socket_path: String },
|
||||
}
|
||||
|
||||
/// Configuration for connecting to a remote MCP server.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct McpServerConfig {
|
||||
@@ -21,6 +39,14 @@ pub struct McpServerConfig {
|
||||
/// Server URL (must be HTTPS for remote servers).
|
||||
pub url: String,
|
||||
|
||||
/// Transport configuration. If `None`, defaults to Http using `url`.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub transport: Option<McpTransportConfig>,
|
||||
|
||||
/// Custom headers to include in every HTTP request.
|
||||
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
|
||||
pub headers: HashMap<String, String>,
|
||||
|
||||
/// OAuth configuration (if server requires authentication).
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub oauth: Option<OAuthConfig>,
|
||||
@@ -44,6 +70,45 @@ impl McpServerConfig {
|
||||
Self {
|
||||
name: name.into(),
|
||||
url: url.into(),
|
||||
transport: None,
|
||||
headers: HashMap::new(),
|
||||
oauth: None,
|
||||
enabled: true,
|
||||
description: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a new stdio transport MCP server configuration.
|
||||
pub fn new_stdio(
|
||||
name: impl Into<String>,
|
||||
command: impl Into<String>,
|
||||
args: Vec<String>,
|
||||
env: HashMap<String, String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
name: name.into(),
|
||||
url: String::new(),
|
||||
transport: Some(McpTransportConfig::Stdio {
|
||||
command: command.into(),
|
||||
args,
|
||||
env,
|
||||
}),
|
||||
headers: HashMap::new(),
|
||||
oauth: None,
|
||||
enabled: true,
|
||||
description: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a new Unix socket transport MCP server configuration.
|
||||
pub fn new_unix(name: impl Into<String>, socket_path: impl Into<String>) -> Self {
|
||||
Self {
|
||||
name: name.into(),
|
||||
url: String::new(),
|
||||
transport: Some(McpTransportConfig::Unix {
|
||||
socket_path: socket_path.into(),
|
||||
}),
|
||||
headers: HashMap::new(),
|
||||
oauth: None,
|
||||
enabled: true,
|
||||
description: None,
|
||||
@@ -62,6 +127,25 @@ impl McpServerConfig {
|
||||
self
|
||||
}
|
||||
|
||||
/// Set custom headers.
|
||||
pub fn with_headers(mut self, headers: HashMap<String, String>) -> Self {
|
||||
self.headers = headers;
|
||||
self
|
||||
}
|
||||
|
||||
/// Get the effective transport type.
|
||||
pub fn effective_transport(&self) -> EffectiveTransport<'_> {
|
||||
match &self.transport {
|
||||
Some(McpTransportConfig::Http) | None => EffectiveTransport::Http,
|
||||
Some(McpTransportConfig::Stdio { command, args, env }) => {
|
||||
EffectiveTransport::Stdio { command, args, env }
|
||||
}
|
||||
Some(McpTransportConfig::Unix { socket_path }) => {
|
||||
EffectiveTransport::Unix { socket_path }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate the server configuration.
|
||||
pub fn validate(&self) -> Result<(), ConfigError> {
|
||||
if self.name.is_empty() {
|
||||
@@ -70,19 +154,38 @@ impl McpServerConfig {
|
||||
});
|
||||
}
|
||||
|
||||
if self.url.is_empty() {
|
||||
return Err(ConfigError::InvalidConfig {
|
||||
reason: "Server URL cannot be empty".to_string(),
|
||||
});
|
||||
}
|
||||
match self.effective_transport() {
|
||||
EffectiveTransport::Http => {
|
||||
if self.url.is_empty() {
|
||||
return Err(ConfigError::InvalidConfig {
|
||||
reason: "Server URL cannot be empty".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
// Remote servers must use HTTPS (localhost is allowed for development)
|
||||
let url_lower = self.url.to_lowercase();
|
||||
let is_localhost = url_lower.contains("localhost") || url_lower.contains("127.0.0.1");
|
||||
if !is_localhost && !url_lower.starts_with("https://") {
|
||||
return Err(ConfigError::InvalidConfig {
|
||||
reason: "Remote MCP servers must use HTTPS".to_string(),
|
||||
});
|
||||
// Remote servers must use HTTPS (localhost is allowed for development)
|
||||
let url_lower = self.url.to_lowercase();
|
||||
let is_localhost =
|
||||
url_lower.contains("localhost") || url_lower.contains("127.0.0.1");
|
||||
if !is_localhost && !url_lower.starts_with("https://") {
|
||||
return Err(ConfigError::InvalidConfig {
|
||||
reason: "Remote MCP servers must use HTTPS".to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
EffectiveTransport::Stdio { command, .. } => {
|
||||
if command.is_empty() {
|
||||
return Err(ConfigError::InvalidConfig {
|
||||
reason: "Stdio transport command cannot be empty".to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
EffectiveTransport::Unix { socket_path } => {
|
||||
if socket_path.is_empty() {
|
||||
return Err(ConfigError::InvalidConfig {
|
||||
reason: "Unix socket path cannot be empty".to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -92,7 +195,14 @@ impl McpServerConfig {
|
||||
///
|
||||
/// Returns true if OAuth is pre-configured OR if this is a remote HTTPS server
|
||||
/// (which likely supports Dynamic Client Registration even without pre-configured OAuth).
|
||||
///
|
||||
/// Non-HTTP transports (stdio, unix) never require auth.
|
||||
pub fn requires_auth(&self) -> bool {
|
||||
// Non-HTTP transports don't use HTTP auth
|
||||
if !matches!(self.effective_transport(), EffectiveTransport::Http) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if self.oauth.is_some() {
|
||||
return true;
|
||||
}
|
||||
@@ -426,6 +536,20 @@ fn is_localhost_url(url: &str) -> bool {
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolved transport type (borrows from config).
|
||||
#[derive(Debug)]
|
||||
pub enum EffectiveTransport<'a> {
|
||||
Http,
|
||||
Stdio {
|
||||
command: &'a str,
|
||||
args: &'a [String],
|
||||
env: &'a HashMap<String, String>,
|
||||
},
|
||||
Unix {
|
||||
socket_path: &'a str,
|
||||
},
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -593,4 +717,250 @@ mod tests {
|
||||
let config = McpServerConfig::new("bad", "http://mcp.example.com");
|
||||
assert!(!config.requires_auth());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_stdio_config_creation() {
|
||||
let env = HashMap::from([("PATH".to_string(), "/usr/bin".to_string())]);
|
||||
let config = McpServerConfig::new_stdio(
|
||||
"my-server",
|
||||
"npx",
|
||||
vec!["-y".to_string(), "@modelcontextprotocol/server".to_string()],
|
||||
env.clone(),
|
||||
);
|
||||
|
||||
assert_eq!(config.name, "my-server");
|
||||
assert!(config.url.is_empty());
|
||||
assert!(config.enabled);
|
||||
assert!(config.oauth.is_none());
|
||||
assert!(config.headers.is_empty());
|
||||
|
||||
match &config.transport {
|
||||
Some(McpTransportConfig::Stdio {
|
||||
command,
|
||||
args,
|
||||
env: e,
|
||||
}) => {
|
||||
assert_eq!(command, "npx");
|
||||
assert_eq!(
|
||||
args,
|
||||
&["-y".to_string(), "@modelcontextprotocol/server".to_string()]
|
||||
);
|
||||
assert_eq!(e, &env);
|
||||
}
|
||||
other => panic!("Expected Stdio transport, got {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_unix_config_creation() {
|
||||
let config = McpServerConfig::new_unix("local-server", "/tmp/mcp.sock");
|
||||
|
||||
assert_eq!(config.name, "local-server");
|
||||
assert!(config.url.is_empty());
|
||||
assert!(config.enabled);
|
||||
|
||||
match &config.transport {
|
||||
Some(McpTransportConfig::Unix { socket_path }) => {
|
||||
assert_eq!(socket_path, "/tmp/mcp.sock");
|
||||
}
|
||||
other => panic!("Expected Unix transport, got {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_stdio_validation() {
|
||||
// Valid stdio config
|
||||
let config = McpServerConfig::new_stdio("server", "npx", vec![], HashMap::new());
|
||||
assert!(config.validate().is_ok());
|
||||
|
||||
// Invalid: empty command
|
||||
let config = McpServerConfig::new_stdio("server", "", vec![], HashMap::new());
|
||||
assert!(config.validate().is_err());
|
||||
let err = config.validate().unwrap_err().to_string();
|
||||
assert!(
|
||||
err.contains("command"),
|
||||
"Error should mention command: {}",
|
||||
err
|
||||
);
|
||||
|
||||
// Invalid: empty name
|
||||
let config = McpServerConfig::new_stdio("", "npx", vec![], HashMap::new());
|
||||
assert!(config.validate().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_unix_validation() {
|
||||
// Valid unix config
|
||||
let config = McpServerConfig::new_unix("server", "/tmp/mcp.sock");
|
||||
assert!(config.validate().is_ok());
|
||||
|
||||
// Invalid: empty socket path
|
||||
let config = McpServerConfig::new_unix("server", "");
|
||||
assert!(config.validate().is_err());
|
||||
let err = config.validate().unwrap_err().to_string();
|
||||
assert!(
|
||||
err.contains("socket"),
|
||||
"Error should mention socket: {}",
|
||||
err
|
||||
);
|
||||
|
||||
// Invalid: empty name
|
||||
let config = McpServerConfig::new_unix("", "/tmp/mcp.sock");
|
||||
assert!(config.validate().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_requires_auth_stdio_never() {
|
||||
// Stdio transport should never require auth, even with OAuth configured
|
||||
let mut config = McpServerConfig::new_stdio("server", "npx", vec![], HashMap::new());
|
||||
assert!(!config.requires_auth());
|
||||
|
||||
// Even if OAuth is set, stdio doesn't use HTTP auth
|
||||
config.oauth = Some(OAuthConfig::new("client-123"));
|
||||
assert!(!config.requires_auth());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_requires_auth_unix_never() {
|
||||
// Unix transport should never require auth
|
||||
let mut config = McpServerConfig::new_unix("server", "/tmp/mcp.sock");
|
||||
assert!(!config.requires_auth());
|
||||
|
||||
config.oauth = Some(OAuthConfig::new("client-123"));
|
||||
assert!(!config.requires_auth());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_custom_headers() {
|
||||
let headers = HashMap::from([
|
||||
("X-Api-Key".to_string(), "secret".to_string()),
|
||||
("Authorization".to_string(), "Bearer token".to_string()),
|
||||
]);
|
||||
let config =
|
||||
McpServerConfig::new("server", "https://mcp.example.com").with_headers(headers.clone());
|
||||
|
||||
assert_eq!(config.headers, headers);
|
||||
assert_eq!(config.headers.get("X-Api-Key").unwrap(), "secret");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transport_config_serde_http() {
|
||||
let transport = McpTransportConfig::Http;
|
||||
let json = serde_json::to_string(&transport).unwrap();
|
||||
assert!(json.contains("\"transport\":\"http\""));
|
||||
|
||||
let parsed: McpTransportConfig = serde_json::from_str(&json).unwrap();
|
||||
assert!(matches!(parsed, McpTransportConfig::Http));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transport_config_serde_stdio() {
|
||||
let transport = McpTransportConfig::Stdio {
|
||||
command: "npx".to_string(),
|
||||
args: vec!["-y".to_string(), "server".to_string()],
|
||||
env: HashMap::from([("KEY".to_string(), "val".to_string())]),
|
||||
};
|
||||
let json = serde_json::to_string(&transport).unwrap();
|
||||
assert!(json.contains("\"transport\":\"stdio\""));
|
||||
assert!(json.contains("\"command\":\"npx\""));
|
||||
|
||||
let parsed: McpTransportConfig = serde_json::from_str(&json).unwrap();
|
||||
match parsed {
|
||||
McpTransportConfig::Stdio { command, args, env } => {
|
||||
assert_eq!(command, "npx");
|
||||
assert_eq!(args, vec!["-y".to_string(), "server".to_string()]);
|
||||
assert_eq!(env.get("KEY").unwrap(), "val");
|
||||
}
|
||||
other => panic!("Expected Stdio, got {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transport_config_serde_unix() {
|
||||
let transport = McpTransportConfig::Unix {
|
||||
socket_path: "/tmp/mcp.sock".to_string(),
|
||||
};
|
||||
let json = serde_json::to_string(&transport).unwrap();
|
||||
assert!(json.contains("\"transport\":\"unix\""));
|
||||
assert!(json.contains("\"socket_path\":\"/tmp/mcp.sock\""));
|
||||
|
||||
let parsed: McpTransportConfig = serde_json::from_str(&json).unwrap();
|
||||
match parsed {
|
||||
McpTransportConfig::Unix { socket_path } => {
|
||||
assert_eq!(socket_path, "/tmp/mcp.sock");
|
||||
}
|
||||
other => panic!("Expected Unix, got {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_backward_compat_no_transport_field() {
|
||||
// Existing configs without transport field should still deserialize
|
||||
let json = r#"{
|
||||
"name": "notion",
|
||||
"url": "https://mcp.notion.com",
|
||||
"enabled": true
|
||||
}"#;
|
||||
let config: McpServerConfig = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(config.name, "notion");
|
||||
assert_eq!(config.url, "https://mcp.notion.com");
|
||||
assert!(config.transport.is_none());
|
||||
assert!(config.headers.is_empty());
|
||||
assert!(matches!(
|
||||
config.effective_transport(),
|
||||
EffectiveTransport::Http
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_config_roundtrip_with_transport() {
|
||||
// Test full roundtrip with stdio transport
|
||||
let config = McpServerConfig::new_stdio(
|
||||
"test-server",
|
||||
"node",
|
||||
vec!["server.js".to_string()],
|
||||
HashMap::from([("NODE_ENV".to_string(), "production".to_string())]),
|
||||
)
|
||||
.with_description("A test server");
|
||||
|
||||
let json = serde_json::to_string_pretty(&config).unwrap();
|
||||
let parsed: McpServerConfig = serde_json::from_str(&json).unwrap();
|
||||
|
||||
assert_eq!(parsed.name, "test-server");
|
||||
assert!(parsed.url.is_empty());
|
||||
assert_eq!(parsed.description.as_deref(), Some("A test server"));
|
||||
|
||||
match &parsed.transport {
|
||||
Some(McpTransportConfig::Stdio { command, args, env }) => {
|
||||
assert_eq!(command, "node");
|
||||
assert_eq!(args, &["server.js".to_string()]);
|
||||
assert_eq!(env.get("NODE_ENV").unwrap(), "production");
|
||||
}
|
||||
other => panic!("Expected Stdio transport, got {:?}", other),
|
||||
}
|
||||
|
||||
// Test full roundtrip with unix transport
|
||||
let config = McpServerConfig::new_unix("unix-server", "/var/run/mcp.sock");
|
||||
let json = serde_json::to_string_pretty(&config).unwrap();
|
||||
let parsed: McpServerConfig = serde_json::from_str(&json).unwrap();
|
||||
|
||||
assert_eq!(parsed.name, "unix-server");
|
||||
match &parsed.transport {
|
||||
Some(McpTransportConfig::Unix { socket_path }) => {
|
||||
assert_eq!(socket_path, "/var/run/mcp.sock");
|
||||
}
|
||||
other => panic!("Expected Unix transport, got {:?}", other),
|
||||
}
|
||||
|
||||
// Test roundtrip with HTTP + headers
|
||||
let headers = HashMap::from([("X-Custom".to_string(), "value".to_string())]);
|
||||
let config =
|
||||
McpServerConfig::new("http-server", "https://mcp.example.com").with_headers(headers);
|
||||
let json = serde_json::to_string_pretty(&config).unwrap();
|
||||
let parsed: McpServerConfig = serde_json::from_str(&json).unwrap();
|
||||
|
||||
assert_eq!(parsed.name, "http-server");
|
||||
assert!(parsed.transport.is_none());
|
||||
assert_eq!(parsed.headers.get("X-Custom").unwrap(), "value");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,386 @@
|
||||
//! HTTP transport for MCP servers.
|
||||
//!
|
||||
//! Implements the Streamable HTTP transport, communicating with MCP servers
|
||||
//! over HTTP POST with JSON and SSE response support.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::tools::mcp::protocol::{McpRequest, McpResponse};
|
||||
use crate::tools::mcp::session::McpSessionManager;
|
||||
use crate::tools::mcp::transport::McpTransport;
|
||||
use crate::tools::tool::ToolError;
|
||||
|
||||
/// MCP transport that communicates with a server over HTTP.
|
||||
///
|
||||
/// Sends JSON-RPC requests as HTTP POST with `Content-Type: application/json`
|
||||
/// and accepts either JSON or SSE (`text/event-stream`) responses. Optionally
|
||||
/// manages session IDs via [`McpSessionManager`] and supports custom headers.
|
||||
pub struct HttpMcpTransport {
|
||||
server_url: String,
|
||||
server_name: String,
|
||||
http_client: reqwest::Client,
|
||||
session_manager: Option<Arc<McpSessionManager>>,
|
||||
custom_headers: HashMap<String, String>,
|
||||
}
|
||||
|
||||
impl HttpMcpTransport {
|
||||
/// Create a new HTTP transport for the given server URL.
|
||||
pub fn new(server_url: impl Into<String>, server_name: impl Into<String>) -> Self {
|
||||
Self {
|
||||
server_url: server_url.into(),
|
||||
server_name: server_name.into(),
|
||||
// reqwest::Client::builder().build() only fails if the TLS backend
|
||||
// cannot initialize, which does not happen with the default rustls
|
||||
// feature set. Panic is acceptable here (same as reqwest's own
|
||||
// `Client::new()`).
|
||||
http_client: reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(30))
|
||||
.build()
|
||||
.expect("Failed to create HTTP client"),
|
||||
session_manager: None,
|
||||
custom_headers: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Attach a session manager for Mcp-Session-Id tracking.
|
||||
pub fn with_session_manager(mut self, session_manager: Arc<McpSessionManager>) -> Self {
|
||||
self.session_manager = Some(session_manager);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set custom headers that will be sent with every request.
|
||||
#[cfg(test)]
|
||||
pub fn with_custom_headers(mut self, headers: HashMap<String, String>) -> Self {
|
||||
self.custom_headers = headers;
|
||||
self
|
||||
}
|
||||
|
||||
/// Get the server URL.
|
||||
#[cfg(test)]
|
||||
pub(crate) fn server_url(&self) -> &str {
|
||||
&self.server_url
|
||||
}
|
||||
|
||||
/// Get the session manager, if one is configured.
|
||||
#[cfg(test)]
|
||||
pub(crate) fn session_manager(&self) -> Option<&Arc<McpSessionManager>> {
|
||||
self.session_manager.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl McpTransport for HttpMcpTransport {
|
||||
async fn send(
|
||||
&self,
|
||||
request: &McpRequest,
|
||||
headers: &HashMap<String, String>,
|
||||
) -> Result<McpResponse, ToolError> {
|
||||
// Build the HTTP request.
|
||||
let mut req_builder = self
|
||||
.http_client
|
||||
.post(&self.server_url)
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "application/json, text/event-stream")
|
||||
.json(request);
|
||||
|
||||
// Apply custom headers configured on the transport.
|
||||
for (key, value) in &self.custom_headers {
|
||||
req_builder = req_builder.header(key.as_str(), value.as_str());
|
||||
}
|
||||
|
||||
// Apply per-request headers (e.g. Authorization, Mcp-Session-Id).
|
||||
for (key, value) in headers {
|
||||
req_builder = req_builder.header(key.as_str(), value.as_str());
|
||||
}
|
||||
|
||||
// Send the request.
|
||||
let response = req_builder.send().await.map_err(|e| {
|
||||
let mut chain = format!("[{}] MCP HTTP request failed: {}", self.server_name, e);
|
||||
let mut source = std::error::Error::source(&e);
|
||||
while let Some(cause) = source {
|
||||
chain.push_str(&format!(" -> {}", cause));
|
||||
source = cause.source();
|
||||
}
|
||||
ToolError::ExternalService(chain)
|
||||
})?;
|
||||
|
||||
// Extract session ID from response headers before consuming the body.
|
||||
if let Some(ref session_manager) = self.session_manager
|
||||
&& let Some(session_id) = response
|
||||
.headers()
|
||||
.get("Mcp-Session-Id")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
{
|
||||
session_manager
|
||||
.update_session_id(&self.server_name, Some(session_id.to_string()))
|
||||
.await;
|
||||
}
|
||||
|
||||
// Handle error status codes.
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
let sanitized = sanitize_error_body(&body);
|
||||
return Err(ToolError::ExternalService(format!(
|
||||
"[{}] MCP server returned status: {} - {}",
|
||||
self.server_name, status, sanitized
|
||||
)));
|
||||
}
|
||||
|
||||
// Determine response format from Content-Type.
|
||||
let content_type = response
|
||||
.headers()
|
||||
.get("content-type")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
|
||||
if content_type.contains("text/event-stream") {
|
||||
self.parse_sse_response(response).await
|
||||
} else {
|
||||
response.json().await.map_err(|e| {
|
||||
ToolError::ExternalService(format!(
|
||||
"[{}] Failed to parse MCP response: {}",
|
||||
self.server_name, e
|
||||
))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
async fn shutdown(&self) -> Result<(), ToolError> {
|
||||
// HTTP transport is stateless; nothing to shut down.
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn supports_http_features(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
impl HttpMcpTransport {
|
||||
/// Parse a Server-Sent Events response, returning the first valid JSON-RPC
|
||||
/// `data:` line as an [`McpResponse`].
|
||||
async fn parse_sse_response(
|
||||
&self,
|
||||
response: reqwest::Response,
|
||||
) -> Result<McpResponse, ToolError> {
|
||||
use futures::StreamExt;
|
||||
|
||||
const MAX_SSE_BUFFER: usize = 10 * 1024 * 1024; // 10 MB
|
||||
|
||||
let mut stream = response.bytes_stream();
|
||||
let mut buffer = String::new();
|
||||
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = chunk.map_err(|e| {
|
||||
ToolError::ExternalService(format!(
|
||||
"[{}] Failed to read SSE chunk: {}",
|
||||
self.server_name, e
|
||||
))
|
||||
})?;
|
||||
|
||||
buffer.push_str(&String::from_utf8_lossy(&chunk));
|
||||
|
||||
if buffer.len() > MAX_SSE_BUFFER {
|
||||
return Err(ToolError::ExternalService(format!(
|
||||
"[{}] SSE response exceeded {} byte limit",
|
||||
self.server_name, MAX_SSE_BUFFER
|
||||
)));
|
||||
}
|
||||
|
||||
// Process only complete lines (terminated by \n). The last
|
||||
// element of split('\n') may be an incomplete line; keep it
|
||||
// in the buffer for the next chunk.
|
||||
let mut remaining_start = 0;
|
||||
let bytes = buffer.as_bytes();
|
||||
for (i, &b) in bytes.iter().enumerate() {
|
||||
if b == b'\n' {
|
||||
let line = &buffer[remaining_start..i];
|
||||
remaining_start = i + 1;
|
||||
|
||||
if let Some(json_str) = line.strip_prefix("data: ")
|
||||
&& let Ok(response) = serde_json::from_str::<McpResponse>(json_str)
|
||||
{
|
||||
return Ok(response);
|
||||
}
|
||||
}
|
||||
}
|
||||
// Keep only the unprocessed trailing fragment.
|
||||
if remaining_start > 0 {
|
||||
buffer = buffer[remaining_start..].to_string();
|
||||
}
|
||||
}
|
||||
|
||||
// Process any remaining data without a trailing newline.
|
||||
if let Some(json_str) = buffer.strip_prefix("data: ")
|
||||
&& let Ok(response) = serde_json::from_str::<McpResponse>(json_str.trim())
|
||||
{
|
||||
return Ok(response);
|
||||
}
|
||||
|
||||
Err(ToolError::ExternalService(format!(
|
||||
"[{}] No valid data in SSE response: {}",
|
||||
self.server_name, buffer
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
/// Sanitize an HTTP error body for safe inclusion in error messages.
|
||||
///
|
||||
/// When the body looks like a full HTML document (`<html` or `<!doctype`),
|
||||
/// strips all tags, collapsing whitespace. Non-HTML bodies are left
|
||||
/// intact. In both cases the result is truncated to 200 *characters*
|
||||
/// (char-boundary safe) so that large payloads don't bloat error messages.
|
||||
///
|
||||
/// See #263 — raw HTML error pages were propagating through the error
|
||||
/// chain into the web UI, causing a white screen.
|
||||
pub(crate) fn sanitize_error_body(body: &str) -> String {
|
||||
const MAX_CHARS: usize = 200;
|
||||
|
||||
// Only strip tags when the body looks like a full HTML document.
|
||||
// Plain text that happens to contain `<` / `>` (e.g. log lines,
|
||||
// comparison expressions) is left untouched.
|
||||
let lower = body.to_ascii_lowercase();
|
||||
let is_html_document = lower.contains("<html") || lower.contains("<!doctype");
|
||||
|
||||
let text = if is_html_document {
|
||||
let stripped = body
|
||||
.chars()
|
||||
.fold((String::new(), false), |(mut out, in_tag), c| {
|
||||
if c == '<' {
|
||||
(out, true)
|
||||
} else if c == '>' {
|
||||
(out, false)
|
||||
} else if !in_tag {
|
||||
out.push(c);
|
||||
(out, false)
|
||||
} else {
|
||||
(out, true)
|
||||
}
|
||||
})
|
||||
.0;
|
||||
stripped.split_whitespace().collect::<Vec<_>>().join(" ")
|
||||
} else {
|
||||
body.to_string()
|
||||
};
|
||||
|
||||
// Truncate at a char boundary (safe for multi-byte UTF-8).
|
||||
if text.chars().count() > MAX_CHARS {
|
||||
let byte_offset = text
|
||||
.char_indices()
|
||||
.nth(MAX_CHARS)
|
||||
.map(|(i, _)| i)
|
||||
.unwrap_or(text.len());
|
||||
format!("{}... ({} bytes total)", &text[..byte_offset], body.len())
|
||||
} else {
|
||||
text
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_sanitize_error_body_strips_html_tags() {
|
||||
let html =
|
||||
r#"<!DOCTYPE html><html><body><h1>422 Error</h1><p>Invalid token</p></body></html>"#;
|
||||
let result = sanitize_error_body(html);
|
||||
assert!(!result.contains('<'), "HTML tags must be stripped");
|
||||
assert!(!result.contains('>'), "HTML tags must be stripped");
|
||||
assert!(result.contains("422 Error"));
|
||||
assert!(result.contains("Invalid token"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sanitize_error_body_truncates_large_html_page() {
|
||||
let html = format!(
|
||||
"<html><body><p>{}</p></body></html>",
|
||||
"error detail ".repeat(50)
|
||||
);
|
||||
let result = sanitize_error_body(&html);
|
||||
assert!(result.contains("..."));
|
||||
assert!(result.contains("bytes total)"));
|
||||
assert!(!result.contains('<'));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sanitize_error_body_passes_short_plain_text() {
|
||||
assert_eq!(sanitize_error_body("Not Found"), "Not Found");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sanitize_error_body_truncates_long_plain_text() {
|
||||
let long = "x".repeat(300);
|
||||
let result = sanitize_error_body(&long);
|
||||
assert!(result.contains("..."));
|
||||
assert!(result.contains("300 bytes total)"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sanitize_error_body_multibyte_no_panic() {
|
||||
// 300 CJK characters = 900 bytes; truncation must land on a
|
||||
// char boundary, not in the middle of a multi-byte sequence.
|
||||
let cjk = "错误".repeat(150);
|
||||
let result = sanitize_error_body(&cjk);
|
||||
assert!(result.contains("..."));
|
||||
// Must be valid UTF-8 (would have panicked otherwise).
|
||||
assert!(result.is_char_boundary(result.len()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sanitize_error_body_strips_uppercase_html() {
|
||||
let html = "<HTML><BODY><H1>500 Internal Server Error</H1></BODY></HTML>";
|
||||
let result = sanitize_error_body(html);
|
||||
assert!(
|
||||
!result.contains('<'),
|
||||
"uppercase HTML tags must be stripped"
|
||||
);
|
||||
assert!(result.contains("500 Internal Server Error"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sanitize_error_body_preserves_angle_brackets_in_non_html() {
|
||||
let text = "value < 10 and value > 0";
|
||||
assert_eq!(sanitize_error_body(text), text);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sanitize_error_body_empty_string() {
|
||||
assert_eq!(sanitize_error_body(""), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_new_creates_transport() {
|
||||
let transport = HttpMcpTransport::new("http://localhost:8080", "test");
|
||||
assert_eq!(transport.server_url(), "http://localhost:8080");
|
||||
assert!(transport.session_manager().is_none());
|
||||
assert!(transport.custom_headers.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_supports_http_features() {
|
||||
let http_transport = HttpMcpTransport::new("http://localhost:8080", "test");
|
||||
assert!(http_transport.supports_http_features());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_with_session_manager() {
|
||||
let session_manager = Arc::new(McpSessionManager::new());
|
||||
let transport = HttpMcpTransport::new("http://localhost:8080", "test")
|
||||
.with_session_manager(session_manager.clone());
|
||||
assert!(transport.session_manager().is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_with_custom_headers() {
|
||||
let mut headers = HashMap::new();
|
||||
headers.insert("X-Custom".to_string(), "value".to_string());
|
||||
let transport =
|
||||
HttpMcpTransport::new("http://localhost:8080", "test").with_custom_headers(headers);
|
||||
assert_eq!(transport.custom_headers.get("X-Custom").unwrap(), "value");
|
||||
}
|
||||
}
|
||||
@@ -4,6 +4,8 @@
|
||||
//! additional capabilities through a standardized protocol.
|
||||
//!
|
||||
//! Supports both local (unauthenticated) and hosted (OAuth-authenticated) servers.
|
||||
//! Transport options include HTTP (Streamable HTTP / SSE), stdio (subprocess),
|
||||
//! and Unix domain sockets.
|
||||
//!
|
||||
//! ## Usage
|
||||
//!
|
||||
@@ -29,11 +31,19 @@
|
||||
pub mod auth;
|
||||
mod client;
|
||||
pub mod config;
|
||||
pub(crate) mod http_transport;
|
||||
pub(crate) mod process;
|
||||
mod protocol;
|
||||
pub mod session;
|
||||
pub(crate) mod stdio_transport;
|
||||
pub(crate) mod transport;
|
||||
#[cfg(unix)]
|
||||
pub(crate) mod unix_transport;
|
||||
|
||||
pub use auth::{is_authenticated, refresh_access_token};
|
||||
pub use client::McpClient;
|
||||
pub use config::{McpServerConfig, McpServersFile, OAuthConfig};
|
||||
pub use process::McpProcessManager;
|
||||
pub use protocol::{InitializeResult, McpRequest, McpResponse, McpTool};
|
||||
pub use session::McpSessionManager;
|
||||
pub use transport::McpTransport;
|
||||
|
||||
@@ -0,0 +1,206 @@
|
||||
//! MCP stdio process manager.
|
||||
//!
|
||||
//! Manages the lifecycle of MCP servers running as child processes.
|
||||
//! Handles spawning, shutdown, and crash recovery with exponential backoff.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::tools::mcp::stdio_transport::StdioMcpTransport;
|
||||
use crate::tools::mcp::transport::McpTransport;
|
||||
use crate::tools::tool::ToolError;
|
||||
|
||||
/// Configuration for spawning a stdio MCP server.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct StdioSpawnConfig {
|
||||
pub command: String,
|
||||
pub args: Vec<String>,
|
||||
pub env: HashMap<String, String>,
|
||||
}
|
||||
|
||||
/// Manages stdio MCP server processes.
|
||||
///
|
||||
/// Handles spawning, tracking, and shutdown of child processes.
|
||||
pub struct McpProcessManager {
|
||||
transports: RwLock<HashMap<String, Arc<StdioMcpTransport>>>,
|
||||
configs: RwLock<HashMap<String, StdioSpawnConfig>>,
|
||||
}
|
||||
|
||||
impl McpProcessManager {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
transports: RwLock::new(HashMap::new()),
|
||||
configs: RwLock::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Spawn a new stdio MCP server process.
|
||||
pub async fn spawn_stdio(
|
||||
&self,
|
||||
name: impl Into<String>,
|
||||
command: impl Into<String>,
|
||||
args: Vec<String>,
|
||||
env: HashMap<String, String>,
|
||||
) -> Result<Arc<StdioMcpTransport>, ToolError> {
|
||||
let name = name.into();
|
||||
let command = command.into();
|
||||
|
||||
// Store config for potential restart
|
||||
self.configs.write().await.insert(
|
||||
name.clone(),
|
||||
StdioSpawnConfig {
|
||||
command: command.clone(),
|
||||
args: args.clone(),
|
||||
env: env.clone(),
|
||||
},
|
||||
);
|
||||
|
||||
let transport = Arc::new(StdioMcpTransport::spawn(&name, &command, args, env).await?);
|
||||
|
||||
self.transports
|
||||
.write()
|
||||
.await
|
||||
.insert(name, Arc::clone(&transport));
|
||||
|
||||
Ok(transport)
|
||||
}
|
||||
|
||||
/// Get a transport by server name.
|
||||
pub async fn get(&self, name: &str) -> Option<Arc<StdioMcpTransport>> {
|
||||
self.transports.read().await.get(name).cloned()
|
||||
}
|
||||
|
||||
/// Shut down all managed transports.
|
||||
pub async fn shutdown_all(&self) {
|
||||
let transports: Vec<(String, Arc<StdioMcpTransport>)> = {
|
||||
let mut map = self.transports.write().await;
|
||||
map.drain().collect()
|
||||
};
|
||||
|
||||
for (name, transport) in transports {
|
||||
if let Err(e) = transport.shutdown().await {
|
||||
tracing::warn!("Failed to shut down MCP stdio server '{}': {}", name, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Shut down a specific transport by name.
|
||||
pub async fn shutdown(&self, name: &str) -> Result<(), ToolError> {
|
||||
let transport = self.transports.write().await.remove(name);
|
||||
|
||||
if let Some(transport) = transport {
|
||||
transport.shutdown().await?;
|
||||
}
|
||||
|
||||
self.configs.write().await.remove(name);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Attempt to restart a crashed transport with exponential backoff.
|
||||
///
|
||||
/// Tries up to 5 times with delays of 1s, 2s, 4s, 8s, 16s (total: 31s max wait).
|
||||
pub async fn try_restart(&self, name: &str) -> Result<Arc<StdioMcpTransport>, ToolError> {
|
||||
let config = self
|
||||
.configs
|
||||
.read()
|
||||
.await
|
||||
.get(name)
|
||||
.cloned()
|
||||
.ok_or_else(|| {
|
||||
ToolError::ExternalService(format!(
|
||||
"No spawn config for MCP server '{}', cannot restart",
|
||||
name
|
||||
))
|
||||
})?;
|
||||
|
||||
// Shut down and remove old transport to avoid orphaning a wedged process.
|
||||
if let Some(old_transport) = self.transports.write().await.remove(name) {
|
||||
let _ = old_transport.shutdown().await;
|
||||
}
|
||||
|
||||
let max_retries = 5;
|
||||
let mut last_err = None;
|
||||
|
||||
for attempt in 0..max_retries {
|
||||
let delay = Duration::from_secs(1 << attempt);
|
||||
tokio::time::sleep(delay).await;
|
||||
|
||||
match StdioMcpTransport::spawn(
|
||||
name,
|
||||
&config.command,
|
||||
config.args.clone(),
|
||||
config.env.clone(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(transport) => {
|
||||
let transport = Arc::new(transport);
|
||||
self.transports
|
||||
.write()
|
||||
.await
|
||||
.insert(name.to_string(), Arc::clone(&transport));
|
||||
tracing::info!(
|
||||
"MCP stdio server '{}' restarted after {} attempt(s)",
|
||||
name,
|
||||
attempt + 1
|
||||
);
|
||||
return Ok(transport);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Restart attempt {}/{} for MCP server '{}' failed: {}",
|
||||
attempt + 1,
|
||||
max_retries,
|
||||
name,
|
||||
e
|
||||
);
|
||||
last_err = Some(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Err(last_err.unwrap_or_else(|| {
|
||||
ToolError::ExternalService(format!(
|
||||
"Failed to restart MCP server '{}' after {} attempts",
|
||||
name, max_retries
|
||||
))
|
||||
}))
|
||||
}
|
||||
|
||||
/// Get names of all managed transports.
|
||||
pub async fn managed_servers(&self) -> Vec<String> {
|
||||
self.transports.read().await.keys().cloned().collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for McpProcessManager {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_new_creates_empty_manager() {
|
||||
let _manager = McpProcessManager::new();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_managed_servers_returns_empty_list_initially() {
|
||||
let manager = McpProcessManager::new();
|
||||
let servers = manager.managed_servers().await;
|
||||
assert!(servers.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_shutdown_all_on_empty_manager_does_not_panic() {
|
||||
let manager = McpProcessManager::new();
|
||||
manager.shutdown_all().await;
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user